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

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

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

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

@@ -1,3 +1,4 @@
+import warnings
 import math
 from typing import Optional, Tuple
 
@@ -20,6 +21,12 @@ except ImportError:
     FLASH_ATTN = False
 
 
+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(
     x, drop_prob: float = 0.0, training: bool = False, scale_by_keep: bool = True
 ):
@@ -78,9 +85,11 @@ class MHAttention(nn.Module):
         seq_len=None,
         linear_module: nn.Module = nn.Linear,
         utility_tokens=0,
+        talking_heads=False,
         rotary_embedding=None,
         source_size=None,
         scaling="d",
+        beta=1.0,
     ):
         """
         Args:
@@ -96,10 +105,20 @@ class MHAttention(nn.Module):
         if causal:
             assert seq_len is not None
 
+        self.talking_heads = talking_heads
+
+        if self.talking_heads:
+            self.head_projection = nn.Linear(n_heads, n_heads, bias=False)
+            self.sample_projection = nn.Linear(n_heads, n_heads, bias=False)
+        else:
+            self.head_projection = None
+            self.sample_projection = None
+
         self.embed_dim = embed_dim
         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
 
@@ -181,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,
         )
 
@@ -202,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(
@@ -243,7 +274,7 @@ class MHAttention(nn.Module):
 
         q, k, v = self.project_qkv(q, k, v)
 
-        if FLASH_ATTN:
+        if FLASH_ATTN and not self.talking_heads:
             # Divide Q/K/V into heads
             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)
@@ -271,12 +302,22 @@ class MHAttention(nn.Module):
 
             qk_scores *= self.scaling_factor
 
+            if self.talking_heads:
+                qk_scores = torch.einsum(
+                    "b h i j, o h -> b o i j", qk_scores, self.head_projection.weight
+                )
+
             # Apply mask if causal (must come before softmax)
             if self.causal:
                 qk_scores.masked_fill_(self.mask, float("-inf"))
 
             qk_scores = F.softmax(qk_scores, dim=-1)
 
+            if self.talking_heads:
+                qk_scores = torch.einsum(
+                    "b h i j, o h -> b o i j", qk_scores, self.sample_projection.weight
+                )
+
             qk_scores = self.dropout(qk_scores)
 
             output_with_heads = qk_scores @ v
@@ -309,7 +350,14 @@ 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)
+            nn.init.eye_(self.sample_projection.weight)
 
 
 class FeedforwardBlock(nn.Module):
@@ -329,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):
@@ -365,19 +410,26 @@ class FeedforwardBlock(nn.Module):
         )
 
         self.max_features = (
-            2 * ratio * output_features if self.xglu else ratio * output_features
+            2 * int(ratio * output_features)
+            if self.xglu
+            else int(ratio * output_features)
         )
 
         self.linear_in = linear_module_up(input_features, self.max_features)
-        self.linear_out = linear_module_down(ratio * output_features, output_features)
+        self.linear_out = linear_module_down(
+            int(ratio * output_features), output_features
+        )
 
         self.process = nn.Sequential(
             *[
-                nn.LayerNorm(input_features) if pre_norm else nn.Identity(),
                 self.linear_in,
                 self.activation,
                 self.inner_dropout,
-                nn.LayerNorm(ratio * output_features) if normformer else nn.Identity(),
+                (
+                    nn.LayerNorm(int(ratio * output_features))
+                    if normformer
+                    else nn.Identity()
+                ),
                 self.linear_out,
                 self.outer_dropout,
             ]
@@ -419,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:
@@ -435,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.
@@ -453,6 +503,7 @@ class TransformerBlock(nn.Module):
         relative_position_embedding=False,
         source_size=None,
         utility_tokens=0,
+        talking_heads=False,
         mlp_ratio=4,
         activation: nn.Module = nn.ReLU,
         activation_kwargs: Optional[dict] = None,
@@ -470,6 +521,8 @@ class TransformerBlock(nn.Module):
         post_norm=False,
         normformer=False,
         checkpoint_ff=True,
+        alpha=1.0,
+        beta=1.0,
     ):
         """
         Args:
@@ -481,15 +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 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!
@@ -513,7 +580,9 @@ class TransformerBlock(nn.Module):
             rotary_embedding=self.rotary_embedding,
             source_size=source_size,
             utility_tokens=utility_tokens,
+            talking_heads=talking_heads,
             scaling=msa_scaling,
+            beta=beta,
         )
 
         # Submodule for the feedforward process
