broccoli-ml 6.0.0__py3-none-any.whl → 13.0.1__py3-none-any.whl

broccoli/activation.py

@@ -46,10 +46,7 @@ class GELU(nn.Module):
 
 class Swish(nn.Module):
     """
-    Implementation of (beta) SwiGLU, as introduced in "GLU Variants Improve Transformer"
-        (https://arxiv.org/abs/2002.05202v1) and used to great effect in LLaMa 2.0.
-
-    Halves the incoming parameter count, which should be scaled up before input.
+    Implementation of (beta) Swish
     """
 
     def __init__(self) -> None:

broccoli/linear.py

@@ -1,4 +1,8 @@
 import math
+import random
+import warnings
+from typing import Union, List, Iterable
+
 import torch
 from torch import nn
 from torch.nn import functional as F
@@ -136,3 +140,229 @@ class WeightNormedLinear(nn.Module):
             f"WeightNormedLinear(in_features={self.in_features},"
             f"out_features={self.out_features}, bias={self.use_bias})"
         )
+
+
+class RecyclingLinear(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = True,
+        row_recycling_rate: float = 0.0,
+        column_recycling_rate: float = 0.0,
+        adaptive=False,
+        xglu=False,
+    ):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.bias = bias
+        self.xglu = xglu
+        self.linear = nn.Linear(in_features, out_features, bias=bias)
+        self.row_recycling_rate = row_recycling_rate
+        self.column_recycling_rate = column_recycling_rate
+        self.adaptive = adaptive
+        self.optimisers = []
+        self.initial_learning_rates = []
+        self._warned_about_registration = False
+
+    def register_optimiser(self, optimiser: torch.optim.Optimizer):
+        self.optimisers.append(optimiser)
+        self.initial_learning_rates.append(self._get_learning_rate(optimiser))
+        if self.initial_learning_rates[-1] == 0.0:
+            warnings.warn(
+                "Learning rate of registered optimiser was 0.0 - make sure "
+                "you haven't initialised a scheduler before registering the "
+                "optimiser",
+                stacklevel=2,
+            )
+
+    def _get_learning_rate(self, optimiser: torch.optim.Optimizer):
+        for group in optimiser.param_groups:
+            for param in group["params"]:
+                if param is self.linear.weight:
+                    return group["lr"]
+
+    def _get_multiplier(self):
+        if not self.adaptive or not self.optimisers:
+            return 1.0
+        else:
+            init = self.initial_learning_rates
+            current = [self._get_learning_rate(o) for o in self.optimisers]
+            pairs = zip(current, init, strict=True)
+            multipliers = [a / b for a, b in pairs if b != 0.0]
+            return min(multipliers) if multipliers else 0.0
+
+    def reset_rows(self, indices, incoming_data=None):
+        """
+        Resets rows.
+        If incoming_data is provided, resets to the centroid (mean) of that data.
+        If not, resets to the mean of existing weights.
+        """
+        if not torch.is_tensor(indices):
+            idx_tensor = torch.as_tensor(
+                list(indices), dtype=torch.long, device=self.linear.weight.device
+            )
+        else:
+            idx_tensor = indices
+
+        if idx_tensor.numel() == 0:
+            return
+
+        if incoming_data is not None:
+            target_center = self._mean_input_weights(incoming_data)
+        else:
+            target_center = self._mean_value_weights()
+
+        target_center = target_center.expand(idx_tensor.size(0), -1)
+
+        if self.xglu:
+            gate_indices = idx_tensor
+            value_indices = idx_tensor + (self.linear.out_features // 2)
+            self._update_weights(gate_indices, 0, target_center, self.optimisers)
+            self._update_weights(value_indices, 0, target_center, self.optimisers)
+        else:
+            self._update_weights(idx_tensor, 0, target_center, self.optimisers)
+
+    def reset_columns(self, indices):
+        if not torch.is_tensor(indices):
+            idx_tensor = torch.as_tensor(
+                list(indices), dtype=torch.long, device=self.linear.weight.device
+            )
+        else:
+            idx_tensor = indices
+
+        if idx_tensor.size(0):
+            random_weights = self._random_weights(
+                self.linear.weight.size(0), indices.size(0)
+            )
+            # Make random col weights quiet so they don't introduce loud noise...
+            # ...but not so quiet that FP16 zeros them and ruins symmetry breaking!
+            random_weights *= 0.1
+            self._update_weights(indices, 1, random_weights, self.optimisers)  # dim
+        else:
+            return
+
+    def forward(self, x):
+        if self.training and self.optimisers:
+            self.reset_rows(self.get_reset_indices(0))
+            self.reset_columns(self.get_reset_indices(1))
+        elif self.training and not self._warned_about_registration:
+            warnings.warn(
+                "RecyclingLinear: No optimiser registered. Recycling disabled.",
+                stacklevel=2,
+            )
+            self._warned_about_registration = True
+
+        return self.linear(x)
+
+    def get_reset_indices(self, dim):
+        base_rate = self.row_recycling_rate if dim == 0 else self.column_recycling_rate
+        p = base_rate * self._get_multiplier()
+        if dim == 0:
+            if self.xglu:
+                sample_space = self.linear.out_features // 2
+            else:
+                sample_space = self.linear.out_features
+        elif dim == 1:
+            sample_space = self.linear.in_features
+        else:
+            raise ValueError("`dim` must be 0 or 1")
+
+        # Sample the indices
+        probs = torch.rand(sample_space, device=self.linear.weight.device)
+        mask = probs < p
+        if mask.any():
+            return torch.nonzero(mask).squeeze(-1)
+        else:
+            return torch.tensor([], dtype=torch.long, device=self.linear.weight.device)
+
+    def _random_weights(self, rows, columns):
+        device = self.linear.weight.device
+        weights = self.linear.weight.data
+        stdv = 1.0 / math.sqrt(weights.size(1))
+        random_weights = torch.rand(rows, columns, device=device)
+        random_weights -= 0.5  # Range [-0.5, +0.5]
+        random_weights *= 2.0 * stdv  # Range [-stdv, +stdv]
+        return random_weights
+
+    def _mean_input_weights(self, input):
+        reduce_dims = list(range(input.ndim - 1))
+        data_mean = input.detach().mean(dim=reduce_dims, keepdim=True)
+
+        weights = self.linear.weight.data
+        stdv = 1.0 / math.sqrt(weights.size(1))
+        data_norm = data_mean.std() + 1e-6
+        scale_factor = stdv / data_norm
+
+        return data_mean * scale_factor
+
+    def _mean_value_weights(self):
+        """
+        Only used when self.xglu
+        """
+        weights = self.linear.weight.data
+        rows = weights.size(0)
+        if self.xglu:
+            return self.linear.weight[int(rows / 2) :].data.mean(dim=0, keepdim=True)
+        else:
+            return self.linear.weight.data.mean(dim=0, keepdim=True)
+
+    def _mean_gate_weights(self):
+        """
+        Only used when self.xglu
+        """
+        weights = self.linear.weight.data
+        rows = weights.size(0)
+        return self.linear.weight[: int(rows / 2)].data.mean(dim=0, keepdim=True)
+
+    def _update_weights(
+        self,
+        indices: Iterable[int],
+        dim: int,
+        data: torch.Tensor,
+        optimisers: Union[
+            List[torch.optim.Optimizer], torch.optim.Optimizer, None
+        ] = None,
+    ):
+        if optimisers is None:
+            optimisers = []
+        if not isinstance(optimisers, list):
+            optimisers = [optimisers]
+
+        if not torch.is_tensor(indices):
+            idx_tensor = torch.as_tensor(
+                list(indices), dtype=torch.long, device=self.linear.weight.device
+            )
+        else:
+            idx_tensor = indices
+
+        if idx_tensor.numel() == 0:
+            return
+
+        with torch.no_grad():
+            if dim == 0:
+                self.linear.weight.data[idx_tensor] = data
+            elif dim == 1:
+                self.linear.weight.data[:, idx_tensor] = data
+            else:
+                raise ValueError("`dim` must be 0 or 1")
+
+            self._reset_optim_state(self.linear.weight, idx_tensor, optimisers, dim=dim)
+
+    def _reset_optim_state(self, param, idx_tensor, optimisers, dim):
+        """
+        Zeroes out the optimizer state for the given indices in a single operation.
+        """
+        for optimiser in optimisers:
+            if param not in optimiser.state:
+                continue
+            state = optimiser.state[param]
+
+            for _, buffer in state.items():
+                if torch.is_tensor(buffer) and buffer.shape == param.shape:
+                    # Vectorized zeroing
+                    if dim == 0:
+                        buffer[idx_tensor] = 0.0
+                    else:
+                        buffer[:, idx_tensor] = 0.0

broccoli/transformer.py

@@ -1,3 +1,4 @@
+import warnings
 import math
 from typing import Optional, Tuple
 
@@ -13,12 +14,19 @@ from .rope import RotaryEmbedding, apply_rotary_emb
 try:
     from flash_attn import flash_attn_func
 
+    print("Using flash-attn.")
     FLASH_ATTN = True
 except ImportError:
     pass
     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
 ):
