broccoli-ml 5.1.1__tar.gz → 9.5.1__tar.gz

{broccoli_ml-5.1.1 → broccoli_ml-9.5.1}/PKG-INFO

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

{broccoli_ml-5.1.1 → broccoli_ml-9.5.1}/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_ml-9.5.1/broccoli/linear.py

@@ -0,0 +1,352 @@
+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
+
+from .tensor import SigmaReparamTensor, AnchoredReparamTensor, NormReparamTensor
+
+
+class SpectralNormLinear(nn.Module):
+    """
+    Inspired by Apple's Spectral Normed Linear Layers
+        (https://github.com/apple/ml-sigma-reparam)
+    """
+
+    def __init__(self, in_features: int, out_features: int, bias: bool = True):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.use_bias = bias
+
+        self.weights = None
+
+        # Define the bias vector as a learnable parameter if required.
+        if self.use_bias:
+            self.bias = nn.Parameter(torch.empty(out_features))
+        else:
+            # If no bias, register it as None.
+            # This is important so that PyTorch doesn't complain when saving/loading the model.
+            self.register_parameter("bias", None)
+
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        weights = torch.empty(self.out_features, self.in_features)
+        stdv = 1.0 / math.sqrt(self.in_features)
+        nn.init.uniform_(weights, a=-stdv, b=stdv)
+        if self.use_bias:
+            fan_in, _ = nn.init._calculate_fan_in_and_fan_out(weights)
+            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
+            nn.init.uniform_(self.bias, -bound, bound)
+        self.weights = SigmaReparamTensor(weights)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return F.linear(x, self.weights(), self.bias)
+
+    def __repr__(self) -> str:
+        # Optional: A nice representation for printing the module.
+        return (
+            f"SpectralNormFeedForward(in_features={self.in_features},"
+            f"out_features={self.out_features}, bias={self.use_bias})"
+        )
+
+
+class AnchoredLinear(nn.Module):
+    """
+    ...
+    """
+
+    def __init__(self, in_features: int, out_features: int, bias: bool = True):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.use_bias = bias
+
+        self.weights = None
+
+        # Define the bias vector as a learnable parameter if required.
+        if self.use_bias:
+            self.bias = nn.Parameter(torch.empty(out_features))
+        else:
+            # If no bias, register it as None.
+            # This is important so that PyTorch doesn't complain when saving/loading the model.
+            self.register_parameter("bias", None)
+
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        weights = torch.empty(self.out_features, self.in_features)
+        stdv = 1.0 / math.sqrt(self.in_features)
+        nn.init.uniform_(weights, a=-stdv, b=stdv)
+        if self.use_bias:
+            fan_in, _ = nn.init._calculate_fan_in_and_fan_out(weights)
+            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
+            nn.init.uniform_(self.bias, -bound, bound)
+        self.weights = AnchoredReparamTensor(weights)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return F.linear(x, self.weights(), self.bias)
+
+    def __repr__(self) -> str:
+        # Optional: A nice representation for printing the module.
+        return (
+            f"AnchoredLinear(in_features={self.in_features},"
+            f"out_features={self.out_features}, bias={self.use_bias})"
+        )
+
+
+class WeightNormedLinear(nn.Module):
+    """
+    ...
+    """
+
+    def __init__(self, in_features: int, out_features: int, bias: bool = True):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.use_bias = bias
+
+        self.weights = None
+
+        # Define the bias vector as a learnable parameter if required.
+        if self.use_bias:
+            self.bias = nn.Parameter(torch.empty(out_features))
+        else:
+            # If no bias, register it as None.
+            # This is important so that PyTorch doesn't complain when saving/loading the model.
+            self.register_parameter("bias", None)
+
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        weights = torch.empty(self.out_features, self.in_features)
+        stdv = 1.0 / math.sqrt(self.in_features)
+        nn.init.uniform_(weights, a=-stdv, b=stdv)
+        if self.use_bias:
+            fan_in, _ = nn.init._calculate_fan_in_and_fan_out(weights)
+            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
+            nn.init.uniform_(self.bias, -bound, bound)
+        self.weights = NormReparamTensor(weights)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return F.linear(x, self.weights(), self.bias)
+
+    def __repr__(self) -> str:
+        return (
+            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):
+        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):
+            value_indices = indices
+            centred_value_weights = self._mean_value_weights()
+            centred_value_weights = centred_value_weights.expand(indices.size(0), -1)
+            if self.xglu:
+                gate_indices = indices
+                value_indices = indices + (self.linear.out_features // 2)
+                centred_gate_weights = self._mean_gate_weights()
+                centred_gate_weights = centred_gate_weights.expand(indices.size(0), -1)
+                self._update_weights(
+                    gate_indices, 0, centred_gate_weights, self.optimisers  # dim
+                )
+            self._update_weights(
+                value_indices, 0, centred_value_weights, self.optimisers
+            )
+        else:
+            return
+
+    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_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_ml-5.1.1 → broccoli_ml-9.5.1}/broccoli/transformer.py

