broccoli-ml 15.1.0__py3-none-any.whl → 15.3.0__py3-none-any.whl

broccoli/transformer.py

@@ -354,9 +354,17 @@ class MHAttention(nn.Module):
         self.q_proj.reset_parameters()
         self.k_proj.reset_parameters()
         self.v_proj.reset_parameters()
-        scale_parameters(self.v_proj, self.beta)  # per Microsoft DeepNet
+        scale_parameters(
+            self.v_proj,
+            math.sqrt(6)
+            * self.beta,  # sqrt(6) to compensate for PyTorch tiny default init
+        )
         self.out_proj.reset_parameters()
-        scale_parameters(self.out_proj, self.beta)  # per Microsoft DeepNet
+        scale_parameters(
+            self.out_proj,
+            math.sqrt(6)
+            * self.beta,  # sqrt(6) to compensate for PyTorch tiny default init
+        )
 
         if self.talking_heads:
             # Initialize close to identity
@@ -473,8 +481,16 @@ class FeedforwardBlock(nn.Module):
             if hasattr(module, "reset_parameters"):
                 module.reset_parameters()
 
-        scale_parameters(self.linear_in, self.beta)  # per Microsoft DeepNet
-        scale_parameters(self.linear_out, self.beta)
+        scale_parameters(
+            self.linear_in,
+            math.sqrt(6)
+            * self.beta,  # sqrt(6) to compensate for PyTorch tiny default init
+        )
+        scale_parameters(
+            self.linear_out,
+            math.sqrt(6)
+            * self.beta,  # sqrt(6) to compensate for PyTorch tiny default init
+        )
 
 
 class EncoderBlock(nn.Module):

{broccoli_ml-15.1.0.dist-info → broccoli_ml-15.3.0.dist-info}/METADATA

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

{broccoli_ml-15.1.0.dist-info → broccoli_ml-15.3.0.dist-info}/RECORD

@@ -4,10 +4,10 @@ broccoli/cnn.py,sha256=WjoPDSpe3ttwxCBNfCVRdaCHvbeZ7G-a5_i8fUsK_d8,4889
 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=Uuz_jCMZRg6GY2DmW3-Tn47gV9a-xkGVN3xQ5BYFM5w,27784
+broccoli/transformer.py,sha256=XCr3fkDrYvNEA3m5rdVDHU5qbf1IYJQS1Cjgmn5RX-o,28177
 broccoli/utils.py,sha256=oOWzn6dJ5nC_9r4zq0emmfmaYACJXJNFS48AOpW2jqc,358
 broccoli/vit.py,sha256=9z1MJce38dppX35OGJvmnNIOnt-euMVV4n160W_aQYU,22799
-broccoli_ml-15.1.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
-broccoli_ml-15.1.0.dist-info/METADATA,sha256=lQjyxZ5c98_f4DlTJ_0b0dIbbsa2FfWTx56EG3SXGws,1369
-broccoli_ml-15.1.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-broccoli_ml-15.1.0.dist-info/RECORD,,
+broccoli_ml-15.3.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
+broccoli_ml-15.3.0.dist-info/METADATA,sha256=StOW7vZd1c4xNddKcA8XJ-zbXgUnAsrQlg9G97wT-2w,1369
+broccoli_ml-15.3.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+broccoli_ml-15.3.0.dist-info/RECORD,,