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

作者/机构: Yongyi Yang, Jianyang Gao

A1 主要贡献

本文探讨了深度神经网络中残差连接的演进。传统的残差连接对于稳定训练至关重要，但其单流和固定恒等映射的设计存在局限性。Hyper-Connections (HC) 通过引入多残差流和动态残差矩阵$H_{res}^l$来混合信息，从而扩展了残差连接，加速了收敛。然而，不受约束的$H_{res}^l$可能导致训练不稳定。

为了解决这个问题，Manifold-Constrained Hyper-Connections (mHC) 提出将$H_{res}^l$约束在Birkhoff多胞体（即双随机矩阵集合）上。双随机矩阵的光谱范数有界且在矩阵乘法下是封闭的，这有助于防止梯度爆炸。mHC通过迭代式的Sinkhorn–Knopp (SK) 算法来近似实现这一约束。

本文指出了mHC方法的两个核心局限性：
1. 近似差距：有限次数（如20次）的SK迭代无法保证矩阵完全满足双随机性。本文通过一个具体例子（$\alpha = 10^{-13}$）说明，20次迭代后矩阵的列和（1.92, 0.59, 0.59）与目标值1相差甚远。这种近似误差会随着网络深度累积，可能破坏mHC旨在实现的稳定性。

$$ \begin{pmatrix} \frac{1}{2}, & \alpha, & \alpha \\ \frac{1}{2}, & \alpha, & \alpha \\ \alpha, & 1, & 1 \end{pmatrix} \stackrel{\text{SK (20 iters)}}{\Longrightarrow}\begin{pmatrix} 0.91, & 0.045, & 0.045 \\ 0.91, & 0.045, & 0.045 \\ 0., & 0.5, & 0.5 \end{pmatrix} $$

1. 工程壁垒：高效的SK迭代实现依赖于高度定制的CUDA核函数，这增加了工程复杂性并降低了跨软件栈的可移植性。此外，为了控制内存占用，mHC的实现需要在反向传播中重新计算SK迭代的中间结果。

受Birkhoff–von Neumann定理（任何双随机矩阵都是置换矩阵的凸组合）的启发，本文提出了mHC-lite。该方法通过置换矩阵的凸组合直接参数化双随机矩阵，完全绕过了SK迭代。

mHC-lite的优势：
1. 构造保证：通过构造保证了矩阵的精确双随机性，消除了近似误差。
2. 实现简单：仅使用原生的矩阵乘法即可高效实现，无需依赖专门的CUDA核。

本文的主要贡献总结如下：
1. 提出了mHC-lite，一种mHC的简单重参数化方法，通过显式构造双随机残差矩阵，消除了SK迭代的需求，完全弥合了近似差距，并实现了仅基于原生矩阵运算的简单快速实现。
2. 从理论和实证两方面证明了mHC中有限的SK迭代会留下不可忽略的近似差距，表明尽管有流形约束，mHC中仍然存在稳定性问题。
3. 通过大量实验表明，mHC-lite在下游任务性能上与mHC相当或更优，同时实现了更高的训练吞吐量，并消除了在mHC和HC中观察到的残差矩阵不稳定性。

A3 背景知识

残差连接范式的重要性与局限性。由ResNet【7，Deep residual learning for image recognition，He et al.，2016a，CVPR】首创的残差连接范式，已成为现代深度学习的基石。它通过构建恒等映射路径来缓解梯度消失问题，使得训练极深的网络成为可能【8，Identity mappings in deep residual networks，He et al.，2016b，ECCV】。此设计后来被Transformer架构【20，Attention is all you need，Vaswani et al.，2017，NeurIPS】采用，并被证明对大型语言模型（LLMs）的可扩展性至关重要，例如GPT-3【3，Language models are few-shot learners，Brown et al.，2020，NeurIPS】和Llama【19，Llama: Open and efficient foundation language models，Touvron et al.，2023】。尽管取得了巨大成功，标准的残差连接仍有其固有的局限性。单流设计限制了信息流向单一路径，可能限制了极深网络的表示能力【9，Densely connected convolutional networks，Huang et al.，2017，CVPR】。此外，固定的恒等映射虽然稳定了训练，但无法适应不同层或输入上下文之间变化的计算需求【18，Highway networks，Srivastava et al.，2015】。这些观察激发了近期对更灵活、更具表现力的连接机制的研究，这些机制超越了简单的恒等捷径，同时保持了训练的稳定性【23，Xie et al.，2024; 26，Zhu et al.，2024; 16，Mak & Flanigan，2025; 1，Bhendawade et al.，2025; 24，Xie et al.，2025; 14，Liu et al.，2025】。

