x-transformers 2.11.20__tar.gz → 2.11.23__tar.gz

{x_transformers-2.11.20 → x_transformers-2.11.23}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.11.20
+Version: 2.11.23
 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 → x_transformers-2.11.23}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "x-transformers"
-version = "2.11.20"
+version = "2.11.23"
 description = "X-Transformers"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{x_transformers-2.11.20 → x_transformers-2.11.23}/tests/test_x_transformers.py

@@ -1463,7 +1463,12 @@ def test_kv_input_residual():
 
     assert tokens.shape == out.shape
 
-def test_belief_attn():
+@param('orthog_project', (False, True))
+@param('orthog_project_per_head', (False, True))
+def test_belief_attn(
+    orthog_project,
+    orthog_project_per_head
+):
     from x_transformers import TransformerWrapper, Decoder
 
     model = TransformerWrapper(
@@ -1473,8 +1478,10 @@ def test_belief_attn():
             dim = 512,
             depth = 6,
             heads = 8,
+            attn_kv_heads = 4,
             rotary_pos_emb = True,
-            attn_orthog_projected_values = True
+            attn_orthog_projected_values = orthog_project,
+            attn_orthog_projected_values_per_head = orthog_project_per_head
         )
     )
 

{x_transformers-2.11.20 → x_transformers-2.11.23}/train_enwik8.py

@@ -52,7 +52,9 @@ model = TransformerWrapper(
         dim = 512,
         depth = 6,
         heads = 8,
-        rotary_pos_emb = True
+        rotary_pos_emb = True,
+        attn_orthog_projected_values = True,
+        attn_orthog_projected_values_per_head = True
     )
 )
 

{x_transformers-2.11.20 → x_transformers-2.11.23}/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,
@@ -1430,6 +1431,7 @@ class Attention(Module):
         assert divisible_by(heads, kv_heads)
 
         self.kv_heads = kv_heads
+        self.groups = heads // kv_heads
 
         q_dim = dim_head * heads
         k_dim = dim_head * kv_heads
@@ -1627,6 +1629,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 +2074,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 = repeat(orig_values, 'b h n d -> b n (g h d)', g = self.groups)
+
+            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