broccoli-ml 1.0.0__py3-none-any.whl → 1.2.0__py3-none-any.whl

broccoli/transformer.py

@@ -73,7 +73,15 @@ class MHAttention(nn.Module):
         bos_tokens=0,
         rotary_embedding=None,
         source_size=None,
+        scaling="d",
     ):
+        """
+        Args:
+            scaling: how should the attention logits be scaled? Can be "sqrtd"
+                to mimic the original Attention is All You Need approach of
+                dividing by the sqrt of the embedding Dimension or "d" per
+                "Tensor Programs V...". Default "d"
+        """
         super().__init__()
 
         if rotary_embedding is not None:
@@ -84,6 +92,7 @@ class MHAttention(nn.Module):
         self.embed_dim = embed_dim
         self.n_heads = n_heads
         assert embed_dim % n_heads == 0
+        self.scaling = scaling
 
         self.head_dim = self.embed_dim // self.n_heads
 
@@ -207,7 +216,12 @@ class MHAttention(nn.Module):
 
         qk_scores = q @ k.transpose(-1, -2)
 
-        qk_scores /= math.sqrt(self.head_dim)
+        if self.scaling == "sqrtd":
+            qk_scores /= math.sqrt(self.head_dim)
+        elif self.scaling == "d":
+            qk_scores /= self.head_dim
+        else:
+            raise ValueError('`scaling` argument to MHAttention must be "d" or "sqrtd"')
 
         # Apply mask if causal (must come before softmax)
         if self.causal:
@@ -305,6 +319,7 @@ class TransformerBlock(nn.Module):
         activation_kwargs: Optional[dict] = None,
         ff_linear_module_up=None,
         ff_linear_module_down=None,
+        msa_scaling="d",
         mlp_dropout=0.0,
         msa_dropout=0.0,
         identity_probability=0.0,
@@ -314,6 +329,14 @@ class TransformerBlock(nn.Module):
         post_norm=False,
         normformer=False,
     ):
+        """
+        Args:
+            msa_scaling: how should the attention logits be scaled? Can be "sqrtd"
+                to mimic the original Attention is All You Need approach of
+                dividing by the sqrt of the embedding Dimension or "d" per
+                "Tensor Programs V...". Default "d"
+        """
+
         super().__init__()
 
         self.pre_norm = pre_norm
@@ -348,6 +371,7 @@ class TransformerBlock(nn.Module):
             rotary_embedding=self.rotary_embedding,
             source_size=source_size,
             bos_tokens=bos_tokens,
+            scaling=msa_scaling,
         )
 
         # Submodule for the feedforward process
@@ -429,9 +453,21 @@ class TransformerEncoder(nn.Module):
         pre_norm=True,
         post_norm=False,
         normformer=False,
+        msa_scaling="d",
     ):
-        if position_embedding_type == "relative":
-            assert source_size is not None  # TODO: make this a proper exception
+        """
+        Args:
+            msa_scaling: how should the attention logits be scaled? Can be "sqrtd"
+                to mimic the original Attention is All You Need approach of
+                dividing by the sqrt of the embedding Dimension or "d" per
+                "Tensor Programs V...". Default "d"
+        """
+
+        if (position_embedding_type == "relative") and (source_size is None):
+            raise ValueError(
+                "`source_size` for TransformerEncoder cannot be None if"
+                " `position_embedding_type` is relative"
+            )
 
         super().__init__()
         self.seq_len = seq_len
@@ -461,7 +497,8 @@ class TransformerEncoder(nn.Module):
         self.msa_dropout = msa_dropout
         self.stochastic_depth = stochastic_depth
 
-        assert isinstance(n_layers, int)  # XXX: make this a proper Exception
+        assert isinstance(n_layers, int)
+
         if n_layers == 1:
             self.stochastic_depth_probabilities = [0.0]
         else:
@@ -484,6 +521,7 @@ class TransformerEncoder(nn.Module):
                     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,
                     msa_dropout=msa_dropout,
                     identity_probability=self.stochastic_depth_probabilities[i],

broccoli/vit.py