Hyper-Connections (HC)。Hyper-Connections (HC)【26，Hyper-connections，Zhu et al.，2024，ArXiv】通过将单个残差流扩展为多个流，并引入这些流之间的动态连接，从而推广了残差连接。这种广义残差连接丰富了模型的连通性，并据报道能在几乎不增加额外计算量的情况下加速收敛。设$x_l \in \mathbb{R}^{n \times C}$表示第$l$层的输入特征，其中$n$是残差流的数量，$C$是维度。该架构的公式如下。

$$\boldsymbol{x}_{l+1}=\boldsymbol{H}_{l}^{\text{res}} \boldsymbol{x}_{l}+\boldsymbol{H}_{l}^{\text{post}} f(\boldsymbol{H}_{l}^{\text{pre}} \boldsymbol{x}_{l} ; \mathcal{W}_{l})$$

其中残差矩阵$H_{res}^l \in \mathbb{R}^{n \times n}$由可学习参数和$x_l$动态决定，用于混合残差流。$H_{pre}^l, H_{post}^l \in \mathbb{R}^{1 \times n}$由可学习参数和$x_l$决定，分别用于聚合输入和扩展输出。$f(\cdot; W_l)$代表由权重$W_l$参数化的可学习函数。关于HC中$H_{res}^l$和$H_{pre}^l, H_{post}^l$的详细计算，我们建议读者参考原始论文【26，Hyper-connections，Zhu et al.，2024，ArXiv】。

流形约束的超连接 (mHC)。流形约束的超连接修改了$H_{pre}^l, H_{post}^l$和$H_{res}^l$的计算方式，特别是尝试将$H_{res}^l$约束在Birkhoff多胞体$B_n$上，即双随机矩阵的集合，其定义如下。

$$\mathcal{B}_{n}=\left\{\boldsymbol{X} \in \mathbb{R}^{n \times n} \mid \boldsymbol{X}^{\top} \mathbf{1}_{n}=\boldsymbol{X} \mathbf{1}_{n}=\mathbf{1}_{n}, \boldsymbol{X} \geq 0\right\}$$

其中$1_n$表示全1向量，$X \ge 0$是逐元素的。双随机矩阵表现出类似恒等矩阵的稳定性，因为它们的光谱范数有界为1，并且该集合在矩阵乘法下是封闭的：双随机矩阵的重复组合仍然是双随机的。设$x_l \in \mathbb{R}^{n \times C}$表示第$l$层的输入特征，$\hat{x}_l \in \mathbb{R}^{1 \times nC}$表示展平的输入特征。mHC的计算详情如下。

$$\begin{aligned} \begin{aligned} \hat{\boldsymbol{x}}_{l}^{\prime} & =\operatorname{RMSNorm}\left(\hat{\boldsymbol{x}}_{l}\right) \\ \boldsymbol{H}_{l}^{\text {pre }} & =\operatorname{sigmoid}\left(\alpha_{l}^{\text {pre }} \hat{\boldsymbol{x}}_{l}^{\prime} \boldsymbol{W}_{l}^{\text {pre }}+\boldsymbol{b}_{l}^{\text {pre }}\right) \\ \boldsymbol{H}_{l}^{\text {post }} & =2 \cdot \operatorname{sigmoid}\left(\alpha_{l}^{\text {post }} \hat{\boldsymbol{x}}_{l}^{\prime} \boldsymbol{W}_{l}^{\text {post }}+\boldsymbol{b}_{l}^{\text {post }}\right) \\ \boldsymbol{H}_{l}^{\text {res }} & =\operatorname{SK}\left(\exp \left(\operatorname{mat}\left(\alpha_{l}^{\text {res }} \hat{\boldsymbol{x}}_{l}^{\prime} \boldsymbol{W}_{l}^{\text {res }}+\boldsymbol{b}_{l}^{\text {res }}\right)\right)\right) \end{aligned} \end{aligned}$$

