mHC:流形约束超连接
mHC: Manifold-Constrained Hyper-Connections
摘要
近年来,以 Hyper‑Connections(HC)为代表的研究通过扩展残差流(residual stream)宽度并丰富连接模式,拓展了过去十年广泛采用的残差连接范式。尽管这种多样化带来了显著性能收益,但它从根本上破坏了残差连接所固有的恒等映射性质,从而导致严重的训练不稳定与可扩展性受限,并额外带来明显的内存访问开销。为解决这些问题,我们提出 Manifold‑Constrained Hyper‑Connections(mHC):一个通用框架,将 HC 的残差连接空间投影到特定流形上以恢复恒等映射性质,同时配合严格的基础设施级优化以保证效率。实证结果表明,mHC 适用于大规模训练,能够带来可观的性能提升与更强的可扩展性。我们期待 mHC 作为 HC 灵活且实用的扩展,能帮助更深入理解拓扑结构化的架构设计,并为基础模型的演进提供有前景的方向。
1. 引言
自 ResNet[1] 提出以来,深度神经网络架构经历了快速演进。如图 1(a) 所示,单层结构可写为:
其中 $\mathbf{x}_l$ 与 $\mathbf{x}_{l+1}$ 分别表示第 $l$ 层的 $C$ 维输入与输出,$\mathcal{F}$ 表示残差函数。尽管在过去十年中残差函数 $\mathcal{F}$ 已从卷积发展到注意力机制与前馈网络等多种形式,但残差连接这一范式仍保持其原始形态。随着 Transformer[2] 架构的普及,残差连接已成为大语言模型(LLM)[3–5] 的基础设计要素之一。
这种成功一方面源于残差连接的简洁形式,更重要的是,早期研究[6] 指出:残差连接的恒等映射性质能够在大规模训练中同时带来稳定性与效率。将残差连接在多层上递归展开,式 (1) 可推得:
其中 $L$ 与 $l$ 分别对应更深与更浅的层。“恒等映射”指的正是式 (2) 中的 $\mathbf{x}_l$ 项:来自浅层的信号无需任何修改即可直接映射到深层。
近期,以 Hyper‑Connections(HC)[7] 为代表的研究为残差连接引入了新的维度,并在实证上展示了其性能潜力。HC 的单层结构如图 1(b) 所示。通过扩展残差流宽度并增加连接复杂度,HC 在不改变单元计算 FLOPs 的前提下显著提升了拓扑复杂度。形式化地,HC 的单层传播定义为:
其中 $\mathbf{x}_{l}$ 与 $\mathbf{x}_{l+1}$ 分别表示第 $l$ 层的输入与输出。与式 (1) 不同,$\mathbf{x}_{l}$ 与 $\mathbf{x}_{l+1}$ 的特征维度由 $C$ 扩展为 $n\times C$,其中 $n$ 为扩展率。$\mathcal{H}_{l}^{\mathrm{res}}\in\mathbb{R}^{n\times n}$ 是一个可学习映射,用于在残差流内混合各路特征;同样是可学习映射的 $\mathcal{H}_{l}^{\mathrm{pre}}\in\mathbb{R}^{1\times n}$ 将 $nC$ 维残差流聚合为 $C$ 维的层输入,而 $\mathcal{H}_{l}^{\mathrm{post}}\in\mathbb{R}^{1\times n}$ 则将层输出写回到残差流。
然而,随着训练规模提升,HC 会引入潜在的不稳定风险。核心问题在于:当网络跨多层扩展时,HC 的无约束特性会破坏恒等映射性质。在包含多条并行残差流的架构中,理想的恒等映射可作为一种“守恒机制”,确保前向与反向传播时各路残差流的平均信号强度保持不变。将式 (3) 递归扩展到多层,可得:
其中 $L$ 与 $l$ 分别代表更深与更浅的层。与式 (2) 不同,HC 中的复合映射 $\prod_{i=1}^{L-l}\mathcal{H}_{L-i}^{\mathrm{res}}$ 无法保持特征的全局均值不变,从而导致信号无界放大或衰减,最终在大规模训练中引发不稳定。另一方面,尽管 HC 在 FLOPs 维度维持了计算效率,但其对加宽残差流带来的硬件层面内存访问成本并未给出有效应对。这些因素共同限制了 HC 的实用可扩展性,阻碍其在大规模训练中的应用。
为解决上述挑战,我们提出如图 1(c) 所示的 mHC:一个通用框架,将 HC 的残差连接空间投影到特定流形上以恢复恒等映射性质,并配合严格的基础设施优化以保证效率。具体而言,mHC 使用 Sinkhorn‑Knopp 算法[8] 将 $\mathcal{H}_{l}^{\mathrm{res}}$ 以熵正则方式投影到 Birkhoff 多面体上,从而将残差连接矩阵约束到由双随机矩阵构成的流形中。由于双随机矩阵的行和与列和均为 1,$\mathcal{H}_{l}^{\mathrm{res}}\mathbf{x}_l$ 等价于对输入特征做凸组合:它能实现良条件的信号传播,既保持特征均值守恒,又对信号范数施加强正则,从而有效缓解梯度/信号消失或爆炸风险。进一步地,双随机矩阵集合在乘法下封闭,因此复合映射 $\prod_{i=1}^{L-l}\mathcal{H}_{L-i}^{\mathrm{res}}$ 也会保留这种守恒性质,使得 mHC 能在任意深度间维持恒等映射的稳定性。为保证效率,我们采用 kernel 融合,并利用 TileLang[9] 开发混合精度 kernels;同时通过选择性重计算降低内存占用,并在 DualPipe 调度[5] 中精细重叠通信以进一步提升吞吐。
大量语言模型预训练实验表明,mHC 在保持 HC 性能优势的同时具备更强的稳定性与可扩展性。我们的内部大规模训练进一步显示:当扩展率 $n=4$ 时,mHC 的额外训练时间开销仅为 6.7%。
3. 预备知识
我们首先给出本文使用的记号。在 HC 的表述中,第 $l$ 层输入 $\mathbf{x}_l\in\mathbb{R}^{1\times C}$ 会按扩展率 $n$ 被展开为隐藏矩阵 $\mathbf{x}_l=(\mathbf{x}^\top_{l,0},\ldots,\mathbf{x}^\top_{l,n-1})^\top\in\mathbb{R}^{n\times C}$,可视为 $n$ 路残差流($n$‑stream residual),从而有效增大残差流宽度。为控制该残差流的读出、写入与更新过程,HC 引入三类可学习线性映射:$\mathcal{H}^{\mathrm{pre}}_{l},\mathcal{H}^{\mathrm{post}}_{l}\in\mathbb{R}^{1\times n}$ 与 $\mathcal{H}^{\mathrm{res}}_{l}\in\mathbb{R}^{n\times n}$。这些映射将式 (1) 的标准残差连接修改为式 (3) 的形式。
在 HC 中,可学习映射的系数由两部分组成:随输入变化的部分与全局部分,分别称为动态映射与静态映射。形式化地,HC 按如下方式计算系数:
其中 $\text{RMSNorm}(\cdot)$[30] 作用于最后一维;标量 $\alpha_l^\mathrm{pre},\alpha_l^\mathrm{post},\alpha_l^\mathrm{res}\in\mathbb{R}$ 为可学习门控因子,通常以较小值初始化。动态映射由线性投影得到:$\theta^\mathrm{pre}_l,\theta^\mathrm{post}_l\in\mathbb{R}^{1\times C}$ 与 $\theta^\mathrm{res}_l\in\mathbb{R}^{n\times C}$;静态映射则由可学习偏置 $\mathbf{b}_l^\mathrm{pre},\mathbf{b}_l^\mathrm{post}\in\mathbb{R}^{1\times n}$ 与 $\mathbf{b}_l^\mathrm{res}\in\mathbb{R}^{n\times n}$ 表示。
需要指出的是,引入 $\mathcal{H}^{\mathrm{pre}}_{l}$、$\mathcal{H}^{\mathrm{post}}_{l}$ 与 $\mathcal{H}^{\mathrm{res}}_{l}$ 的额外计算开销通常可以忽略,因为典型扩展率 $n$(例如 4)远小于输入维度 $C$。借助该设计,HC 将残差流的信息容量与层输入维度解耦,而后者与模型计算复杂度(FLOPs)强相关。因此,HC 提供了一条通过调节残差流宽度进行扩展的新路径,补充了预训练 scaling law[31] 中更传统的 FLOPs 与训练数据规模等扩展维度。
尽管 HC 需要三类映射来处理残差流与层输入之间的维度不匹配,但表 1 的初步消融实验显示:残差映射 $\mathcal{H}^{\mathrm{res}}_{l}$ 带来最显著的性能收益,这强调了残差流内部有效信息交换的重要性。
表 1:HC 组件消融。当某个映射($\mathcal{H}^{\mathrm{pre}}_{l}$、$\mathcal{H}^{\mathrm{post}}_{l}$ 或 $\mathcal{H}^{\mathrm{res}}_{l}$)被禁用时,我们用固定映射来保持维度一致:对 $\mathcal{H}^{\mathrm{pre}}_{l}$ 使用 $1/n$ 的均匀权重,对 $\mathcal{H}^{\mathrm{post}}_{l}$ 使用全 1 均匀权重,对 $\mathcal{H}^{\mathrm{res}}_{l}$ 使用单位矩阵。
| $\mathcal{H}^{\mathrm{res}}_{l}$ | $\mathcal{H}^{\mathrm{pre}}_{l}$ | $\mathcal{H}^{\mathrm{post}}_{l}$ | 绝对损失差(Absolute Loss Gap) |
|---|---|---|---|
| 0.0 | |||
| ✓ | -0.022 | ||
| ✓ | ✓ | -0.025 | |
| ✓ | ✓ | ✓ | -0.027 |
3.1 数值不稳定性
尽管残差映射 $\mathcal{H}^{\mathrm{res}}_{l}$ 对性能至关重要,但其在多层上的连续应用会显著增加数值不稳定风险。正如式 (4) 所示,当 HC 扩展到多层时,层 $l$ 到层 $L$ 的有效信号传播由复合映射 $\prod_{i=1}^{L-l}\mathcal{H}_{L-i}^{\mathrm{res}}$ 控制。由于 $\mathcal{H}^{\mathrm{res}}_{l}$ 无任何约束,该复合映射不可避免地偏离恒等映射,从而导致前向与反向传播中的信号幅度容易爆炸或消失。这直接破坏了残差学习依赖的“无阻碍信号流”前提,进而在更深或更大规模模型训练中引发不稳定。
我们在大规模实验中观察到的现象与上述分析一致。如图 2 所示,当以 mHC 作为基线时,HC 在约 12k step 处出现异常的损失激增,并且与梯度范数的不稳定高度相关。对 $\mathcal{H}^{\mathrm{res}}_{l}$ 的进一步分析也验证了该不稳定机制。
为量化复合映射 $\prod_{i=1}^{L-l}\mathcal{H}_{L-i}^{\mathrm{res}}$ 沿残差流对信号的放大程度,我们采用两项指标:其一为复合映射行和的最大绝对值,用以刻画前向传播的最坏放大;其二为列和的最大绝对值,对应反向传播。我们将这两项统称为复合映射的 Amax 增益幅度(Amax Gain Magnitude)。如图 3(b) 所示,Amax 增益幅度出现峰值达 3000 的极端数值,远离理想值 1,清晰表明残差流发生爆炸。
3.2 系统开销
尽管 HC 由于额外映射的线性形式而保持了可控的计算复杂度,但系统层面的开销仍是不可忽视的挑战。尤其是,在现代模型架构中,内存访问(I/O)成本往往构成主要瓶颈之一,这被广泛称为“内存墙”(memory wall)[32]。该瓶颈常在架构设计中被忽视,但它会决定性地影响实际运行效率。
表 2:每 token 的内存访问成本对比。该分析仅统计前向传播中为维护残差流所引入的额外开销,不包含层函数 $\mathcal{F}$ 内部的 I/O。
| 方法 | 操作 | 读(Elements) | 写(Elements) |
|---|---|---|---|
| Residual Connection | Residual Merge | $2C$ | $C$ |
| Total I/O | $2C$ | $C$ | |
| Hyper‑Connections | 计算 $\mathcal{H}^{\mathrm{pre}}_{l}$、$\mathcal{H}^{\mathrm{post}}_{l}$、$\mathcal{H}^{\mathrm{res}}_{l}$ | $nC$ | $n^2+2n$ |
| $\mathcal{H}^{\mathrm{pre}}_{l}$ | $nC+n$ | $C$ | |
| $\mathcal{H}^{\mathrm{post}}_{l}$ | $C+n$ | $nC$ | |
| $\mathcal{H}^{\mathrm{res}}_{l}$ | $nC+n^2$ | $nC$ | |
| Residual Merge | $2nC$ | $nC$ | |
| Total I/O | $(5n+1)C+n^2+2n$ | $(3n+1)C+n^2+2n$ |
聚焦于广泛采用的 pre‑norm Transformer[2] 架构,我们分析了 HC 的 I/O 模式。表 2 总结了 $n$ 路残差流设计在单个残差层中引入的 per token 内存访问开销。分析表明,HC 使内存访问成本按近似与 $n$ 成正比的比例增长;若不使用融合 kernel 进行缓解,这种过度 I/O 会显著降低训练吞吐。除此之外,由于 $\mathcal{H}^{\mathrm{pre}}_{l}$、$\mathcal{H}^{\mathrm{post}}_{l}$、$\mathcal{H}^{\mathrm{res}}_{l}$ 含有可学习参数,其反向传播需要保存中间激活,从而大幅增加 GPU 显存占用,往往需要使用梯度检查点(gradient checkpointing)才能维持可训练的内存规模。进一步地,在流水线并行(pipeline parallelism)[33] 中,HC 还会引入约 $n$ 倍的通信成本,导致更大的 bubble 并降低整体吞吐。
4. 方法
4.1 mHC:流形约束超连接
受恒等映射原则[6] 启发,mHC 的核心思想是:将残差映射 $\mathcal{H}^{\mathrm{res}}_{l}$ 约束到某个特定流形上。尽管原始恒等映射通过强制 $\mathcal{H}^{\mathrm{res}}_{l}=\mathbf{I}$ 可保证稳定性,但它从根本上禁止了残差流内部的信息交换,而后者对发挥多路残差流架构潜力至关重要。因此,我们希望将残差映射投影到一个既能维持跨层信号传播稳定、又能促进残差流之间交互以保留模型表达性的流形上。
为此,我们将 $\mathcal{H}^{\mathrm{res}}_{l}$ 约束为双随机矩阵:其元素非负,且每行与每列的和均为 1。令 $\mathcal{M}^{\mathrm{res}}$ 表示双随机矩阵构成的流形(亦称 Birkhoff 多面体),我们将 $\mathcal{H}^{\mathrm{res}}_{l}$ 约束到 $\mathcal{P}_{\mathcal{M}^{\mathrm{res}}}(\mathcal{H}^{\mathrm{res}}_{l})$,其定义为:
其中 $\mathbf{1}_n$ 是 $n$ 维全 1 向量。需要注意的是,当 $n=1$ 时,双随机条件退化为标量 1,从而恢复原始恒等映射。选择双随机约束带来了若干有利于大规模训练的严格性质:
- 范数保持:双随机矩阵的谱范数有上界 1(即 $\|\mathcal{H}^{\mathrm{res}}_{l}\|_2\le 1$)。这意味着该可学习映射是非扩张(non‑expansive)的,可有效缓解梯度爆炸。
- 复合封闭:双随机矩阵集合在矩阵乘法下封闭。这保证了跨多层的复合残差映射 $\prod_{i=1}^{L-l}\mathcal{H}_{L-i}^{\mathrm{res}}$ 仍为双随机,从而在整个深度范围内维持稳定性。
- Birkhoff 多面体的几何解释:$\mathcal{M}^{\mathrm{res}}$ 是置换矩阵集合的凸包,因此残差映射可解释为置换的凸组合。从数学上看,反复应用此类矩阵会单调增强跨流的信息混合,可视为一种稳健的特征融合机制。
此外,我们还对输入映射 $\mathcal{H}^{\mathrm{pre}}_{l}$ 与输出映射 $\mathcal{H}^{\mathrm{post}}_{l}$ 施加非负约束,以避免正负系数组合造成的信号抵消;这一约束也可视为一种特殊的流形投影。
4.2 参数化与流形投影
本节给出 mHC 中 $\mathcal{H}^{\mathrm{pre}}_{l}$、$\mathcal{H}^{\mathrm{post}}_{l}$ 与 $\mathcal{H}^{\mathrm{res}}_{l}$ 的计算过程。给定第 $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 的思路,得到动态映射与静态映射:
其中 $\phi^\mathrm{pre}_l,\phi^\mathrm{post}_l\in\mathbb{R}^{nC\times n}$ 与 $\phi^\mathrm{res}_l\in\mathbb{R}^{nC\times n^2}$ 为动态映射的线性投影参数,$\text{mat}(\cdot)$ 将 $\mathbb{R}^{1\times n^2}$ 重排为 $\mathbb{R}^{n\times n}$。
随后,通过如下方式得到最终满足约束的映射:
其中 $\sigma(\cdot)$ 为 Sigmoid 函数。$\text{Sinkhorn‑Knopp}(\cdot)$ 算子先用指数运算将所有元素变为正值,然后进行交替归一化迭代,使得行和与列和依次被重标到 1。具体而言,给定初始正矩阵 $\mathbf{M}^{(0)}=\exp(\tilde{\mathcal{H}}^{\mathrm{res}}_{l})$,迭代过程为:
其中 $\mathcal{T}_r$ 与 $\mathcal{T}_c$ 分别表示行归一化与列归一化。该过程在 $t_{\text{max}}\to\infty$ 时收敛到双随机矩阵 $\mathcal{H}^{\mathrm{res}}_{l}=\mathbf{M}^{(t_{\text{max}})}$。在我们的实验中取 $t_{\text{max}}=20$ 作为实用设置。
4.3 高效基础设施设计
本节给出面向 mHC 的基础设施设计。通过严格优化,我们在大规模模型中实现了 mHC($n=4$),其训练额外开销仅为 6.7%。
4.3.1 Kernel 融合
我们观察到:在 mHC 中,RMSNorm 作用于高维隐藏状态 $\vec{\mathbf{x}}_l\in\mathbb{R}^{1\times nC}$ 会引入显著延迟。因此,我们将“除以范数”的操作重排到矩阵乘之后,从而在保持数学等价的同时提升效率。我们进一步采用混合精度策略,在不牺牲速度的前提下尽可能提高数值精度,并将多个共享内存访问的操作融合到统一计算 kernel 中,以缓解带宽瓶颈。基于式 (10)–(13) 中给出的输入与参数,我们实现了三个 mHC 专用 kernel 来计算 $\mathcal{H}^{\mathrm{pre}}_{l}$、$\mathcal{H}^{\mathrm{post}}_{l}$ 与 $\mathcal{H}^{\mathrm{res}}_{l}$;在这些 kernels 中,偏置与线性投影分别被整合到 $\mathbf{b}_l$ 与 $\phi_l$ 中,RMSNorm 的权重也被吸收到 $\phi_l$ 内。
- 式 (14)–(15):我们开发了一个统一 kernel,将对 $\vec{\mathbf{x}}_l$ 的两次扫描融合,并利用矩阵乘法单元最大化带宽利用率。反向传播(两次矩阵乘)也被整合为单 kernel,避免重复加载 $\vec{\mathbf{x}}_l$。两个 kernels 都采用精细调度的流水线(load、cast、compute、store)以高效处理混合精度。
- 式 (16)–(18):这些针对小系数的轻量操作被机会性融合为单 kernel,从而显著降低 kernel 启动开销。
- 式 (19):我们在单个 kernel 内实现 Sinkhorn‑Knopp 迭代;其反向则使用自定义 backward kernel,在片上重计算中间结果并遍历完整迭代过程。
利用上述 kernels 得到的系数,我们再引入两个 kernels 来应用这些映射:其一用于 $\mathcal{F}_{\mathrm{pre}}\mathrel{:=} \mathcal{H}^{\mathrm{pre}}_{l}\mathbf{x}_l$,其二用于 $\mathcal{F}_{\mathrm{post,res}}\mathrel{:=} \mathcal{H}^{\mathrm{res}}_{l}\mathbf{x}_l + \mathcal{H}_{l}^{\mathrm{post}\, \top}\mathcal{F}(\cdot,\cdot)$。通过将 $\mathcal{H}^{\mathrm{post}}_{l}$ 与 $\mathcal{H}^{\mathrm{res}}_{l}$ 的应用与残差合并融合,我们将该 kernel 的读取元素数从 $(3n+1)C$ 降至 $(n+1)C$,写入元素数从 $3nC$ 降至 $nC$。除式 (14)–(15) 涉及的大规模矩阵乘外,我们使用 TileLang[9] 高效实现了大多数 kernels;该框架能够简化复杂计算流程的 kernel 实现,并以较小工程代价充分利用内存带宽。
4.3.2 重计算
$n$ 路残差流设计会在训练中引入显著的激活内存开销。为缓解这一点,我们在前向传播后丢弃 mHC kernels 的中间激活,并在反向传播中通过重新执行 mHC kernels(不包含重的层函数 $\mathcal{F}$)按需重计算这些中间量。因此,对于连续 $L_r$ 层构成的一个块,我们仅需存储该块首层的输入 $\mathbf{x}_{l_0}$。在不计轻量系数、并考虑 $\mathcal{F}$ 中 pre‑norm 的情况下,表 3 总结了反向传播所需保留/重计算的中间激活。
表 3:保存与重计算的中间激活。按 token 统计:列出反向传播需要持久保存的激活,以及在连续 $L_r$ 层内会被临时重计算的激活。$l_0$ 表示该 $L_r$ 层块的首层,$l\in[l_0, l_0+L_r-1]$。
| 激活 | $\mathbf{x}_{l_0}$ | $\mathcal{F}(\mathcal{H}^{\mathrm{pre}}_{l}\mathbf{x}_l, \mathcal{W}_l)$ | $\mathbf{x}_{l}$ | $\mathcal{H}^{\mathrm{pre}}_{l}\mathbf{x}_l$ | $\text{RMSNorm}(\mathcal{H}^{\mathrm{pre}}_{l}\mathbf{x}_l)$ |
|---|---|---|---|---|---|
| 大小(Elements) | $nC$ | $C$ | $nC$ | $C$ | $C$ |
| 存储方式 | 每 $L_r$ 层存一次 | 每层都存 | 在 $L_r$ 层块内为临时量 | ||
由于 mHC 的重计算以连续 $L_r$ 层为块进行,在总层数为 $L$ 时,我们需要为反向传播持久存储所有 $\lceil\tfrac{L}{L_r}\rceil$ 个块的首层输入 $\mathbf{x}_{l_0}$。除这一常驻内存外,重计算过程还会为当前活跃块引入 $(n+2)C\times L_r$ 个元素的瞬时内存开销,从而决定反向传播的峰值显存。于是我们通过最小化随 $L_r$ 变化的总内存占用来确定最优块大小 $L_r^*$:
此外,大规模训练中的流水线并行还带来约束:重计算块不能跨越流水线 stage 的边界。我们观察到理论最优 $L_r^*$ 通常与每个 stage 的层数相匹配,因此选择将重计算边界与流水线 stage 对齐。
4.3.3 DualPipe 中的通信重叠
在大规模训练中,流水线并行是降低参数与梯度内存占用的标准做法。我们采用 DualPipe 调度[5],它能有效重叠跨节点互联通信(例如专家并行与流水线并行中的通信流量)。然而,相比单路残差流设计,mHC 的 $n$ 路残差流会显著增加流水线 stage 之间的通信延迟;同时在 stage 边界处,对所有 $L_r$ 层执行 mHC kernels 的重计算也会带来不可忽视的计算开销。为解决这些瓶颈,我们对 DualPipe 调度进行扩展(见图 4),以增强 stage 边界处通信与计算的重叠效果。
需要特别指出的是:为避免阻塞通信流,我们在专用的高优先级计算流上执行 MLP(即 FFN)层的 $\mathcal{F}_{\mathrm{post,res}}$ kernels;同时我们避免在注意力层的长耗时操作中使用持久化 kernels,以防长时间 stall。该设计允许对被重叠的注意力计算进行抢占,从而实现更灵活的调度,并维持对计算设备处理单元的高利用率。另一方面,由于每个 stage 的初始激活 $\mathbf{x}_{l_0}$ 已经在本地缓存,重计算过程也与流水线通信依赖解耦。
5. 实验
5.1 实验设置
我们通过语言模型预训练来验证所提方法,并在 Baseline、HC 与 mHC 之间进行对比分析。基于 DeepSeek‑V3[5] 启发的 MoE 架构,我们训练了四个不同的模型变体以覆盖不同评估场景。HC 与 mHC 的扩展率均设为 $n=4$。我们主要关注一个 27B 模型:其训练数据规模与参数规模成比例,并作为系统层面主要结果的研究对象。进一步地,我们还加入更小的 3B 与 9B 模型(同样采用与参数规模成比例的数据规模)来分析计算扩展(compute scaling)行为,从而观察不同计算预算下的趋势。除此之外,为专门研究 token 扩展(token scaling)行为,我们还在固定的 1T token 语料上训练了一个独立的 3B 模型。更详细的模型配置与训练超参数见附录 A。
5.2 主要结果
表 4:27B 模型的系统级基准结果。比较 Baseline、HC 与 mHC 在 8 个多样下游基准上的 zero‑shot / few‑shot 表现。mHC 在多数基准上优于 HC,并稳定超过 Baseline,表明其对大规模预训练有效。
| 基准(指标) | BBH(EM) | DROP(F1) | GSM8K(EM) | HellaSwag(Acc.) | MATH(EM) | MMLU(Acc.) | PIQA(Acc.) | TriviaQA(EM) |
|---|---|---|---|---|---|---|---|---|
| # Shots | 3‑shot | 3‑shot | 8‑shot | 10‑shot | 4‑shot | 5‑shot | 0‑shot | 5‑shot |
| 27B Baseline | 43.8 | 47.0 | 46.7 | 73.7 | 22.0 | 59.0 | 78.5 | 54.3 |
| 27B + HC | 48.9 | 51.6 | 53.2 | 74.3 | 26.4 | 63.0 | 79.9 | 56.3 |
| 27B + mHC | 51.0 | 53.9 | 53.8 | 74.7 | 26.0 | 63.4 | 80.5 | 57.6 |
我们首先考察 27B 模型的训练稳定性与收敛性。如图 5(a) 所示,mHC 有效缓解了 HC 的训练不稳定问题,并在训练末端相对基线实现 0.021 的损失下降。图 5(b) 的梯度范数分析也进一步印证了这一点:mHC 的梯度行为显著优于 HC,整体轮廓与基线相近且更加稳定。
表 4 展示了多个下游基准上的表现[34–39]。mHC 在整体上带来全面提升:它稳定优于基线,并在多数任务上超过 HC。值得注意的是,相比 HC,mHC 进一步增强了模型推理能力:在 BBH[40] 上提升 2.1%,在 DROP[41] 上提升 2.3%。
5.3 扩展性实验
为评估方法的可扩展性,我们报告了 mHC 相对基线在不同规模下的相对损失改进。如图 6(a) 所示,跨越 3B、9B 与 27B 参数规模的计算扩展曲线表明:即使在更高计算预算下,mHC 的性能优势也能稳健保持,仅出现轻微衰减。我们还在图 6(b) 中进一步考察了单次训练过程内的动态,给出了 3B 模型的 token 扩展曲线。综合这些发现,mHC 在大规模场景中表现出有效性,这一点也被我们的内部大规模训练实验进一步佐证。
5.4 稳定性分析
与图 3 类似,图 7 展示了 mHC 的传播稳定性。理想情况下,单层映射应满足双随机约束,这意味着前向信号增益与反向梯度增益都应等于 1。然而,实际实现中为保证计算效率,Sinkhorn‑Knopp 迭代次数必须受限;在我们的设置中取 20 次以获得近似解。因此,如图 7(a) 所示,反向梯度增益会轻微偏离 1;在图 7(b) 的复合情形下,该偏离有所增大但仍被严格限制,最大约为 1.6。值得注意的是,相比 HC 中接近 3000 的最大增益幅度,mHC 将其降低了约三个数量级。这些结果表明:mHC 能显著增强传播稳定性,从而保证前向信号与反向梯度的稳定流动。
此外,图 8 给出了代表性映射的可视化。我们观察到:在 HC 中,当最大增益很大时,其它值往往也会显著增大,这表明几乎所有传播路径都处于普遍不稳定状态;相比之下,mHC 始终给出稳定结果。
6. 结论与展望
本文指出:尽管 Hyper‑Connections(HC)通过扩展残差流宽度并丰富连接模式带来了性能收益,但其无约束连接会导致信号发散。这种破坏削弱了跨层信号能量守恒,诱发训练不稳定并限制深层网络的可扩展性。为解决这些问题,我们提出 Manifold‑Constrained Hyper‑Connections(mHC):一个将残差连接空间投影到特定流形上的通用框架。通过使用 Sinkhorn‑Knopp 算法对残差映射施加双随机约束,mHC 将信号传播转化为特征的凸组合。实证结果证实,mHC 能有效恢复恒等映射性质,使得相较传统 HC 具备更好的稳定性与可扩展性,从而支持稳定的大规模训练。更关键的是,借助基础设施级优化,mHC 能在几乎可忽略的计算开销下实现上述收益。
作为 HC 范式的广义扩展,mHC 打开了若干有前景的研究方向。尽管本文使用双随机矩阵来保证稳定性,但该框架同样支持探索面向不同学习目标的多样流形约束。我们预期:对不同几何约束的进一步研究有望产生新的方法,更好地优化可塑性(plasticity)与稳定性(stability)之间的权衡。我们也希望 mHC 能重新激发社区对宏观架构设计的兴趣:通过加深对拓扑结构如何影响优化与表示学习的理解,mHC 有望帮助突破当前限制,并为下一代基础架构的演进指明新的路径。
参考文献
- Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun. “Deep residual learning for image recognition.” Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770--778, 2016.
- Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, Illia Polosukhin. “Attention is all you need.” Advances in neural information processing systems, vol. 30, 2017.
- Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al.. “Language models are few-shot learners.” Advances in neural information processing systems, vol. 33, pp. 1877--1901, 2020.
- Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al.. “Llama: Open and efficient foundation language models.” arXiv preprint arXiv:2302.13971, 2023. https://arxiv.org/abs/2302.13971
- Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, et al.. “Deepseek-v3 technical report.” arXiv preprint arXiv:2412.19437, 2024. https://arxiv.org/abs/2412.19437
- Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun. “Identity mappings in deep residual networks.” European conference on computer vision, pp. 630--645, 2016.
- Defa Zhu, Hongzhi Huang, Zihao Huang, Yutao Zeng, Yunyao Mao, Banggu Wu, Qiyang Min, Xun Zhou. “Hyper-connections.” arXiv preprint arXiv:2409.19606, 2024. https://arxiv.org/abs/2409.19606
- Richard Sinkhorn, Paul Knopp. “Concerning nonnegative matrices and doubly stochastic matrices.” Pacific Journal of Mathematics, vol. 21, no. 2, pp. 343--348, 1967.
- Lei Wang, Yu Cheng, Yining Shi, Zhengju Tang, Zhiwen Mo, Wenhao Xie, Lingxiao Ma, Yuqing Xia, Jilong Xue, Fan Yang, et al.. “TileLang: A Composable Tiled Programming Model for AI Systems.” arXiv preprint arXiv:2504.17577, 2025. https://arxiv.org/abs/2504.17577
- François Chollet. “Xception: Deep learning with depthwise separable convolutions.” Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 1251--1258, 2017.
- Saining Xie, Ross Girshick, Piotr Dollár, Zhuowen Tu, Kaiming He. “Aggregated residual transformations for deep neural networks.” Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 1492--1500, 2017.
- Noam Shazeer. “Fast transformer decoding: One write-head is all you need.” arXiv preprint arXiv:1911.02150, 2019. https://arxiv.org/abs/1911.02150
- Joshua Ainslie, James Lee-Thorp, Michiel De Jong, Yury Zemlyanskiy, Federico Lebrón, Sumit Sanghai. “Gqa: Training generalized multi-query transformer models from multi-head checkpoints.” arXiv preprint arXiv:2305.13245, 2023. https://arxiv.org/abs/2305.13245
- Aixin Liu, Bei Feng, Bin Wang, Bingxuan Wang, Bo Liu, Chenggang Zhao, Chengqi Dengr, Chong Ruan, Damai Dai, Daya Guo, et al.. “Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model.” arXiv preprint arXiv:2405.04434, 2024. https://arxiv.org/abs/2405.04434
- Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, Jeff Dean. “Outrageously large neural networks: The sparsely-gated mixture-of-experts layer.” arXiv preprint arXiv:1701.06538, 2017. https://arxiv.org/abs/1701.06538
- Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, Zhifeng Chen. “Gshard: Scaling giant models with conditional computation and automatic sharding.” arXiv preprint arXiv:2006.16668, 2020. https://arxiv.org/abs/2006.16668
- William Fedus, Barret Zoph, Noam Shazeer. “Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity.” Journal of Machine Learning Research, vol. 23, no. 120, pp. 1--39, 2022.
- Rupesh K Srivastava, Klaus Greff, Jürgen Schmidhuber. “Training Very Deep Networks.” Advances in Neural Information Processing Systems, vol. 28, 2015. https://proceedings.neurips.cc/paper_files/paper/2015/file/215a71a12769b056c3c32e7299f1c5ed-Paper.pdf
- Gao Huang, Zhuang Liu, Laurens Van Der Maaten, Kilian Q Weinberger. “Densely connected convolutional networks.” Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 4700--4708, 2017.
- Gustav Larsson, Michael Maire, Gregory Shakhnarovich. “Fractalnet: Ultra-deep neural networks without residuals.” arXiv preprint arXiv:1605.07648, 2016. https://arxiv.org/abs/1605.07648
- Fisher Yu, Dequan Wang, Evan Shelhamer, Trevor Darrell. “Deep layer aggregation.” Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 2403--2412, 2018.
- Yekun Chai, Shuo Jin, Xinwen Hou. “Highway Transformer: Self-Gating Enhanced Self-Attentive Networks.” Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 6887--6900, 2020. https://aclanthology.org/2020.acl-main.616/
- Yanwen Fang, Yuxi CAI, Jintai Chen, Jingyu Zhao, Guangjian Tian, Guodong Li. “Cross-Layer Retrospective Retrieving via Layer Attention.” The Eleventh International Conference on Learning Representations, 2023. https://openreview.net/forum?id=pvgEL1yS3Ql
- Shufang Xie, Huishuai Zhang, Junliang Guo, Xu Tan, Jiang Bian, Hany Hassan Awadalla, Arul Menezes, Tao Qin, Rui Yan. “ResiDual: Transformer with Dual Residual Connections.” 2023. https://arxiv.org/abs/2304.14802
- Matteo Pagliardini, Amirkeivan Mohtashami, François Fleuret, Martin Jaggi. “DenseFormer: Enhancing Information Flow in Transformers via Depth Weighted Averaging.” The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. https://openreview.net/forum?id=kMnoh7CXrq
- Gaurav Menghani, Ravi Kumar, Sanjiv Kumar. “LAuReL: Learned Augmented Residual Layer.” Forty-second International Conference on Machine Learning, 2025. https://openreview.net/forum?id=rUDRWP9WvZ
- Mike Heddes, Adel Javanmard, Kyriakos Axiotis, Gang Fu, Mohammadhossein Bateni, Vahab Mirrokni. “DeepCrossAttention: Supercharging Transformer Residual Connections.” Forty-second International Conference on Machine Learning, 2025. https://openreview.net/forum?id=j3JBfFnGYh
- Brian Mak, Jeffrey Flanigan. “Residual Matrix Transformers: Scaling the Size of the Residual Stream.” arXiv preprint arXiv:2506.22696, 2025. https://arxiv.org/abs/2506.22696
- Da Xiao, Qingye Meng, Shengping Li, Xingyuan Yuan. “Muddformer: Breaking residual bottlenecks in transformers via multiway dynamic dense connections.” arXiv preprint arXiv:2502.12170, 2025. https://arxiv.org/abs/2502.12170
- Biao Zhang, Rico Sennrich. “Root mean square layer normalization.” Advances in neural information processing systems, vol. 32, 2019.
- Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Thomas Hennigan, Eric Noland, Katherine Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karén Simonyan, Erich Elsen, Oriol Vinyals, Jack Rae, Laurent Sifre. “An empirical analysis of compute-optimal large language model training.” Advances in Neural Information Processing Systems, vol. 35, pp. 30016--30030, 2022. https://proceedings.neurips.cc/paper_files/paper/2022/file/c1e2faff6f588870935f114ebe04a3e5-Paper-Conference.pdf
- Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré. “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.” Advances in Neural Information Processing Systems (NeurIPS), 2022.
- Penghui Qi, Xinyi Wan, Guangxing Huang, Min Lin. “Zero Bubble (Almost) Pipeline Parallelism.” The Twelfth International Conference on Learning Representations, 2024. https://openreview.net/forum?id=tuzTN0eIO5
- Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, Jacob Steinhardt. “Measuring massive multitask language understanding.” arXiv preprint arXiv:2009.03300, 2020. https://arxiv.org/abs/2009.03300
- Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al.. “Training verifiers to solve math word problems.” arXiv preprint arXiv:2110.14168, 2021. https://arxiv.org/abs/2110.14168
- Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, Yejin Choi. “HellaSwag: Can a Machine Really Finish Your Sentence?.” Proceedings of the 57th Conference of the Association for Computational Linguistics, ACL 2019, Florence, Italy, July 28- August 2, 2019, Volume 1: Long Papers, pp. 4791--4800, 2019. https://doi.org/10.18653/v1/p19-1472
- Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, Jacob Steinhardt. “Measuring mathematical problem solving with the math dataset.” arXiv preprint arXiv:2103.03874, 2021. https://arxiv.org/abs/2103.03874
- Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, Yejin Choi. “PIQA: Reasoning about Physical Commonsense in Natural Language.” The Thirty-Fourth AAAI Conference on Artificial Intelligence, AAAI 2020, The Thirty-Second Innovative Applications of Artificial Intelligence Conference, IAAI 2020, The Tenth AAAI Symposium on Educational Advances in Artificial Intelligence, EAAI 2020, New York, NY, USA, February 7-12, 2020, pp. 7432--7439, 2020. https://doi.org/10.1609/aaai.v34i05.6239
- Mandar Joshi, Eunsol Choi, Daniel Weld, Luke Zettlemoyer. “TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehension.” Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 1601--1611, 2017. https://aclanthology.org/P17-1147
- Mirac Suzgun, Nathan Scales, Nathanael Schärli, Sebastian Gehrmann, Yi Tay, Hyung Won Chung, Aakanksha Chowdhery, Quoc V Le, Ed H Chi, Denny Zhou, et al.. “Challenging big-bench tasks and whether chain-of-thought can solve them.” arXiv preprint arXiv:2210.09261, 2022. https://arxiv.org/abs/2210.09261
- Dheeru Dua, Yizhong Wang, Pradeep Dasigi, Gabriel Stanovsky, Sameer Singh, Matt Gardner. “DROP: A Reading Comprehension Benchmark Requiring Discrete Reasoning Over Paragraphs.” Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume 1 (Long and Short Papers), pp. 2368--2378, 2019. https://doi.org/10.18653/v1/n19-1246
- Lean Wang, Huazuo Gao, Chenggang Zhao, Xu Sun, Damai Dai. “Auxiliary-loss-free load balancing strategy for mixture-of-experts.” arXiv preprint arXiv:2408.15664, 2024. https://arxiv.org/abs/2408.15664
- Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, Yunfeng Liu. “Roformer: Enhanced transformer with rotary position embedding.” Neurocomputing, vol. 568, pp. 127063, 2024.
- Ilya Loshchilov, Frank Hutter. “Decoupled weight decay regularization.” arXiv preprint arXiv:1711.05101, 2017. https://arxiv.org/abs/1711.05101
注:参考文献按本文首次引用顺序编号,对应正文中的方括号引用。
附录
A. 模型规格与超参数
表 A.1:详细模型规格与超参数。该表给出了基于 DeepSeek‑V3[5] 架构的 3B、9B 与 27B 模型配置,并列出 mHC/HC 的关键超参数(如残差流扩展率与 Sinkhorn‑Knopp 设置),以及实验中采用的优化与训练协议。
| 属性 | 3B | 9B | 27B | 3B |
|---|---|---|---|---|
| 1T Tokens | ||||
| Vocab Params | 331M | 496M | 662M | 331M |
| Active Params | 612M | 1.66B | 4.14B | 612M |
| Total Params | 2.97B | 9.18B | 27.0B | 2.97B |
| 架构 | ||||
| Layers | 12 | 18 | 30 | 12 |
| Leading Dense Layers | 1 | 1 | ||
| Routed Experts | 64 | 64 | 72 | 64 |
| Active Experts | 6 | 6 | ||
| Shared Experts | 2 | 2 | ||
| Dimension | 1280 | 1920 | 2560 | 1280 |
| FFN Dimension | 896 | 1280 | 1536 | 896 |
| Load Balancing Method | Loss‑Free[42] | Loss‑Free | ||
| Attention Heads | 16 | 24 | 32 | 16 |
| Attention Dimension | 128 | 128 | ||
| Attention Variant | MLA[14] | MLA | ||
| KV Rank | 512 | 512 | ||
| Position Embedding | RoPE[43] | RoPE | ||
| RoPE Dimension | 64 | 64 | ||
| RoPE $\theta$ | 10000 | 10000 | ||
| Layer Norm Type | RMSNorm[30] | RMSNorm | ||
| Layer Norm $\epsilon$ | 1e‑20 | 1e‑20 | ||
| mHC / HC | ||||
| mHC/HC Expansion Rate $n$ | 4 | 4 | ||
| mHC/HC Gating Factor Init $\alpha$ | 0.01 | 0.01 | ||
| mHC Sinkhorn‑Knopp $t_{\text{max}}$ | 20 | 20 | ||
| 训练 | ||||
| Sequence Length | 4096 | 4096 | ||
| Vocab Size | 129280 | 129280 | ||
| Batch Size | 320 | 512 | 1280 | 2560 |
| Training Steps | 30000 | 50000 | 50000 | 100000 |
| Training Tokens | 39.3B | 105B | 262B | 1.05T |
| Warmup Steps | 2000 | 2000 | ||
| Optimizer | AdamW[44] | AdamW | ||
| AdamW Betas | (0.9, 0.95) | (0.9, 0.95) | ||
| AdamW $\epsilon$ | 1e‑20 | 1e‑20 | ||
| Base Learning Rate | 8.6e‑4 | 5.9e‑4 | 4.0e‑4 | 9.0e‑4 |
| Lr Scheduler | Step | Step | ||
| Lr Decay Step Ratio | [0.8×, 0.9×] | [0.8×, 0.9×] | ||
| Lr Decay Rate | [0.316, 0.1] | [0.316, 0.1] | ||
| Weight Decay | 0.1 | 0.1 | ||