broccoli-ml 7.0.0__py3-none-any.whl → 9.0.0__py3-none-any.whl

broccoli/linear.py

@@ -1,4 +1,7 @@
 import math
+import random
+from typing import Union, List, Iterable
+
 import torch
 from torch import nn
 from torch.nn import functional as F
@@ -136,3 +139,136 @@ class WeightNormedLinear(nn.Module):
             f"WeightNormedLinear(in_features={self.in_features},"
             f"out_features={self.out_features}, bias={self.use_bias})"
         )
+
+
+class RecyclingLinear(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = True,
+        row_recycling_rate: float = 0.0,
+        column_recycling_rate: float = 0.0,
+    ):
+        super().__init__()
+        self.linear = nn.Linear(in_features, out_features, bias=bias)
+        self.row_recycling_rate = row_recycling_rate
+        self.column_recycling_rate = column_recycling_rate
+        self.optimisers = []
+
+    def register_optimiser(self, optimiser: torch.optim.Optimizer):
+        self.optimisers.append(optimiser)
+
+    def forward(self, x):
+        if self.training and self.optimisers:
+
+            if self.row_recycling_rate > 0:
+                probs = torch.rand(self.linear.out_features, device=x.device)
+                mask = probs < self.row_recycling_rate
+                if mask.any():
+                    # nonzero returns [N, 1], squeeze to get [N]
+                    indices = torch.nonzero(mask).squeeze(-1)
+                    self.reset_rows(indices, self.optimisers)
+
+            if self.column_recycling_rate > 0:
+                probs = torch.rand(self.linear.in_features, device=x.device)
+                mask = probs < self.column_recycling_rate
+                if mask.any():
+                    indices = torch.nonzero(mask).squeeze(-1)
+                    self.reset_columns(indices, self.optimisers)
+
+        return self.linear(x)
+
+    def reset_rows(
+        self,
+        indices: Iterable[int],
+        optimisers: Union[
+            List[torch.optim.Optimizer], torch.optim.Optimizer, None
+        ] = None,
+    ):
+        """
+        Update some of the weight rows to be equal to the mean of all weight rows.
+        """
+        if optimisers is None:
+            optimisers = []
+        if not isinstance(optimisers, list):
+            optimisers = [optimisers]
+
+        device = self.linear.weight.device
+        idx_tensor = torch.as_tensor(list(indices), dtype=torch.long, device=device)
+
+        if idx_tensor.numel() == 0:
+            return
+
+        with torch.no_grad():
+            # Calculate mean of all rows including the rows to be reset
+            mean_vector = self.linear.weight.data.mean(
+                dim=0, keepdim=True
+            )  # [1, in_features]
+            update_data = mean_vector.expand(idx_tensor.size(0), -1)
+            self.linear.weight.data[idx_tensor] = update_data
+
+            if self.linear.bias is not None:
+                self.linear.bias.data[idx_tensor] = 0.0
+
+            self._reset_optim_state(self.linear.weight, idx_tensor, optimisers, dim=0)
+            if self.linear.bias is not None:
+                self._reset_optim_state(self.linear.bias, idx_tensor, optimisers, dim=0)
+
+    def reset_columns(
+        self,
+        indices: Iterable[int],
+        optimisers: Union[
+            List[torch.optim.Optimizer], torch.optim.Optimizer, None
+        ] = None,
+    ):
+        """
+        Update some of the weight columns to be random as though reinitialised.
+        """
+        if optimisers is None:
+            optimisers = []
+        if not isinstance(optimisers, list):
+            optimisers = [optimisers]
+
+        device = self.linear.weight.device
+        idx_tensor = torch.as_tensor(list(indices), dtype=torch.long, device=device)
+
+        if idx_tensor.numel() == 0:
+            return
+
+        with torch.no_grad():
+            # 1. Generate Random Columns
+            # Shape: [out_features, N_indices]
+            weights = self.linear.weight.data
+            stdv = 1.0 / math.sqrt(weights.size(1))
+
+            # Generate [Rows, N] block
+            random_weights = torch.rand(
+                weights.size(0), idx_tensor.size(0), device=device
+            )
+            random_weights = (random_weights - 0.5) * 2.0 * stdv
+
+            # 2. Update Weights (One-shot)
+            # We assign into the columns specified by idx_tensor
+            self.linear.weight.data[:, idx_tensor] = random_weights
+
+            # 3. Update Optimizers
+            # Bias is untouched by column resets (bias is shape [Out], cols are [In])
+            self._reset_optim_state(self.linear.weight, idx_tensor, optimisers, dim=1)
+
+    def _reset_optim_state(self, param, idx_tensor, optimisers, dim):
+        """
+        Zeroes out the optimizer state for the given indices in a single operation.
+        """
+        for optimiser in optimisers:
+            if param not in optimiser.state:
+                continue
+            state = optimiser.state[param]
+
+            for _, buffer in state.items():
+                if torch.is_tensor(buffer) and buffer.shape == param.shape:
+                    # Vectorized zeroing
+                    if dim == 0:
+                        buffer[idx_tensor] = 0.0
+                    else:
+                        buffer[:, idx_tensor] = 0.0

