broccoli-ml 14.0.1__tar.gz → 15.0.0__tar.gz

{broccoli_ml-14.0.1 → broccoli_ml-15.0.0}/PKG-INFO

@@ -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 → broccoli_ml-15.0.0}/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_ml-14.0.1 → broccoli_ml-15.0.0}/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 → broccoli_ml-15.0.0}/pyproject.toml

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