@@ -77,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:
@@ -95,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
 
@@ -180,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,
         )
 
@@ -201,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(
@@ -242,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)
@@ -270,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
@@ -308,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):
@@ -324,19 +373,19 @@ class FeedforwardBlock(nn.Module):
         activation=nn.ReLU,
         activation_kwargs=None,
         dropout=0.0,
+        inner_dropout=None,
+        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):
             raise ValueError(
@@ -353,40 +402,76 @@ class FeedforwardBlock(nn.Module):
         else:
             self.activation = activation()
 
-        self.dropout = nn.Dropout(dropout)
+        self.inner_dropout = nn.Dropout(
+            inner_dropout if inner_dropout is not None else dropout
+        )
+        self.outer_dropout = nn.Dropout(
+            outer_dropout if outer_dropout is not None else dropout
+        )
 
         self.max_features = (
-            2 * ratio * output_features
-            if activation.__name__.endswith("GLU")
-            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(
+            int(ratio * output_features), output_features
         )
 
         self.process = nn.Sequential(
             *[
-                nn.LayerNorm(input_features) if pre_norm else nn.Identity(),
-                linear_module_up(input_features, self.max_features),
+                self.linear_in,
                 self.activation,
-                nn.LayerNorm(ratio * output_features) if normformer else nn.Identity(),
-                linear_module_down(ratio * output_features, output_features),
-                self.dropout,
+                self.inner_dropout,
+                (
+                    nn.LayerNorm(int(ratio * output_features))
+                    if normformer
+                    else nn.Identity()
+                ),
+                self.linear_out,
+                self.outer_dropout,
             ]
         )
 
+        self.recycling_enabled = False
+        if hasattr(self.linear_in, "row_recycling_rate") and hasattr(
+            self.linear_out, "column_recycling_rate"
+        ):
+            self.recycling_enabled = True
+            self.master_recycling_rate = self.linear_in.row_recycling_rate
+            self.linear_in.row_recycling_rate = 0.0
+            self.linear_out.column_recycling_rate = 0.0
+            if (
+                hasattr(self.linear_in, "column_recycling_rate")
+                and self.linear_in.column_recycling_rate > 0
+            ) or (
+                hasattr(self.linear_out, "row_recycling_rate")
+                and self.linear_out.row_recycling_rate > 0
+            ):
+                raise NotImplementedError(
+                    "At the moment this layer can only support recycling linear "
+                    "layers if the in layer resets only rows and the out layer "
+                    "resets only columns."
+                )
+
         self.reset_parameters()
 
     def forward(self, x):
 
+        # Recycle weights if using recycling linear layers
+        if self.training and self.recycling_enabled:
+            indices = self.linear_out.get_reset_indices(1)
+            self.linear_in.reset_rows(indices, incoming_data=x)
+            self.linear_out.reset_columns(indices)
+
         if self.checkpoint:
             processed = checkpoint(self.process, x, use_reentrant=False)
         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:
@@ -397,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.
@@ -415,13 +503,16 @@ 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,
         ff_linear_module_up=None,
         ff_linear_module_down=None,
         msa_scaling="d",
-        mlp_dropout=0.0,
+        ff_dropout=0.0,
+        ff_inner_dropout=0.0,
+        ff_outer_dropout=0.0,
         msa_dropout=0.0,
         identity_probability=0.0,
         causal=False,
@@ -430,6 +521,8 @@ class TransformerBlock(nn.Module):
         post_norm=False,
         normformer=False,
         checkpoint_ff=True,
+        alpha=1.0,
+        beta=1.0,
     ):
         """
         Args:
@@ -441,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!
@@ -473,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
@@ -483,7 +592,9 @@ class TransformerBlock(nn.Module):
             d_model,
             activation=activation,
             activation_kwargs=activation_kwargs,
-            dropout=mlp_dropout,
+            dropout=ff_dropout,
+            inner_dropout=ff_inner_dropout,
+            outer_dropout=ff_outer_dropout,
             linear_module_up=(
                 ff_linear_module_up
                 if ff_linear_module_up is not None
@@ -494,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()
@@ -508,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
 
@@ -531,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()
@@ -566,18 +697,23 @@ class TransformerEncoder(nn.Module):
         activation_kwargs: Optional[dict] = None,
         ff_linear_module_up=None,
         ff_linear_module_down=None,
-        mlp_dropout=0.0,
+        ff_dropout=0.0,
+        ff_inner_dropout=0.0,
+        ff_outer_dropout=0.0,
         msa_dropout=0.0,
         stochastic_depth=0.0,
         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:
@@ -590,10 +726,23 @@ class TransformerEncoder(nn.Module):
         if relative_position_embedding and (source_size is None):
             raise ValueError(
                 "`source_size` for TransformerEncoder cannot be None if"
-                " `position_embedding_type` is relative"
+                " `relative_position_embedding` is True"
+            )
+
+        if absolute_position_embedding and (seq_len is None):
+            raise ValueError(
+                "`seq_len` for TransformerEncoder cannot be None if"
+                " `absolute_position_embedding` is True"
             )
 
         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
@@ -605,9 +754,12 @@ class TransformerEncoder(nn.Module):
                 torch.empty(self._utility_tokens, d_model)
             )
             nn.init.normal_(self._utility_token_embedding, mean=0.0, std=1.0)
-            self.full_sequence_length = self.seq_len + self._utility_tokens
         else:
             self._utility_token_embedding = None
+
+        if self._utility_tokens and (self.seq_len is not None):
+            self.full_sequence_length = self.seq_len + self._utility_tokens
+        else:
             self.full_sequence_length = self.seq_len
 
         self.d_model = d_model
@@ -619,7 +771,7 @@ class TransformerEncoder(nn.Module):
         else:
             self.absolute_position_embedding = None
 
-        self.mlp_dropout = mlp_dropout
+        self.mlp_dropout = ff_dropout
         self.msa_dropout = msa_dropout
         self.stochastic_depth = stochastic_depth
 
@@ -635,20 +787,23 @@ 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,
                     ff_linear_module_up=ff_linear_module_up,
                     ff_linear_module_down=ff_linear_module_down,
                     msa_scaling=msa_scaling,
-                    mlp_dropout=mlp_dropout,
+                    ff_dropout=ff_dropout,
+                    ff_inner_dropout=ff_inner_dropout,
+                    ff_outer_dropout=ff_outer_dropout,
                     msa_dropout=msa_dropout,
                     identity_probability=self.stochastic_depth_probabilities[i],
                     causal=causal,
@@ -657,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)
             ]
@@ -677,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/vit.py

@@ -158,10 +158,11 @@ 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_mlp_dropout=None,
+        transformer_initial_ff_dropout=None,
+        transformer_initial_ff_inner_dropout=None,
+        transformer_initial_ff_outer_dropout=None,
         transformer_pre_norm=True,
         transformer_normformer=False,
         transformer_post_norm=False,
@@ -172,20 +173,28 @@ 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,
         transformer_ff_linear_module_up=None,
         transformer_ff_linear_module_down=None,
         transformer_msa_scaling="d",
-        transformer_mlp_dropout=0.0,
+        transformer_ff_dropout=0.0,
+        transformer_ff_inner_dropout=0.0,
+        transformer_ff_outer_dropout=0.0,
         transformer_msa_dropout=0.1,
         transformer_stochastic_depth=0.1,
         transformer_checkpoint_ff=True,
         linear_module=nn.Linear,
+        alpha=1.0,
+        beta=1.0,
     ):
         super().__init__()
 
+        self.alpha = alpha
+        self.beta = beta
+
         if cnn_activation_kwargs is not None:
             self.cnn_activation = cnn_activation(**cnn_activation_kwargs)
         else:
@@ -333,17 +342,22 @@ class ViTEncoder(nn.Module):
                 ff_linear_module_up=transformer_ff_linear_module_up,
                 ff_linear_module_down=transformer_ff_linear_module_down,
                 msa_scaling=transformer_msa_scaling,
-                mlp_dropout=transformer_mlp_dropout,
+                ff_dropout=transformer_ff_dropout,
+                ff_inner_dropout=transformer_ff_inner_dropout,
+                ff_outer_dropout=transformer_ff_outer_dropout,
                 msa_dropout=transformer_msa_dropout,
                 stochastic_depth=transformer_stochastic_depth,
                 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()
@@ -357,9 +371,21 @@ class ViTEncoder(nn.Module):
                 activation_kwargs=transformer_activation_kwargs,
                 dropout=(
                     # First truthy assigned value
-                    transformer_initial_ff_mlp_dropout
-                    if transformer_initial_ff_mlp_dropout is not None
-                    else transformer_mlp_dropout
+                    transformer_initial_ff_dropout
+                    if transformer_initial_ff_dropout is not None
+                    else transformer_ff_dropout
+                ),
+                inner_dropout=(
+                    # First truthy assigned value
+                    transformer_initial_ff_inner_dropout
+                    if transformer_initial_ff_inner_dropout is not None
+                    else transformer_ff_inner_dropout
+                ),
+                outer_dropout=(
+                    # First truthy assigned value
+                    transformer_initial_ff_outer_dropout
+                    if transformer_initial_ff_outer_dropout is not None
+                    else transformer_ff_outer_dropout
                 ),
                 linear_module_up=(
                     # First truthy assigned value
@@ -373,16 +399,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,
@@ -392,19 +416,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:
@@ -438,10 +464,11 @@ 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_mlp_dropout=None,
+        transformer_initial_ff_dropout=None,
+        transformer_initial_ff_inner_dropout=None,
+        transformer_initial_ff_outer_dropout=None,
         transformer_pre_norm=True,
         transformer_normformer=False,
         transformer_post_norm=False,
@@ -452,13 +479,16 @@ 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,
         transformer_ff_linear_module_up=None,
         transformer_ff_linear_module_down=None,
         transformer_msa_scaling="d",
-        transformer_mlp_dropout=0.0,
+        transformer_ff_dropout=0.0,
+        transformer_ff_inner_dropout=0.0,
+        transformer_ff_outer_dropout=0.0,
         transformer_msa_dropout=0.1,
         transformer_stochastic_depth=0.1,
         transformer_checkpoint_ff=True,
@@ -466,6 +496,8 @@ 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__()
@@ -486,6 +518,9 @@ class ViT(nn.Module):
                 "SwiGLU": SwiGLU,
             }[transformer_activation]
 
+        self.alpha = alpha
+        self.beta = beta
+
         self.encoder = ViTEncoder(
             input_size=input_size,
             initial_batch_norm=initial_batch_norm,
@@ -505,10 +540,11 @@ 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_mlp_dropout=transformer_initial_ff_mlp_dropout,
+            transformer_initial_ff_dropout=transformer_initial_ff_dropout,
+            transformer_initial_ff_inner_dropout=transformer_initial_ff_inner_dropout,
+            transformer_initial_ff_outer_dropout=transformer_initial_ff_outer_dropout,
             transformer_pre_norm=transformer_pre_norm,
             transformer_normformer=transformer_normformer,
             transformer_post_norm=transformer_post_norm,
@@ -519,17 +555,22 @@ 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,
             transformer_ff_linear_module_up=transformer_ff_linear_module_up,
             transformer_ff_linear_module_down=transformer_ff_linear_module_down,
             transformer_msa_scaling=transformer_msa_scaling,
-            transformer_mlp_dropout=transformer_mlp_dropout,
+            transformer_ff_dropout=transformer_ff_dropout,
+            transformer_ff_inner_dropout=transformer_ff_inner_dropout,
+            transformer_ff_outer_dropout=transformer_ff_outer_dropout,
             transformer_msa_dropout=transformer_msa_dropout,
             transformer_stochastic_depth=transformer_stochastic_depth,
             transformer_checkpoint_ff=transformer_checkpoint_ff,
             linear_module=linear_module,
+            alpha=alpha,
+            beta=beta,
         )
 
         self.pool = head(

{broccoli_ml-6.0.0.dist-info → broccoli_ml-13.0.1.dist-info}/METADATA

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

broccoli_ml-13.0.1.dist-info/RECORD

@@ -0,0 +1,13 @@
+broccoli/__init__.py,sha256=tmyspsVxqPZHRQCY_NRwpW4SMNBbtE8E_8z7l-SAzSo,127
+broccoli/activation.py,sha256=nrpTOrpg9k23_E4AJWy7VlXXAJCtCJCOR-TonEWJr04,3218
+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=3vAQQ75SAyr4-m3e7vSru8M-RpUy2Enp5cVUafaVYMU,28410
+broccoli/utils.py,sha256=oOWzn6dJ5nC_9r4zq0emmfmaYACJXJNFS48AOpW2jqc,358
+broccoli/vit.py,sha256=jd4e6MjL2JKB8ynSQssWRh6Hs36RuLj4uWyUNVhIMUY,22472
+broccoli_ml-13.0.1.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
+broccoli_ml-13.0.1.dist-info/METADATA,sha256=gN9cKQDpwRr8JG_Ilj0ZknlPdxhq62I8xEBExVspKGw,1369
+broccoli_ml-13.0.1.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+broccoli_ml-13.0.1.dist-info/RECORD,,

broccoli_ml-6.0.0.dist-info/RECORD

@@ -1,13 +0,0 @@
-broccoli/__init__.py,sha256=tmyspsVxqPZHRQCY_NRwpW4SMNBbtE8E_8z7l-SAzSo,127
-broccoli/activation.py,sha256=-Jf30C6iGqWCorC9HEGn2oduWwjeaCAxGLUUYIy1zX8,3438
-broccoli/cnn.py,sha256=WjoPDSpe3ttwxCBNfCVRdaCHvbeZ7G-a5_i8fUsK_d8,4889
-broccoli/linear.py,sha256=Y7s-DzcwsOipRboNHc4HTScw4mJRalNoVFsNcxOB6a4,4872
-broccoli/rope.py,sha256=GRqApBNmYCFaDak0WL1xE_BC5CTTYKQU_PBdeTcQcjc,12557
-broccoli/tensor.py,sha256=um8mrxkYbvNDo-QvHlmJm8Aw6qcngOlUZPoAk_PMReA,4480
-broccoli/transformer.py,sha256=MxIdzoxoWx_IWcq86vDZJIV4tk-dMNivhopZu8zJk90,23293
-broccoli/utils.py,sha256=oOWzn6dJ5nC_9r4zq0emmfmaYACJXJNFS48AOpW2jqc,358
-broccoli/vit.py,sha256=9oyh76ulmX5lDPMCDicQhhqm8RYCvJIgAJkDbYRVdi4,20873
-broccoli_ml-6.0.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
-broccoli_ml-6.0.0.dist-info/METADATA,sha256=Sv8nRPb7oCAeoMe3AAHIYDewAETvb0ZDxN8IKFniVHk,1368
-broccoli_ml-6.0.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-broccoli_ml-6.0.0.dist-info/RECORD,,