其中$W_{pre}, W_{post} \in \mathbb{R}^{nC \times n}$和$W_{res} \in \mathbb{R}^{nC \times n^2}$是第$l$层的可学习权重矩阵。$b_{pre}^l, b_{post}^l \in \mathbb{R}^{1 \times n}$和$b_{res}^l \in \mathbb{R}^{1 \times n^2}$是可学习的偏置。$\alpha_{pre}^l, \alpha_{post}^l$和$\alpha_{res}^l$是可学习的标量。函数mat(·)将矩阵从$\mathbb{R}^{1 \times n^2}$重塑为$\mathbb{R}^{n \times n}$。RMSNorm(·)指的是RMSNorm【25，Root mean square layer normalization，Zhang & Sennrich，2019】。exp(·)函数是逐元素的。SK(·)迭代交替地重新缩放所有列和行，使其和等于1。在mHC的设置中，SK迭代重复20次。

A2 方法细节

正如第1节所讨论的，mHC对有限次数SK迭代的依赖引发了关于可移植性和稳定性的担忧。从系统角度看，实现SK迭代的竞争性效率通常依赖于专门的、融合的CUDA核，这使得该组件难以作为标准残差连接在不同框架间的即插即用替代品。除了可移植性，一个更根本的问题在于残差矩阵的稳定性。特别是，有限步的近似可能导致与精确双随机性的不可忽略的偏差，这些偏差可能随深度累积，从而破坏mHC旨在实现的稳定性。我们在3.1节详细分析了这一稳定性问题。这些观察共同促使我们在3.2节提出一种重参数化方法，该方法通过构造确保了精确的双随机性，并避免了对CUDA核的重度定制。对于任意$X \in B_n$，存在一个权重$a = (a_1, ..., a_{n!}) \in \mathbb{R}^{1 \times n!}$，其中$a_k \ge 0, \forall k \in [n!], \|a\|_{\ell_1} = 1$，使得

3.1. 稳定性分析

SK算法的收敛性问题。在mHC中，固定次数的SK迭代（例如mHC中的20次）在收敛缓慢时并不能保证高质量的近似。关于矩阵缩放的经典研究表明，SK算法通常并非一致快速的【13，Linial et al.，1998; 10，Knight，2008; 4，Chakrabarty & Khanna，2021】。对于一般的非负矩阵，SK算法只有一个最坏情况下的迭代界：要获得一个$\ell_1$误差最多为$\epsilon$的双随机性近似，可能需要多达$O(\frac{n^2 \log(n/\nu)}{\epsilon^2})$次迭代，其中相对范围$\nu$定义为

$$\nu := \frac{\min_{i,j : x_{i,j} > 0} x_{i,j}}{\max_{i,j} x_{i,j}},$$

其中$x_{i,j}$是$X$的第$(i, j)$个元素。即使对于严格为正的矩阵，收敛性仍然对$1/\nu$敏感，当$1/\nu$很大时，收敛可能极其缓慢【13，A deterministic strongly polynomial algorithm for matrix scaling and approximate permanents，Linial et al.，1998，STOC】（参见第1节中的例子）。

mHC中的实际影响。这个问题在mHC中具有实际意义。如公式(2)所示，SK的输入是通过对特征的仿射函数进行指数运算得到的，这可能产生具有非常大相对范围的病态矩阵。在我们的测量中（图4），大约27.9%的SK输入满足$1/\nu \ge 10^{13}$。对于这样的输入，固定的SK预算可能无法产生一个接近双随机的矩阵。图3显示，mHC中单个残差矩阵的列和可能偏离1高达100%。更重要的是，这些每层的偏差会随深度累积：图3显示，在一个24层网络中，$Q H_{res}^l$的列和可能偏离1高达220%，这意味着当模型进一步扩大时存在不稳定的风险。在实践中，最新的一个模型基于经典的恒等残差连接【7，Deep residual learning for image recognition，He et al.，2016a，CVPR】构建了一个1000层的网络用于自监督强化学习【22，1000 layer networks for self-supervised RL: Scaling depth can enable new goal-reaching capabilities，Wang et al.，2025，NeurIPS】。这一经验趋势表明了具有理论保证的稳定残差矩阵的重要性。

