broccoli-ml 15.4.1__py3-none-any.whl → 15.5.0__py3-none-any.whl

broccoli/transformer.py

@@ -122,9 +122,6 @@ class MHAttention(nn.Module):
 
         self.head_dim = self.embed_dim // self.n_heads
 
-        self.query_norm = nn.RMSNorm(self.head_dim)
-        self.key_norm = nn.RMSNorm(self.head_dim)
-
         if self.scaling == "sqrtd":
             self.scaling_factor = 1 / math.sqrt(self.head_dim)
         elif self.scaling == "d":
@@ -229,8 +226,8 @@ class MHAttention(nn.Module):
         freqs = self.rotary_embedding.get_axial_freqs(*self.source_size)
 
         # norm Qs/Ks to protect axial rope, like https://arxiv.org/abs/2302.05442
-        q_img = apply_rotary_emb(freqs, self.query_norm(q_img))
-        k_img = apply_rotary_emb(freqs, self.key_norm(k_img))
+        q_img = apply_rotary_emb(freqs, q_img)
+        k_img = apply_rotary_emb(freqs, k_img)
 
         q_img = rearrange(
             q_img,
@@ -354,17 +351,9 @@ class MHAttention(nn.Module):
         self.q_proj.reset_parameters()
         self.k_proj.reset_parameters()
         self.v_proj.reset_parameters()
-        scale_parameters(
-            self.v_proj,
-            math.sqrt(6)
-            * self.beta,  # sqrt(6) to compensate for PyTorch tiny default init
-        )
+        scale_parameters(self.v_proj, self.beta)
         self.out_proj.reset_parameters()
-        scale_parameters(
-            self.out_proj,
-            math.sqrt(6)
-            * self.beta,  # sqrt(6) to compensate for PyTorch tiny default init
-        )
+        scale_parameters(self.out_proj, self.beta)
 
         if self.talking_heads:
             # Initialize close to identity
@@ -481,16 +470,8 @@ class FeedforwardBlock(nn.Module):
             if hasattr(module, "reset_parameters"):
                 module.reset_parameters()
 
-        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
-        )
+        scale_parameters(self.linear_in, self.beta)
+        scale_parameters(self.linear_out, self.beta)
 
 
 class EncoderBlock(nn.Module):

broccoli/vit.py

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

{broccoli_ml-15.4.1.dist-info → broccoli_ml-15.5.0.dist-info}/METADATA

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

{broccoli_ml-15.4.1.dist-info → broccoli_ml-15.5.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=XCr3fkDrYvNEA3m5rdVDHU5qbf1IYJQS1Cjgmn5RX-o,28177
+broccoli/transformer.py,sha256=xhMKGWgQqSMhCpN-cqM6Fv_MfyKU9-Gq1t9nGpUAmzE,27574
 broccoli/utils.py,sha256=oOWzn6dJ5nC_9r4zq0emmfmaYACJXJNFS48AOpW2jqc,358
-broccoli/vit.py,sha256=lhWOcVg4GfTHJp3ABflKjQQzWGKMqrOU6QKqRWU11wo,23012
-broccoli_ml-15.4.1.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
-broccoli_ml-15.4.1.dist-info/METADATA,sha256=Vg2KdoHSlbV1IDYuNk8h07sVT1UTW69bt49JLiEr0os,1369
-broccoli_ml-15.4.1.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-broccoli_ml-15.4.1.dist-info/RECORD,,
+broccoli/vit.py,sha256=v3U_UVIZd2t3Nt60K6KGJcI5ci9t9S8h2ENwklnHg8M,22735
+broccoli_ml-15.5.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
+broccoli_ml-15.5.0.dist-info/METADATA,sha256=ANXSYDDts212i3b0rySkKT71_2ZSpcmpHSloNayfNns,1369
+broccoli_ml-15.5.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+broccoli_ml-15.5.0.dist-info/RECORD,,