x-transformers 2.4.12__tar.gz → 2.5.0__tar.gz

{x_transformers-2.4.12 → x_transformers-2.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.4.12
+Version: 2.5.0
 Summary: X-Transformers
 Project-URL: Homepage, https://pypi.org/project/x-transformers/
 Project-URL: Repository, https://github.com/lucidrains/x-transformers

{x_transformers-2.4.12 → x_transformers-2.5.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "x-transformers"
-version = "2.4.12"
+version = "2.5.0"
 description = "X-Transformers"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{x_transformers-2.4.12 → x_transformers-2.5.0}/tests/test_x_transformers.py

@@ -1086,7 +1086,7 @@ def add_attn_pool():
         num_tokens = 256,
         max_seq_len = 1024,
         attn_pool = True,
-        num_attn_pool_queries =  3,
+        num_pooled_tokens =  3,
         attn_layers = Decoder(
             dim = 512,
             depth = 12,
@@ -1110,7 +1110,6 @@ def test_up(
         num_tokens = 256,
         max_seq_len = 1024,
         attn_pool = True,
-        num_attn_pool_queries =  3,
         attn_layers = Decoder(
             dim = 512,
             depth = 12,
@@ -1153,3 +1152,32 @@ def test_beam_search(stochastic):
 
     assert beams.shape == (4, 2, 10)
     assert scores.shape == (4, 2)
+
+
+@pytest.mark.parametrize('num_pooled_tokens', (1, 3))
+@pytest.mark.parametrize('attn_pool_depth', (1, 3))
+def test_attn_pooler(
+    num_pooled_tokens,
+    attn_pool_depth
+):
+
+    model = TransformerWrapper(
+        num_tokens = 256,
+        max_seq_len = 1024,
+        attn_pool = True,
+        num_pooled_tokens =  num_pooled_tokens,
+        attn_pool_depth = attn_pool_depth,
+        dim_pooled_tokens = 77,
+        attn_layers = Encoder(
+            dim = 512,
+            depth = 12,
+            heads = 8,
+            attn_value_rmsnorm = True
+        ),
+    )
+
+    x = torch.randint(0, 256, (2, 10))
+
+    out = model(x)
+
+    assert out.shape == (2, num_pooled_tokens, 77)

{x_transformers-2.4.12 → x_transformers-2.5.0}/x_transformers/__init__.py

@@ -4,12 +4,13 @@ from x_transformers.x_transformers import (
     Decoder,
     PrefixDecoder,
     CrossAttender,
+    AttentionPool,
     Attention,
     FeedForward,
     RMSNorm,
     AdaptiveRMSNorm,
     TransformerWrapper,
-    ViTransformerWrapper
+    ViTransformerWrapper,
 )
 
 from x_transformers.autoregressive_wrapper import AutoregressiveWrapper

{x_transformers-2.4.12 → x_transformers-2.5.0}/x_transformers/x_transformers.py

@@ -51,6 +51,7 @@ class LayerIntermediates:
     attn_pooled_tokens: Tensor | None = None
     memory_tokens:      Tensor | None = None
     logit_entropies:    Tensor | None = None
+    logits:             Tensor | None = None
     cache_length:       int = 0
 
 LinearNoBias = partial(nn.Linear, bias = False)
@@ -1304,6 +1305,7 @@ class Attention(Module):
         qk_norm_groups = 1,
         qk_norm_scale = 10,
         qk_norm_dim_scale = False,
+        value_rmsnorm = False,      # used in alphagenome and bytedance's GR3 for further stability
         l2_distance = False,
         sigmoid = False,
         selective = False,
@@ -1458,6 +1460,10 @@ class Attention(Module):
         assert (not qk_norm) or divisible_by(dim_head, qk_norm_groups), 'dimension per attention head must be divisible by the qk norm groups'
         assert not (qk_norm and (dim_head // qk_norm_groups) <= 2), 'the group dimension may be too small (2 was too small in my tests, but 4 still works, surprisingly)'
 
+        # value rms norm
+
+        self.value_rmsnorm = MultiheadRMSNorm(dim_head, heads = heads) if value_rmsnorm else None
+
         # contextual positional encoding
         # https://arxiv.org/html/2405.18719v2
 
@@ -1697,6 +1703,10 @@ class Attention(Module):
             q = q * self.qk_norm_q_scale
             k = k * self.qk_norm_k_scale
 
+        # maybe value rmsnorm
+
+        v = maybe(self.value_rmsnorm)(v)
+
         # take care of caching
 
         if not is_multi_latent_attn:
@@ -2740,6 +2750,45 @@ class CrossAttender(AttentionLayers):
     def __init__(self, **kwargs):
         super().__init__(cross_attend = True, only_cross = True, **kwargs)
 
+class AttentionPool(Module):
+    def __init__(
+        self,
+        dim,
+        num_pooled_tokens = 1,
+        dim_context = None,
+        add_residual = False,
+        depth = 1,
+        squeeze_output = None,
+        attn_kwargs: dict = dict()
+    ):
+        super().__init__()
+        dim_context = default(dim_context, dim)
+
+        squeeze_output = default(squeeze_output, False)
+        assert not (squeeze_output and num_pooled_tokens > 1)
+
+        self.queries = nn.Parameter(torch.randn(num_pooled_tokens, dim) * 1e-2)
+
+        if depth > 1:
+            assert not add_residual, 'residual already in effect when doing a full cross attention based transformer for pooling'
+            self.pooler = CrossAttender(dim = dim, cross_attn_dim_context = dim_context, depth = depth, **attn_kwargs)
+        else:
+            self.pooler = Attention(dim = dim, dim_context = dim_context, **attn_kwargs)
+
+        self.add_residual = add_residual
+
+    def forward(self, context, mask = None):
+        batch = context.shape[0]
+
+        queries = repeat(self.queries, 'n d -> b n d', b = batch)
+
+        pooled = self.pooler(queries, context, context_mask = mask)
+
+        if self.add_residual:
+            pooled = pooled + queries
+
+        return pooled
+
 class ViTransformerWrapper(Module):
     def __init__(
         self,
@@ -2851,8 +2900,9 @@ class TransformerWrapper(Module):
         use_cls_token = False,
         num_cls_tokens = 1,
         attn_pool = False,
-        num_attn_pool_queries = 1,
-        dim_attn_pool_query = None,
+        num_pooled_tokens = 1,
+        attn_pool_depth = 1,
+        dim_pooled_tokens = None,
         squeeze_out_last_dim = False,
         token_emb: TokenEmbedding | None = None,
         mixture_of_softmax = False,
@@ -2949,10 +2999,7 @@ class TransformerWrapper(Module):
         self.attn_pool = None
 
         if attn_pool:
-            self.attn_pool = Attention(dim = default(dim_attn_pool_query, dim), dim_context = dim)
-
-            self.attn_pool_queries = nn.Parameter(torch.zeros(num_attn_pool_queries, dim))
-            nn.init.normal_(self.attn_pool_queries, std = 0.02)
+            self.attn_pool = AttentionPool(dim = default(dim_pooled_tokens, dim), dim_context = dim, num_pooled_tokens = num_pooled_tokens, depth = attn_pool_depth)
 
         # whether to average pool the embed (`global average pool`)
 
@@ -3250,7 +3297,6 @@ class TransformerWrapper(Module):
         if self.average_pool_embed:
             x = masked_mean(x, mask = orig_mask, dim = 1)
 
-
         # cls token(s)
 
         if exists(self.cls_token):
@@ -3263,13 +3309,15 @@ class TransformerWrapper(Module):
 
         # attention pool
 
-        if exists(self.attn_pool) and return_intermediates:
-            queries = repeat(self.attn_pool_queries, 'n d -> b n d', b = x.shape[0])
+        is_encoder = not self.attn_layers.causal
+        return_pooled_tokens = exists(self.attn_pool) and is_encoder
 
-            attn_pooled_tokens = self.attn_pool(queries, context = x, context_mask = mask)
+        if (
+            exists(self.attn_pool) and
+            (return_intermediates or is_encoder) # in a new paper, they use attention pooling on decoder - so we'll default to returning pooled tokens if encoder, but for decoder, they must set `return_intermediates`
+        ):
 
-            if attn_pooled_tokens.shape[1] == 1:
-                attn_pooled_tokens = rearrange(attn_pooled_tokens, 'b 1 d -> b d')
+            attn_pooled_tokens = self.attn_pool(x, mask = mask)
 
             intermediates.attn_pooled_tokens = attn_pooled_tokens
 
@@ -3318,6 +3366,9 @@ class TransformerWrapper(Module):
             out = (x, intermediates)
         elif return_embeddings:
             out = x
+        elif return_pooled_tokens:
+            intermediates.logits = logits
+            out = attn_pooled_tokens
         else:
             out = logits