mHC: Manifold-Constrained Hyper-Connections

作者/机构: Zhenda Xie*†, Yixuan Wei*, Huanqi Cao*, Chenggang Zhao, Chengqi Deng, Jiashi Li, Damai Dai, Huazuo Gao, Jiang Chang, Liang Zhao, Shangyan Zhou, Zhean Xu, Zhengyan Zhang, Wangding Zeng, Shengding Hu, Yuqing Wang, Jingyang Yuan, Lean Wang, Wenfeng Liang, DeepSeek-AI

A1 主要贡献

本文探讨了深度神经网络中残差连接的演进。传统的残差连接（图 1a）在过去十年中已成为基础设计，尤其是在大型语言模型中，其成功关键在于简洁的形式和“恒等映射”特性，这保证了大规模训练的稳定性和效率。然而，以超连接（Hyper-Connections, HC）为代表的新研究（图 1b）通过扩展残差流的宽度和多样化连接模式来提升性能，但这种无约束的设计从根本上破坏了恒等映射属性。

核心问题：
1. 训练不稳定性：HC 中无约束的可学习映射矩阵 $H_{res}$ 在多层堆叠后，其复合映射 $\prod H_{res}$ 无法保持特征的全局均值，导致信号无限放大或衰减，从而在大规模训练中引发不稳定性。
2. 系统开销：虽然 HC 在浮点运算（FLOPs）方面计算开销不大，但其加宽的残差流带来了显著的内存访问成本和通信开销，限制了其实际可扩展性。

研究目标与创新点：
为解决上述挑战，本文提出了流形约束超连接（Manifold-Constrained Hyper-Connections, mHC），一个通用框架（图 1c），旨在恢复 HC 的恒等映射属性，同时通过严格的基础设施优化来保证效率。

核心创新如下：
1. 恢复恒等映射以保证稳定性：mHC 的核心思想是将 HC 中的残差连接空间投影到一个特定的流形上。具体地，它使用 Sinkhorn-Knopp 算法将残差连接矩阵 $H_{res}$ 投影到由双随机矩阵构成的伯克霍夫多胞体（Birkhoff polytope）上。由于双随机矩阵的行和与列和均为 1，这使得残差映射操作等效于对输入特征的凸组合，从而保持了特征均值的守恒，并严格正则化了信号范数，有效缓解了信号消失或爆炸的风险。由于双随机矩阵在矩阵乘法下是封闭的，这种稳定性可以贯穿任意深度的网络。
2. 高效的基础设施优化：为了解决系统开销问题，本文进行了一系列工程优化。
* Kernel Fusion（核函数融合）：利用 TileLang 开发了混合精度核函数，将多个共享内存访问的操作融合成统一的计算核，以减少内存带宽瓶颈。
* Recomputing（重计算）：通过选择性重计算来减少内存占用，即在反向传播过程中重新计算前向传播中的中间激活值。
* Overlapping Communication（通信重叠）：在 DualPipe 调度中精心设计了通信与计算的重叠，以缓解流水线并行中的通信延迟。

主要成果：
大量实验证明，mHC 在保持 HC 性能优势的同时，展现出卓越的稳定性和可扩展性。在内部大规模训练中，当扩展率 $n=4$ 时，mHC 仅引入了 6.7% 的额外时间开销，证明了其在大规模训练中的有效性和实用性。

图 1 | 残差连接范式图示。该图比较了 (a) 标准残差连接, (b) 超连接 (HC), 以及 (c) 我们提出的流形约束超连接 (mHC) 的结构设计。与无约束的 HC 不同，mHC 专注于优化残差连接空间，通过将矩阵投影到受约束的流形上来确保稳定性。
图 1 | 残差连接范式图示。该图比较了 (a) 标准残差连接, (b) 超连接 (HC), 以及 (c) 我们提出的流形约束超连接 (mHC) 的结构设计。与无约束的 HC 不同，mHC 专注于优化残差连接空间，通过将矩阵投影到受约束的流形上来确保稳定性。

A3 背景知识与关键观察

初步介绍

HC 公式化符号定义
首先，我们确定本文中使用的符号。在 HC 的公式中，第 $l$ 层的输入 $x_l \in \mathbb{R}^{1 \times C}$ 被扩展一个因子 $n$，构建出一个隐藏矩阵 $\mathbf{x}_l = (\mathbf{x}_{l,0}^\top, \dots, \mathbf{x}_{l,n-1}^\top)^\top \in \mathbb{R}^{n \times C}$，这可以看作是 $n$ 流的残差。这个操作有效地加宽了残差流的宽度。为了管理这个流的读出、写入和更新过程，HC 引入了三个可学习的线性映射——$H_l^{\text{pre}}, H_l^{\text{post}} \in \mathbb{R}^{1 \times n}$ 和 $H_l^{\text{res}} \in \mathbb{R}^{n \times n}$。这些映射修改了方程 (1) 中所示的标准残差连接，从而得到了方程 (3) 中的公式。

