1. 首页
  2. 精选文章
  3. 从 DeepSeek mHC 到 xHC:如何高效 scale 残差流?

从 DeepSeek mHC 到 xHC:如何高效 scale 残差流?

  • 发布于 2026-07-23
  • ·
  • 31 次阅读
  • ·
  • ·

作者:null,上海交通大学AI博一
https://zhuanlan.zhihu.com/p/2063300859472221420

上周五我们发布了 xHC: Expanded Hyper-Connections,这段时间的工作也算告一段落。先用一句话概括我们做了什么:xHC 首次把语言模型的残差流高效扩展到 16 流。

在 18B 和 28B MoE 语言模型的预训练实验中,xHC 在 DeepSeek Manifold-constrained Hyper-Connections (mHC) 的基础上,把收益又提升了近一倍(如下图)。据我们所知,这可能是目前同 FLOPs 下最强的残差流改进方法之一(包括 Kimi Attention Residual)。

xHC、mHC 与 Vanilla Baseline 在 18B MoE 上的性能对比

这项工作的探索过程,远比 paper 里呈现的更加曲折。受篇幅和叙事限制,论文最终只保留了相对干净的方法主线和关键实验,但很多我认为同样有价值的探索并没有展开叙述:哪些看似合理的方案没有奏效,我们为什么误判了某些问题,以及哪些观察最终改变了研究方向。

所以这篇 blog 更像是一份研究复盘。我想把 xHC 从复现、试错到最终成形的过程尽量完整地记录下来,也希望这些论文之外的细节,能给想复现、使用或继续改进 xHC 的 researcher 提供一些不同的视角。

Paper 链接:https://arxiv.org/pdf/2607.14530
Github 链接:https://github.com/aHapBean/xHC 
(开源日期待定,但是最终会release。欢迎各位 star,感谢支持!)

从复现 DeepSeek mHC 开始

故事开始于 mHC 发布。第一次看到它的实验结果时,我想到的第一句话是:vanilla residual reaches its end。原因很简单,mHC 带来的并不是一个很小的涨点,而是在大规模预训练中明显超过了传统 residual connection。

对一个做 LLM pre-training 的 researcher 来说,这种差距很难不让人兴奋:我们使用了这么多年的 residual pathway,可能仍然存在很大的结构改进空间。

mHC论文中展示的结果

在研究 mHC 的过程中,我回头仔细阅读了 Hyper-Connections。HC 将 Transformer 原本单一的 residual stream 扩展为多条并行 streams,让模型能够保存不同的跨层历史;mHC 则通过 manifold constraint 解决了多流残差在大规模训练中的稳定性问题,它也是公开结果中真正经过大规模训练验证的最强 HC 变体。

我们最初并没有打算提出一个新的 Hyper-Connections 方法,目标只是把 mHC 完整复现出来,scaling!因为对于 LLM 预训练,最重要的不是一个方法能否在小规模实验中涨点,而是这种收益能否在模型和训练规模增大后继续成立。

复现结果和 mHC 原论文大致吻合。在 18B MoE backbone 和训练设置下,mHC 像原论文展示的那样,明显优于 Vanilla residual baseline:

mHC复现结果

平均分从 40.6 提升到 44.8,显著提高了 4.2 points。更具体地看,MMLU、ARC-Challenge、GSM8K 和 C3 分别提升了 5.8、10.6、6.8 和 5.6 points。这也让我确信 residual-stream expansion 确实是一个值得继续投入的方向。

复现之后,我又仔细复盘了实验和 mHC 的方法细节。我注意到,无论是 HC 还是 mHC,大多数公开的实验都把 residual streams 的数量(记作 N ,即 expansion rate) 设为 4。HC 的结果显示,从 N=1 扩展到 N=4 时收益明显,但从 4 继续增加到 8,改善已经非常有限。这让我产生了一个很直接的疑问:

N 一定要取 4 吗,还能不能继续 scale?如果 residual-stream expansion 只能停在一个固定的配置上,它是否只能被看作为一个受限的 architecture choice,而不是一个真正可以继续 scaling 的 axis?

