broccoli-ml 10.1.0__tar.gz → 13.0.0__tar.gz

{broccoli_ml-10.1.0 → broccoli_ml-13.0.0}/PKG-INFO

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

{broccoli_ml-10.1.0 → broccoli_ml-13.0.0}/broccoli/transformer.py

@@ -21,13 +21,10 @@ except ImportError:
     FLASH_ATTN = False
 
 
-class LayerScale(nn.Module):
-    def __init__(self, dim, init_values=1e-4):
-        super().__init__()
-        self.nondecay_scale = nn.Parameter(init_values * torch.ones(dim))
-
-    def forward(self, x):
-        return x * self.nondecay_scale
+def scale_parameters(torch_module: nn.Module, factor: float):
+    with torch.no_grad():
+        for param in torch_module.parameters():
+            param.mul_(factor)
 
 
 def drop_path(
@@ -92,6 +89,7 @@ class MHAttention(nn.Module):
         rotary_embedding=None,
         source_size=None,
         scaling="d",
+        beta=1.0,
     ):
         """
         Args:
@@ -120,6 +118,7 @@ class MHAttention(nn.Module):
         self.n_heads = n_heads
         assert embed_dim % n_heads == 0
         self.scaling = scaling
+        self.beta = beta
 
         self.head_dim = self.embed_dim // self.n_heads
 
@@ -201,17 +200,26 @@ class MHAttention(nn.Module):
                 "`source_size` must be a tuple of 1, 2 or 3 integers"
             )
 
-        q_util, q_img = q[:, : self.utility_tokens, :], q[:, self.utility_tokens :, :]
-        k_util, k_img = k[:, : self.utility_tokens, :], k[:, self.utility_tokens :, :]
+        q = rearrange(q, "b t (h d) -> b t h d", h=self.n_heads)
+        k = rearrange(k, "b t (h d) -> b t h d", h=self.n_heads)
+
+        q_util, q_img = (
+            q[:, : self.utility_tokens, :, :],
+            q[:, self.utility_tokens :, :, :],
+        )
+        k_util, k_img = (
+            k[:, : self.utility_tokens, :, :],
+            k[:, self.utility_tokens :, :, :],
+        )
 
         q_img = rearrange(
             q_img,
-            f"b ({spatial_dimension_names}) d -> b {spatial_dimension_names} d",
+            f"b ({spatial_dimension_names}) h d -> b {spatial_dimension_names} h d",
             **spatial_dimension_values,
         )
         k_img = rearrange(
             k_img,
-            f"b ({spatial_dimension_names}) d -> b {spatial_dimension_names} d",
+            f"b ({spatial_dimension_names}) h d -> b {spatial_dimension_names} h d",
             **spatial_dimension_values,
         )
 
@@ -222,17 +230,20 @@ class MHAttention(nn.Module):
 
         q_img = rearrange(
             q_img,
-            f"b {spatial_dimension_names} d -> b ({spatial_dimension_names}) d",
+            f"b {spatial_dimension_names} h d -> b ({spatial_dimension_names}) h d",
         )
         k_img = rearrange(
             k_img,
-            f"b {spatial_dimension_names} d -> b ({spatial_dimension_names}) d",
+            f"b {spatial_dimension_names} h d -> b ({spatial_dimension_names}) h d",
         )
 
         # Re-combine the utility tokens and the RoPE-enhanced sequence tokens
         q = torch.cat([q_util, q_img], dim=1)
         k = torch.cat([k_util, k_img], dim=1)
 
+        q = rearrange(q, "b t h d -> b t (h d)")
+        k = rearrange(k, "b t h d -> b t (h d)")
+
         return q, k
 
     def project_qkv(
@@ -339,7 +350,10 @@ 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
         self.out_proj.reset_parameters()
+        scale_parameters(self.out_proj, self.beta)  # per Microsoft DeepNet
+
         if self.talking_heads:
             # Initialize close to identity
             nn.init.eye_(self.head_projection.weight)
@@ -363,17 +377,14 @@ class FeedforwardBlock(nn.Module):
         outer_dropout=None,
         linear_module_up=nn.Linear,
         linear_module_down=nn.Linear,
-        pre_norm=True,
         normformer=False,
-        post_norm=True,
-        residual_path=True,
         checkpoint=True,
+        beta=1.0,
     ):
         super().__init__()
 
         self.checkpoint = checkpoint
-        self.residual_path = residual_path
-        self.post_norm = post_norm
+        self.beta = beta
         self.xglu = activation.__name__.endswith("GLU")
 
         if self.residual_path and (output_features < input_features):
@@ -411,7 +422,6 @@ class FeedforwardBlock(nn.Module):
 
         self.process = nn.Sequential(
             *[
-                nn.LayerNorm(input_features) if pre_norm else nn.Identity(),
                 self.linear_in,
                 self.activation,
                 self.inner_dropout,
@@ -461,12 +471,7 @@ class FeedforwardBlock(nn.Module):
         else:
             processed = self.process(x)
 
-        if self.residual_path and self.post_norm:
-            return self.layernorm(x + processed)
-        elif self.residual_path:
-            return x + processed
-        else:
-            return processed
+        return processed
 
     def reset_parameters(self):
         if self.post_norm:
@@ -477,8 +482,11 @@ 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)
+
 
-class TransformerBlock(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.
@@ -513,7 +521,8 @@ class TransformerBlock(nn.Module):
         post_norm=False,
         normformer=False,
         checkpoint_ff=True,
-        layerscale=True,
+        alpha=1.0,
+        beta=1.0,
     ):
         """
         Args:
@@ -525,22 +534,29 @@ class TransformerBlock(nn.Module):
 
         super().__init__()
 
+        if pre_norm and post_norm:
+            raise ValueError("A transformer cannot be both prenorm and postnorm.")
+
         self.pre_norm = pre_norm
         self.post_norm = post_norm
         self.normformer = normformer
 
+        self.alpha = alpha
+        self.beta = beta
+
         self.drop_path = DropPath(drop_prob=identity_probability, scale_by_keep=True)
 
-        self.layer_norm_1 = nn.LayerNorm(d_model)
-        self.layer_norm_2 = nn.LayerNorm(d_model)
-        self.layer_norm_3 = nn.LayerNorm(d_model)
+        if self.pre_norm:
+            self.pre_attention_norm = nn.LayerNorm(d_model)
+            self.pre_mlp_norm = nn.LayerNorm(d_model)
 
-        if layerscale:
-            self.layerscale1 = LayerScale(d_model)
-            self.layerscale2 = LayerScale(d_model)
-        else:
-            self.layerscale1 = nn.Identity()
-            self.layerscale2 = nn.Identity()
+        if normformer:
+            self.normformer_norm = nn.LayerNorm(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)
 
         if relative_position_embedding:
             max_freq = int(max(source_size) / 2)  # Suggested by Gemini!
@@ -566,6 +582,7 @@ class TransformerBlock(nn.Module):
             utility_tokens=utility_tokens,
             talking_heads=talking_heads,
             scaling=msa_scaling,
+            beta=beta,
         )
 
         # Submodule for the feedforward process
@@ -588,11 +605,9 @@ class TransformerBlock(nn.Module):
                 if ff_linear_module_down is not None
                 else linear_module
             ),
-            pre_norm=False,  # Handled outside the block
             normformer=normformer,
-            post_norm=False,  # Handled outside the block
-            residual_path=False,  # Handled outside the block
             checkpoint=checkpoint_ff,
+            beta=beta,
         )
 
         self.reset_parameters()
@@ -602,22 +617,34 @@ class TransformerBlock(nn.Module):
         return self.attn._kv_distance
 
     def forward(self, x):
+        if self.post_norm:
+            x = self.input_norm(x)
 
         if self.pre_norm:
