broccoli-ml 9.5.1__py3-none-any.whl → 9.6.0__py3-none-any.whl

broccoli/linear.py

@@ -193,7 +193,12 @@ class RecyclingLinear(nn.Module):
             multipliers = [a / b for a, b in pairs if b != 0.0]
             return min(multipliers) if multipliers else 0.0
 
-    def reset_rows(self, indices):
+    def reset_rows(self, indices, incoming_data=None):
+        """
+        Resets rows.
+        If incoming_data is provided, resets to the centroid (mean) of that data.
+        If not, resets to the mean of existing weights.
+        """
         if not torch.is_tensor(indices):
             idx_tensor = torch.as_tensor(
                 list(indices), dtype=torch.long, device=self.linear.weight.device
@@ -201,24 +206,24 @@ class RecyclingLinear(nn.Module):
         else:
             idx_tensor = indices
 
-        if idx_tensor.size(0):
-            value_indices = indices
-            centred_value_weights = self._mean_value_weights()
-            centred_value_weights = centred_value_weights.expand(indices.size(0), -1)
-            if self.xglu:
-                gate_indices = indices
-                value_indices = indices + (self.linear.out_features // 2)
-                centred_gate_weights = self._mean_gate_weights()
-                centred_gate_weights = centred_gate_weights.expand(indices.size(0), -1)
-                self._update_weights(
-                    gate_indices, 0, centred_gate_weights, self.optimisers  # dim
-                )
-            self._update_weights(
-                value_indices, 0, centred_value_weights, self.optimisers
-            )
-        else:
+        if idx_tensor.numel() == 0:
             return
 
+        if incoming_data is not None:
+            target_center = self._mean_input_weights(incoming_data)
+        else:
+            target_center = self._mean_value_weights()
+
+        target_center = target_center.expand(idx_tensor.size(0), -1)
+
+        if self.xglu:
+            gate_indices = idx_tensor
+            value_indices = idx_tensor + (self.linear.out_features // 2)
+            self._update_weights(gate_indices, 0, target_center, self.optimisers)
+            self._update_weights(value_indices, 0, target_center, self.optimisers)
+        else:
+            self._update_weights(idx_tensor, 0, target_center, self.optimisers)
+
     def reset_columns(self, indices):
         if not torch.is_tensor(indices):
             idx_tensor = torch.as_tensor(
@@ -281,6 +286,17 @@ class RecyclingLinear(nn.Module):
         random_weights *= 2.0 * stdv  # Range [-stdv, +stdv]
         return random_weights
 
+    def _mean_input_weights(self, input):
+        reduce_dims = list(range(input.ndim - 1))
+        data_mean = input.detach().mean(dim=reduce_dims, keepdim=True)
+
+        weights = self.linear.weight.data
+        stdv = 1.0 / math.sqrt(weights.size(1))
+        data_norm = data_mean.std() + 1e-6
+        scale_factor = stdv / data_norm
+
+        return data_mean * scale_factor
+
     def _mean_value_weights(self):
         """
         Only used when self.xglu

broccoli/transformer.py

@@ -411,7 +411,7 @@ class FeedforwardBlock(nn.Module):
         # Recycle weights if using recycling linear layers
         if self.training and self.recycling_enabled:
             indices = self.linear_out.get_reset_indices(1)
-            self.linear_in.reset_rows(indices)
+            self.linear_in.reset_rows(indices, incoming_data=x)
             self.linear_out.reset_columns(indices)
 
         if self.checkpoint:

{broccoli_ml-9.5.1.dist-info → broccoli_ml-9.6.0.dist-info}/METADATA

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

{broccoli_ml-9.5.1.dist-info → broccoli_ml-9.6.0.dist-info}/RECORD

@@ -1,13 +1,13 @@
 broccoli/__init__.py,sha256=tmyspsVxqPZHRQCY_NRwpW4SMNBbtE8E_8z7l-SAzSo,127
 broccoli/activation.py,sha256=nrpTOrpg9k23_E4AJWy7VlXXAJCtCJCOR-TonEWJr04,3218
 broccoli/cnn.py,sha256=WjoPDSpe3ttwxCBNfCVRdaCHvbeZ7G-a5_i8fUsK_d8,4889
-broccoli/linear.py,sha256=i4U7ZC4ZWEH82YpDasx0Qs1pc3gkyL-3ajuyKCbsGTM,12649
+broccoli/linear.py,sha256=W-3aNpBjd_0xRyzbCKkmg4H1qmslQOIQhB-WDDay2nM,13125
 broccoli/rope.py,sha256=GRqApBNmYCFaDak0WL1xE_BC5CTTYKQU_PBdeTcQcjc,12557
 broccoli/tensor.py,sha256=um8mrxkYbvNDo-QvHlmJm8Aw6qcngOlUZPoAk_PMReA,4480
-broccoli/transformer.py,sha256=ULk-QQX3hAI14-aCKhp9QSebzX4KUjlisEGup2Eycck,25565
+broccoli/transformer.py,sha256=vmAgFuevFn9ekRqq2Fbiblzr5DOrucPvlHF87l0vUuM,25582
 broccoli/utils.py,sha256=oOWzn6dJ5nC_9r4zq0emmfmaYACJXJNFS48AOpW2jqc,358
 broccoli/vit.py,sha256=sC6K3FK3a8ojOgvNWSWhuZHBtnFrrTQbsDdlagcKJH4,22224
-broccoli_ml-9.5.1.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
-broccoli_ml-9.5.1.dist-info/METADATA,sha256=HXRWnuc_-Gs_g37_RP3-POTLmi7sZamlzYv5SJEun1Y,1368
-broccoli_ml-9.5.1.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-broccoli_ml-9.5.1.dist-info/RECORD,,
+broccoli_ml-9.6.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
+broccoli_ml-9.6.0.dist-info/METADATA,sha256=zzGvnF60IZx7Htq3-R91wsQ0ey3VmSNde2XQZY0H11U,1368
+broccoli_ml-9.6.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+broccoli_ml-9.6.0.dist-info/RECORD,,