@@ -13,6 +13,7 @@ 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
@@ -76,7 +77,7 @@ class MHAttention(nn.Module):
         causal=False,
         seq_len=None,
         linear_module: nn.Module = nn.Linear,
-        bos_tokens=0,
+        utility_tokens=0,
         rotary_embedding=None,
         source_size=None,
         scaling="d",
@@ -129,7 +130,7 @@ class MHAttention(nn.Module):
             )
         self.rotary_embedding = rotary_embedding
         self.source_size = source_size
-        self.bos_tokens = bos_tokens
+        self.utility_tokens = utility_tokens
 
         self.reset_parameters()
 
@@ -156,7 +157,7 @@ class MHAttention(nn.Module):
         self, q: torch.Tensor, k: torch.Tensor
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         """
-        Apply Axial RoPE to all tokens except BOS tokens
+        Apply Axial RoPE to all tokens except utility tokens
         """
 
         if len(self.source_size) == 1:
@@ -180,8 +181,8 @@ class MHAttention(nn.Module):
                 "`source_size` must be a tuple of 1, 2 or 3 integers"
             )
 
-        q_bos, q_img = q[:, : self.bos_tokens, :], q[:, self.bos_tokens :, :]
-        k_bos, k_img = k[:, : self.bos_tokens, :], k[:, self.bos_tokens :, :]
+        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,
@@ -208,9 +209,9 @@ class MHAttention(nn.Module):
             f"b {spatial_dimension_names} d -> b ({spatial_dimension_names}) d",
         )
 
-        # Re-combine the BOS tokens and the RoPE-enhanced image tokens
-        q = torch.cat([q_bos, q_img], dim=1)
-        k = torch.cat([k_bos, k_img], dim=1)
+        # 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)
 
         return q, k
 
@@ -284,7 +285,7 @@ class MHAttention(nn.Module):
 
             return self.out_proj(output_without_heads)
 
-    def attention_scores(self, q, k, v):
+    def attention_logits(self, q, k, v):
 
         q, k, v = self.project_qkv(q, k, v)
 
@@ -301,8 +302,6 @@ class MHAttention(nn.Module):
         if self.causal:
             qk_scores.masked_fill_(self.mask, float("-inf"))
 
-        qk_scores = F.softmax(qk_scores, dim=-1)
-
         return qk_scores  # (batch, head, seq_len, seq_len)
 
     def reset_parameters(self):
@@ -326,6 +325,8 @@ 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,
@@ -339,6 +340,7 @@ class FeedforwardBlock(nn.Module):
         self.checkpoint = checkpoint
         self.residual_path = residual_path
         self.post_norm = post_norm
+        self.xglu = activation.__name__.endswith("GLU")
 
         if self.residual_path and (output_features < input_features):
             raise ValueError(
@@ -355,29 +357,63 @@ 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 * ratio * output_features if self.xglu else 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.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,
+                self.inner_dropout,
                 nn.LayerNorm(ratio * output_features) if normformer else nn.Identity(),
-                linear_module_down(ratio * output_features, output_features),
-                self.dropout,
+                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)
+            self.linear_out.reset_columns(indices)
+
         if self.checkpoint:
             processed = checkpoint(self.process, x, use_reentrant=False)
         else:
@@ -416,14 +452,16 @@ class TransformerBlock(nn.Module):
         n_heads,
         relative_position_embedding=False,
         source_size=None,
-        bos_tokens=0,
+        utility_tokens=0,
         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,
@@ -474,7 +512,7 @@ class TransformerBlock(nn.Module):
             linear_module=linear_module,
             rotary_embedding=self.rotary_embedding,
             source_size=source_size,
-            bos_tokens=bos_tokens,
+            utility_tokens=utility_tokens,
             scaling=msa_scaling,
         )
 
@@ -485,7 +523,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
@@ -529,15 +569,15 @@ class TransformerBlock(nn.Module):
 
         return x
 