-            x = self.layer_norm_1(x)
-            x = x + self.drop_path(self.layerscale1(self.attn(x, x, x)))
-            x = self.layer_norm_2(x)
-            x = x + self.drop_path(self.layerscale2(self.ff(x)))
-            if self.post_norm:  # i.e. in addition! Pre and post.
-                x = self.layer_norm_3(x)
-        elif self.post_norm:  # i.e. only, not prenorm, just post
-            x = x + self.drop_path(self.layerscale1(self.attn(x, x, x)))
-            x = self.layer_norm_1(x)
-            x = x + self.drop_path(self.layerscale2(self.ff(x)))
-            x = self.layer_norm_2(x)
-        else:  # Not pre or post norm. Stand well back.
-            x = x + self.drop_path(self.layerscale1(self.attn(x, x, x)))
-            x = x + self.drop_path(self.layerscale2(self.ff(x)))
+            process_x = self.pre_attention_norm(x)
+        else:
+            process_x = x
+
+        processed = self.drop_path(self.attn(process_x, process_x, process_x))
+
+        if self.normformer:
+            processed = self.normformer_norm(processed)
+
+        x = self.alpha * x + processed
+
+        if self.post_norm:
+            x = self.post_attention_norm(x)
+        elif self.pre_norm:
+            process_x = self.pre_mlp_norm(x)
+        else:
+            process_x = x
+
+        processed = self.drop_path(self.ff(process_x))
+
+        x = self.alpha * x + processed
+
+        if self.post_norm:
+            x = self.post_mlp_norm(x)
 
         return x
 
@@ -625,16 +652,26 @@ class TransformerBlock(nn.Module):
         """
         Give back the attention scores used in this layer.
         """
+        # Fix: Use the correct attribute name 'pre_attention_norm'
         if self.pre_norm:
-            x = self.layer_norm_1(x)
+            # We must normalize the input before measuring attention logits
+            # to match what the model actually sees during forward()
+            x = self.pre_attention_norm(x)
             return self.attn.attention_logits(x, x, x)
         else:
             return self.attn.attention_logits(x, x, x)
 
     def reset_parameters(self):
-        self.layer_norm_1.reset_parameters()
-        self.layer_norm_2.reset_parameters()
-        self.layer_norm_3.reset_parameters()
+        if self.pre_norm:
+            self.pre_attention_norm.reset_parameters()
+            self.pre_mlp_norm.reset_parameters()
+
+        if self.post_norm:
+            self.post_attention_norm.reset_parameters()
+            self.post_mlp_norm.reset_parameters()
+
+        if self.normformer:
+            self.normformer_norm.reset_parameters()
 
         self.attn.reset_parameters()
         self.ff.reset_parameters()
@@ -675,7 +712,8 @@ class TransformerEncoder(nn.Module):
         normformer=False,
         msa_scaling="d",
         checkpoint_ff=True,
