x-transformers 1.42.24__tar.gz → 1.42.25__tar.gz

{x_transformers-1.42.24/x_transformers.egg-info → x_transformers-1.42.25}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: x-transformers
-Version: 1.42.24
+Version: 1.42.25
 Summary: X-Transformers - Pytorch
 Home-page: https://github.com/lucidrains/x-transformers
 Author: Phil Wang

{x_transformers-1.42.24 → x_transformers-1.42.25}/README.md

@@ -317,6 +317,8 @@ model = TransformerWrapper(
 
 Update: MetaAI researchers <a href="https://arxiv.org/abs/2309.16588">have found</a> that adding memory tokens (they call them register tokens), alleviates outliers (which is suspected now to be a pathology of attention networks unable to <a href="https://arxiv.org/abs/2306.12929">attend to nothing</a>).
 
+Update 2: a hybrid architecture out of Nvidia named <a href="https://openreview.net/forum?id=A1ztozypga">Hymba</a> used memory tokens successfully in the autoregressive case, termed meta tokens in their paper
+
 ### Transformers Without Tears
 
 <img src="./images/scalenorm.png"></img>

{x_transformers-1.42.24 → x_transformers-1.42.25}/setup.py

@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
 setup(
   name = 'x-transformers',
   packages = find_packages(exclude=['examples']),
-  version = '1.42.24',
+  version = '1.42.25',
   license='MIT',
   description = 'X-Transformers - Pytorch',
   author = 'Phil Wang',

{x_transformers-1.42.24 → x_transformers-1.42.25}/tests/test_x_transformers.py

@@ -557,3 +557,32 @@ def test_laser():
     x = torch.randint(0, 20000, (2, 1024))
 
     model(x)
+
+@pytest.mark.parametrize('cross_attn_rotary', (True, False))
+def test_cross_attn_rotary(
+    cross_attn_rotary: bool
+):
+
+    x = torch.randn((1, 64, 256))
+    mask = torch.ones((1, 64)).bool()
+    context = torch.randn((1, 128, 512))
+    context_mask = torch.ones((1, 128)).bool()
+
+    model = Encoder(
+        dim = 256,
+        depth = 4,
+        heads = 4,
+        rotary_pos_emb = True,
+        cross_attend = True,
+        cross_attn_dim_context = 512
+    )
+
+    context_pos = torch.arange(128)
+
+    embed = model(
+      x = x,
+      mask = mask,
+      context = context,
+      context_pos = context_pos if cross_attn_rotary else None,
+      context_mask = context_mask
+    )

{x_transformers-1.42.24 → x_transformers-1.42.25}/x_transformers/x_transformers.py

@@ -1284,6 +1284,7 @@ class Attention(Module):
         rel_pos = None,
         attn_bias = None,
         rotary_pos_emb = None,
+        context_rotary_pos_emb = None,
         pos = None, # for custom alibi positions
         prev_attn = None,
         mem = None,
@@ -1355,11 +1356,19 @@ class Attention(Module):
             q = q * self.qk_norm_q_scale
             k = k * self.qk_norm_k_scale
 
-        if exists(rotary_pos_emb) and not has_context:
+        if exists(rotary_pos_emb):
+
             freqs, xpos_scale = rotary_pos_emb
             q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale ** -1.) if exists(xpos_scale) else (1., 1.)
 
             q = apply_rotary_pos_emb(q, freqs, q_xpos_scale)
+
+            if has_context:
+                # override with `context_rotary_pos_emb` if provided
+
+                freqs, xpos_scale = context_rotary_pos_emb
+                _, k_xpos_scale = (xpos_scale, xpos_scale ** -1.) if exists(xpos_scale) else (1., 1.)
+
             k = apply_rotary_pos_emb(k, freqs, k_xpos_scale)
 
             if self.rotary_embed_values:
@@ -1848,7 +1857,6 @@ class AttentionLayers(Module):
                 layer = Attention(dim, heads = heads, causal = causal, learned_value_residual_mix = self_attn_learned_value_residual, **attn_kwargs)
                 is_first_self_attn = False
             elif layer_type == 'c':
-                cross_attn_learned_value_residual = learned_value_residual_mix and not is_first_cross_attn
                 layer = Attention(dim, heads = heads, **{**attn_kwargs, **cross_attn_kwargs})
                 is_first_cross_attn = False
             elif layer_type == 'f':
@@ -1917,6 +1925,7 @@ class AttentionLayers(Module):
         return_hiddens = False,
         rotary_pos_emb = None,
         pos = None,
+        context_pos = None,
         attn_bias = None,
         condition = None,
         in_attn_cond = None, # https://arxiv.org/abs/2105.04090
@@ -1976,14 +1985,28 @@ class AttentionLayers(Module):
 
         # rotary positions
 
-        if not exists(rotary_pos_emb) and exists(self.rotary_pos_emb):
-            maybe_mem = mems[0] # todo - handle edge case where different layers get different memory lengths. don't think this will ever come up but who knows
-            mem_len = maybe_mem.shape[1] if exists(maybe_mem) else 0
+        cross_attn_rotary_pos_emb = dict()
+
+        if exists(self.rotary_pos_emb):
+            if not exists(rotary_pos_emb):
+                maybe_mem = mems[0] # todo - handle edge case where different layers get different memory lengths. don't think this will ever come up but who knows
+                mem_len = maybe_mem.shape[1] if exists(maybe_mem) else 0
+
+                if not exists(pos):
+                    pos = torch.arange(x.shape[1] + mem_len, device = x.device) - mem_len
+
+                rotary_pos_emb = self.rotary_pos_emb(pos)
+
+            # allow for rotary positions for context if provided
 
-            if not exists(pos):
-                pos = torch.arange(x.shape[1] + mem_len, device = x.device) - mem_len
+            if exists(context_pos):
+                assert self.cross_attend
+                context_rotary_pos_emb = self.rotary_pos_emb(context_pos)
 
-            rotary_pos_emb = self.rotary_pos_emb(pos)
+                cross_attn_rotary_pos_emb.update(
+                    rotary_pos_emb = rotary_pos_emb,
+                    context_rotary_pos_emb = context_rotary_pos_emb
+                )
 
         # assume cached key / values
 
@@ -2108,7 +2131,7 @@ class AttentionLayers(Module):
             if layer_type == 'a':
                 out, inter = block(x, mask = mask, context_mask = self_attn_kv_mask, attn_mask = attn_mask, rel_pos = self.rel_pos, pos = pos, rotary_pos_emb = rotary_pos_emb, prev_attn = prev_attn, cache = next(iter_attn_cache, None), mem = layer_mem, mem_mask = layer_mem_mask, attn_bias = attn_bias, value_residual = maybe_self_attn_value_residual, return_intermediates = True)
             elif layer_type == 'c':
-                out, inter = block(x, context = context, mask = mask, context_mask = context_mask, prev_attn = prev_cross_attn, cache = next(iter_attn_cache, None), value_residual = maybe_cross_attn_value_residual, return_intermediates = True)
+                out, inter = block(x, context = context, mask = mask, context_mask = context_mask, prev_attn = prev_cross_attn, cache = next(iter_attn_cache, None), value_residual = maybe_cross_attn_value_residual, **cross_attn_rotary_pos_emb, return_intermediates = True)
             elif layer_type == 'f':
                 out = block(x)
 

{x_transformers-1.42.24 → x_transformers-1.42.25/x_transformers.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: x-transformers
-Version: 1.42.24
+Version: 1.42.25
 Summary: X-Transformers - Pytorch
 Home-page: https://github.com/lucidrains/x-transformers
 Author: Phil Wang