3.2. 重参数化与mHC-lite

理论基础：Birkhoff-von Neumann定理。我们的方法基于Birkhoff-von Neumann定理【2，Birkhoff，1946; 21，von Neumann，1953】，mHC的论文【24，mhc: Manifold-constrained hyper-connections，Xie et al.，2025】也强调了这一点。为使本文内容完整，我们重述该定理如下。
定理 3.1 (Birkhoff-von Neumann定理)。对于任意$X \in B_n$，存在一个权重$a = (a_1, ..., a_{n!}) \in \mathbb{R}^{1 \times n!}$，其中$a_k \ge 0, \forall k \in [n!], |a|_{\ell_1} = 1$，使得

$$\boldsymbol{X}=\sum_{i=1}^{n!} a_{k} \boldsymbol{P}_{k}$$

其中 ${P_k}_{k=1}^{n!}$ 是 $n \times n$ 置换矩阵的序列。

参数化方法。基于Birkhoff-von Neumann定理，我们直接将双随机矩阵表示为置换矩阵的凸组合。这种参数化保证了矩阵是精确双随机的。此外，通过消除迭代近似，该参数化移除了它们在训练和推理中的计算开销，避免了对高度专业化基础设施的重度依赖。

mHC-lite的动态矩阵构建。在mHC-lite中，为了控制混杂因素，我们保持了mHC的结构不变，除了$H_{res}^l$。设$x_l \in \mathbb{R}^{n \times C}$为第$l$层的输入特征，$\hat{x}l \in \mathbb{R}^{1 \times nC}$为展平的输入特征。然后我们根据$x_l$动态构建映射$H^l$如下。}^l, H_{pre}^l$和$H_{post

$$\begin{aligned} \begin{aligned} \hat{\boldsymbol{x}}_{l}^{\prime} & =\operatorname{RMSNorm}\left(\hat{\boldsymbol{x}}_{l}\right) \\ \boldsymbol{H}_{l}^{\text {pre }} & =\operatorname{sigmoid}\left(\alpha_{l}^{\text {pre }} \hat{\boldsymbol{x}}_{l}^{\prime} \boldsymbol{W}_{l}^{\text {pre }}+\boldsymbol{b}_{l}^{\text {pre }}\right) \\ \boldsymbol{H}_{l}^{\text {post }} & =2 \cdot \operatorname{sigmoid}\left(\alpha_{l}^{\text {post }} \hat{\boldsymbol{x}}_{l}^{\prime} \boldsymbol{W}_{l}^{\text {post }}+\boldsymbol{b}_{l}^{\text {post }}\right) \\ \boldsymbol{a}_{l} & =\operatorname{softmax}\left(\alpha_{l}^{\text {res }} \hat{\boldsymbol{x}}_{l}^{\prime} \boldsymbol{W}_{l}^{\text {res }}+\boldsymbol{b}_{l}^{\text {res }}\right) \end{aligned} \end{aligned}$$ $$H_{l}^{\mathrm{res}}=\sum_{k=1}^{n!} a_{l, k} \boldsymbol{P}_{k}$$

其中$W_{pre}^l, W_{post}^l \in \mathbb{R}^{nC \times n}$和$W_{res}^l \in \mathbb{R}^{nC \times n!}$是第$l$层的可学习权重矩阵。这里$b_{pre}, b_{post} \in \mathbb{R}^{1 \times n}$和$b_{res}^l \in \mathbb{R}^{1 \times n!}$是可学习的偏置。$\alpha_{pre}, \alpha_{post}$和$\alpha_{res}$是可学习的标量。RMSNorm(·)指的是RMSNorm【25，Root mean square layer normalization，Zhang & Sennrich，2019】。

具体实现。在实践中，我们首先通过一个带softmax激活的线性层计算一个动态权重向量$a_l = (a_{l,1}, ..., a_{l,n!}) \in \mathbb{R}^{n!}$。回想一下，$n$表示残差流的数量，在HC和mHC中$n = 4$【26，Zhu et al.，2024; 24，Xie et al.，2025】，所以$n! = 24$是一个很小的常数。为了生成$H_{res}^l$，公式5通过$a_{res}$和一个常数0/1矩阵（由所有置换矩阵拼接并重塑而成，维度为$\mathbb{R}^{n! \times n^2}$）之间的矩阵乘法来实现。

计算优势。与HC和mHC【23，Xie et al.，2024; 26，Zhu et al.，2024】类似，由残差连接引入的额外FLOPs与主要变换$f(\cdot; W_l)$相比通常可以忽略不计。例如，在Transformer架构【20，Attention is all you need，Vaswani et al.，2017，NeurIPS】中，$f(\cdot; W_l)$对应于注意力（attention）和多层感知机（MLP）操作，它们占据了主要的计算量。我们在计算上的关键优势反而是面向工程的：该构造可以完全用标准算子实现，避免了对重复迭代的专用核的依赖，因此在不同框架间更具通用可移植性。

A4 实验环境

实现框架：实验基于nanoGPT框架【17，nanogpt，nanoGPT，2022，GitHub repository】实现。
数据集：
- OpenWebText
- FineWeb-Edu
模型架构：
- 采用三种不同规模的模型：
- S (Small)：6层，约4500万参数
- M (Medium)：12层，约1.2亿参数
- L (Large)：24层，约3.6亿参数

• 所有实验中，残差流的数量$n$均设为4。
  硬件配置：
• GPU：8块 NVIDIA A100 80GB (SXM4) GPU，单节点运行。
  软件配置：
• 实现：使用PyTorch的DistributedDataParallel (DDP) 与NCCL后端进行分布式训练。

• 训练细节：

  • 使用AdamW优化器【15，Decoupled weight decay regularization，Loshchilov & Hutter，2017】从头开始训练。
  • 学习率调度采用带线性预热的余弦衰减策略。
  • 使用bfloat16进行混合精度训练，并设置梯度裁剪。
  • 总训练步数为10,000步（约处理13亿个token）。

• 初始化策略：HC/mHC/mHC-lite模块中的参数被特殊初始化，使得在训练开始时，每个模块等效于一个普通的残差连接。具体来说，$W_{pre}, W_{post}, W_{res}^l$初始化为零，$\alpha$标量初始化为0.01。偏置向量$b_{pre}^l$和$b_{post}^l$除一个元素设为1外其余均为-1。对于mHC，$b_{res}^l$除对角线元素设为0外其余为-8，以在指数化后近似恒等矩阵。对于mHC-lite，$b_{res}$除对应恒等矩阵的条目设为0外其余为-8，以使softmax后的权重集中在恒等矩阵上。

A4 实验结果

性能与训练稳定性

模型损失：为了验证mHC-lite是否能达到与mHC相当的模型损失改进，实验比较了不同残差连接组件在训练结束时的训练和验证损失。如表1所示，mHC-lite在所有数据集和模型规模上都取得了与mHC相当甚至略优的性能。

训练稳定性：图2展示了在FineWeb-Edu数据集上训练L模型时的梯度范数曲线。结果表明，mHC-lite与mHC一样对训练具有稳定作用。从图2右侧的放大图可以看出，mHC-lite的梯度范数均值略低于mHC，且波动更小，进一步证实了其在稳定训练动态方面的有效性。

效率

训练吞吐量：实验通过测量在OpenWebText数据集上使用M模型的平均训练吞吐量（每秒token数）来比较mHC-lite与HC的计算效率。如图5所示，mHC-lite在没有任何系统级优化的情况下，实现了比HC更高的吞吐量。作为对比，基于PyTorch重新实现的mHC吞吐量较低。作者指出，尽管mHC的原始论文【24，mhc: Manifold-constrained hyper-connections，Xie et al.，2025】声称使用专用核的优化实现仅比HC慢6.7%，但mHC-lite的优异表现凸显了其高度的实现友好性，易于集成到现有训练代码和实际系统中。

稳定性分析

mHC稳定性问题的实证：本节通过分析每个token级别的$H_{res}^l$矩阵（而非mHC论文中的token平均矩阵）来评估mHC的真实稳定性。实验首先分析了SK算法输入的相对范围$1/\nu$（定义见公式3）。如图4所示，在应用SK之前，mHC中有相当一部分（约27.9%）输入的$\log(1/\nu)$大于30（即$1/\nu > 10^{13}$），在这种情况下20次SK迭代可能无法很好地收敛到双随机矩阵。此外，更深的模型（L模型）通常具有更大的$1/\nu$，这表明固定的20次SK迭代对于更深的网络可能不足够。

列和偏差分析：实验进一步直接检查了mHC中$H_{res}$的列和分布。如图3所示，虽然单个$H_{res}^l$矩阵的列和中位数通常接近1，但存在许多显著偏离1的离群值。更重要的是，当考虑跨层矩阵的乘积$Q H_{res}^l$时，即使是中位数也可能远离1。这种偏差在更深的模型中更为严重，暗示了当模型进一步扩展时潜在的不稳定风险。

mHC-lite的稳定性保证：相比之下，mHC-lite不依赖迭代归一化，因此避免了与收敛相关的失败。对于mHC-lite，$H_{res}^l$及其组合$Q H_{res}^l$的完美双随机性是通过Birkhoff-von Neumann定理从构造上保证的。

A5 结论

本文从稳定性和系统可移植性的角度重新审视了mHC的残差连接设计。研究发现，mHC依赖的迭代式SK算法需要专门的核函数才能高效执行，这为通用采纳设置了工程障碍。此外，通过理论分析和实证评估，我们发现由于mHC依赖有限步数的SK迭代，当SK算法未能收敛时，其残差矩阵可能显著偏离双随机性，从而引入潜在的稳定性风险。

为解决这些局限性，本文提出了mHC-lite，一个基于Birkhoff–von Neumann定理重参数化双随机矩阵的简单、强大且高效的替代方案。该重参数化方法使得我们可以完全跳过SK迭代，消除了近似差距，并仅用基本算子即可支持计算。这使得我们的方法成为经典残差架构的即插即用替代品，提供了有保证的鲁棒性，而无需牺牲部署的简便性。

mHC-lite的设计验证了一个简单而强大的原则：当可实现时，精确性通常是最高效的近似形式。这种从“投影”到“重参数化”的转变，确保了约束通过构造得到满足，消除了近似差距（例如由有限次Sinkhorn–Knopp迭代引起的差距），同时可能实现更高效的实现。

关于mHC-lite在更大n值下的计算效率讨论：尽管mHC-lite在$n=4$时表现良好，但其时空复杂度随$n$呈指数增长。对此，作者提出两点看法：
1. 原始HC论文【26，Hyper-connections，Zhu et al.，2024，ArXiv】的消融研究表明，$n=4$在实践中是一个优越的选择。
2. 即使需要更大的$n$，也可以通过采样一部分置换矩阵而非全部来降低计算成本。这相当于将可行域限制在Birkhoff多胞体的一个子集上。由此产生的残差矩阵仍然保证是双随机的，同时可以通过控制采样置换的数量来调整计算预算。

A6 附录

实现细节。我们的实现基于nanoGPT【17，nanogpt，nanoGPT，2022，GitHub repository】，除非另有说明，所有参数均设为默认值。所有模型均使用AdamW优化器【15，Decoupled weight decay regularization，Loshchilov & Hutter，2017】从头开始训练，采用余弦学习率调度和线性预热。我们使用bfloat16进行混合精度训练并进行梯度裁剪。所有实验均在8个NVIDIA A100 80GB GPU上使用PyTorch的DistributedDataParallel (DDP)和NCCL后端进行。

共享超参数。所有实验中使用的共享超参数总结在表2中。

模型规模特定超参数。对于三种模型规模（S、M和L），其特定规模的超参数列于表3中。