broccoli-ml 4.0.1__py3-none-any.whl → 5.1.0__py3-none-any.whl

broccoli/transformer.py

@@ -1,5 +1,5 @@
 import math
-from typing import Optional
+from typing import Optional, Tuple
 
 import torch
 import torch.nn as nn
@@ -102,6 +102,15 @@ class MHAttention(nn.Module):
 
         self.head_dim = self.embed_dim // self.n_heads
 
+        if self.scaling == "sqrtd":
+            self.scaling_factor = 1 / math.sqrt(self.head_dim)
+        elif self.scaling == "d":
+            # 8/d_model for backwards compatibility,
+            #     per https://github.com/microsoft/mup
+            self.scaling_factor = 8 / self.head_dim
+        else:
+            raise ValueError('`scaling` argument to MHAttention must be "d" or "sqrtd"')
+
         self.q_proj = linear_module(self.embed_dim, self.embed_dim, bias=False)
         self.k_proj = linear_module(self.embed_dim, self.embed_dim, bias=False)
         self.v_proj = linear_module(self.embed_dim, self.embed_dim, bias=False)
@@ -122,6 +131,8 @@ class MHAttention(nn.Module):
         self.source_size = source_size
         self.bos_tokens = bos_tokens
 
+        self.reset_parameters()
+
     @property
     def _kv_distance(self) -> float:
         """
@@ -141,7 +152,71 @@ class MHAttention(nn.Module):
 
         return 1 - similarity
 
-    def forward(self, q, k, v):
+    def add_axial_rope(
+        self, q: torch.Tensor, k: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Apply Axial RoPE to all tokens except BOS tokens
+        """
+
+        if len(self.source_size) == 1:
+            spatial_dimension_names = "D1"
+            spatial_dimension_values = {"D1": self.source_size[0]}
+        elif len(self.source_size) == 2:
+            spatial_dimension_names = "D1 D2"
+            spatial_dimension_values = {
+                "D1": self.source_size[0],
+                "D2": self.source_size[1],
+            }
+        elif len(self.source_size) == 3:
+            spatial_dimension_names = "D1 D2 D3"
+            spatial_dimension_values = {
+                "D1": self.source_size[0],
+                "D2": self.source_size[1],
+                "D3": self.source_size[2],
+            }
+        else:
+            raise NotImplementedError(
+                "`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_img = rearrange(
+            q_img,
+            f"b ({spatial_dimension_names}) d -> b {spatial_dimension_names} d",
+            **spatial_dimension_values,
+        )
+        k_img = rearrange(
+            k_img,
+            f"b ({spatial_dimension_names}) d -> b {spatial_dimension_names} d",
+            **spatial_dimension_values,
+        )
+
+        freqs = self.rotary_embedding.get_axial_freqs(*self.source_size)
+
+        q_img = apply_rotary_emb(freqs, q_img)
+        k_img = apply_rotary_emb(freqs, k_img)
+
+        q_img = rearrange(
+            q_img,
+            f"b {spatial_dimension_names} d -> b ({spatial_dimension_names}) d",
+        )
+        k_img = rearrange(
+            k_img,
+            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)
+
+        return q, k
+
+    def project_qkv(
+        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         query_batch_size, query_tokens, query_features = q.size()
         key_batch_size, key_tokens, key_features = k.size()
 
@@ -154,74 +229,18 @@ class MHAttention(nn.Module):
 
         if self.causal:
             assert query_tokens == key_tokens
-            assert query_tokens == self.sequence_length
+            assert query_tokens == self.seq_len
 
-        # Project q, k and v
-        q = self.q_proj(q)
-        k = self.k_proj(k)
-        v = self.v_proj(v)
+        q, k, v = self.q_proj(q), self.k_proj(k), self.v_proj(v)
 
-        # Rearrange dimensions and add RoPE if needed
         if self.rotary_embedding is not None:
+            q, k = self.add_axial_rope(q, k)
 
-            if len(self.source_size) == 1:
-                spatial_dimension_names = "D1"
-                spatial_dimension_values = {"D1": self.source_size[0]}
-            elif len(self.source_size) == 2:
-                spatial_dimension_names = "D1 D2"
-                spatial_dimension_values = {
-                    "D1": self.source_size[0],
-                    "D2": self.source_size[1],
-                }
-            elif len(self.source_size) == 3:
-                spatial_dimension_names = "D1 D2 D3"
-                spatial_dimension_values = {
-                    "D1": self.source_size[0],
-                    "D2": self.source_size[1],
-                    "D3": self.source_size[2],
-                }
-            else:
-                raise NotImplementedError(
-                    "`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_img = rearrange(
-                q_img,
-                f"b ({spatial_dimension_names}) d -> b {spatial_dimension_names} d",
-                **spatial_dimension_values,
-            )
-            k_img = rearrange(
-                k_img,
-                f"b ({spatial_dimension_names}) d -> b {spatial_dimension_names} d",
-                **spatial_dimension_values,
-            )
-            freqs = self.rotary_embedding.get_axial_freqs(*self.source_size)
-            q_img = apply_rotary_emb(freqs, q_img)
-            k_img = apply_rotary_emb(freqs, k_img)
-
-            q_img = rearrange(
-                q_img,
-                f"b {spatial_dimension_names} d -> b ({spatial_dimension_names}) d",
-            )
-            k_img = rearrange(
-                k_img,
-                f"b {spatial_dimension_names} d -> b ({spatial_dimension_names}) d",
-            )
+        return q, k, v
 
-            # 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)
+    def forward(self, q, k, v):
 
-        if self.scaling == "sqrtd":
-            scaling_factor = 1 / math.sqrt(self.head_dim)
-        elif self.scaling == "d":
-            # for backwards compatibility, per https://github.com/microsoft/mup
-            scaling_factor = 8 / self.head_dim
-        else:
-            raise ValueError('`scaling` argument to MHAttention must be "d" or "sqrtd"')
+        q, k, v = self.project_qkv(q, k, v)
 
         if FLASH_ATTN:
             # Divide Q/K/V into heads
@@ -234,7 +253,7 @@ class MHAttention(nn.Module):
                 k,
                 v,
                 dropout_p=self.dropout.p if self.training else 0.0,
-                softmax_scale=scaling_factor,
+                softmax_scale=self.scaling_factor,
                 causal=self.causal,
             )
 
