x-transformers 1.37.10__py3-none-any.whl → 1.38.1__py3-none-any.whl

x_transformers/attend.py

@@ -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):
@@ -109,6 +138,7 @@ class Attend(Module):
         softclamp_logits = False,
         logit_softclamp_value = 50.,
         add_zero_kv = False,
+        selective = False,
         sigsoftmax = False,
         cope = None,
         onnxable = False,
@@ -155,6 +185,12 @@ 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 (flash and selective), 'selective attention cannot work on flash 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'
@@ -427,6 +463,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/x_transformers.py

@@ -925,6 +925,7 @@ 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,
@@ -1042,6 +1043,7 @@ 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,

{x_transformers-1.37.10.dist-info → x_transformers-1.38.1.dist-info}/METADATA

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

{x_transformers-1.37.10.dist-info → x_transformers-1.38.1.dist-info}/RECORD

@@ -1,15 +1,15 @@
 x_transformers/__init__.py,sha256=-MkQrSc37cTVDX7AOykxunYnqVtFlQ7lb0Cse5dsGWU,793
-x_transformers/attend.py,sha256=MJmMLIt0rHx-JNNsc2auUsCjsB-69NewufaRV32ADmA,14012
+x_transformers/attend.py,sha256=Ysh8fFQ-BW0aKl7Ls3-gBJQ0Y-QGahMCuvTRJZ-C1rM,15246
 x_transformers/autoregressive_wrapper.py,sha256=DOJJCMMDOqDYKWy_IaG5IyKsXD3AW6amzfUgdAADOLY,10500
 x_transformers/continuous.py,sha256=cIVEdhfei258__ziV7kQBrJMxCel54bExBTDrO9rfCI,6450
 x_transformers/dpo.py,sha256=LjvWgCkqTl-UuehrzQ8nkX5guLr4whYwsmm7SKSwdls,3450
 x_transformers/multi_input.py,sha256=tCh-fTJDj2ib4SMGtsa-AM8MxKzJAQSwqAXOu3HU2mg,9252
 x_transformers/nonautoregressive_wrapper.py,sha256=2NU58hYMgn-4Jzg3mie-mXb0XH_dCN7fjlzd3K1rLUY,10510
-x_transformers/x_transformers.py,sha256=q6I6rvyYUWLgwtKOxPwF12UL1HzcIlauI8YrM8gvZac,83901
+x_transformers/x_transformers.py,sha256=Dol6GMZOoHGOFdHwe21o2SbJp6b3YKCUHoIs_AjfvTo,83963
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=QE1ltYZTR_eGgIHPP2BrMWVWVLqMW-OpDZh87BSmQEg,8563
-x_transformers-1.37.10.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.37.10.dist-info/METADATA,sha256=qA7YAj5ZeaesnGamBR-cPSR_0HSwquwPBptaUmi7c3c,662
-x_transformers-1.37.10.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-x_transformers-1.37.10.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.37.10.dist-info/RECORD,,
+x_transformers-1.38.1.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.38.1.dist-info/METADATA,sha256=mdClNuSZCRyILMjlRi0WUuRFRnRoRSDsXZEF5A4NbpA,661
+x_transformers-1.38.1.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+x_transformers-1.38.1.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.38.1.dist-info/RECORD,,