Table of Contents
Fetching ...

mHC-lite: You Don't Need 20 Sinkhorn-Knopp Iterations

Yongyi Yang, Jianyang Gao

TL;DR

This work tackles instability and portability issues in hyper-connected residuals by replacing SK-based projection with a Birkhoff–von Neumann reparameterization, yielding mHC-lite that expresses each residual matrix as a convex combination of permutation matrices to guarantee exact doubly stochasticity. By using only native matrix operations, mHC-lite removes the need for specialized SK kernels and reduces implementation barriers, while matching or surpassing mHC in downstream performance and improving training stability. Empirical results on nanoGPT-scale models across OpenWebText and FineWeb-Edu show that mHC-lite offers higher throughput and eliminates the residual instabilities observed in HC and mHC, with code available for reproduction. The approach highlights that exactness can be more efficient than approximate projection in practice and points to scalable strategies (e.g., subset permutation sampling) for larger residual dimensions.

Abstract

Hyper-Connections (HC) generalizes residual connections by introducing dynamic residual matrices that mix information across multiple residual streams, accelerating convergence in deep neural networks. However, unconstrained residual matrices can compromise training stability. To address this, DeepSeek's Manifold-Constrained Hyper-Connections (mHC) approximately projects these matrices onto the Birkhoff polytope via iterative Sinkhorn--Knopp (SK) normalization. We identify two limitations of this approach: (i) finite SK iterations do not guarantee exact doubly stochasticity, leaving an approximation gap that can accumulate through network depth and undermine stability; (ii) efficient SK implementation requires highly specialized CUDA kernels, raising engineering barriers and reducing portability. Motivated by the Birkhoff--von Neumann theorem, we propose mHC-lite, a simple reparameterization that explicitly constructs doubly stochastic matrices as convex combinations of permutation matrices. This approach guarantees exact doubly stochasticity by construction and can be implemented using only native matrix operations. Extensive experiments demonstrate that mHC-lite matches or exceeds mHC in performance while achieving higher training throughput with a naive implementation and eliminating the residual instabilities observed in both HC and mHC. The code is publicly available at https://github.com/FFTYYY/mhc-lite.

mHC-lite: You Don't Need 20 Sinkhorn-Knopp Iterations

TL;DR

This work tackles instability and portability issues in hyper-connected residuals by replacing SK-based projection with a Birkhoff–von Neumann reparameterization, yielding mHC-lite that expresses each residual matrix as a convex combination of permutation matrices to guarantee exact doubly stochasticity. By using only native matrix operations, mHC-lite removes the need for specialized SK kernels and reduces implementation barriers, while matching or surpassing mHC in downstream performance and improving training stability. Empirical results on nanoGPT-scale models across OpenWebText and FineWeb-Edu show that mHC-lite offers higher throughput and eliminates the residual instabilities observed in HC and mHC, with code available for reproduction. The approach highlights that exactness can be more efficient than approximate projection in practice and points to scalable strategies (e.g., subset permutation sampling) for larger residual dimensions.

Abstract

Hyper-Connections (HC) generalizes residual connections by introducing dynamic residual matrices that mix information across multiple residual streams, accelerating convergence in deep neural networks. However, unconstrained residual matrices can compromise training stability. To address this, DeepSeek's Manifold-Constrained Hyper-Connections (mHC) approximately projects these matrices onto the Birkhoff polytope via iterative Sinkhorn--Knopp (SK) normalization. We identify two limitations of this approach: (i) finite SK iterations do not guarantee exact doubly stochasticity, leaving an approximation gap that can accumulate through network depth and undermine stability; (ii) efficient SK implementation requires highly specialized CUDA kernels, raising engineering barriers and reducing portability. Motivated by the Birkhoff--von Neumann theorem, we propose mHC-lite, a simple reparameterization that explicitly constructs doubly stochastic matrices as convex combinations of permutation matrices. This approach guarantees exact doubly stochasticity by construction and can be implemented using only native matrix operations. Extensive experiments demonstrate that mHC-lite matches or exceeds mHC in performance while achieving higher training throughput with a naive implementation and eliminating the residual instabilities observed in both HC and mHC. The code is publicly available at https://github.com/FFTYYY/mhc-lite.
Paper Structure (15 sections, 1 theorem, 7 equations, 5 figures, 3 tables)

This paper contains 15 sections, 1 theorem, 7 equations, 5 figures, 3 tables.

Key Result

Theorem 3.1

For any $\boldsymbol{X}\in\mathcal{B}_n$, there exists a weight $\mathbf{a}=(a_1,...,a_{n!}) \in \mathbb{R}^{1\times n!}$, where $a_k \ge 0,\forall k\in [n!], \|\mathbf{a}\|_{\ell_1}=1$, such that where $\left\{\boldsymbol{P}_k\right\}_{k=1}^{n!}$ is the sequence of $n\times n$ permutation matrices.

Figures (5)

  • Figure 1: Residual matrix construction in mHC vs. mHC-lite. The method mHC relies on repeated Sinkhorn–Knopp iterations to approximate doubly stochastic matrices, whereas mHC-lite directly computes the matrix via a convex combination of permutation matrices, achieving exact doubly stochasticity.
  • Figure 2: Gradient-norm dynamics during training. We compare the evolution of gradient norms over the course of training. Left: overall trajectories, showing that both mHC and mHC-lite exhibit substantially smaller gradient norms (and improved stability) than HC. Right: a zoomed-in view of mHC and mHC-lite; curves are smoothed using a 200-step moving average, and the shaded region indicates the standard deviation within the same window. From the zoomed-in view, it is clear that mHC-lite yields a smaller mean gradient norm and reduced fluctuations compared to mHC. Results are obtained with the L model on the FineWeb-Edu dataset.
  • Figure 3: Column sums of ${\boldsymbol{H}}^{\text{res}}$. We compute column sums for token-level ${\boldsymbol{H}}^{\text{res}}$ matrices and summarize their distribution with standard boxplots (points indicate outliers). per-matrix: statistics for individual ${\boldsymbol{H}}^{\text{res}}$ matrices. prod: statistics for the layer-wise product of ${\boldsymbol{H}}^{\text{res}}$ across all layers.
  • Figure 4: Distribution of $\log(1/\nu)$. Distribution of the relative range $\log (1/\nu)$ (defined in \ref{['eq:nu']}) for mHC before applying SK. Large values (e.g., $\log(1/\nu)>30$) suggest that 20 SK iterations may not converge well to a doubly stochastic matrix.
  • Figure 5: Token throughput during training. We report training throughput in tokens/s, computed as the number of tokens per batch divided by the wall-clock time of each optimizer update and averaged over the entire training run. All experiments are run on a single node with 8$\times$ NVIDIA A100 80GB (SXM4) GPUs. Notice that the mHC result is based on our PyTorch re-implementation and may underestimate the throughput of the specialized-kernel implementation in xie25mhc, which is reported to incur only a $6.7\%$ overhead relative to HC.

Theorems & Definitions (1)

  • Theorem 3.1: The Birkhoff-von Neumann theorem