这个念头后来逐渐变成了一种想要 scale N 的信念。因为xHC的研究过程比较曲折,没有这个信念,可能也不会有接下来的故事。

Expanded Hyper-Connections: Make N scaling meaningful and practical

把 N 调大以后,事情开始不太对了

复现完mHC以后,我们做了进一步的实验:其他设置都不变,只把 residual streams 的数量 N 往上调,看看它能不能像 hidden dimension、层数或者专家数一样,成为一个可以持续 scaling 的维度。

我们首先回忆一下 mHC 的原理。mHC 的更新过程:设第 l 层的多流残差状态为

X_l=(x_{l,1},\ldots,x_{l,N})^\top\in\mathbb{R}^{N\times C},

其中 N 是 residual streams 的数量(expansion rate),C 是 hidden dimension。每个 sublayer 的更新可以写成

X_{l+1}=H_l^{\mathrm{res}}X_l+H_l^{\mathrm{post}}F(H_l^{\mathrm{pre}}X_l,W_l).

H^{\mathrm{pre}} 负责从多条 streams 中读取 sublayer 的输入,H^{\mathrm{post}} 负责把 sublayer 输出写回不同 streams,H^{\mathrm{res}} 则在 streams 之间做 residual mixing。

mHC 使用 Sinkhorn normalization 将 H^{\mathrm{res}} 约束为 doubly stochastic matrix,从而保留深层训练所需要的稳定信息传播。从 N=2N=4,mHC 的收益很明显;但一旦越过 N=4,曲线很快开始变平,训练开销却还在快速增加。

不同N下的mHC实验结果

在我们的 2.5B MoE 实验中,mHC 从 N=4 扩展到 N=16,validation loss 只下降了 0.006,训练 FLOPs 却增加了 32%。这并不意味着更大的 N 完全没有用,而是它的边际收益和额外成本已经明显不成比例。

这个结果也从另一个角度解释了,为什么已有的大规模 mHC 实验通常停在 N=4:继续增加 streams,至少按照原始形式,并不是一个划算的操作。

以上图中的现象其实对应着两个不同的问题。

  • 第一个问题在收益侧:为什么从 N=1N=4 收益明显,继续增加 streams 后,边际收益却迅速衰减?
  • 第二个问题在成本侧:增加 streams,为什么方法的计算代价增长得这么快?

这两个问题后来分别导向了 xHC 的两个核心设计:Temporal Feature Augmentation,以及稀疏的的 residual-stream architecture。

接下来,我们分别沿着收益和成本这两条线展开。xHC 的 formulation 是在解决这两个问题的过程中逐渐形成的。我们先在第一个问题上绕了相当长的一段路。

更多 streams,为什么仍然只写入一份信息?

mHC 对第 i 条 stream 的新增写入可以简化为

\Delta x_{l,i}=h^{\mathrm{post}}_{l,i}\cdot \mathrm{out},

其中 \mathrm{out}\in\mathbb{R}^{C} 是当前 sublayer 的输出,h^{\mathrm{post}}_{l,i} 是一项动态生成的标量。把所有 streams 放在一起,一次 write-back 可以写成

\Delta X_l=H_l^{\mathrm{post}}\mathrm{out}^{\top}.

其中 H_l^{\mathrm{post}}\in\mathbb{R}^{N\times 1}\mathrm{out}^{\top}\in\mathbb{R}^{1\times C},因此 \Delta X_l\in\mathbb{R}^{N\times C}

在 mHC 中,每条 stream 会积累历史 layer outputs 的不同加权组合,因此能够形成不同的历史信息“视角”。

h^{\mathrm{post}}_{l,i} 也是 input-dependent 且 stream-specific 的,所以不同 streams 可以决定以多大权重接收当前层的输出。

但无论这些权重如何变化,在一次 sublayer update 中,新注入到所有 streams 的信息都来自同一个 write-back component,也就是 \mathrm{out}。不同 streams 可以为它分配不同权重,却无法接收彼此不同的 components。

N 较小时,这未必构成严重限制。即使每层只提供一个新的 component,不同 streams 仍然可以通过跨层积累不同的加权组合,形成有意义的历史表示。这与从 N=1 扩展到 N=4 时观察到的明显收益是一致的。

但随着 N 继续增大,模型拥有了更多用于保存历史状态的 streams,每层提供的基础 write-back components 却没有同步增加。新增 streams 仍然只能围绕同一份 layer output 学习更多加权方式,因此更容易形成彼此重叠的历史组合。

这让我们开始怀疑:限制大 N 扩展的可能是 write-back information 的多样性没有随 N 一起增长。

最直接的解决办法,是为每条 stream 分别计算一份 layer output,但这会把 layer FLOPs 近似放大 N 倍,失去 residual-stream expansion 原本低成本的意义。

我们最初没有立刻转向寻找其他信息源,而是先怀疑 mHC 的 residual update 过于线性:H^{\mathrm{res}} 只对已有 streams 做线性混合,H^{\mathrm{post}} 也只是用 scalar coefficient 缩放当前 layer output,因此每条 stream 最终保存的仍是历史 layer outputs 的不同线性组合。

一个很自然的想法是:引入非线性,能否在不重复额外进行 sublayer 计算的前提下,产生更加丰富的 write-back components?

沿着这个思路,我们围绕如何引入更细粒度的非线性做了几类尝试:

这些尝试的具体形式不同,但最终指向了同一个问题:非线性可以产生不同的 feature directions,却不一定能提供有意义的新信息。

如果所有变换仍然只作用于当前 token 的同一份 layer output,模型可能只是在更复杂地重组已有内容,而没有真正增加信息含量。我们需要的不是更复杂的写入方式,而是一些计算成本足够低、同时又来自不同信息来源的 write-back components。

最便宜的新信息,可能就在相邻 token 里

于是问题变得很具体:既不能为每条 stream 重新计算一份 layer output,又需要为当前 token 提供更多有意义的 write-back components,还能从哪里获得这些信息?我们最后把目光转向了 sequence dimension。

对于当前位置,autoregressive prediction 本来就建立在当前及此前 token 的上下文之上。因此,相邻位置的 hidden outputs 并不是无关的外部信号,而是与当前位置的 next-token prediction 语义兼容的局部上下文;

更重要的是,它们已经在同一次 sublayer computation 中被计算出来了。如果能利用这些表征增加当前 token 的 write-back diversity,就可以在不重复运行完整 sublayer 的前提下,为 write-back 提供更多候选 components。

我们一开始也试过更便宜的 average pooling,结果并不好。这一点其实很重要,因为它说明“看到其他 token”本身还不够。模型还需要学习应该从局部上下文中提取什么,以及不同时间范围的信息应该如何进入 residual memory。

最终有效的形式是 causal depthwise convolution。Depthwise convolution 让每个 channel 独立地沿 sequence dimension 聚合局部信息,避免昂贵的 cross-channel mixing。

具体来说,我们先恢复 MLP/MoE output 的 sequence dimension,记作 \mathrm{out}\in\mathbb{R}^{S\times C},其中 S 是 sequence length。然后使用三个 kernel size 分别为 4、8、12 的 causal depthwise conv 分支,从不同的局部范围提取 multi-granularity temporal features,再把它们与原始输出放在一起:

\mathrm{out}_{\mathrm{aug}}= [\mathrm{out};\mathrm{DWConv}_4(\mathrm{out}); \mathrm{DWConv}_8(\mathrm{out}); \mathrm{DWConv}_{12}(\mathrm{out})].

这里的四组输出沿新增的 component dimension 堆叠。这样,post-mapping 面对的不再只有一份 layer output,而是四组来自不同局部时间范围的候选 components,可以为不同 active streams 学习差异化的 write-back signals。我们把这个模块称为 Temporal Feature Augmentation(TFA)。

Temporal Feature Augmentation示意图

左侧的 mHC 只能将同一份 layer output 以不同权重写入各条 streams;右侧的 TFA 则提供四组具有不同时间范围的候选 components,使不同 streams 可以得到多样的 write-back signals。