-    def attention_scores(self, x):
+    def attention_logits(self, x):
         """
         Give back the attention scores used in this layer.
         """
         if self.pre_norm:
             x = self.layer_norm_1(x)
-            return self.attn.attention_scores(x, x, x)
+            return self.attn.attention_logits(x, x, x)
         else:
-            return self.attn.attention_scores(x, x, x)
+            return self.attn.attention_logits(x, x, x)
 
     def reset_parameters(self):
         self.layer_norm_1.reset_parameters()
@@ -568,13 +608,15 @@ 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,
-        bos_tokens=0,
-        return_bos_tokens=False,
+        utility_tokens=0,
+        return_utility_tokens=False,
         pre_norm=True,
         post_norm=False,
         normformer=False,
@@ -592,22 +634,33 @@ 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__()
         self.seq_len = seq_len
         self.n_heads = n_heads
-        self._bos_tokens = bos_tokens
-        self.return_bos_tokens = return_bos_tokens
-
-        # Initialise BOS tokens with normal init, like usual Pytorch embeddings
-        if self._bos_tokens:
-            self._bos_embedding = nn.Parameter(torch.empty(self._bos_tokens, d_model))
-            nn.init.normal_(self._bos_embedding, mean=0.0, std=1.0)
-            self.full_sequence_length = self.seq_len + self._bos_tokens
+        self._utility_tokens = utility_tokens
+        self.return_utility_tokens = return_utility_tokens
+
+        # Initialise utility tokens with normal init, like usual Pytorch embeddings
+        if self._utility_tokens:
+            self._utility_token_embedding = nn.Parameter(
+                torch.empty(self._utility_tokens, d_model)
+            )
+            nn.init.normal_(self._utility_token_embedding, mean=0.0, std=1.0)
+        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._bos_embedding = None
             self.full_sequence_length = self.seq_len
 
         self.d_model = d_model
@@ -619,7 +672,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
 
@@ -641,14 +694,16 @@ class TransformerEncoder(nn.Module):
                     n_heads,
                     relative_position_embedding=relative_position_embedding,
                     source_size=source_size,
-                    bos_tokens=bos_tokens,
+                    utility_tokens=utility_tokens,
                     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,
@@ -669,8 +724,10 @@ class TransformerEncoder(nn.Module):
         return ",".join([str(block._kv_distance) for block in self.blocks])
 
     def preprocess(self, x):
-        if self._bos_tokens:
-            x = torch.cat([self._bos_embedding.expand(x.size(0), -1, -1), x], dim=1)
+        if self._utility_tokens:
+            x = torch.cat(
+                [self._utility_token_embedding.expand(x.size(0), -1, -1), x], dim=1
+            )
         else:
             x = x
 
@@ -692,12 +749,12 @@ class TransformerEncoder(nn.Module):
         for block in self.blocks:
             x = block(x)
 
-        if self._bos_tokens and not self.return_bos_tokens:
-            return x[:, self._bos_tokens :, :]
+        if self._utility_tokens and not self.return_utility_tokens:
+            return x[:, self._utility_tokens :, :]
         else:
             return x
 
-    def attention_scores(self, x):
+    def attention_logits(self, x):
 
         x = self.preprocess(x)
 
@@ -705,15 +762,15 @@ class TransformerEncoder(nn.Module):
 
         for block in self.blocks:
             # Get attention scores with shape (batch, 1, head, seq_len, seq_len)
-            layer_attention_scores = block.attention_scores(x).unsqueeze(1)
-            layer_scores.append(layer_attention_scores)
+            layer_attention_logits = block.attention_logits(x).unsqueeze(1)
+            layer_scores.append(layer_attention_logits)
             x = block(x)
 
         return torch.cat(layer_scores, dim=1)  # (batch, layer, head, seq_len, seq_len)
 
     def reset_parameters(self):
-        if self._bos_embedding is not None:
-            nn.init.normal_(self._bos_embedding, mean=0.0, std=1.0)
+        if self._utility_token_embedding is not None:
+            nn.init.normal_(self._utility_token_embedding, mean=0.0, std=1.0)
 
         if self.absolute_position_embedding is not None:
             self.absolute_position_embedding.reset_parameters()

{broccoli_ml-5.1.1 → broccoli_ml-9.5.1}/broccoli/vit.py

@@ -11,7 +11,6 @@ from einops.layers.torch import Rearrange
 
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 
 
 class GetCLSToken(nn.Module):
@@ -39,6 +38,9 @@ class SequencePool(nn.Module):
         weights = self.attention(x)
         return einsum(weights, x, "batch seq, batch seq d_model -> batch d_model")
 
+    def attention_scores(self, x):
+        return self.attention(x)
+
     def reset_parameters(self):
         # Iterate over modules in the sequential block
         for module in self.attention:
@@ -159,7 +161,9 @@ class ViTEncoder(nn.Module):
         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,
@@ -169,14 +173,16 @@ class ViTEncoder(nn.Module):
         transformer_layers=7,
         transformer_heads=4,
         transformer_mlp_ratio=2,
-        transformer_bos_tokens=0,
-        transformer_return_bos_tokens=False,
+        transformer_utility_tokens=0,
+        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,
@@ -331,13 +337,15 @@ 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,
-                bos_tokens=transformer_bos_tokens,
-                return_bos_tokens=transformer_return_bos_tokens,
+                utility_tokens=transformer_utility_tokens,
+                return_utility_tokens=transformer_return_utility_tokens,
                 pre_norm=transformer_pre_norm,
                 normformer=transformer_normformer,
                 post_norm=transformer_post_norm,
@@ -355,9 +363,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
@@ -400,9 +420,9 @@ class ViTEncoder(nn.Module):
     def forward(self, x):
         return self.encoder(x)
 
-    def attention_scores(self, x):
+    def attention_logits(self, x):
         x = self.encoder[:-1](x)
-        return self.encoder[-1].attention_scores(x)
+        return self.encoder[-1].attention_logits(x)
 
     def reset_parameters(self):
         for module in self.encoder:
@@ -439,7 +459,9 @@ class ViT(nn.Module):
         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,
@@ -449,14 +471,16 @@ class ViT(nn.Module):
         transformer_layers=7,
         transformer_heads=4,
         transformer_mlp_ratio=2,
-        transformer_bos_tokens=0,
-        transformer_return_bos_tokens=False,
+        transformer_utility_tokens=0,
+        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,
@@ -506,7 +530,9 @@ class ViT(nn.Module):
             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,
@@ -516,14 +542,16 @@ class ViT(nn.Module):
             transformer_layers=transformer_layers,
             transformer_heads=transformer_heads,
             transformer_mlp_ratio=transformer_mlp_ratio,
-            transformer_bos_tokens=transformer_bos_tokens,
-            transformer_return_bos_tokens=transformer_return_bos_tokens,
+            transformer_utility_tokens=transformer_utility_tokens,
+            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,
@@ -546,16 +574,26 @@ class ViT(nn.Module):
     def forward(self, x):
         return self.pool(self.encoder(x))
 
-    def attention_scores(self, x):
-        return self.encoder.attention_scores(x)
+    def attention_logits(self, x):
+        return self.encoder.attention_logits(x)
+
+    def pool_attention(self, x):
+        if hasattr(self.pool.summarize, "attention"):
+            return self.pool.summarize.attention(self.encoder(x))
+        else:
+            raise NotImplementedError(
+                "`pool_attention` is currently only implemented where"
+                " head class is SequencePoolClassificationHead"
+            )
 
-    def head_to_bos_token_attention(self, x):
-        all_attention = self.attention_scores(x)
+    def head_to_utility_token_attention_logits(self, x):
+        all_attention = self.attention_logits(x)
         batch_averages = torch.mean(all_attention, dim=0, keepdim=False)
         sequence_averages = torch.mean(batch_averages, dim=-1, keepdim=False)
-        n_bos_tokens = self.encoder.encoder._bos_tokens
-        just_bos = sequence_averages[:, :, :n_bos_tokens]
-        return F.softmax(just_bos, dim=-1)  # (layer, head, bos_token)
+        n_utility_tokens = self.encoder.encoder[-1]._utility_tokens
+        return sequence_averages[
+            :, :, :n_utility_tokens
+        ]  # (layer, head, utility_tokens)
 
     def reset_parameters(self):
         self.encoder.reset_parameters()

{broccoli_ml-5.1.1 → broccoli_ml-9.5.1}/pyproject.toml

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

broccoli_ml-5.1.1/broccoli/linear.py

