x-transformers 2.9.2__py3-none-any.whl → 2.10.1__py3-none-any.whl

x_transformers/attend.py

@@ -183,6 +183,7 @@ class Attend(Module):
         gumbel_softmax = False,
         gumbel_softmax_temp = 1.,
         gumbel_softmax_hard = True,
+        cog_signed = False,
         custom_attn_fn: Callable | None = None,
         flash = False,
         softclamp_logits = False,
@@ -260,6 +261,12 @@ class Attend(Module):
         assert not (selective and not causal), 'selective attention is designed for autoregressive'
         self.selective = selective
 
+        # cog attention - negative weights for expressiveness
+        # https://openreview.net/forum?id=ezRrwwbxd0
+
+        assert not (flash and cog_signed), 'cog attention not available for flash'
+        self.cog_signed = cog_signed
+
         # l2 distance attention
 
         self.l2_distance = l2_distance
@@ -509,6 +516,14 @@ class Attend(Module):
         if self.softclamp_logits:
             sim = softclamp(sim, self.logit_softclamp_value)
 
+        # pre-masking - handle cog by storing sign
+
+        if self.cog_signed:
+            sim_sign = sim.sign()
+            sim = sim.abs()
+
+        # masking
+
         i, j, dtype = *sim.shape[-2:], sim.dtype
 
         mask_value = -torch.finfo(sim.dtype).max
@@ -542,6 +557,11 @@ class Attend(Module):
 
         attn = attn.type(dtype)
 
+        # add back the sign
+
+        if self.cog_signed:
+            attn = attn * sim_sign
+
         post_softmax_attn = attn
 
         if self.head_learned_sink:

x_transformers/x_transformers.py

@@ -1340,6 +1340,7 @@ class Attention(Module):
         gumbel_softmax_temp = 1.,
         gumbel_softmax_hard = True,
         selective = False,
+        cog_signed = False,
         custom_attn_fn: Callable | None = None,
         hybrid_module: Module | None = None,
         hybrid_mask_kwarg: str | None = None,
@@ -1548,6 +1549,7 @@ class Attention(Module):
             gumbel_softmax_temp = gumbel_softmax_temp,
             gumbel_softmax_hard = gumbel_softmax_hard,
             selective = selective,
+            cog_signed = cog_signed,
             custom_attn_fn = custom_attn_fn,
             add_zero_kv = add_zero_kv,
             head_learned_sink = head_learned_sink,

{x_transformers-2.9.2.dist-info → x_transformers-2.10.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.9.2
+Version: 2.10.1
 Summary: X-Transformers
 Project-URL: Homepage, https://pypi.org/project/x-transformers/
 Project-URL: Repository, https://github.com/lucidrains/x-transformers
@@ -2586,4 +2586,16 @@ ids_out, num_out, is_number_mask = model.generate(start_ids, start_nums, 17)
 }
 ```
 
+```bibtex
+@misc{lv2025expressiveattentionnegativeweights,
+    title   = {More Expressive Attention with Negative Weights}, 
+    author  = {Ang Lv and Ruobing Xie and Shuaipeng Li and Jiayi Liao and Xingwu Sun and Zhanhui Kang and Di Wang and Rui Yan},
+    year    = {2025},
+    eprint  = {2411.07176},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.CL},
+    url     = {https://arxiv.org/abs/2411.07176}, 
+}
+```
+
 *solve intelligence... then use that to solve everything else.* - Demis Hassabis

{x_transformers-2.9.2.dist-info → x_transformers-2.10.1.dist-info}/RECORD

@@ -1,5 +1,5 @@
 x_transformers/__init__.py,sha256=aVuhUU0572TJHW88BVc4yA2tla0Zb8l3NH7W4RZ1AEs,1005
-x_transformers/attend.py,sha256=RZJT9pPlpqSG3nOUqQHNRR6jOeJ2r-Fvvar2wdu9HLw,18687
+x_transformers/attend.py,sha256=uu4lIEfiwzZLeuBY2dJLG9709DZbWK8-on4ds8SCCJ0,19207
 x_transformers/autoregressive_wrapper.py,sha256=BsGO9xfVYkvynqbU1__tu_S_cxl7gss0YwnkhIa2baY,18401
 x_transformers/belief_state_wrapper.py,sha256=YLUMk6t2MhFBEw5lHDDHJHcoCxTIkHvxTNY__GGZEKU,13374
 x_transformers/continuous.py,sha256=WwpQCjyVY4PtuEAOFY68zqgklbF9I7AL5w6874YlDe8,13249
@@ -10,10 +10,10 @@ x_transformers/multi_input.py,sha256=tCh-fTJDj2ib4SMGtsa-AM8MxKzJAQSwqAXOu3HU2mg
 x_transformers/neo_mlp.py,sha256=XCNnnop9WLarcxap1kGuYc1x8GHvwkZiDRnXOxSl3Po,3452
 x_transformers/nonautoregressive_wrapper.py,sha256=hMQqNimGtchNIe13cR5LZule1V7I1qM5LmY8VQfVdnA,11698
 x_transformers/up_wrapper.py,sha256=YC2LN14_7Xx9Wtiek2rtEJ_qHqdfSmKlh3d7Cgxwd80,7073
-x_transformers/x_transformers.py,sha256=o6B10urcC7MRUrmoHOgYJgkrVDzHhX-jt6zZY3pZEgA,125700
+x_transformers/x_transformers.py,sha256=ADr83Fz2cehj_F7N1bMwxhAg-r48fGhlaZqw3hxoxMQ,125765
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=AwwYUm8yDAtKQyKJDIhYMsiLTJ_skh3scUFMjp5sda8,8597
-x_transformers-2.9.2.dist-info/METADATA,sha256=3JsbSIp9fsGpuXopeIaIq4ffjYTJIHyqdRLxM21cfUM,95381
-x_transformers-2.9.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.9.2.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.9.2.dist-info/RECORD,,
+x_transformers-2.10.1.dist-info/METADATA,sha256=sEfcxJr3l0W4Yga0NLHq1sMk90Zr5-Lpr-9fIlmG9H4,95799
+x_transformers-2.10.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.10.1.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.10.1.dist-info/RECORD,,