x-transformers 2.11.18__py3-none-any.whl → 2.11.20__py3-none-any.whl

x_transformers/x_transformers.py

@@ -161,6 +161,21 @@ def or_reduce(masks):
         head = head | rest
     return head
 
+def orthog_project(x, y):
+    x, packed_shape = pack([x], 'b *')
+    y, _ = pack([y], 'b *')
+
+    dtype = x.dtype
+    x, y = x.double(), y.double()
+    unit = F.normalize(y, dim = -1)
+
+    parallel = (x * unit).sum(dim = -1, keepdim = True) * unit
+    orthog = x - parallel
+
+    orthog, = unpack(orthog, packed_shape, 'b *')
+
+    return orthog.to(dtype)
+
 # cache helpers
 
 def get_cached_kvs(
@@ -276,8 +291,13 @@ class ReluSquared(Module):
         return F.relu(x) ** 2
 
 class SoLU(Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.norm = LayerNorm(dim)
+
     def forward(self, x):
-        return x.softmax(dim = -1) * x
+        activated = x.softmax(dim = -1) * x
+        return self.norm(activated)
 
 # embedding
 
@@ -1262,7 +1282,7 @@ class FeedForward(Module):
         elif relu_squared:
             activation = ReluSquared()
         elif solu:
-            activation = SoLU()
+            activation = SoLU(inner_dim)
         elif swish:
             activation = nn.SiLU()
         else:
@@ -1376,7 +1396,8 @@ class Attention(Module):
         softclamp_logits = False,
         logit_softclamp_value = 50.,
         learned_value_residual_mix = False,
-        laser = False,                # https://arxiv.org/abs/2411.03493v1
+        orthog_projected_values = False,  # https://openreview.net/forum?id=Ard2QzPAUK
+        laser = False,                    # https://arxiv.org/abs/2411.03493v1
         laser_softclamp_value = 15.,
         qkv_receive_diff_residuals = False,
         use_latent_q = False,
@@ -1602,6 +1623,11 @@ class Attention(Module):
 
         self.attn_on_attn = on_attn
 
+        # return orthogonal projected weighted values on original values
+        # "belief attention" - iclr 2026
+
+        self.orthog_projected_values = orthog_projected_values
+
         # hybrid module, in same vein as hymba https://www.arxiv.org/abs/2411.13676
 
         hybrid_mix = None
@@ -2043,6 +2069,12 @@ class Attention(Module):
             gates = self.to_v_gate(x)
             out = out * self.to_v_gate_activation(gates)
 
+        # maybe return orthogonal projected - "belief" attention
+
+        if self.orthog_projected_values:
+            merged_v = self.merge_heads(orig_values)
+            out = orthog_project(out, merged_v)
+
         # combine the heads
 
         out = self.to_out(out)

{x_transformers-2.11.18.dist-info → x_transformers-2.11.20.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.11.18
+Version: 2.11.20
 Summary: X-Transformers
 Project-URL: Homepage, https://pypi.org/project/x-transformers/
 Project-URL: Repository, https://github.com/lucidrains/x-transformers
@@ -2608,12 +2608,13 @@ ids_out, num_out, is_number_mask = model.generate(start_ids, start_nums, 17)
 ```
 
 ```bibtex
-@article{elhage2022solu,
-    title   = {Softmax Linear Units},
-    author  = {Elhage, Nelson and Hume, Tristan and Olsson, Catherine and Nanda, Neel and Henighan, Tom and Johnston, Scott and ElShowk, Sheer and Joseph, Nicholas and DasSarma, Nova and Mann, Ben and Hernandez, Danny and Askell, Amanda and Ndousse, Kamal and Jones, Andy and Drain, Dawn and Chen, Anna and Bai, Yuntao and Ganguli, Deep and Lovitt, Liane and Hatfield-Dodds, Zac and Kernion, Jackson and Conerly, Tom and Kravec, Shauna and Fort, Stanislav and Kadavath, Saurav and Jacobson, Josh and Tran-Johnson, Eli and Kaplan, Jared and Clark, Jack and Brown, Tom and McCandlish, Sam and Amodei, Dario and Olah, Christopher},
-    year    = {2022},
-    journal = {Transformer Circuits Thread},
-    note    = {https://transformer-circuits.pub/2022/solu/index.html}
+@inproceedings{anonymous2025beliefformer,
+    title   = {BeliefFormer: Belief Attention in Transformer},
+    author  = {Anonymous},
+    booktitle = {Submitted to The Fourteenth International Conference on Learning Representations},
+    year    = {2025},
+    url     = {https://openreview.net/forum?id=Ard2QzPAUK},
+    note    = {under review}
 }
 ```
 

{x_transformers-2.11.18.dist-info → x_transformers-2.11.20.dist-info}/RECORD

@@ -11,10 +11,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=pIUxQmj_wLHIMOxyqjy4hKww6NdYtzxtRMWROovHoDA,127212
+x_transformers/x_transformers.py,sha256=igiWHF93dDCRGaiLL-YVtAYcyhxNQr0Yw_Ievvs1cD4,128188
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=AwwYUm8yDAtKQyKJDIhYMsiLTJ_skh3scUFMjp5sda8,8597
-x_transformers-2.11.18.dist-info/METADATA,sha256=9VPaNWK5WVVltDqRqkb_4OtPEmJFzfkYkln3aJpKdfQ,96858
-x_transformers-2.11.18.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.11.18.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.11.18.dist-info/RECORD,,
+x_transformers-2.11.20.dist-info/METADATA,sha256=Ii1ODqNuQz8ge5SrePZ33Sahr4izJEcpZH2SrAltNcA,96372
+x_transformers-2.11.20.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.11.20.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.11.20.dist-info/RECORD,,