x-transformers 2.4.14__tar.gz → 2.5.1__tar.gz

{x_transformers-2.4.14 → x_transformers-2.5.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.4.14
+Version: 2.5.1
 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.14 → x_transformers-2.5.1}/pyproject.toml

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

{x_transformers-2.4.14 → x_transformers-2.5.1}/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.14 → x_transformers-2.5.1}/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.14 → x_transformers-2.5.1}/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)
@@ -2037,6 +2038,9 @@ class AttentionLayers(Module):
         self.causal = causal
         self.layers = ModuleList([])
 
+        self.attn_heads = heads
+        self.attn_dim_head = dim_head
+
         # routing related
         # 1. greater than one residual stream, proposed in Hyper-Connections paper https://arxiv.org/abs/2409.19606
         # 2. integrating more than one past layer, from LIMe paper https://arxiv.org/abs/2502.09245
@@ -2749,6 +2753,49 @@ 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,
+        heads = 8,
+        dim_head = 64,
+        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'
+            attn_kwargs = {f'attn_{k}': v for k, v in attn_kwargs.items()}
+
+            self.pooler = CrossAttender(dim = dim, cross_attn_dim_context = dim_context, depth = depth, heads = heads, attn_dim_head = dim_head, )
+        else:
+            self.pooler = Attention(dim = dim, dim_context = dim_context, heads = heads, dim_head = dim_head, **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,
@@ -2860,8 +2907,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,
@@ -2958,10 +3006,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, heads = self.attn_layers.attn_heads, dim_head = self.attn_layers.attn_dim_head)
 
         # whether to average pool the embed (`global average pool`)
 
@@ -3259,7 +3304,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):
@@ -3272,13 +3316,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
 
@@ -3327,6 +3373,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