broccoli-ml 14.0.1__py3-none-any.whl → 15.0.0__py3-none-any.whl

broccoli/transformer.py

@@ -604,8 +604,6 @@ class EncoderBlock(nn.Module):
         return self.attn._kv_distance
 
     def forward(self, x):
-        if self.post_norm:
-            x = self.input_norm(x)
 
         if self.pre_norm:
             process_x = self.pre_attention_norm(x)

broccoli/vit.py

@@ -403,8 +403,9 @@ class ViTEncoder(nn.Module):
                 checkpoint=transformer_checkpoint_ff,
                 beta=self.beta,
             )
+            self.layer_norm = nn.LayerNorm(transformer_embedding_size)
         else:
-            self.initial_ff = nn.Identity()
+            self.initial_ff = None
 
         self.preprocess = nn.Sequential(
             *[
@@ -424,7 +425,8 @@ class ViTEncoder(nn.Module):
 
     def forward(self, x):
         x = self.preprocess(x)
-        x = x + self.initial_ff(x)
+        if self.initial_ff is not None:
+            x = self.layer_norm(x + self.initial_ff(x))
         return self.transformer(x)
 
     def attention_logits(self, x):
@@ -498,8 +500,6 @@ class ViT(nn.Module):
         batch_norm_logits=True,
         logit_projection_layer=nn.Linear,
         linear_module=nn.Linear,
-        alpha=1.0,
-        beta=1.0,
     ):
 
         super().__init__()
@@ -520,8 +520,9 @@ class ViT(nn.Module):
                 "SwiGLU": SwiGLU,
             }[transformer_activation]
 
-        self.alpha = alpha
-        self.beta = beta
+        # Set alpha and beta according to Microsoft's DeepNorm
+        self.alpha = (2 * transformer_layers) ** 0.25
+        self.beta = (8 * transformer_layers) ** 0.25
 
         self.encoder = ViTEncoder(
             input_size=input_size,
@@ -571,8 +572,8 @@ class ViT(nn.Module):
             transformer_stochastic_depth=transformer_stochastic_depth,
             transformer_checkpoint_ff=transformer_checkpoint_ff,
             linear_module=linear_module,
-            alpha=alpha,
-            beta=beta,
+            alpha=self.alpha,
+            beta=self.beta,
         )
 
         self.pool = head(

{broccoli_ml-14.0.1.dist-info → broccoli_ml-15.0.0.dist-info}/METADATA

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

{broccoli_ml-14.0.1.dist-info → broccoli_ml-15.0.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=h1pvwdp3UntZC336ax9BPFWNetd6gfaGUe5o9nVWU1M,27906
+broccoli/transformer.py,sha256=jvlSfMXFlz3oB3tXX-LK_HTLFbXTTzFl9NdwWyTQAxY,27844
 broccoli/utils.py,sha256=oOWzn6dJ5nC_9r4zq0emmfmaYACJXJNFS48AOpW2jqc,358
-broccoli/vit.py,sha256=SYAHGuGVZRa-zYCLK5siKdUauDgHtMYUmQ6W2tUIqKA,22588
-broccoli_ml-14.0.1.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
-broccoli_ml-14.0.1.dist-info/METADATA,sha256=6Zl9CoeEG2phu5OwuAlJ8NSlwAnKJ1PRdmlgBH8au6g,1369
-broccoli_ml-14.0.1.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-broccoli_ml-14.0.1.dist-info/RECORD,,
+broccoli/vit.py,sha256=RZZNaEN-DbCENyCRyYpmVzxS-bl0wSaRfYICuAt4mv4,22802
+broccoli_ml-15.0.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
+broccoli_ml-15.0.0.dist-info/METADATA,sha256=527mYp49whcP9FHdLJTz20dUsgxAFc_cCvdn0jj_oPg,1369
+broccoli_ml-15.0.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+broccoli_ml-15.0.0.dist-info/RECORD,,