broccoli/transformer.py

@@ -325,6 +325,8 @@ class FeedforwardBlock(nn.Module):
         activation=nn.ReLU,
         activation_kwargs=None,
         dropout=0.0,
+        inner_dropout=None,
+        outer_dropout=None,
         linear_module_up=nn.Linear,
         linear_module_down=nn.Linear,
         pre_norm=True,
@@ -354,7 +356,12 @@ class FeedforwardBlock(nn.Module):
         else:
             self.activation = activation()
 
-        self.dropout = nn.Dropout(dropout)
+        self.inner_dropout = nn.Dropout(
+            inner_dropout if inner_dropout is not None else dropout
+        )
+        self.outer_dropout = nn.Dropout(
+            outer_dropout if outer_dropout is not None else dropout
+        )
 
         self.max_features = (
             2 * ratio * output_features
@@ -367,9 +374,10 @@ class FeedforwardBlock(nn.Module):
                 nn.LayerNorm(input_features) if pre_norm else nn.Identity(),
                 linear_module_up(input_features, self.max_features),
                 self.activation,
+                self.inner_dropout,
                 nn.LayerNorm(ratio * output_features) if normformer else nn.Identity(),
                 linear_module_down(ratio * output_features, output_features),
-                self.dropout,
+                self.outer_dropout,
             ]
         )
 
@@ -422,7 +430,9 @@ class TransformerBlock(nn.Module):
         ff_linear_module_up=None,
         ff_linear_module_down=None,
         msa_scaling="d",
-        mlp_dropout=0.0,
+        ff_dropout=0.0,
+        ff_inner_dropout=0.0,
+        ff_outer_dropout=0.0,
         msa_dropout=0.0,
         identity_probability=0.0,
         causal=False,
@@ -484,7 +494,9 @@ class TransformerBlock(nn.Module):
             d_model,
             activation=activation,
             activation_kwargs=activation_kwargs,
