broccoli-ml 15.0.1__py3-none-any.whl → 15.2.0__py3-none-any.whl

broccoli/transformer.py

@@ -122,6 +122,9 @@ 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":
@@ -225,8 +228,9 @@ class MHAttention(nn.Module):
 
         freqs = self.rotary_embedding.get_axial_freqs(*self.source_size)
 
-        q_img = apply_rotary_emb(freqs, q_img)
-        k_img = apply_rotary_emb(freqs, k_img)
+        # 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 = rearrange(
             q_img,
@@ -416,7 +420,7 @@ class FeedforwardBlock(nn.Module):
                 self.activation,
                 self.inner_dropout,
                 (
-                    nn.LayerNorm(int(ratio * output_features))
+                    nn.RMSNorm(int(ratio * output_features))
                     if normformer
                     else nn.Identity()
                 ),
@@ -474,13 +478,7 @@ class FeedforwardBlock(nn.Module):
 
 
 class EncoderBlock(nn.Module):
-    """
-    Performs LayerNorms first (as in PyTorch Transformers when norm_first=True),
-        which is also what is seen in e.g.
-        https://github.com/karpathy/minGPT/blob/master/mingpt/model.py
-        and is recommended by https://arxiv.org/abs/2002.04745
-
-    """
+    """ """
 
     def __init__(
         self,
@@ -534,16 +532,16 @@ class EncoderBlock(nn.Module):
         self.drop_path = DropPath(drop_prob=identity_probability, scale_by_keep=True)
 
         if self.pre_norm:
-            self.pre_attention_norm = nn.LayerNorm(d_model)
-            self.pre_mlp_norm = nn.LayerNorm(d_model)
+            self.pre_attention_norm = nn.RMSNorm(d_model)
+            self.pre_mlp_norm = nn.RMSNorm(d_model)
 
         if normformer:
-            self.normformer_norm = nn.LayerNorm(d_model)
+            self.normformer_norm = nn.RMSNorm(d_model)
 
         if self.post_norm:
-            self.input_norm = nn.LayerNorm(d_model)
-            self.post_attention_norm = nn.LayerNorm(d_model)
-            self.post_mlp_norm = nn.LayerNorm(d_model)
+            self.input_norm = nn.RMSNorm(d_model)
+            self.post_attention_norm = nn.RMSNorm(d_model)
+            self.post_mlp_norm = nn.RMSNorm(d_model)
 
         if relative_position_embedding:
             max_freq = int(max(source_size) / 2)  # Suggested by Gemini!

broccoli/vit.py

@@ -403,7 +403,7 @@ class ViTEncoder(nn.Module):
                 checkpoint=transformer_checkpoint_ff,
                 beta=self.beta,
             )
-            self.layer_norm = nn.LayerNorm(transformer_embedding_size)
+            self.layer_norm = nn.RMSNorm(transformer_embedding_size)
         else:
             self.initial_ff = None
 
@@ -417,7 +417,7 @@ class ViTEncoder(nn.Module):
                     f"N C {spatial_dim_names} -> N ({spatial_dim_names}) C"
                 ),
                 self.pooling_channels_padding,
-                nn.LayerNorm(transformer_embedding_size),
+                nn.RMSNorm(transformer_embedding_size),
             ]
         )
 
@@ -522,7 +522,8 @@ class ViT(nn.Module):
 
         # Set alpha and beta according to Microsoft's DeepNorm
         self.alpha = (2 * transformer_layers) ** 0.25
-        self.beta = (8 * 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
 
         self.encoder = ViTEncoder(
             input_size=input_size,

{broccoli_ml-15.0.1.dist-info → broccoli_ml-15.2.0.dist-info}/METADATA

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

{broccoli_ml-15.0.1.dist-info → broccoli_ml-15.2.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=jvlSfMXFlz3oB3tXX-LK_HTLFbXTTzFl9NdwWyTQAxY,27844
+broccoli/transformer.py,sha256=Uuz_jCMZRg6GY2DmW3-Tn47gV9a-xkGVN3xQ5BYFM5w,27784
 broccoli/utils.py,sha256=oOWzn6dJ5nC_9r4zq0emmfmaYACJXJNFS48AOpW2jqc,358
-broccoli/vit.py,sha256=3YSKOBcgVeeGaED69x1h3MSmAl5a7zrkBwGOX8cYozM,22803
-broccoli_ml-15.0.1.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
-broccoli_ml-15.0.1.dist-info/METADATA,sha256=ARxRoey2lg2wHXbpl33L4CIIr7-jmIDGr_s5tokPGPE,1369
-broccoli_ml-15.0.1.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-broccoli_ml-15.0.1.dist-info/RECORD,,
+broccoli/vit.py,sha256=fMpdUDUARGs8nPXR6l-SLSlVMNyeSrHIH8Fst4-e5wU,22884
+broccoli_ml-15.2.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
+broccoli_ml-15.2.0.dist-info/METADATA,sha256=4Pqm4a8K-40gZxm0mMRsvBAUkapsBPXJZRZpzqQAoyY,1369
+broccoli_ml-15.2.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+broccoli_ml-15.2.0.dist-info/RECORD,,