三个 depthwise conv 分支在实际实现中采用 zero initialization。训练刚开始时,新增的 temporal branches 不会立刻扰乱原有的 residual update;随着优化进行,它们再逐渐学会从相邻 token 中提取有用信息。相比一开始就向多流状态注入几组未经训练的局部特征,这个起点更加可控。

我们最终只在 MLP/MoE 之后使用 TFA。Attention 本身已经进行了 content-dependent token mixing,再加入相同的 temporal convolution,不仅功能上有些重复,实验中也没有带来额外收益。

在 10B MoE 消融实验中,只使用 MLP-side TFA 的默认 xHC 达到 1.983 validation loss;加入 Attention-side TFA 后变为 1.985,同时这个操作还会增加训练不稳定性。

还有一个比单点 loss 更能说明问题的现象:我们在 N\in\{4,8,16\} 的 mHC 上分别加入相同的 TFA,发现随着 N 增大,TFA 相对原始 mHC 的收益也越来越明显。

这个趋势与 information supply bottleneck 的判断是一致的:streams 越多,增加 write-back component diversity 带来的价值越大。

不同N下的Loss

Conv 带来了新信息,也带来了新的麻烦

承接上一节,把 TFA 得到的四组 component 记为

u_0=\mathrm{out},\quad u_1=\mathrm{DWConv}_4(\mathrm{out}),\quad u_2=\mathrm{DWConv}_8(\mathrm{out}),\quad u_3=\mathrm{DWConv}_{12}(\mathrm{out}),

其中每个 u_c\in\mathbb{R}^{C}。问题在于,三条卷积分支与原始输出方向高度相关。在 18B 实验中,我们观察到 u_1,u_2,u_3u_0 的 cosine similarity 可以超过 0.7。

如果直接让 H^{\mathrm{post}} 组合这些方向相近的 components,模型可能在同一个方向上反复放大 write-back signal,进一步引发 activation growth 和 gradient-scale spikes。

我们先尝试分析和降低这种相关性,最后采用了 modified Gram-Schmidt orthogonalization。它保留原始输出、逐个移除 convolutional component 在已有 components 上的投影:

\tilde{u}_0=u_0,\qquad \tilde{u}_j=u_j-\sum_{i<j}\frac{\langle u_j,\tilde{u}_i\rangle}{\langle \tilde{u}_i,\tilde{u}_i\rangle}\,\tilde{u}_i,\quad j=1,2,3,

使四组 write-back features \{\tilde{u}_c\}_{c=0}^{3} 保持相对独立。在 10B ablation 上,加上 GS 后 loss 从 1.984 轻微变化到 1.983,说明它没有通过牺牲性能来换取稳定性;

但在 18B 规模上,移除 GS 会导致训练不稳定。它更像是一个随着规模增大才显现必要性的 conditioning step。

不过还有一个需要注意的点。在使用 Muon 训练 backbone 时,我们移除了 GS。Muon 对 backbone matrix updates 进行 Newton-Schulz orthogonalization,更新谱更加受控;此时再在 forward 中投影掉平行分量,会显得有些多余,甚至对性能略有限制。

到这里,TFA 基本解决了大 N 下 write-back information 不足的问题。但它并没有改变 mHC 原有的计算方式,接下来还要处理另一个问题:如何控制增加 streams 带来的额外开销?

N 从 4 扩展到 16,成本为什么涨得这么快?

mHC 的 residual mapping H^{\mathrm{res}}\in\mathbb{R}^{N\times N} 是从 NC 维的完整多流状态动态生成的。输入维度随 N 线性增长,需要预测的 mixing coefficients 又有 N^2 个,因此这部分 projection 的参数量和逐 token 计算量都是

O(N^3C).

N=16 时,仅这一 projection 就达到 4096C 的量级。对于中等 hidden dimension(例如 1024),这已经不是可以忽略的附加开销。

即使新增 streams 变得更有信息,如果每个 sublayer 仍然为全部 streams 生成并执行完整的 residual mixing,N 也很难成为一个真正实用的 scaling axis。

一个直接的想法是做稀疏化:模型可以拥有 N 条 streams,但每一层是否真的需要同时更新全部 N 条?