HC 映射的计算方式
在 HC 的公式中，可学习的映射由两部分系数组成：依赖于输入的部分和全局部分，分别称为动态映射和静态映射。形式上，HC 如下计算系数：

$$\begin{aligned} \begin{cases} \tilde{\mathbf{x}}_{l} = \mathrm{RMSNorm}(\mathbf{x}_{l}) \\ \mathcal{H}_{l}^{\mathrm{pre}} = \alpha_{l}^{\mathrm{pre}} \cdot \tanh(\theta_{l}^{\mathrm{pre}} \tilde{\mathbf{x}}_{l}^{\top}) + \mathbf{b}_{l}^{\mathrm{pre}} \\ \mathcal{H}_{l}^{\mathrm{post}} = \alpha_{l}^{\mathrm{post}} \cdot \tanh(\theta_{l}^{\mathrm{post}} \tilde{\mathbf{x}}_{l}^{\top}) + \mathbf{b}_{l}^{\mathrm{post}} \\ \mathcal{H}_{l}^{\mathrm{res}} = \alpha_{l}^{\mathrm{res}} \cdot \tanh(\theta_{l}^{\mathrm{res}} \tilde{\mathbf{x}}_{l}^{\top}) + \mathbf{b}_{l}^{\mathrm{res}}, \end{cases} \end{aligned}$$


其中，RMSNorm(·) 【45, Root mean square layer normalization, 2019, Advances in neural information processing systems】 应用于最后一个维度，标量 $\alpha_l^{\text{pre}}, \alpha_l^{\text{post}}$ 和 $\alpha_l^{\text{res}} \in \mathbb{R}$ 是初始化为小值的可学习门控因子。动态映射通过由 $\theta_l^{\text{pre}}, \theta_l^{\text{post}} \in \mathbb{R}^{1 \times n}$ 和 $\theta_l^{\text{res}} \in \mathbb{R}^{n \times n}$ 参数化的线性投影导出，而静态映射由可学习的偏置 $b_l^{\text{pre}}, b_l^{\text{post}} \in \mathbb{R}^{1 \times n}$ 和 $b_l^{\text{res}} \in \mathbb{R}^{n \times n}$ 表示。

HC 的计算开销与扩展维度
值得注意的是，引入这些映射——$H_l^{\text{pre}}, H_l^{\text{post}}$ 和 $H_l^{\text{res}}$——所带来的计算开销可以忽略不计，因为典型的扩展率 $n$（例如 4）远小于输入维度 $C$。通过这种设计，HC 有效地将残差流的信息容量与层的输入维度解耦，而后者与模型的计算复杂度（FLOPs）密切相关。因此，HC 通过调整残差流宽度提供了一个新的扩展途径，补充了预训练缩放定律【16, An empirical analysis of compute-optimal large language model training, 2022, Advances in Neural Information Processing Systems】中讨论的模型 FLOPs 和训练数据大小等传统扩展维度。

HC 组件的消融研究
虽然 HC 需要三个映射来管理残差流和层输入之间的维度不匹配，但表 1 中展示的初步实验表明，残差映射 $H_l^{\text{res}}$ 产生了最显著的性能增益。这一发现强调了在残差流内进行有效信息交换的至关重要性。

表 1 | HC 组件的消融研究。当禁用特定的映射（$H_l^{\text{pre}}$, $H_l^{\text{post}}$, 或 $H_l^{\text{res}}$）时，我们采用一个固定的映射来维持维度一致性：对于 $H_l^{\text{pre}}$ 使用权重为 $1/n$ 的均匀权重，对于 $H_l^{\text{post}}$ 使用权重为 1 的均匀权重，对于 $H_l^{\text{res}}$ 使用单位矩阵。
表 1

数值不稳定性

不稳定性根源分析
虽然残差映射 $H_l^{\text{res}}$ 对性能至关重要，但其顺序应用对数值稳定性构成了重大风险。如方程 (4) 中详述，当 HC 扩展到多层时，从层 $l$ 到 $L$ 的有效信号传播由复合映射 $\prod_{i=1}^{L-l} H_{L-i}^{\text{res}}$ 控制。由于可学习的映射 $H_l^{\text{res}}$ 是无约束的，这个复合映射不可避免地偏离了恒等映射。因此，信号幅度在前向传播和反向传播过程中容易发生爆炸或消失。这种现象破坏了残差学习的基本前提，即依赖于无阻碍的信号流，从而在更深或更大规模的模型中破坏了训练过程的稳定性。