-        layerscale=True,
+        alpha=1.0,
+        beta=1.0,
     ):
         """
         Args:
@@ -710,11 +748,6 @@ class TransformerEncoder(nn.Module):
         self._utility_tokens = utility_tokens
         self.return_utility_tokens = return_utility_tokens
 
-        if layerscale:
-            self.layerscale = LayerScale(d_model)
-        else:
-            self.layerscale = None
-
         # Initialise utility tokens with normal init, like usual Pytorch embeddings
         if self._utility_tokens:
             self._utility_token_embedding = nn.Parameter(
@@ -754,7 +787,7 @@ class TransformerEncoder(nn.Module):
 
         self.blocks = nn.ModuleList(
             [
-                TransformerBlock(
+                EncoderBlock(
                     self.full_sequence_length,
                     d_model,
                     n_heads,
@@ -779,7 +812,8 @@ class TransformerEncoder(nn.Module):
                     post_norm=post_norm,
                     normformer=normformer,
                     checkpoint_ff=checkpoint_ff,
-                    layerscale=layerscale,
+                    alpha=alpha,
+                    beta=beta,
                 )
                 for i in range(n_layers)
             ]
@@ -807,10 +841,9 @@ class TransformerEncoder(nn.Module):
                     0
                 )  # to shape (1, seq_len) to broadcast over batch
             )
-            if self.layerscale is not None:
-                position_embedding = self.layerscale(position_embedding)
+            x += position_embedding
 
-        return x + position_embedding
+        return x
 
     def forward(self, x):
 

{broccoli_ml-10.1.0 → broccoli_ml-13.0.0}/broccoli/vit.py

@@ -158,7 +158,6 @@ class ViTEncoder(nn.Module):
         pooling_kernel_stride=2,
         pooling_padding=1,
         transformer_feedforward_first=True,
-        transformer_initial_ff_residual_path=True,
         transformer_initial_ff_linear_module_up=None,
         transformer_initial_ff_linear_module_down=None,
         transformer_initial_ff_dropout=None,
@@ -187,7 +186,6 @@ class ViTEncoder(nn.Module):
         transformer_msa_dropout=0.1,
         transformer_stochastic_depth=0.1,
         transformer_checkpoint_ff=True,
-        transformer_layerscale=True,
         linear_module=nn.Linear,
     ):
         super().__init__()
@@ -353,7 +351,8 @@ class ViTEncoder(nn.Module):
                 normformer=transformer_normformer,
                 post_norm=transformer_post_norm,
                 checkpoint_ff=transformer_checkpoint_ff,
-                layerscale=transformer_layerscale,
+                alpha=self.alpha,
+                beta=self.beta,
             )
         else:
             self.transformer = nn.Identity()
@@ -395,16 +394,14 @@ class ViTEncoder(nn.Module):
                     or transformer_ff_linear_module_down
                     or linear_module
                 ),
-                pre_norm=transformer_pre_norm,
                 normformer=transformer_normformer,
-                post_norm=transformer_post_norm,
-                residual_path=transformer_initial_ff_residual_path,
                 checkpoint=transformer_checkpoint_ff,
+                beta=self.beta,
             )
         else:
             self.initial_ff = nn.Identity()
 
-        self.encoder = nn.Sequential(
+        self.preprocess = nn.Sequential(
             *[
                 batchnormxd(in_channels) if initial_batch_norm else nn.Identity(),
                 self.cnn,
@@ -414,19 +411,21 @@ class ViTEncoder(nn.Module):
                     f"N C {spatial_dim_names} -> N ({spatial_dim_names}) C"
                 ),
                 self.pooling_channels_padding,
-                self.initial_ff,
-                self.transformer,
+                nn.LayerNorm(),
             ]
         )
 
         self.reset_parameters()
 
     def forward(self, x):
-        return self.encoder(x)
+        x = self.preprocess(x)
+        x = x + self.initial_ff(x)
+        return self.transformer(x)
 
     def attention_logits(self, x):
-        x = self.encoder[:-1](x)
-        return self.encoder[-1].attention_logits(x)
+        x = self.preprocess(x)
+        x = x + self.initial_ff(x)
+        return self.transformer.attention_logits(x)
 
     def reset_parameters(self):
         for module in self.encoder:
@@ -460,7 +459,6 @@ class ViT(nn.Module):
         pooling_kernel_stride=2,
         pooling_padding=1,
         transformer_feedforward_first=True,
-        transformer_initial_ff_residual_path=True,
         transformer_initial_ff_linear_module_up=None,
         transformer_initial_ff_linear_module_down=None,
         transformer_initial_ff_dropout=None,
@@ -489,7 +487,6 @@ class ViT(nn.Module):
         transformer_msa_dropout=0.1,
         transformer_stochastic_depth=0.1,
         transformer_checkpoint_ff=True,
-        transformer_layerscale=True,
         head=SequencePoolClassificationHead,
         batch_norm_logits=True,
         logit_projection_layer=nn.Linear,
@@ -533,7 +530,6 @@ class ViT(nn.Module):
             pooling_kernel_stride=pooling_kernel_stride,
             pooling_padding=pooling_padding,
             transformer_feedforward_first=transformer_feedforward_first,
-            transformer_initial_ff_residual_path=transformer_initial_ff_residual_path,
             transformer_initial_ff_linear_module_up=transformer_initial_ff_linear_module_up,
             transformer_initial_ff_linear_module_down=transformer_initial_ff_linear_module_down,
             transformer_initial_ff_dropout=transformer_initial_ff_dropout,
@@ -562,7 +558,6 @@ class ViT(nn.Module):
             transformer_msa_dropout=transformer_msa_dropout,
             transformer_stochastic_depth=transformer_stochastic_depth,
             transformer_checkpoint_ff=transformer_checkpoint_ff,
-            transformer_layerscale=transformer_layerscale,
             linear_module=linear_module,
         )
 

{broccoli_ml-10.1.0 → broccoli_ml-13.0.0}/pyproject.toml

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