broccoli-ml 15.2.0__tar.gz → 15.4.0__tar.gz

{broccoli_ml-15.2.0 → broccoli_ml-15.4.0}/PKG-INFO

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

{broccoli_ml-15.2.0 → broccoli_ml-15.4.0}/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.2.0 → broccoli_ml-15.4.0}/broccoli/vit.py

@@ -520,10 +520,10 @@ class ViT(nn.Module):
                 "SwiGLU": SwiGLU,
             }[transformer_activation]
 
-        # Set alpha and beta according to Microsoft's DeepNorm
-        self.alpha = (2 * transformer_layers) ** 0.25
-        # beta is only needed for very deep models
-        self.beta = 1 if transformer_layers < 50 else (8 * transformer_layers) ** -0.25
+        # Set alpha according to Microsoft's DeepNorm if layers > 50
+        if transformer_layers > 50:
+            self.alpha = (2 * transformer_layers) ** 0.25
+            # beta is not needed as we norm the Q and K vectors in MSA!
 
         self.encoder = ViTEncoder(
             input_size=input_size,

{broccoli_ml-15.2.0 → broccoli_ml-15.4.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "broccoli-ml"
-version = "15.2.0"
+version = "15.4.0"
 description = "Some useful Pytorch models, circa 2025"
 authors = [
     {name = "Nicholas Bailey"}