这个问题看起来很像 MoE routing,但 residual streams 与 experts 有一个重要区别:MoE expert 本身不保存跨层持续存在的状态,residual stream 却会携带历史信息穿过后续层。

如果一层把信息写进 stream A,下一层既不选择也不读取 A,这部分信息就暂时失去了继续影响计算的路径。直接把 read 和 write 一起做成 top-k,虽然节省了计算,却可能破坏 Hyper-Connections 最重要的跨层信息流。

这就是 xHC 最终采用非对称结构的原因:Dense Read,Sparse Write。

模型可以拥有 16 条 streams,但每层只更新 4 条

xHC的架构图

xHC 的每个 sublayer 仍然从全部 N 条 streams 生成输入:

\mathrm{input}_l=\sum_{i=1}^{N}h^{\mathrm{pre}}_{l,i}x_{l,i}.

因此,无论某条 stream 最近是否被 router 选中,它保存的信息都仍然对后续层可见。Dense Read 保留了完整 residual memory 的跨层连接性。

写入侧则只更新 k 条 active streams。我们的主设置是 N=16,k=4:其中两条 streams 始终激活,另外两条由 router 动态选择。

保留 fixed streams 也是在训练过程中逐渐确定下来的。完全动态的 routing 看起来更灵活,但 residual memory 仍然需要一些持续存在的写入路径:fixed streams 提供稳定的 active targets,routed streams 再负责动态选择。

移除 fixed streams 后,validation loss 从 1.983 变为 1.986;如果同时移除 Dense Read 和 fixed streams,loss 会进一步退化到 1.997。

Router 读取 LayerNorm 后的完整 NC 状态,并为 N 条 streams 输出 sigmoid scores。我们没有使用 softmax,因为 streams 之间并不需要竞争一份固定的总权重。

一个可能的解释是,softmax 更容易形成 winner-take-all routing;实验上,sigmoid router 的 loss 为 1.983,softmax 则为 1.988。

Top-k 的索引选择本身是离散的。对于被选中的 routed streams,sigmoid score 会作为 routing weight 乘在新增 write-back 上,因此 router 仍然可以通过 write-back 的幅度获得梯度;

没有被选中的 logits 在当前 sublayer 不接收这部分梯度。我们也尝试过额外的 balancing 设计,但没有观察到足够稳定的收益,最终选择保持 router 简洁,去掉 balance loss。

把这些放在一起,xHC 对第 i 条 stream 的写入可以写成

\Delta x_{l,i}= \begin{cases} s_{l,i}\displaystyle\sum_{c=0}^{3} h^{\mathrm{post}}_{l,i,c}\,\tilde{u}_c, & i\in \mathcal{A}_l,\\[6pt] 0, & i\notin \mathcal{A}_l, \end{cases}

其中 \mathcal{A}_l 是当前 sublayer 的 active streams(|\mathcal{A}_l|=k),s_{l,i} 是 router 给出的 sigmoid routing weight(fixed streams 恒为激活),\tilde{u}_c 是上一节 GS 之后的四组 write-back components,h^{\mathrm{post}}_{l,i,c} 则决定第 i 条 stream 如何组合它们。

对照 mHC 的 \Delta x_{l,i}=h^{\mathrm{post}}_{l,i}\cdot\mathrm{out} 可以看出三处改动:单一 component \mathrm{out} 变成四组 \tilde{u}_c(TFA),标量系数变成对 component 的组合,以及只有 \mathcal{A}_l 中的 streams 才被更新(Sparse Write)。

因此,Sparse Write 只让 k 条 active streams 参与动态 residual mixing,H^{\mathrm{res}} mapping generation 中最昂贵的 cubic term 从 O(N^3C) 降到了 O(k^3C);inactive streams 保持不变,等待后续 Dense Read,或者在未来某层再次被激活。

这里需要说明,xHC 的整体开销并没有因此完全与 N 无关。Router 仍然读取 NC 状态并输出 N 个 scores,对应 O(N^2C) 的计算;Dense Read 也仍然需要访问全部 streams。Sparse Write 解耦的是最昂贵的动态更新成本,而不是抹掉所有随 N 增长的开销。

从这个角度看,N=16,k=4 不只是两个超参数。N 表示模型拥有多少 residual state,k 表示一个 sublayer 实际修改多少 state。在 xHC 中,状态容量和活跃更新带宽被明确拆开了。

实验结果

到这里,我们已经从两个问题导出了 xHC 的两个核心设计:TFA 用来增加大 N 下的 write-back information,Sparse Write 用来控制 residual update 的成本。接下来通过实验验证:这两个设计是否分别解决了前文提到的 bottleneck,以及 xHC 能否在大规模语言模型预训练下,提升模型的性能。

Main Results:xHC 相比 mHC 提升 4.0 和 3.1 分

在公平的训练 FLOPs 下,我们对使用 Vanilla、mHC 和 xHC 的 18B、28B MoE 模型进行了评测,覆盖 language understanding、reasoning、code 和中文任务。

在 18B MoE 上,Vanilla、mHC 和 xHC 的平均分分别为 40.6、44.8 和 48.8,xHC 相比 mHC 提升了 4.0 points。在 28B MoE 上,三者的平均分分别为 47.8、50.5 和 53.6,xHC 相比 mHC 提升了 3.1 points。

以 18B 为例,xHC 相比 mHC 在 ARC-Challenge、BBH、C3 和 HumanEval 上分别提升 5.9、5.8、5.6 和 6.1 points;相比 Vanilla,MMLU 从 48.9 提升到 57.2,MMLU-Pro 从 21.1 提升到 29.7,分别增加 8.3 和 8.6 points。这些提升都可以称之为显著。

前面的主实验均使用 AdamW。考虑到 Muon 在 large-scale Transformer pre-training 中已经展现出很强的效果,我们进一步测试了 xHC 与 Muon 的兼容性。

实验保持 18B MoE 的模型结构和其余 training recipe 不变,在 backbone 的二维权重上使用 Muon,同时让 xHC-specific parameters 继续使用 AdamW。

这个对比可以直观地看出 xHC 的增益量级:把 Vanilla backbone 的 optimizer 从 AdamW 换成 Muon,平均分从 40.6 提升到 43.1,增加了 2.5 points;

在 Muon baseline 上加入 xHC 后,平均分进一步提升到 49.9,又增加了 6.8 points。这说明 xHC 的收益在 Muon 下依然成立,并不依赖 AdamW 特有的 optimization dynamics。

Muon 也改变了 TFA 对 GS 的需求。由于 Muon 通过 Newton–Schulz iteration 对 backbone matrix updates 进行正交化,我们观察到额外的 forward-side projection 不再必要,甚至略有限制。因此,Muon + xHC 的实验移除了 GS,其余 xHC 组件保持不变。

从 N-sweep 到 Scaling Law

18B 和 28B 的结果说明 xHC 在两个 scale 上都有效。回到我们最初的想法:N 能否真正成为一个可以继续扩展的维度。为此,我们先在 2.5B MoE 上对 N\in\{2,4,8,16\} 做了完整的 N-sweep 实验。

当 mHC 从 N=4 扩展到 N=16 时,validation loss 只下降了 0.006,训练 FLOPs 却增加了 32%。在相同范围内,xHC 将 loss 降低了 0.012,而训练 FLOPs 只增加约 4%。

对我来说,这张图最重要的并不是 N=16 这个单点的绝对 loss,而是扩大 N 时的 marginal return 发生了变化。mHC 越过 N=4 后很快趋于饱和;xHC 则在我们测试的 N\leq16 范围内,仍然能从更大的 residual memory 中持续获益。这当然不能说明 N 可以无限扩大,但至少在当前实验范围内,它不再天然停在 4。

N-sweep 单独考察了 expansion rate,Scaling Law 则进一步检验这种优势能否跨 compute budgets 延续。我们为 Vanilla、mHC 和 xHC 分别训练了四个有不同训练 FLOPs 的模型,拟合 scaling curves。

scaling law曲线

Scaling Law 的结果进一步展现了 xHC 很强的 scaling 潜力。在我们测试的 compute range 内,xHC 始终保持最优的 scaling curve;达到相同 loss,Vanilla 和 mHC 分别需要消耗 xHC 的 1.50 倍和 1.19 倍计算量。

