adopt mHC from deepseek

Author: lucidrains

README.md

@@ -2593,4 +2593,16 @@ ids_out, num_out, is_number_mask = model.generate(start_ids, start_nums, 17)
 }
 ```
 
+```bibtex
+@misc{xie2025mhcmanifoldconstrainedhyperconnections,
+    title   = {mHC: Manifold-Constrained Hyper-Connections}, 
+    author  = {Zhenda Xie and Yixuan Wei and Huanqi Cao and Chenggang Zhao and Chengqi Deng and Jiashi Li and Damai Dai and Huazuo Gao and Jiang Chang and Liang Zhao and Shangyan Zhou and Zhean Xu and Zhengyan Zhang and Wangding Zeng and Shengding Hu and Yuqing Wang and Jingyang Yuan and Lean Wang and Wenfeng Liang},
+    year    = {2025},
+    eprint  = {2512.24880},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.CL},
+    url     = {https://arxiv.org/abs/2512.24880}, 
+}
+```
+
 *solve intelligence... then use that to solve everything else.* - Demis Hassabis

pyproject.toml

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

tests/test_x_transformers.py

@@ -594,8 +594,8 @@ def test_cross_attn_rotary(
         context_mask = context_mask
     )
 
-@param('tanh', (True, False))
-def test_hyper_connections(tanh):
+@param('qkv_receive_diff_residuals', (False, True))
+def test_hyper_connections(qkv_receive_diff_residuals):
 
     model = TransformerWrapper(
         num_tokens = 20000,
@@ -605,9 +605,7 @@ def test_hyper_connections(tanh):
             depth = 6,
             heads = 8,
             num_residual_streams = 8, # 8 dynamic hyper connection residual streams
-            residual_fn_kwargs = dict(
-                tanh = tanh
-            )
+            qkv_receive_diff_residuals = qkv_receive_diff_residuals
         )
     )
 

x_transformers/x_transformers.py

@@ -1052,6 +1052,18 @@ def forward(self, x, residual, **kwargs):
 
 # hyper connections
 
+def sinkhorn_knopps(t, iters = 20):
+    dtype = t.dtype
+    t = t.float()
+
+    t = t.softmax(dim = -2)
+
+    for _ in range(iters):
+        t = F.normalize(t, p = 1, dim = -1)
+        t = F.normalize(t, p = 1, dim = -2)
+
+    return t.to(dtype)
+
 class HyperConnection(Module):
     def __init__(
         self,
@@ -1066,11 +1078,12 @@ def __init__(
         """
         https://arxiv.org/abs/2409.19606
         Appendix J - Algorithm 2, Dynamic only
+
+        https://arxiv.org/abs/2512.24880
+        "Manifold constrained" mixing matrices
         """
         super().__init__()
 
-        self.act = nn.Tanh() if tanh else nn.Identity()
-
         self.norm = nn.LayerNorm(dim, bias = False)
 
         self.num_residual_streams = num_residual_streams
@@ -1092,25 +1105,39 @@ def __init__(
         self.dynamic_beta_scale = nn.Parameter(torch.ones(()) * 1e-2)
 
     def prepare(self, residuals):
+        views = self.num_input_views
+        streams = self.num_residual_streams
 
         residuals = rearrange(residuals, '(b s) n d -> b n s d', s = self.num_residual_streams)
 
         normed = self.norm(residuals)
 
-        wc_weight = self.act(normed @ self.dynamic_alpha_fn)
+        wc_weight = normed @ self.dynamic_alpha_fn
         dynamic_alpha = wc_weight * self.dynamic_alpha_scale
         alpha = dynamic_alpha + self.static_alpha
 
-        dc_weight = self.act(normed @ self.dynamic_beta_fn)
+        alpha_input, alpha_residual = alpha[..., :views], alpha[..., views:]
+
+        alpha_input = alpha_input.sigmoid() # constraint Hpre
+
+        # the sinkhorn knopps constraint for the residual mixing
+
+        alpha_residual = rearrange(alpha_residual, '... (s1 s2) -> ... s1 s2', s2 = streams)
+        alpha_residual = sinkhorn_knopps(alpha_residual)
+        alpha_residual = rearrange(alpha_residual, '... s1 s2 -> ... (s1 s2)')
+
+        alpha = cat((alpha_input, alpha_residual), dim = -1)
+
+        dc_weight = (normed @ self.dynamic_beta_fn).sigmoid() * 2
         dynamic_beta = dc_weight * self.dynamic_beta_scale
         beta = dynamic_beta + self.static_beta
 
+        beta = beta.sigmoid() * 2 # constraint Hpost
+
         # width connection
 
         mix_h = einsum('... s t, ... s d -> ... t d', alpha, residuals)
 
-        views = self.num_input_views
-
         if views == 1:
             branch_input, residuals = mix_h[..., 0, :], mix_h[..., 1:, :]
         else: