[paper] mHC: Manifold-Constrained Hyper-Connections

Based on Manifold-Constrained Hyper-Connections arXiv:2512.24880 and the Megatron-LM integration tracked in issue #2919, PyTorch-path PR #2943, and cuTile kernel-fusion PR #3828. And the latest DeepSeek-V4 released @today (2026/04/24) also harnesses this hyper-connection structure, as expected.

In this post, we're focusing more on the computation flow and dedicated optimizations instead of the theoretical analysis.

Parameterization and manifold projection

Hyper-Connections (HC) widen the residual from a single stream of width $C$ to $n$ parallel streams with learnable mixing:

with $\mathbf{H}_l^{\mathrm{res}} \in \mathbb{R}^{n \times n}$ and $\mathbf{H}_l^{\mathrm{pre}}, \mathbf{H}_l^{\mathrm{post}} \in \mathbb{R}^{1 \times n}$. mHC parameterizes these three mappings so that each is projected onto a well-behaved manifold, keeping the composite stable across depth. Given the input hidden matrix $\mathbf{x}_l \in \mathbb{R}^{n \times C}$ at layer $l$, the computation proceeds in two phases.

Phase 1: initial mapping computation

The flattened input $\vec{\mathbf{x}}_l = \mathrm{vec}(\mathbf{x}_l) \in \mathbb{R}^{1 \times nC}$ preserves full cross-stream context. After an RMSNorm, the dynamic and static mappings are computed as:

where $\boldsymbol{\phi}_l^{\mathrm{pre}}, \boldsymbol{\phi}_l^{\mathrm{post}} \in \mathbb{R}^{nC \times n}$ and $\boldsymbol{\phi}_l^{\mathrm{res}} \in \mathbb{R}^{nC \times n^2}$ are learnable linear projections, and $\mathrm{mat}(\cdot)$ reshapes the output from $\mathbb{R}^{1 \times n^2}$ back to $\mathbb{R}^{n \times n}$. In practice the three projections are packed into a single $nC \to n^2 + 2n$ linear. $\boldsymbol{\phi}$, $\mathbf{b}$ and $\alpha$ are all learnable parameters.

Phase 2: manifold projection

Each raw mapping is then pushed onto its target manifold:

where $\sigma(\cdot)$ is the sigmoid and $\mathrm{SK}(\cdot)$ denotes Sinkhorn-Knopp. The non-negativity enforced on $\mathbf{H}_l^{\mathrm{pre}}$ and $\mathbf{H}_l^{\mathrm{post}}$ prevents signal cancellation from positive-negative coefficient composition, acting as a supplementary manifold projection alongside the Birkhoff-polytope constraint on $\mathbf{H}_l^{\mathrm{res}}$.

Forward compute workflow

Each transformer layer applies the mHC pattern twice: once around self-attention, once around the MLP. The per-sublayer flow, with $s$ tokens, batch $b$, expansion $n$, and hidden size $C$, is:

The block entry replicates the single input stream $n$ times, and the block exit averages the $n$ streams back.

Compute and memory overhead

The extra work that mHC adds per sublayer falls into three buckets; everything else (sigmoids, the 20 Sinkhorn-Knopp iterations) is $O(n^2)$ and negligible at scale.

For typical $n = 4$ (in the DS-V4 tech report), per sublayer the total is $\approx 236\, sbC$ FLOPs. Against $8\, sb C^2$ for attention QKVO and $16\, sb C^2$ for a $4\times$ MLP at $C=7168$, the ratio is roughly $0.4\%$ and $0.2\%$ respectively.

Yet the paper reports 6.7% end-to-end training overhead — well above what the FLOP count predicts. The gap is dominated by activation memory: the residual stream widens from $sbC$ to $sb\,nC$, so every op along it — mix, bias-dropout-add, layernorm input, the residual add itself — runs on an $n\times$ larger tensor.

Parameter overhead is small: each HC module adds $nC (n^2 + 2n)$ weights from the packed $\boldsymbol{\phi}$ projection plus biases and three scalar gates. For $n = 4, C = 7168$ that is $\sim$688k parameters per sublayer, on the order of 0.1% of the sublayer's own GEMM weights which is negligible.

Megatron integration

PRs #2943 and #3828 target on dev branch. There is another mHC PR #3430 being merged into main.

mHC integration

PR #2943 is the PyTorch-path landing, correctness-first and without custom kernels. It adds a HyperConnectionModule (one per sublayer, i.e. two per transformer layer) and a HyperConnectionTransformerLayer subclass that wraps the attention and MLP sublayers with aggregate on the input side and the residual-mix on the output side.

Parallelism

• TP: supported. The $nC \to n^2 + 2n$ projection is small and stays replicated; only the sublayer weights are sharded as before.
• SP: HC parameters are marked sequence_parallel=True so their gradients participate in the standard SP allreduce.
• PP: validated at PP=4 on Qwen3-30B-A3B. DualPipe is extended so the first stream copy $\mathbf{x}_{l_0}$ is cached locally per stage and recompute windows do not cross stage boundaries.
• CP: unaffected, since mHC is per-token.

Block-level recomputation

The wide residual costs $n$ times the activation memory, so the non-fused path checkpoints aggregate and groups layers into recompute blocks. Given $L$ total layers and block size $L_r$, the peak-activation memory is $nC \cdot \lceil L/L_r \rceil + (n+2)C \cdot L_r$, minimized at

Only the first-layer input of each block is persisted; the rest is transient inside the recompute bubble.

Kernel fusion

PR #3828 is the GB200-targeted follow-up: four cuTile kernels that collapse the per-layer mHC forward and backward into four launches, enabled with --use-fused-mhc.

Rough $\sim 35\%$ speedups measured in the PR. End-to-end, the paper reports only 6.7% training-time overhead versus baseline dense residuals at 27B with $n=4$, which is what makes the manifold constraint affordable in production.

Results

From the paper (dense 27B, matched tokens):