@@ -249,7 +268,7 @@ class MHAttention(nn.Module):
 
             qk_scores = q @ k.transpose(-1, -2)
 
-            qk_scores *= scaling_factor
+            qk_scores *= self.scaling_factor
 
             # Apply mask if causal (must come before softmax)
             if self.causal:
@@ -265,6 +284,34 @@ class MHAttention(nn.Module):
 
             return self.out_proj(output_without_heads)
 
+    def attention_scores(self, q, k, v):
+
+        q, k, v = self.project_qkv(q, k, v)
+
+        # Divide Q/K/V into heads
+        q = rearrange(q, "b t (h d) -> b h t d", h=self.n_heads)
+        k = rearrange(k, "b t (h d) -> b h t d", h=self.n_heads)
+        v = rearrange(v, "b t (h d) -> b h t d", h=self.n_heads)
+
+        qk_scores = q @ k.transpose(-1, -2)
+
+        qk_scores *= self.scaling_factor
+
+        # 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)
+
+        return qk_scores  # (batch, head, seq_len, seq_len)
+
+    def reset_parameters(self):
+        # Default nn.Linear init is kaiming_uniform, which is fine
+        self.q_proj.reset_parameters()
+        self.k_proj.reset_parameters()
+        self.v_proj.reset_parameters()
+        self.out_proj.reset_parameters()
+
 
 class FeedforwardBlock(nn.Module):
     """
@@ -285,9 +332,11 @@ class FeedforwardBlock(nn.Module):
         normformer=False,
         post_norm=True,
         residual_path=True,
+        checkpoint=True,
     ):
         super().__init__()
 
+        self.checkpoint = checkpoint
         self.residual_path = residual_path
         self.post_norm = post_norm
 
@@ -325,13 +374,30 @@ class FeedforwardBlock(nn.Module):
             ]
         )
 
+        self.reset_parameters()
+
     def forward(self, x):
+
+        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 + self.process(x))
+            return self.layernorm(x + processed)
         elif self.residual_path:
-            return x + self.process(x)
+            return x + processed
         else:
-            return self.process(x)
+            return processed
+
+    def reset_parameters(self):
+        if self.post_norm:
+            self.layernorm.reset_parameters()
+
+        # Iterate over the sequential block to reset parameters
+        for module in self.process:
+            if hasattr(module, "reset_parameters"):
+                module.reset_parameters()
 
 
 class TransformerBlock(nn.Module):
@@ -365,6 +431,7 @@ class TransformerBlock(nn.Module):
         pre_norm=True,
         post_norm=False,
         normformer=False,
+        checkpoint_ff=True,
     ):
         """
         Args:
@@ -433,8 +500,11 @@ class TransformerBlock(nn.Module):
             normformer=normformer,
             post_norm=False,  # Handled outside the block
             residual_path=False,  # Handled outside the block
+            checkpoint=checkpoint_ff,
         )
 
+        self.reset_parameters()
+
     @property
     def _kv_distance(self) -> float:
         return self.attn._kv_distance
@@ -445,25 +515,43 @@ class TransformerBlock(nn.Module):
             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(checkpoint(self.ff, x, use_reentrant=False))
+            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(checkpoint(self.ff, x, use_reentrant=False))
+            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(checkpoint(self.ff, x, use_reentrant=False))
+            x = x + self.drop_path(self.ff(x))
 
         return x
 
+    def attention_scores(self, x):
+        """
+        Give back the attention scores used in this layer.
+        """
+        if self.pre_norm:
+            x = self.layer_norm_1(x)
+            return self.attn(x, x, x)
+        else:
+            return self.attn(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()
+
+        self.attn.reset_parameters()
+        self.ff.reset_parameters()
+
 
 class TransformerEncoder(nn.Module):
     """
     This assumes we already get a sequence of embeddings (e.g. word or image
-        patch embeddings). It uses learned positional embeddings.
+        patch embeddings).
     """
 
     def __init__(
@@ -491,6 +579,7 @@ class TransformerEncoder(nn.Module):
         post_norm=False,
         normformer=False,
         msa_scaling="d",
+        checkpoint_ff=True,
     ):
         """
         Args:
@@ -567,16 +656,19 @@ class TransformerEncoder(nn.Module):
                     pre_norm=pre_norm,
                     post_norm=post_norm,
                     normformer=normformer,
+                    checkpoint_ff=checkpoint_ff,
                 )
                 for i in range(n_layers)
             ]
         )
 
+        self.reset_parameters()
+
     @property
     def _kv_distances(self) -> float:
         return ",".join([str(block._kv_distance) for block in self.blocks])
 
-    def forward(self, x):
+    def preprocess(self, x):
         if self._bos_tokens:
             x = torch.cat([self._bos_embedding.expand(x.size(0), -1, -1), x], dim=1)
         else:
@@ -591,6 +683,10 @@ class TransformerEncoder(nn.Module):
                 )  # to shape (1, seq_len) to broadcast over batch
             )
 
+    def forward(self, x):
+
+        x = self.preprocess(x)
+
         for block in self.blocks:
             x = block(x)
 
@@ -598,3 +694,27 @@ class TransformerEncoder(nn.Module):
             return x[:, self._bos_tokens :, :]
         else:
             return x
+
+    def attention_scores(self, x):
+
+        x = self.preprocess(x)
+
+        layer_scores = []
+
+        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)
+            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.absolute_position_embedding is not None:
+            self.absolute_position_embedding.reset_parameters()
+
+        for block in self.blocks:
+            block.reset_parameters()

broccoli/vit.py

@@ -9,7 +9,9 @@ from .utils import PadTensor
 from einops import einsum
 from einops.layers.torch import Rearrange
 
+import torch
 import torch.nn as nn
+import torch.nn.functional as F
 
 
 class GetCLSToken(nn.Module):
@@ -31,10 +33,18 @@ class SequencePool(nn.Module):
             ]
         )
 
+        self.reset_parameters()
+
     def forward(self, x):
         weights = self.attention(x)
         return einsum(weights, x, "batch seq, batch seq d_model -> batch d_model")
 
+    def reset_parameters(self):
+        # Iterate over modules in the sequential block
+        for module in self.attention:
+            if hasattr(module, "reset_parameters"):
+                module.reset_parameters()
+
 
 class ClassificationHead(nn.Module):
     """
@@ -71,9 +81,16 @@ class ClassificationHead(nn.Module):
             ]
         )
 
+        self.reset_parameters()
+
     def forward(self, x):
         return self.classification_process(x)
 
+    def reset_parameters(self):
+        for module in self.classification_process:
+            if hasattr(module, "reset_parameters"):
+                module.reset_parameters()
+
 
 class SequencePoolClassificationHead(ClassificationHead):
     """
@@ -106,6 +123,8 @@ class SequencePoolClassificationHead(ClassificationHead):
             ]
         )
 
+        self.reset_parameters()
+
 
 class ViTEncoder(nn.Module):
     """
@@ -160,6 +179,7 @@ class ViTEncoder(nn.Module):
         transformer_mlp_dropout=0.0,
         transformer_msa_dropout=0.1,
         transformer_stochastic_depth=0.1,
+        transformer_checkpoint_ff=True,
         linear_module=nn.Linear,
     ):
         super().__init__()
@@ -321,6 +341,7 @@ class ViTEncoder(nn.Module):
                 pre_norm=transformer_pre_norm,
                 normformer=transformer_normformer,
                 post_norm=transformer_post_norm,
+                checkpoint_ff=transformer_checkpoint_ff,
             )
         else:
             self.transformer = nn.Identity()
@@ -354,6 +375,7 @@ class ViTEncoder(nn.Module):
                 normformer=transformer_normformer,
                 post_norm=transformer_post_norm,
                 residual_path=transformer_initial_ff_residual_path,
+                checkpoint=transformer_checkpoint_ff,
             )
         else:
             self.initial_ff = nn.Identity()
@@ -373,9 +395,20 @@ class ViTEncoder(nn.Module):
             ]
         )
 
+        self.reset_parameters()
+
     def forward(self, x):
         return self.encoder(x)
 
+    def attention_scores(self, x):
+        x = self.encoder[:-1](x)
+        return self.encoder[-1].attention_scores(x)
+
+    def reset_parameters(self):
+        for module in self.encoder:
+            if hasattr(module, "reset_parameters"):
+                module.reset_parameters()
+
 
 class ViT(nn.Module):
     """
@@ -426,6 +459,7 @@ class ViT(nn.Module):
         transformer_mlp_dropout=0.0,
         transformer_msa_dropout=0.1,
         transformer_stochastic_depth=0.1,
+        transformer_checkpoint_ff=True,
         head=SequencePoolClassificationHead,
         batch_norm_logits=True,
         logit_projection_layer=nn.Linear,
@@ -492,6 +526,7 @@ class ViT(nn.Module):
             transformer_mlp_dropout=transformer_mlp_dropout,
             transformer_msa_dropout=transformer_msa_dropout,
             transformer_stochastic_depth=transformer_stochastic_depth,
+            transformer_checkpoint_ff=transformer_checkpoint_ff,
             linear_module=linear_module,
         )
 
@@ -502,9 +537,26 @@ class ViT(nn.Module):
             batch_norm_logits=batch_norm_logits,
         )
 
+        self.reset_parameters()
+
     @property
     def sequence_length(self):
         return self.encoder.sequence_length
 
     def forward(self, x):
         return self.pool(self.encoder(x))
+
+    def attention_scores(self, x):
+        return self.encoder.attention_scores(x)
+
+    def head_to_bos_token_attention(self, x):
+        all_attention = self.attention_scores(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)
+
+    def reset_parameters(self):
+        self.encoder.reset_parameters()
+        self.pool.reset_parameters()

{broccoli_ml-4.0.1.dist-info → broccoli_ml-5.1.0.dist-info}/METADATA

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

{broccoli_ml-4.0.1.dist-info → broccoli_ml-5.1.0.dist-info}/RECORD

@@ -4,10 +4,10 @@ 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=uqSf8q30MF7Ds7LfqW8Pr206NXpSlf7o6770KZu81Ew,19883
+broccoli/transformer.py,sha256=x3Mo6_1x6fGG6lPDPx9srxn6UdwKEpvjFAO8zoMwAMI,23052
 broccoli/utils.py,sha256=oOWzn6dJ5nC_9r4zq0emmfmaYACJXJNFS48AOpW2jqc,358
-broccoli/vit.py,sha256=_5uLcklmJ1Uoj7V7TkzF0UqroVnl8NCHun5B0mORmOg,18651
-broccoli_ml-4.0.1.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
-broccoli_ml-4.0.1.dist-info/METADATA,sha256=vH_utDdo0-e2q8ReDrRHQ1d6fOzG4nzb9EWlqgyl4XY,1368
-broccoli_ml-4.0.1.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-broccoli_ml-4.0.1.dist-info/RECORD,,
+broccoli/vit.py,sha256=tUYQyoDsBc5ZR_M5_J0huj0T3OAy-vn1f19hCGVDCrM,20425
+broccoli_ml-5.1.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
+broccoli_ml-5.1.0.dist-info/METADATA,sha256=3986lqn1iuWJ53O8ckM9LVU3tTjr32i19SeIXauWDXw,1368
+broccoli_ml-5.1.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+broccoli_ml-5.1.0.dist-info/RECORD,,