@@ -53,12 +53,20 @@ class ClassificationHead(nn.Module):
     A general classification head for a ViT
     """
 
-    def __init__(self, d_model, linear_module, n_classes, batch_norm=True):
+    def __init__(
+        self,
+        d_model,
+        n_classes,
+        linear_module=nn.Linear,
+        logit_projection_layer=nn.Linear,
+        batch_norm_logits=True,
+    ):
         super().__init__()
         self.d_model = d_model
         self.summarize = GetCLSToken()
-        self.projection = linear_module(d_model, n_classes)
-        if batch_norm:
+        self.projection = logit_projection_layer(d_model, n_classes)
+
+        if batch_norm_logits:
             self.batch_norm = nn.BatchNorm1d(n_classes, affine=False)
         else:
             self.batch_norm = nn.Identity()
@@ -83,7 +91,7 @@ class SequencePoolClassificationHead(ClassificationHead):
     """
 
     def __init__(self, d_model, linear_module, out_dim, batch_norm=True):
-        super().__init__(d_model, linear_module, out_dim, batch_norm=batch_norm)
+        super().__init__(d_model, linear_module, out_dim, batch_norm_logits=batch_norm)
         self.summarize = SequencePool(d_model, linear_module)
         # Rebuild the classification process with the correct summary module:
         self.classification_process = nn.Sequential(
@@ -143,6 +151,7 @@ class ViTEncoder(nn.Module):
         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_msa_dropout=0.1,
         transformer_stochastic_depth=0.1,
@@ -295,6 +304,7 @@ class ViTEncoder(nn.Module):
                 activation_kwargs=transformer_activation_kwargs,
                 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,
                 msa_dropout=transformer_msa_dropout,
                 stochastic_depth=transformer_stochastic_depth,
@@ -405,12 +415,14 @@ class ViT(nn.Module):
         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_msa_dropout=0.1,
         transformer_stochastic_depth=0.1,
-        batch_norm_outputs=True,
-        linear_module=SpectralNormLinear,
         head=SequencePoolClassificationHead,
+        batch_norm_logits=True,
+        logit_projection_layer=nn.Linear,
+        linear_module=nn.Linear,
     ):
 
         super().__init__()
@@ -468,6 +480,7 @@ class ViT(nn.Module):
             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_msa_dropout=transformer_msa_dropout,
             transformer_stochastic_depth=transformer_stochastic_depth,
@@ -476,9 +489,10 @@ class ViT(nn.Module):
 
         self.pool = head(
             transformer_embedding_size,
-            linear_module,
             image_classes,
-            batch_norm=batch_norm_outputs,
+            linear_module=linear_module,
+            logit_projection_layer=logit_projection_layer,
+            batch_norm=batch_norm_logits,
         )
 
     @property

{broccoli_ml-1.0.0.dist-info → broccoli_ml-1.2.0.dist-info}/METADATA

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

{broccoli_ml-1.0.0.dist-info → broccoli_ml-1.2.0.dist-info}/RECORD

@@ -8,10 +8,10 @@ broccoli/eigenpatches.py,sha256=J6n2usN1oQuHEHYiBNyYpn_a9eQcHjOBiIlvSei520Y,2413
 broccoli/linear.py,sha256=w021EDzWVDEu9odzrf9QwBZ3G8Ydu3nroV8soIJeRng,4894
 broccoli/rope.py,sha256=hw7kBPNR9GQXj4GxyIAffsGKPfcTPOFh8Bc7oEHtaZY,12108
 broccoli/tensor.py,sha256=um8mrxkYbvNDo-QvHlmJm8Aw6qcngOlUZPoAk_PMReA,4480
-broccoli/transformer.py,sha256=Xw1oBLsvVeHmMqgurhorRa49nrjfooLev5uBPFeK9og,17004
+broccoli/transformer.py,sha256=KhwrV5Nz0utXJQKFFf6MIifmymgOaylc_2_hs9kQt7g,18455
 broccoli/utils.py,sha256=htq_hOsdhUhL0nJi9WkKiEYOjEoWqFpK5X49PtgTf-0,299
-broccoli/vit.py,sha256=AKU8fLyzvMaLYw1V70meyb_t3HLxZOfQ6DKwNOPWd_s,18049
-broccoli_ml-1.0.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
-broccoli_ml-1.0.0.dist-info/METADATA,sha256=7Jg7wN6IPLbzr_RpjpR7IOQxmFq7NliBR52J5YhJZD8,1256
-broccoli_ml-1.0.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-broccoli_ml-1.0.0.dist-info/RECORD,,
+broccoli/vit.py,sha256=9dfOPO7y-o96OYpy_ancH4-Y2_Y62fqgzJvI3-86W_o,18471
+broccoli_ml-1.2.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
+broccoli_ml-1.2.0.dist-info/METADATA,sha256=am67vHH2HDQdq3G336ng56xX6-ys1LyZzLRIWaf8dE8,1256
+broccoli_ml-1.2.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+broccoli_ml-1.2.0.dist-info/RECORD,,