[paper] mHC: Manifold-Constrained Hyper-Connections

Based on Manifold-Constrained Hyper-Connections ¹ 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

Residual connections, introduced in ResNet ², let each layer learn an additive update $\mathbf{x}_{l+1} = \mathbf{x}_l + \mathcal{F}(\mathbf{x}_l, \mathbf{W}_l)$ instead of a full transformation, which stabilizes gradient flow and enables the training of very deep networks. The single residual stream, however, ties every sublayer to the same width-$C$ representation; HC and mHC generalize this skip to $n$ parallel streams with learnable, possibly manifold-constrained mixing.

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}}$.

Intuition: why manifold can help numerical stability?

Think of the $n$ residual streams as buckets that carry signal across depth, and each $\mathbf{H}_l^{*}$ as a per-layer rule for how to refill them. Without constraints, that rule can both scale a bucket up and flip its sign, so a tiny perturbation compounds multiplicatively layer over layer, like an interest rate free to be negative or larger than one. The manifold projections cap exactly the two failure modes:

• The sigmoid on $\mathbf{H}_l^{\mathrm{pre}}$ and $2\sigma$ on $\mathbf{H}_l^{\mathrm{post}}$ enforce non-negativity, so signals can only be attenuated or summed, never sign-cancelled.
• Sinkhorn-Knopp makes $\mathbf{H}_l^{\mathrm{res}}$ doubly stochastic, so the mix becomes a redistribution of mass across buckets rather than a rescaling of the total.

Together these projections turn the deep-network update into a bounded, sign-preserving, roughly mass-conserving map — the regime where pre-norm residual networks are known to train stably across hundreds of layers.

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.

Results

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