不稳定性的实证
经验证据支持了这一分析。我们在大规模实验中观察到不稳定的损失行为，如图 2 所示。以 mHC 为基线，HC 在 12k 步左右表现出意想不到的损失激增，这与梯度范数的不稳定性高度相关。此外，对 $H_l^{\text{res}}$ 的分析验证了这种不稳定性的机制。为了量化复合映射 $\prod_{i=1}^{L-l} H_{L-i}^{\text{res}}$ 如何沿残差流放大信号，我们使用了最大绝对行和，该映射捕捉了前向传播中的最坏情况扩展。第二个是基于最大绝对列和，对应于反向传播。我们将这些度量称为复合映射的 Amax 增益幅度。如图 3 (b) 所示，Amax 增益幅度产生了极值，峰值达到 3000，这与 1 的巨大差异证实了残差流爆炸的存在。

图 2 | 超连接（HC）的训练不稳定性。该图展示了 (a) HC 相对于 mHC 的绝对损失差距，以及 (b) 梯度范数的比较。所有结果均基于 27B 模型。
图 2 | 超连接（HC）的训练不稳定性。该图展示了 (a) HC 相对于 mHC 的绝对损失差距，以及 (b) 梯度范数的比较。所有结果均基于 27B 模型。
图 3 | 超连接（HC）的传播不稳定性。该图展示了 (a) 单层映射 Hres 和 (b) 复合映射 Π Hres 在 27B 模型内的传播动态。层索引 l（x 轴）将每个标准的 Transformer 块展开为两个独立的层（Attention 和 FFN）。Amax 增益幅度（y 轴）计算为最大绝对行和（对于前向信号）和列和（对于后向梯度），并在选定序列中的所有 token 上取平均。
图 3 | 超连接（HC）的传播不稳定性。该图展示了 (a) 单层映射 Hres 和 (b) 复合映射 Π Hres 在 27B 模型内的传播动态。层索引 l（x 轴）将每个标准的 Transformer 块展开为两个独立的层（Attention 和 FFN）。Amax 增益幅度（y 轴）计算为最大绝对行和（对于前向信号）和列和（对于后向梯度），并在选定序列中的所有 token 上取平均。

系统开销

内存墙问题
虽然 HC 的计算复杂度由于附加映射的线性特性而保持可控，但系统级开销构成了不可忽视的挑战。具体来说，内存访问（I/O）成本通常是现代模型架构中的主要瓶颈之一，这被广泛称为“内存墙”【7, FlashAttention: Fast and memory-efficient exact attention with IO-awareness, 2022, Advances in Neural Information Processing Systems (NeurIPS)】。这个瓶颈在架构设计中经常被忽视，但它决定性地影响着运行效率。

HC 的 I/O 开销与内存占用
我们关注广泛采用的 pre-norm Transformer 【36, Attention is all you need, 2017, Advances in neural information processing systems】架构，分析了 HC 固有的 I/O 模式。表 2 总结了由 $n$ 流残差设计引入的单个残差层中每个 token 的内存访问开销。分析显示，HC 将内存访问成本增加了大约与 $n$ 成正比的因子。这种过度的 I/O 需求在没有融合核函数缓解的情况下显著降低了训练吞吐量。此外，由于 $H_l^{\text{pre}}, H_l^{\text{post}}$ 和 $H_l^{\text{res}}$ 涉及可学习的参数，它们的反向传播需要中间激活值。这导致 GPU 内存占用大幅增加，通常需要梯度检查点来维持可行的内存使用。此外，HC 在流水线并行【27, Zero bubble (almost) pipeline parallelism, 2024, The Twelfth International Conference on Learning Representations】中需要 $n$ 倍的通信成本，导致更大的气泡并降低训练吞吐量。

表 2 | 每个 Token 的内存访问成本比较。此分析考虑了前向传播中残差流维护所引入的开销，不包括层函数 F 的内部 I/O。
表 2

A2 方法细节

4.1. 流形约束超连接

核心思想与流形选择
受到恒等映射原理【12, Identity mappings in deep residual networks, 2016, European conference on computer vision】的启发，mHC 的核心前提是将残差映射 $H_l^{\text{res}}$ 约束在一个特定的流形上。虽然原始的恒等映射通过强制 $H_l^{\text{res}} = I$ 来确保稳定性，但它从根本上排除了残差流内部的信息交换，而这对于最大化多流架构的潜力至关重要。因此，我们建议将残差映射投影到一个既能保持跨层信号传播稳定性又能促进残差流之间相互作用以保持模型表达能力的流形上。为此，我们将 $H_l^{\text{res}}$ 限制为一个双随机矩阵，其条目非负，且行和与列和均为 1。形式上，设 $M_{\text{res}}$ 为双随机矩阵的流形（也称为伯克霍夫多胞体）。我们将 $H_l^{\text{res}}$ 约束为 $\mathcal{P}_{M_{\text{res}}}(H_l^{\text{res}})$，定义如下：

