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

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

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

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

@@ -1,3 +1,4 @@
+import warnings
 import math
 from typing import Optional, Tuple
 
@@ -20,6 +21,15 @@ 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 drop_path(
     x, drop_prob: float = 0.0, training: bool = False, scale_by_keep: bool = True
 ):
@@ -78,6 +88,7 @@ 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",
@@ -96,6 +107,15 @@ 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
@@ -243,7 +263,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 +291,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
@@ -310,6 +340,10 @@ class MHAttention(nn.Module):
         self.k_proj.reset_parameters()
         self.v_proj.reset_parameters()
         self.out_proj.reset_parameters()
+        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):
@@ -365,11 +399,15 @@ 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(
             *[
@@ -377,7 +415,11 @@ class FeedforwardBlock(nn.Module):
                 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,
             ]
@@ -453,6 +495,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 +513,7 @@ class TransformerBlock(nn.Module):
         post_norm=False,
         normformer=False,
         checkpoint_ff=True,
+        layerscale=True,
     ):
         """
         Args:
@@ -491,6 +535,13 @@ class TransformerBlock(nn.Module):
         self.layer_norm_2 = nn.LayerNorm(d_model)
         self.layer_norm_3 = 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 relative_position_embedding:
             max_freq = int(max(source_size) / 2)  # Suggested by Gemini!
             if d_model < 16:
@@ -513,6 +564,7 @@ class TransformerBlock(nn.Module):
             rotary_embedding=self.rotary_embedding,
             source_size=source_size,
             utility_tokens=utility_tokens,
+            talking_heads=talking_heads,
             scaling=msa_scaling,
         )
 
@@ -553,19 +605,19 @@ class TransformerBlock(nn.Module):
 
         if self.pre_norm:
             x = self.layer_norm_1(x)
-            x = x + self.drop_path(self.attn(x, x, 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.ff(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.attn(x, x, x))
+            x = x + self.drop_path(self.layerscale1(self.attn(x, x, x)))
             x = self.layer_norm_1(x)
-            x = x + self.drop_path(self.ff(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.attn(x, x, x))
-            x = x + self.drop_path(self.ff(x))
+            x = x + self.drop_path(self.layerscale1(self.attn(x, x, x)))
+            x = x + self.drop_path(self.layerscale2(self.ff(x)))
 
         return x
 
@@ -616,12 +668,14 @@ 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,
+        layerscale=True,
     ):
         """
         Args:
@@ -644,11 +698,23 @@ 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
         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(
@@ -695,6 +761,7 @@ class TransformerEncoder(nn.Module):
                     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 +779,7 @@ class TransformerEncoder(nn.Module):
                     post_norm=post_norm,
                     normformer=normformer,
                     checkpoint_ff=checkpoint_ff,
+                    layerscale=layerscale,
                 )
                 for i in range(n_layers)
             ]
@@ -732,15 +800,17 @@ 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
             )
+            if self.layerscale is not None:
+                position_embedding = self.layerscale(position_embedding)
 
-        return x
+        return x + position_embedding
 
     def forward(self, x):
 

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

@@ -174,6 +174,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,
@@ -186,6 +187,7 @@ 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__()
@@ -345,11 +347,13 @@ 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,
+                layerscale=transformer_layerscale,
             )
         else:
             self.transformer = nn.Identity()
@@ -472,6 +476,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,
@@ -484,6 +489,7 @@ 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,
@@ -543,6 +549,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,
@@ -555,6 +562,7 @@ 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-9.6.0 → broccoli_ml-10.1.0}/pyproject.toml

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