@@ -1,138 +0,0 @@
-import math
-import torch
-from torch import nn
-from torch.nn import functional as F
-
-from .tensor import SigmaReparamTensor, AnchoredReparamTensor, NormReparamTensor
-
-
-class SpectralNormLinear(nn.Module):
-    """
-    Inspired by Apple's Spectral Normed Linear Layers
-        (https://github.com/apple/ml-sigma-reparam)
-    """
-
-    def __init__(self, in_features: int, out_features: int, bias: bool = True):
-        super().__init__()
-        self.in_features = in_features
-        self.out_features = out_features
-        self.use_bias = bias
-
-        self.weights = None
-
-        # Define the bias vector as a learnable parameter if required.
-        if self.use_bias:
-            self.bias = nn.Parameter(torch.empty(out_features))
-        else:
-            # If no bias, register it as None.
-            # This is important so that PyTorch doesn't complain when saving/loading the model.
-            self.register_parameter("bias", None)
-
-        self.reset_parameters()
-
-    def reset_parameters(self) -> None:
-        weights = torch.empty(self.out_features, self.in_features)
-        stdv = 1.0 / math.sqrt(self.in_features)
-        nn.init.uniform_(weights, a=-stdv, b=stdv)
-        if self.use_bias:
-            fan_in, _ = nn.init._calculate_fan_in_and_fan_out(weights)
-            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
-            nn.init.uniform_(self.bias, -bound, bound)
-        self.weights = SigmaReparamTensor(weights)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return F.linear(x, self.weights(), self.bias)
-
-    def __repr__(self) -> str:
-        # Optional: A nice representation for printing the module.
-        return (
-            f"SpectralNormFeedForward(in_features={self.in_features},"
-            f"out_features={self.out_features}, bias={self.use_bias})"
-        )
-
-
-class AnchoredLinear(nn.Module):
-    """
-    ...
-    """
-
-    def __init__(self, in_features: int, out_features: int, bias: bool = True):
-        super().__init__()
-        self.in_features = in_features
-        self.out_features = out_features
-        self.use_bias = bias
-
-        self.weights = None
-
-        # Define the bias vector as a learnable parameter if required.
-        if self.use_bias:
-            self.bias = nn.Parameter(torch.empty(out_features))
-        else:
-            # If no bias, register it as None.
-            # This is important so that PyTorch doesn't complain when saving/loading the model.
-            self.register_parameter("bias", None)
-
-        self.reset_parameters()
-
-    def reset_parameters(self) -> None:
-        weights = torch.empty(self.out_features, self.in_features)
-        stdv = 1.0 / math.sqrt(self.in_features)
-        nn.init.uniform_(weights, a=-stdv, b=stdv)
-        if self.use_bias:
-            fan_in, _ = nn.init._calculate_fan_in_and_fan_out(weights)
-            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
-            nn.init.uniform_(self.bias, -bound, bound)
-        self.weights = AnchoredReparamTensor(weights)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return F.linear(x, self.weights(), self.bias)
-
-    def __repr__(self) -> str:
-        # Optional: A nice representation for printing the module.
-        return (
-            f"AnchoredLinear(in_features={self.in_features},"
-            f"out_features={self.out_features}, bias={self.use_bias})"
-        )
-
-
-class WeightNormedLinear(nn.Module):
-    """
-    ...
-    """
-
-    def __init__(self, in_features: int, out_features: int, bias: bool = True):
-        super().__init__()
-        self.in_features = in_features
-        self.out_features = out_features
-        self.use_bias = bias
-
-        self.weights = None
-
-        # Define the bias vector as a learnable parameter if required.
-        if self.use_bias:
-            self.bias = nn.Parameter(torch.empty(out_features))
-        else:
-            # If no bias, register it as None.
-            # This is important so that PyTorch doesn't complain when saving/loading the model.
-            self.register_parameter("bias", None)
-
-        self.reset_parameters()
-
-    def reset_parameters(self) -> None:
-        weights = torch.empty(self.out_features, self.in_features)
-        stdv = 1.0 / math.sqrt(self.in_features)
-        nn.init.uniform_(weights, a=-stdv, b=stdv)
-        if self.use_bias:
-            fan_in, _ = nn.init._calculate_fan_in_and_fan_out(weights)
-            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
-            nn.init.uniform_(self.bias, -bound, bound)
-        self.weights = NormReparamTensor(weights)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return F.linear(x, self.weights(), self.bias)
-
-    def __repr__(self) -> str:
-        return (
-            f"WeightNormedLinear(in_features={self.in_features},"
-            f"out_features={self.out_features}, bias={self.use_bias})"
-        )