$$\mathcal{P}_{\mathcal{M}^{\text{res}}}(\mathcal{H}_l^{\text{res}}) := \{\mathcal{H}_l^{\text{res}} \in \mathbb{R}^{n \times n} \mid \mathcal{H}_l^{\text{res}} \mathbf{1}_n = \mathbf{1}_n, \mathbf{1}_n^\top \mathcal{H}_l^{\text{res}} = \mathbf{1}_n^\top, \mathcal{H}_l^{\text{res}} \geqslant 0 \}$$


其中 $\mathbf{1}_n$ 代表全为 1 的 $n$ 维向量。

双随机性约束的理论优势
值得注意的是，当 $n=1$ 时，双随机条件退化为标量 1，从而恢复了原始的恒等映射。选择双随机性赋予了几个对大规模模型训练有益的严格理论性质：

  1. 范数保持：双随机矩阵的谱范数有界为 1 (即 $\|H_l^{\text{res}}\|_2 \le 1$ )。这意味着可学习的映射是非扩张的，有效缓解了梯度爆炸问题。
  2. 复合闭包：双随机矩阵集合在矩阵乘法下是封闭的。这确保了跨多层的复合残差映射 $\prod_{i=1}^{L-l} H_{L-i}^{\text{res}}$ 仍然是双随机的，从而在模型的整个深度上保持了稳定性。
  3. 通过伯克霍夫多胞体的几何解释：集合 $M_{\text{res}}$ 构成了伯克霍夫多胞体，它是置换矩阵集合的凸包。这提供了一个清晰的几何解释：残差映射作为置换的凸组合。在数学上，重复应用此类矩阵倾向于单调增加跨流信息的混合，有效地作为一个鲁棒的特征融合机制。

对输入输出映射的约束
此外，我们对输入映射 $H_l^{\text{pre}}$ 和输出映射 $H_l^{\text{post}}$ 施加了非负性约束。这种约束防止了由正负系数组合引起的信号抵消，这也可以被视为一种特殊的流形投影。

4.2. 参数化与流形投影

映射计算过程
在本节中，我们详细介绍 mHC 中 $H_l^{\text{pre}}, H_l^{\text{post}}$ 和 $H_l^{\text{res}}$ 的计算过程。给定第 $l$ 层的输入隐藏矩阵 $\mathbf{x}_l \in \mathbb{R}^{n \times C}$，我们首先将其展平为一个向量 $\vec{\mathbf{x}}_l = \text{vec}(\mathbf{x}_l) \in \mathbb{R}^{1 \times nC}$ 以保留完整的上下文信息。然后，我们遵循原始 HC 的公式来获得动态映射和静态映射，如下所示：

$$\begin{aligned} \begin{cases} \vec{\mathbf{x}}_l' = \text{RMSNorm}(\vec{\mathbf{x}}_l) \\ \tilde{\mathcal{H}}_l^{\text{pre}} = \alpha_l^{\text{pre}} \cdot (\vec{\mathbf{x}}_l' \varphi_l^{\text{pre}}) + \mathbf{b}_l^{\text{pre}} \\ \tilde{\mathcal{H}}_l^{\text{post}} = \alpha_l^{\text{post}} \cdot (\vec{\mathbf{x}}_l' \varphi_l^{\text{post}}) + \mathbf{b}_l^{\text{post}} \\ \tilde{\mathcal{H}}_l^{\text{res}} = \alpha_l^{\text{res}} \cdot \text{mat}(\vec{\mathbf{x}}_l' \varphi_l^{\text{res}}) + \mathbf{b}_l^{\text{res}}, \end{cases} \end{aligned}$$


其中 $\varphi_l^{\text{pre}}, \varphi_l^{\text{post}} \in \mathbb{R}^{nC \times n}$ 和 $\varphi_l^{\text{res}} \in \mathbb{R}^{nC \times n^2}$ 是用于动态映射的线性投影，而 mat(·) 是一个从 $\mathbb{R}^{1 \times n^2}$ 到 $\mathbb{R}^{n \times n}$ 的重塑函数。

流形投影
然后，通过以下方式获得最终的约束映射：