-            dropout=mlp_dropout,
+            dropout=ff_dropout,
+            inner_dropout=ff_inner_dropout,
+            outer_dropout=ff_outer_dropout,
             linear_module_up=(
                 ff_linear_module_up
                 if ff_linear_module_up is not None
@@ -567,7 +579,9 @@ class TransformerEncoder(nn.Module):
         activation_kwargs: Optional[dict] = None,
         ff_linear_module_up=None,
         ff_linear_module_down=None,
-        mlp_dropout=0.0,
+        ff_dropout=0.0,
+        ff_inner_dropout=0.0,
+        ff_outer_dropout=0.0,
         msa_dropout=0.0,
         stochastic_depth=0.0,
         causal=False,
@@ -629,7 +643,7 @@ class TransformerEncoder(nn.Module):
         else:
             self.absolute_position_embedding = None
 
-        self.mlp_dropout = mlp_dropout
+        self.mlp_dropout = ff_dropout
         self.msa_dropout = msa_dropout
         self.stochastic_depth = stochastic_depth
 
@@ -658,7 +672,9 @@ class TransformerEncoder(nn.Module):
                     ff_linear_module_up=ff_linear_module_up,
                     ff_linear_module_down=ff_linear_module_down,
                     msa_scaling=msa_scaling,
-                    mlp_dropout=mlp_dropout,
+                    ff_dropout=ff_dropout,
+                    ff_inner_dropout=ff_inner_dropout,
+                    ff_outer_dropout=ff_outer_dropout,
                     msa_dropout=msa_dropout,
                     identity_probability=self.stochastic_depth_probabilities[i],
                     causal=causal,

broccoli/vit.py

@@ -161,7 +161,9 @@ class ViTEncoder(nn.Module):
         transformer_initial_ff_residual_path=True,
         transformer_initial_ff_linear_module_up=None,
         transformer_initial_ff_linear_module_down=None,
-        transformer_initial_ff_mlp_dropout=None,
+        transformer_initial_ff_dropout=None,
+        transformer_initial_ff_inner_dropout=None,
+        transformer_initial_ff_outer_dropout=None,
         transformer_pre_norm=True,
         transformer_normformer=False,
         transformer_post_norm=False,
@@ -178,7 +180,9 @@ class ViTEncoder(nn.Module):
         transformer_ff_linear_module_up=None,
         transformer_ff_linear_module_down=None,
         transformer_msa_scaling="d",
-        transformer_mlp_dropout=0.0,
+        transformer_ff_dropout=0.0,
+        transformer_ff_inner_dropout=0.0,
+        transformer_ff_outer_dropout=0.0,
         transformer_msa_dropout=0.1,
         transformer_stochastic_depth=0.1,
         transformer_checkpoint_ff=True,
@@ -333,7 +337,9 @@ class ViTEncoder(nn.Module):
                 ff_linear_module_up=transformer_ff_linear_module_up,
                 ff_linear_module_down=transformer_ff_linear_module_down,
                 msa_scaling=transformer_msa_scaling,
-                mlp_dropout=transformer_mlp_dropout,
+                ff_dropout=transformer_ff_dropout,
+                ff_inner_dropout=transformer_ff_inner_dropout,
+                ff_outer_dropout=transformer_ff_outer_dropout,
                 msa_dropout=transformer_msa_dropout,
                 stochastic_depth=transformer_stochastic_depth,
                 causal=False,
@@ -357,9 +363,21 @@ class ViTEncoder(nn.Module):
                 activation_kwargs=transformer_activation_kwargs,
                 dropout=(
                     # First truthy assigned value
-                    transformer_initial_ff_mlp_dropout
-                    if transformer_initial_ff_mlp_dropout is not None
-                    else transformer_mlp_dropout
+                    transformer_initial_ff_dropout
+                    if transformer_initial_ff_dropout is not None
+                    else transformer_ff_dropout
+                ),
+                inner_dropout=(
+                    # First truthy assigned value
+                    transformer_initial_ff_inner_dropout
+                    if transformer_initial_ff_inner_dropout is not None
+                    else transformer_ff_inner_dropout
+                ),
+                outer_dropout=(
+                    # First truthy assigned value
+                    transformer_initial_ff_outer_dropout
+                    if transformer_initial_ff_outer_dropout is not None
+                    else transformer_ff_outer_dropout
                 ),
                 linear_module_up=(
                     # First truthy assigned value
@@ -441,7 +459,9 @@ class ViT(nn.Module):
         transformer_initial_ff_residual_path=True,
         transformer_initial_ff_linear_module_up=None,
         transformer_initial_ff_linear_module_down=None,
-        transformer_initial_ff_mlp_dropout=None,
+        transformer_initial_ff_dropout=None,
+        transformer_initial_ff_inner_dropout=None,
+        transformer_initial_ff_outer_dropout=None,
         transformer_pre_norm=True,
         transformer_normformer=False,
         transformer_post_norm=False,
@@ -458,7 +478,9 @@ class ViT(nn.Module):
         transformer_ff_linear_module_up=None,
         transformer_ff_linear_module_down=None,
         transformer_msa_scaling="d",
-        transformer_mlp_dropout=0.0,
+        transformer_ff_dropout=0.0,
+        transformer_ff_inner_dropout=0.0,
+        transformer_ff_outer_dropout=0.0,
         transformer_msa_dropout=0.1,
         transformer_stochastic_depth=0.1,
         transformer_checkpoint_ff=True,
@@ -508,7 +530,9 @@ class ViT(nn.Module):
             transformer_initial_ff_residual_path=transformer_initial_ff_residual_path,
             transformer_initial_ff_linear_module_up=transformer_initial_ff_linear_module_up,
             transformer_initial_ff_linear_module_down=transformer_initial_ff_linear_module_down,
-            transformer_initial_ff_mlp_dropout=transformer_initial_ff_mlp_dropout,
+            transformer_initial_ff_dropout=transformer_initial_ff_dropout,
+            transformer_initial_ff_inner_dropout=transformer_initial_ff_inner_dropout,
+            transformer_initial_ff_outer_dropout=transformer_initial_ff_outer_dropout,
             transformer_pre_norm=transformer_pre_norm,
             transformer_normformer=transformer_normformer,
             transformer_post_norm=transformer_post_norm,
@@ -525,7 +549,9 @@ class ViT(nn.Module):
             transformer_ff_linear_module_up=transformer_ff_linear_module_up,
             transformer_ff_linear_module_down=transformer_ff_linear_module_down,
             transformer_msa_scaling=transformer_msa_scaling,
-            transformer_mlp_dropout=transformer_mlp_dropout,
+            transformer_ff_dropout=transformer_ff_dropout,
+            transformer_ff_inner_dropout=transformer_ff_inner_dropout,
+            transformer_ff_outer_dropout=transformer_ff_outer_dropout,
             transformer_msa_dropout=transformer_msa_dropout,
             transformer_stochastic_depth=transformer_stochastic_depth,
             transformer_checkpoint_ff=transformer_checkpoint_ff,

{broccoli_ml-7.0.0.dist-info → broccoli_ml-9.0.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: broccoli-ml
-Version: 7.0.0
+Version: 9.0.0
 Summary: Some useful Pytorch models, circa 2025
 License: MIT
 Author: Nicholas Bailey

{broccoli_ml-7.0.0.dist-info → broccoli_ml-9.0.0.dist-info}/RECORD

@@ -1,13 +1,13 @@
 broccoli/__init__.py,sha256=tmyspsVxqPZHRQCY_NRwpW4SMNBbtE8E_8z7l-SAzSo,127
 broccoli/activation.py,sha256=-Jf30C6iGqWCorC9HEGn2oduWwjeaCAxGLUUYIy1zX8,3438
 broccoli/cnn.py,sha256=WjoPDSpe3ttwxCBNfCVRdaCHvbeZ7G-a5_i8fUsK_d8,4889
-broccoli/linear.py,sha256=Y7s-DzcwsOipRboNHc4HTScw4mJRalNoVFsNcxOB6a4,4872
+broccoli/linear.py,sha256=XaGHZguvK-7hvtIt07zo8uQZBQvS7oMD2K9nPvyYJLE,9769
 broccoli/rope.py,sha256=GRqApBNmYCFaDak0WL1xE_BC5CTTYKQU_PBdeTcQcjc,12557
 broccoli/tensor.py,sha256=um8mrxkYbvNDo-QvHlmJm8Aw6qcngOlUZPoAk_PMReA,4480
-broccoli/transformer.py,sha256=-NRO8mzvOkQPCuTrF6OTDQ8sUIQ3_j4HxP-NfTUxu10,23636
+broccoli/transformer.py,sha256=Rozh0hExHjwGvvKbMeZfLoB95dDKyDn3X6o1Ms26aAI,24241
 broccoli/utils.py,sha256=oOWzn6dJ5nC_9r4zq0emmfmaYACJXJNFS48AOpW2jqc,358
-broccoli/vit.py,sha256=9oyh76ulmX5lDPMCDicQhhqm8RYCvJIgAJkDbYRVdi4,20873
-broccoli_ml-7.0.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
-broccoli_ml-7.0.0.dist-info/METADATA,sha256=OdAph_5ItouyYHeZwVgbrTiyexljFb9uKNLNWlnR4wM,1368
-broccoli_ml-7.0.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-broccoli_ml-7.0.0.dist-info/RECORD,,
+broccoli/vit.py,sha256=sC6K3FK3a8ojOgvNWSWhuZHBtnFrrTQbsDdlagcKJH4,22224
+broccoli_ml-9.0.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
+broccoli_ml-9.0.0.dist-info/METADATA,sha256=ecQ2BRxtzmNSO2CMAp2rcNRq9L37urE_pKdsPf-jJKs,1368
+broccoli_ml-9.0.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+broccoli_ml-9.0.0.dist-info/RECORD,,