x-transformers 1.37.9__tar.gz → 1.38.0__tar.gz

{x_transformers-1.37.9/x_transformers.egg-info → x_transformers-1.38.0}/PKG-INFO

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

{x_transformers-1.37.9 → x_transformers-1.38.0}/README.md

@@ -2261,4 +2261,13 @@ ids_out, num_out, is_number_mask = model.generate(start_ids, start_nums, 17)
 }
 ```
 
+```bibtex
+@inproceedings{Leviathan2024SelectiveAI,
+    title   = {Selective Attention Improves Transformer},
+    author  = {Yaniv Leviathan and Matan Kalman and Yossi Matias},
+    year    = {2024},
+    url     = {https://api.semanticscholar.org/CorpusID:273098114}
+}
+```
+
 *solve intelligence... then use that to solve everything else.* - Demis Hassabis

{x_transformers-1.37.9 → x_transformers-1.38.0}/setup.py

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

{x_transformers-1.37.9 → x_transformers-1.38.0}/x_transformers/attend.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from functools import partial
-from typing import Tuple
+from typing import Tuple, Callable
 
 import torch
 from torch.nn import Module
@@ -65,6 +65,35 @@ def once(fn):
 
 print_once = once(print)
 
+# selective attention
+# https://arxiv.org/abs/2410.02703 - section 3.3
+# it is a parameter-less technique to allow each token to prevent itself from being attended to by future tokens
+# if sim_head_gate not supplied, will use the first head of the attention logits (sim in this framework)
+
+def selective_attn(
+    sim,
+    sim_head_gate = None,
+    no_mask_sos = True
+):
+    i, j, device = *sim.shape[-2:], sim.device
+    sim_head_gate = default(sim_head_gate, sim[:, 0])
+
+    gate = F.relu(sim_head_gate) # only positive
+
+    if no_mask_sos:
+        gate[..., -i] = 0.
+
+    eye = torch.eye(j, device = device)
+
+    if j > i:
+        eye = F.pad(eye, (j - i, 0), value = 1.)
+
+    gate = (1. - eye) * gate
+    gate = F.pad(gate, (0, 0, 1, -1), value = 0.) # only allow for masking the future
+    gate = gate.cumsum(dim = -2)
+
+    return sim - rearrange(gate, 'b i j -> b 1 i j')
+
 # alternative distance functions
 
 def qk_l2_dist_squared(q, k):
@@ -104,10 +133,12 @@ class Attend(Module):
         qk_norm = False,
         l2_distance = False,
         sigmoid = False,
+        custom_attn_fn: Callable | None = None,
         flash = False,
         softclamp_logits = False,
         logit_softclamp_value = 50.,
         add_zero_kv = False,
+        selective = False,
         sigsoftmax = False,
         cope = None,
         onnxable = False,
@@ -132,7 +163,9 @@ class Attend(Module):
 
         self.sigmoid = sigmoid
 
-        if not sigmoid:
+        if exists(custom_attn_fn):
+            self.attn_fn = custom_attn_fn
+        elif not sigmoid:
             softmax_fn = partial(F.softmax, dim = -1)
             self.attn_fn = partial(softmax_fn, dtype = torch.float32) if not qk_norm else softmax_fn
         else:
@@ -152,6 +185,11 @@ class Attend(Module):
             self.pre_softmax_talking_heads = nn.Conv2d(heads, heads, 1, bias = False)
             self.post_softmax_talking_heads = nn.Conv2d(heads, heads, 1, bias = False)
 
+        # selective attention
+
+        assert not (selective and not causal), 'selective attention is designed for autoregressive'
+        self.selective = selective
+
         # sparse topk
 
         assert not (flash and sparse_topk), 'sparse topk not compatible with flash attention'
@@ -404,11 +442,6 @@ class Attend(Module):
 
         mask_value = -torch.finfo(sim.dtype).max
 
-        if exists(self.sparse_topk) and self.sparse_topk < j:
-            top_values, _ = sim.topk(self.sparse_topk, dim = -1)
-            sparse_topk_mask = sim < top_values[..., -1:]
-            mask = (mask & sparse_topk_mask) if exists(mask) else sparse_topk_mask
-
         if exists(mask):
             sim = sim.masked_fill(~mask, mask_value)
 
@@ -416,6 +449,11 @@ class Attend(Module):
             causal_mask = self.create_causal_mask(i, j, device = device)
             sim = sim.masked_fill(causal_mask, mask_value)
 
+        if exists(self.sparse_topk):
+            top_values, _ = sim.topk(self.sparse_topk, dim = -1)
+            sparse_topk_mask = (sim >= top_values[..., -1:]) & (sim > mask_value)
+            sim = sim.masked_fill(~sparse_topk_mask, mask_value)
+
         row_is_entirely_masked = None
 
         if exists(mask):
@@ -424,6 +462,9 @@ class Attend(Module):
         if exists(self.cope):
             sim = sim + self.cope(q, sim)
 
+        if self.selective:
+            sim = selective_attn(sim)
+
         pre_softmax_attn = sim
 
         if self.sigsoftmax:

{x_transformers-1.37.9 → x_transformers-1.38.0}/x_transformers/x_transformers.py

@@ -925,6 +925,8 @@ class Attention(Module):
         qk_norm_dim_scale = False,
         l2_distance = False,
         sigmoid = False,
+        selective = False,
+        custom_attn_fn: Callable | None = None,
         one_kv_head = False,
         kv_heads = None,
         shared_kv = False,
@@ -1041,6 +1043,8 @@ class Attention(Module):
             scale = qk_norm_scale if qk_norm else self.scale,
             l2_distance = l2_distance,
             sigmoid = sigmoid,
+            selective = selective,
+            custom_attn_fn = custom_attn_fn,
             add_zero_kv = add_zero_kv,
             flash = flash,
             softclamp_logits = softclamp_logits,

{x_transformers-1.37.9 → x_transformers-1.38.0/x_transformers.egg-info}/PKG-INFO

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