xHC 核心组件消融

为了分别观察各个组件的作用,我们在 10B MoE 上进行了一系列 controlled ablation:

Ablation实验表格

先看最关键的两步。将 mHC 从 N=4 直接扩展到 N=16,loss 只从 2.004 降到 1.998;而额外训练 FLOPs 从 N=4 时的 0.6% 升到了 N=16 时的 18.8%。

=16 的 mHC 上加入 TFA 后,loss 进一步降到 1.984,说明大 N 下丰富 write-back components 确实能够带来明显收益,这与前面对 information supply bottleneck 的分析一致。

不过,此时 16 条 streams 仍然全部参与 dense residual mixing,额外 FLOPs 也随之达到 20.1%。

将残差流结构替换为 xHC 的 Dense Read、Sparse Write architecture 后,loss 进一步变为 1.983,而额外 FLOPs 从 20.1% 降到了 3.3%。

这组对比很直接:TFA 解决的是大 N 下新增 streams 缺少足够 write-back information 的问题;Sparse Write 则让模型不必在每个 sublayer 更新全部 streams,从而把扩大 N 的成本控制下来。

表中其余结果也与前面的设计动机一致:Dense Read 让所有 residual states 始终对后续层可见,fixed streams 则提供稳定、持续的写入路径;

两者共同维持稀疏更新下的跨层信息流。对 k 的消融表明,更新的 streams 太少会限制效果,继续增大 k 的边际收益又很有限,因此我们最终选择 k=4。Sigmoid router 也优于 softmax,这与 streams 之间不需要进行零和竞争的判断一致。

从 N=16 往后

回到最开始的问题:N 究竟只是一个通常取 4 的 architecture choice,还是一个可以继续扩展的 scaling axis?

我们的答案是后者。在 mHC 的原始形式下,N 越过 4 就迅速饱和,看起来更像一个需要调好的常数;而 xHC 让 N 表现得像 hidden dimension、depth、expert 数一样。在我们测试的范围内,扩大 N 能持续换来收益,而代价被控制在可接受的范围内。

N 成为一个可以 scaling 的 axis 之后,更多的问题也随之而来:

1、为什么不做 N=32 的xHC? 主要还是开销侧的问题,除了前文分析的训练FLOPs的增加,N 增大还带来了超大的 I/O 开销以及显存开销,所以综合考虑 N=32 并不适合现在的 MoE backbone。这不是说 N=16 没有带来这些开销,而是 N=32 带来的已经超出当前最大承受范围了。

2、在真实训练中,xHC对训练吞吐的影响是怎么样的? 这个问题,我们在paper中已经深入分析了,为了让 N=16 真正可用到large scale MoE上,最大的阻碍还是 memory access,也就是 I/O 开销(Hyper-Connections类方法存在大量的激活read/write),其中xHC相比于mHC在N=16下已经减少了一大部分memory access,但是还不足以match mHC的开销。

所以在此基础上,我们进一步提出了xHC-Flash,解决 N=16 下 I/O 开销大的问题,将理论 memory access 从73.5C降低到了40C,与mHC的34C相当。

关于xHC如何从 73.5C 的 memory traffic 一步步降到接近 mHC N=4 的水平,以及为什么 dense-read correction 可以做到精确等价,篇幅原因我不在此赘述。

xHC 的完整实现会在文章上方 GitHub 仓库开源;具体开源时间还未最终确定,感兴趣的可以先 star,开源后会第一时间更新。

写在最后

xHC 当然还不是一个完美的方法,仍然有很多可以继续优化的地方,比如端到端开销就还有压缩空间。

但比起方法本身,这段研究过程带给我更深的一点感受是:做研究,很多时候是起源于一个很小的疑惑,在一个又一个的岔路口进行选择,在此过程中形成并坚持自己的判断。

这也是我想把这些论文之外的过程写下来的原因。相比最后那个相对干净的 formulation,这些走过的岔路以及其中的选择可能更接近研究原本的模样。

目录