$$\begin{aligned} \left\{\begin{aligned} \mathcal{H}_{l}^{\text {pre }} & =\sigma\left(\tilde{\mathcal{H}}_{l}^{\text {pre }}\right) \\ \mathcal{H}_{l}^{\text {post }} & =2 \sigma\left(\tilde{\mathcal{H}}_{l}^{\text {post }}\right) \\ \mathcal{H}_{l}^{\text {res }} & =\text { Sinkhorn-Knopp }\left(\tilde{\mathcal{H}}_{l}^{\text {res }}\right), \end{aligned}\right. \end{aligned}$$


其中 $\sigma(\cdot)$ 表示 Sigmoid 函数。Sinkhorn-Knopp(·) 算子首先通过指数算子使所有元素为正，然后进行迭代归一化过程，交替地重新缩放行和列使其和为 1。

Sinkhorn-Knopp 算法细节
具体来说，给定一个正矩阵 $\mathbf{M}^{(0)} = \exp(\tilde{\mathcal{H}}_l^{\text{res}})$ 作为起点，归一化迭代过程如下：

$$\mathbf{M}^{(t)}=\mathcal{T}_{r}\left(\mathcal{T}_{c}\left(\mathbf{M}^{(t-1)}\right)\right),$$


其中 $\mathcal{T}_r$ 和 $\mathcal{T}_c$ 分别表示行归一化和列归一化。当 $t_{\max} \to \infty$ 时，这个过程收敛到一个双随机矩阵 $\mathcal{H}_l^{\text{res}} = \mathbf{M}^{(t_{\max})}$。在我们的实验中，我们选择 $t_{\max} = 20$ 作为一个实际值。

4.3. 高效的基础设施设计

总体开销
在本节中，我们详细介绍为 mHC 量身定制的基础设施设计。通过严格的优化，我们在大规模模型中实现了 mHC (当 $n=4$ 时)，其边际训练开销仅为 6.7%。

4.3.1. Kernel Fusion (核函数融合)

优化策略
观察到 mHC 中的 RMSNorm 在处理高维隐藏状态 $\vec{\mathbf{x}}_l \in \mathbb{R}^{1 \times nC}$ 时会带来显著的延迟，我们重新排序了除以范数的操作，使其在矩阵乘法之后进行。这个优化在保持数学等价性的同时提高了效率。此外，我们采用混合精度策略以最大化数值精度而不牺牲速度，并将具有共享内存访问的多个操作融合成统一的计算核，以减少内存带宽瓶颈。基于方程 (10) 到 (13) 中详述的输入和参数，我们实现了三个专门的 mHC 核来计算 $H_l^{\text{pre}}, H_l^{\text{post}}$ 和 $H_l^{\text{res}}$。在这些核中，偏置和线性投影被整合到 $\mathbf{b}_l$ 和 $\varphi_l$ 中，RMSNorm 的权重也被吸收到 $\varphi_l$ 中。

具体实现
- 方程 (14) 到 (15)：我们开发了一个统一的核，融合了对 $\vec{\mathbf{x}}_l$ 的两次扫描，利用矩阵乘法单元来最大化内存带宽利用率。反向传播——包括两次矩阵乘法——同样被整合到一个单一的核中，消除了对 $\vec{\mathbf{x}}_l$ 的冗余重载。两个核都具有一个精细调整的流水线（加载、转换、计算、存储）以高效处理混合精度。
- 方程 (16) 到 (18)：这些对小系数的轻量级操作被机会性地融合成一个单一的核，显著减少了核启动开销。
- 方程 (19)：我们在一个单一的核中实现了 Sinkhorn-Knopp 迭代。对于反向传播，我们推导了一个自定义的反向核，它在芯片上重计算中间结果并遍历整个迭代过程。

$$\varphi_{l} : \text{tfloat32} \quad [nC, n^{2} + 2n]$$


$\vec{\mathbf{x}}_{l} : \text{bfloat16}$
$[1, nC]$
$\alpha_{l}^{\text {pre }}, \alpha_{l}^{\text {post }}, \alpha_{l}^{\text {res }}: \text { float32 }$
$\text{Scalars}$
$\mathbf{b}_l : \text{float32} \quad [1, n^2 + 2n]$

$$ [\tilde{\tilde{\mathcal{H}}}_l^\text{pre}, \tilde{\tilde{\mathcal{H}}}_l^\text{post}, \tilde{\tilde{\mathcal{H}}}_l^\text{res}] : \text{float32} \qquad = \vec{\mathbf{x}}_l \varphi_l $$
$r: \text{float}32$
$$ = \|\vec{\mathbf{x}}_l\|_2 / \sqrt{nC} $$
$$ \left[\tilde{\mathcal{H}}_{l}^{\text {pre }}, \tilde{\mathcal{H}}_{l}^{\text {post }}, \tilde{\mathcal{H}}_{l}^{\text {res }}\right]: \text { float32 } \quad=1 / r\left[\alpha_{l}^{\text {pre }} \tilde{\mathcal{H}}_{l}^{\text {pre }}, \alpha_{l}^{\text {post }} \tilde{\mathcal{H}}_{l}^{\text {post }}, \alpha_{l}^{\text {res }} \tilde{\mathcal{H}}_{l}^{\text {res }}\right]+\mathbf{b}_{l} $$
$$\mathcal{H}_{l}^{\text {pre }}: \text { float32 } \quad=\sigma\left(\tilde{\mathcal{H}}_{l}^{\text {pre }}\right)$$
公式 (25)
$$\mathcal{H}_{l}^{\text{res}} : \text{float32} \qquad = \text{Sinkhorn-Knopp} \left( \tilde{\mathcal{H}}_{l}^{\text{res}} \right)$$

应用映射与 TileLang
使用从上述核函数派生的系数，我们引入了另外两个核函数来应用这些映射：一个用于 $\mathcal{F}_{\text{pre}} := \mathcal{H}_l^{\text{pre}}\mathbf{x}_l$，另一个用于 $\mathcal{F}_{\text{post,res}} := \mathcal{H}_l^{\text{res}}\mathbf{x}_l + \mathcal{H}_l^{\text{post}\top}\mathcal{F}(\cdot, \cdot)$。通过将 $\mathcal{H}_l^{\text{post}}$ 和 $\mathcal{H}_l^{\text{res}}$ 的应用与残差合并融合，我们将该核函数读取的元素数量从 $(3n+1)C$ 减少到 $(n+1)C$，写入的元素数量从 $3nC$ 减少到 $nC$。我们使用 TileLang 【38, Tilelang: A composable tiled programming model for ai systems, 2025, arXiv preprint】 高效地实现了大部分核函数（除了方程 (14) 到 (15)）。这个框架简化了具有复杂计算过程的核函数的实现，并使我们能够以最小的工程努力充分利用内存带宽。

4.3.2. Recomputing (重计算)

内存优化策略
$n$ 流残差设计在训练期间引入了大量的内存开销。为了缓解这个问题，我们在前向传播后丢弃 mHC 核的中间激活值，并在反向传播中通过重新执行 mHC 核（不包括繁重的层函数 $\mathcal{F}$）来即时重计算它们。因此，对于一个包含 $L_r$ 个连续层的块，我们只需要存储第一层的输入 $\mathbf{x}_{l_0}$。不包括轻量级系数，同时考虑到 $\mathcal{F}$ 内部的 pre-norm，表 3 总结了为反向传播保留的中间激活值。

表 3 | 存储和重计算的中间激活值。我们列出了为反向传播保留的每个 token 的激活值，以及在 $L_r$ 个连续层中重计算的瞬时激活值。层 $l_0$ 表示 $L_r$ 层中的第一层，层 $l$ 在 $[l_0, l_0 + L_r - 1]$ 范围内。
表 3

最优块大小
由于 mHC 核的重计算是针对 $L_r$ 个连续层的块进行的，给定总共 $L$ 层，我们必须为反向传播持久存储所有 $\lceil \frac{L}{L_r} \rceil$ 个块的第一层输入 $\mathbf{x}_{l_0}$。除了这个常驻内存外，重计算过程还为活动块引入了 $(n+2)C \times L_r$ 个元素的瞬时内存开销，这决定了反向传播期间的峰值内存使用。因此，我们通过最小化与 $L_r$ 相关的总内存占用量来确定最优块大小 $L_r^*$：

$$L_r^* = \arg \min_{L_r} \left[ nC \times \left\lceil \frac{L}{L_r} \right\rceil + (n+2)C \times L_r \right] \approx \sqrt{\frac{nL}{n+2}}.$$

与流水线并行的结合
此外，大规模训练中的流水线并行施加了一个约束：重计算块不能跨越流水线阶段的边界。观察到理论上的最优 $L_r^*$ 通常与每个流水线阶段的层数一致，我们选择将重计算边界与流水线阶段同步。

4.3.3. Overlapping Communication in DualPipe (在 DualPipe 中重叠通信)

背景与挑战
在大规模训练中，流水线并行是减轻参数和梯度内存占用的标准做法。具体来说，我们采用 DualPipe 调度【22, Deepseek-v3 technical report, 2024b, arXiv preprint】，它有效地重叠了专家并行和流水线并行等横向扩展互连通信流量。然而，与单流设计相比，mHC 中提出的 $n$ 流残差在流水线阶段之间产生了大量的通信延迟。此外，在阶段边界处，对所有 $L_r$ 层的 mHC 核的重计算引入了不可忽略的计算开销。

解决方案
为了解决这些瓶颈，我们扩展了 DualPipe 调度（见图 4），以促进在流水线阶段边界处更好地重叠通信和计算。值得注意的是，为了防止阻塞通信流，我们在一个专用的高优先级计算流上执行 MLP（即 FFN）层的 $\mathcal{F}_{\text{post,res}}$ 核。我们进一步避免在注意力层中使用持久性核进行长时间运行的操作，从而防止长时间的停顿。这种设计使得重叠的注意力计算可以被抢占，从而在保持计算设备处理单元高利用率的同时，实现灵活的调度。此外，重计算过程与流水线通信依赖解耦，因为每个阶段的初始激活值 $\mathbf{x}_{l_0}$ 已经本地缓存。

图 4 | mHC 的通信-计算重叠。我们扩展了 DualPipe 调度来处理 mHC 引入的开销。每个块的长度仅为示意，不代表实际持续时间。(F)、(B)、(W) 分别指前向传播、反向传播、权重梯度计算。FA 和 FM 分别代表对应于 Attention 和 MLP 的核函数。
图 4 | mHC 的通信-计算重叠。我们扩展了 DualPipe 调度来处理 mHC 引入的开销。每个块的长度仅为示意，不代表实际持续时间。(F)、(B)、(W) 分别指前向传播、反向传播、权重梯度计算。FA 和 FM 分别代表对应于 Attention 和 MLP 的核函数。

A4 实验

实验环境

  • 模型架构：实验采用受 DeepSeek-V3 【22, Deepseek-v3 technical report, 2024b, arXiv preprint】 启发的 MoE 架构。HC 和 mHC 的扩展率 $n$ 均设置为 4。

  • 模型规模与数据集：

    1. 主要模型：一个 27B 参数的模型，使用与其参数量成比例的数据集进行训练，用于系统级的主要结果分析。
    2. 计算扩展性分析：训练了较小的 3B 和 9B 模型，数据量也与其参数量成比例，用于观察不同计算量下的性能趋势。
    3. Token 扩展性分析：单独训练了一个 3B 模型，使用固定的 1 万亿 token 的语料库。

  • 详细配置：详细的模型配置和训练超参数见附录 A.1。

实验结果

主要结果

训练稳定性与收敛性：
首先分析 27B 模型的训练稳定性和收敛性。如图 5(a) 所示，mHC 有效地缓解了 HC 中观察到的训练不稳定性，与基线相比，最终实现了 0.021 的损失降低。图 5(b) 中的梯度范数分析进一步证实了这种改进的稳定性，其中 mHC 表现出比 HC 好得多的行为，保持了与基线相当的稳定曲线。

图 5 | 流形约束超连接 (mHC) 的训练稳定性。该图展示了 (a) mHC 和 HC 相对于基线的绝对训练损失差距，以及 (b) 三种方法的梯度范数。所有实验均使用 27B 模型。结果表明，mHC 在损失和梯度范数方面均表现出更高的稳定性。
图 5 | 流形约束超连接 (mHC) 的训练稳定性。该图展示了 (a) mHC 和 HC 相对于基线的绝对训练损失差距，以及 (b) 三种方法的梯度范数。所有实验均使用 27B 模型。结果表明，mHC 在损失和梯度范数方面均表现出更高的稳定性。

下游任务性能：
表 4 展示了在一系列多样化的下游基准测试上的性能。mHC 取得了全面的提升，持续优于基线，并在大多数任务上超过了 HC。值得注意的是，与 HC 相比，mHC 进一步增强了模型的推理能力，在 BBH 【34, Challenging big-bench tasks and whether chain-of-thought can solve them, 2022, arXiv preprint】 上实现了 2.1% 的性能提升，在 DROP 【8, DROP: A reading comprehension benchmark requiring discrete reasoning over paragraphs, 2019, Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics】 上实现了 2.3% 的性能提升。

表 4 | 27B 模型的系统级基准测试结果。此表比较了基线、HC 和 mHC 在 8 个不同下游基准测试上的零样本和少样本性能。mHC 始终优于基线，并在大多数基准测试上超过 HC，证明了其在大规模预训练中的有效性。
表 4

扩展性实验

计算扩展性与 Token 扩展性：
为了评估方法的可扩展性，报告了 mHC 相对于基线在不同规模下的相对损失改善。图 6(a) 绘制了跨越 3B、9B 和 27B 参数的计算扩展曲线。轨迹表明，即使在更高的计算预算下，性能优势也能稳健地保持，仅表现出边际衰减。此外，图 6(b) 中展示了 3B 模型的 Token 扩展曲线，显示了其在训练过程中的动态。总的来说，这些发现验证了 mHC 在大规模场景中的有效性。

图 6 | mHC 与基线的扩展性属性比较。(a) 计算扩展曲线。实线描绘了不同计算预算下的性能差距。每个点代表模型大小和数据集大小的特定计算最优配置，从 3B 和 9B 扩展到 27B 参数。(b) Token 扩展曲线。3B 模型在训练期间的轨迹。每个点代表模型在不同训练 token 数下的性能。详细的架构和训练配置在附录 A.1 中提供。
图 6 | mHC 与基线的扩展性属性比较。(a) 计算扩展曲线。实线描绘了不同计算预算下的性能差距。每个点代表模型大小和数据集大小的特定计算最优配置，从 3B 和 9B 扩展到 27B 参数。(b) Token 扩展曲线。3B 模型在训练期间的轨迹。每个点代表模型在不同训练 token 数下的性能。详细的架构和训练配置在附录 A.1 中提供。

稳定性分析

传播稳定性：
与图 3 类似，图 7 展示了 mHC 的传播稳定性。理想情况下，单层映射应满足双随机约束，即前向信号增益和后向梯度增益都应等于 1。然而，使用 Sinkhorn-Knopp 算法的实际实现必须限制迭代次数以实现计算效率。在我们的设置中，我们使用 20 次迭代来获得近似解。因此，如图 7(a) 所示，后向梯度增益略有偏离 1。在图 7(b) 的复合情况下，偏差增加但仍保持有界，最大值约为 1.6。值得注意的是，与 HC 中近 3000 的最大增益幅度相比，mHC 将其显著降低了三个数量级。这些结果表明，mHC 显著增强了传播稳定性，确保了稳定的前向信号和后向梯度流。

图 7 | 流形约束超连接 (mHC) 的传播稳定性。该图展示了 (a) 单层映射 PMres(Hres) 和 (b) 复合映射 Π PMres(Hres) 在 27B 模型内的传播动态。结果表明，与 HC 相比，mHC 显著增强了传播稳定性。
图 7 | 流形约束超连接 (mHC) 的传播稳定性。该图展示了 (a) 单层映射 PMres(Hres) 和 (b) 复合映射 Π PMres(Hres) 在 27B 模型内的传播动态。结果表明，与 HC 相比，mHC 显著增强了传播稳定性。

可学习映射的可视化：
图 8 展示了代表性的映射。我们观察到，对于 HC，当最大增益较大时，其他值也趋于显著，这表明所有传播路径都存在普遍的不稳定性。相比之下，mHC 始终产生稳定的结果。

图 8 | 可学习映射的可视化。该图显示了 HC（第一行）和 mHC（第二行）的代表性单层和复合映射。每个矩阵是通过在选定序列中的所有 token 上取平均值计算得出的。沿 y 轴和 x 轴标注的标签分别表示前向信号增益（行和）和后向梯度增益（列和）。
图 8 | 可学习映射的可视化。该图显示了 HC（第一行）和 mHC（第二行）的代表性单层和复合映射。每个矩阵是通过在选定序列中的所有 token 上取平均值计算得出的。沿 y 轴和 x 轴标注的标签分别表示前向信号增益（行和）和后向梯度增益（列和）。

A5 结论与展望

结论：
本文指出，尽管像超连接（HC）那样扩展残差流宽度和多样化连接能带来性能增益，但这些连接的无约束性会导致信号发散，破坏了信号能量在各层间的守恒，从而引发训练不稳定并阻碍深度网络的可扩展性。为应对这些挑战，我们引入了流形约束超连接（mHC），一个将残差连接空间投影到特定流形上的通用框架。通过采用 Sinkhorn-Knopp 算法对残差映射强制施加双随机约束，mHC 将信号传播转化为特征的凸组合。实验结果证实，mHC 有效地恢复了恒等映射属性，实现了稳定的大规模训练，并具有优于传统 HC 的可扩展性。至关重要的是，通过高效的基础设施级优化，mHC 在实现这些改进的同时，计算开销可以忽略不计。

未来展望：
作为 HC 范式的通用扩展，mHC 为未来的研究开辟了几个有前景的途径。
1. 探索更多流形约束：虽然本文利用双随机矩阵来确保稳定性，但该框架允许探索为特定学习目标量身定制的各种流形约束。我们预计，对不同几何约束的进一步研究可能会产生更好地优化可塑性与稳定性之间权衡的新方法。
2. 推动宏观架构设计：我们希望 mHC 能够重新激发社区对宏观架构设计的兴趣。通过加深对拓扑结构如何影响优化和表示学习的理解，mHC 将有助于解决当前的局限性，并可能为下一代基础架构的演进指明新的道路。

A6 附录

A.1. 详细的模型规格和超参数

表 5 | 详细的模型规格和超参数。该表展示了基于 DeepSeek-V3 架构的 3B、9B 和 27B 模型的架构配置。它概述了 mHC 和 HC 的特定超参数，包括残差流扩展和 Sinkhorn-Knopp 设置，以及实验中使用的优化和训练协议。

表 5
表 5