@@ -536,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()
@@ -550,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.attn(x, x, x))
-            x = self.layer_norm_2(x)
-            x = x + self.drop_path(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.attn(x, x, x))
-            x = self.layer_norm_1(x)
-            x = x + self.drop_path(self.ff(x))
-            x = self.layer_norm_2(x)
-        else:  # Not pre or post norm. Stand well back.
-            x = x + self.drop_path(self.attn(x, x, x))
-            x = x + self.drop_path(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
 
@@ -573,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()
@@ -616,12 +705,15 @@ class TransformerEncoder(nn.Module):
         causal=False,
         linear_module=nn.Linear,
         utility_tokens=0,
+        talking_heads=False,
         return_utility_tokens=False,
         pre_norm=True,
         post_norm=False,
         normformer=False,
         msa_scaling="d",
         checkpoint_ff=True,
+        alpha=1.0,
+        beta=1.0,
     ):
         """
         Args:
@@ -644,6 +736,13 @@ class TransformerEncoder(nn.Module):
             )
 
         super().__init__()
+
+        if FLASH_ATTN and talking_heads:
+            warnings.warn(
+                "Using talking heads currently prevents using flash attention.",
+                stacklevel=2,
+            )
+
         self.seq_len = seq_len
         self.n_heads = n_heads
         self._utility_tokens = utility_tokens
@@ -688,13 +787,14 @@ class TransformerEncoder(nn.Module):
 
         self.blocks = nn.ModuleList(
             [
-                TransformerBlock(
+                EncoderBlock(
                     self.full_sequence_length,
                     d_model,
                     n_heads,
                     relative_position_embedding=relative_position_embedding,
                     source_size=source_size,
                     utility_tokens=utility_tokens,
+                    talking_heads=talking_heads,
                     mlp_ratio=mlp_ratio,
                     activation=activation,
                     activation_kwargs=activation_kwargs,
@@ -712,6 +812,8 @@ class TransformerEncoder(nn.Module):
                     post_norm=post_norm,
                     normformer=normformer,
                     checkpoint_ff=checkpoint_ff,
+                    alpha=alpha,
+                    beta=beta,
                 )
                 for i in range(n_layers)
             ]
@@ -732,13 +834,14 @@ class TransformerEncoder(nn.Module):
             x = x
 
         if self.absolute_position_embedding is not None:
-            x = x + self.absolute_position_embedding(
+            position_embedding = self.absolute_position_embedding(
                 torch.arange(
                     0, self.full_sequence_length, dtype=torch.long, device=x.device
                 ).unsqueeze(
                     0
                 )  # to shape (1, seq_len) to broadcast over batch
             )
+            x += position_embedding
 
         return x
 

{broccoli_ml-9.6.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,
@@ -174,6 +173,7 @@ class ViTEncoder(nn.Module):
         transformer_heads=4,
         transformer_mlp_ratio=2,
         transformer_utility_tokens=0,
+        transformer_talking_heads=False,
         transformer_return_utility_tokens=False,
         transformer_activation: nn.Module = SquaredReLU,
         transformer_activation_kwargs: Optional[dict] = None,
@@ -345,11 +345,14 @@ class ViTEncoder(nn.Module):
                 causal=False,
                 linear_module=linear_module,
                 utility_tokens=transformer_utility_tokens,
+                talking_heads=transformer_talking_heads,
                 return_utility_tokens=transformer_return_utility_tokens,
                 pre_norm=transformer_pre_norm,
                 normformer=transformer_normformer,
                 post_norm=transformer_post_norm,
                 checkpoint_ff=transformer_checkpoint_ff,
+                alpha=self.alpha,
+                beta=self.beta,
             )
         else:
             self.transformer = nn.Identity()
@@ -391,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,
@@ -410,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:
@@ -456,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,
@@ -472,6 +474,7 @@ class ViT(nn.Module):
         transformer_heads=4,
         transformer_mlp_ratio=2,
         transformer_utility_tokens=0,
+        transformer_talking_heads=False,
         transformer_return_utility_tokens=False,
         transformer_activation: nn.Module = SquaredReLU,
         transformer_activation_kwargs: Optional[dict] = None,
@@ -527,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,
@@ -543,6 +545,7 @@ class ViT(nn.Module):
             transformer_heads=transformer_heads,
             transformer_mlp_ratio=transformer_mlp_ratio,
             transformer_utility_tokens=transformer_utility_tokens,
+            transformer_talking_heads=transformer_talking_heads,
             transformer_return_utility_tokens=transformer_return_utility_tokens,
             transformer_activation=transformer_activation,
             transformer_activation_kwargs=transformer_activation_kwargs,

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

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