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

x_transformers/x_transformers.py

@@ -1397,6 +1397,7 @@ class Attention(Module):
         logit_softclamp_value = 50.,
         learned_value_residual_mix = False,
         orthog_projected_values = False,  # https://openreview.net/forum?id=Ard2QzPAUK
+        orthog_projected_values_per_head = False,
         laser = False,                    # https://arxiv.org/abs/2411.03493v1
         laser_softclamp_value = 15.,
         qkv_receive_diff_residuals = False,
@@ -1627,6 +1628,9 @@ class Attention(Module):
         # "belief attention" - iclr 2026
 
         self.orthog_projected_values = orthog_projected_values
+        self.orthog_projected_values_per_head = orthog_projected_values_per_head
+
+        out_dim *= max(1, int(orthog_projected_values) + int(orthog_projected_values_per_head))
 
         # hybrid module, in same vein as hymba https://www.arxiv.org/abs/2411.13676
 
@@ -2069,11 +2073,24 @@ class Attention(Module):
             gates = self.to_v_gate(x)
             out = out * self.to_v_gate_activation(gates)
 
-        # maybe return orthogonal projected - "belief" attention
+        # maybe orthogonal projected weighted values - "belief" attention
+
+        if self.orthog_projected_values or self.orthog_projected_values_per_head:
+            orthog_projected = []
+            v_for_proj = self.merge_heads(orig_values)
+
+            if self.orthog_projected_values:
+                projected = orthog_project(out, v_for_proj)
+                orthog_projected.append(projected)
+
+            if self.orthog_projected_values_per_head:
+                v_for_proj = rearrange(v_for_proj, 'b n (h d) -> b n h d', h = h)
+                out = rearrange(out, 'b n (h d) -> b n h d', h = h)
+                projected = orthog_project(out, v_for_proj)
+                projected = rearrange(projected, 'b n h d -> b n (h d)')
+                orthog_projected.append(projected)
 
-        if self.orthog_projected_values:
-            merged_v = self.merge_heads(orig_values)
-            out = orthog_project(out, merged_v)
+            out = cat(orthog_projected, dim = -1)
 
         # combine the heads
 

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.11.20
+Version: 2.11.22
 Summary: X-Transformers
 Project-URL: Homepage, https://pypi.org/project/x-transformers/
 Project-URL: Repository, https://github.com/lucidrains/x-transformers

{x_transformers-2.11.20.dist-info → x_transformers-2.11.22.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=igiWHF93dDCRGaiLL-YVtAYcyhxNQr0Yw_Ievvs1cD4,128188
+x_transformers/x_transformers.py,sha256=G3_RQ-Z8szUS4WzKcwMCO2lnrFWDRsrLJ0hJL1rv-WM,129051
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=AwwYUm8yDAtKQyKJDIhYMsiLTJ_skh3scUFMjp5sda8,8597
-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,,
+x_transformers-2.11.22.dist-info/METADATA,sha256=vtm9uUCTPJ8K8azq1m6YfyFwEXmRPX1IATTfkv6V_Sc,96372
+x_transformers-2.11.22.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.11.22.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.11.22.dist-info/RECORD,,