LLM跨层信息传递优化:DeepSeek mHC & Kimi AttnRes解析

less than 1 minute read

Published:

LLM沿宽度方向(Attention)的优化已十分充分,但沿深度方向则一直使用简单的残差连接来进行跨层信息传递。残差连接会直接将前层的信息不加筛选地传递到后层,使得深层输入的信息中包含大量无用噪声,可能对模型性能带来不好的影响。因此在2026上半年,DeepSeek和Kimi分别通过mHC和AttnRes来优化了跨层信息传递,使得有用的信息能更多地传递到深层,而无用的信息则受到抑制。

目录

DeepSeek mHC

https://arxiv.org/abs/2512.24880

https://zhuanlan.zhihu.com/p/1997998943259342644

https://www.zhihu.com/question/1990084005744362243/answer/1990253031078515110

alt text

普通残差:

\[x_{l+1}=x_l+F_l(x_l,W_l)\]

其中,$x_l\in\mathbb R^{1\times d}$为第$l$层的网络输入(这里考虑输入为单个向量的情况),$F_l(\cdot)$为第$l$层的网络变换函数,$W_l$为该层的权重参数。残差连接的恒等映射特性可以使得浅层信息无损地传入深层,确保了模型训练的稳定性。多层传播后的结果如下:

\[x_L=x_l+\sum_{i=1}^{L-1}F(x_i,W_i)\]

可见,无论网络多深,短接部分的梯度系数会恒等于1。

字节提出的HC(Hyper- Connections,超连接)通过扩展残差流的宽度和增加残差连接的复杂度,在没有增加额外计算开销的情况下提升了网络的拓扑复杂度。相比于普通残差连接不加区分地传播所有信息,其对于信息进行了更细粒度的区分和筛选,使得传递到深层的信息是经过筛选变换的。单层传播公式如下:

\[x_{l+1}=\textcolor{red}{\mathcal H_l^{res}}x_l+\textcolor{red}{\mathcal H_l^{post^T}}F_l(\textcolor{red}{\mathcal H_l^{pre}}x_l,W_l)\]

在这里,层间流动的隐含向量的形状被扩展为$x_l,x_{l+1}\in\mathbb R^{\textcolor{red}{n}\times d}$,其中$n$为扩展率,也即将隐含向量由$1$个特征通道扩展为了拥有了$n$个特征通道(首层输入时直接复制$n$份)。

  • $\mathcal H_l^{res}\in\mathbb R^{n\times n}$为res mapping,是用于在残差流($x_l$)内部混合其($n$个通道的)特征的信息的可学习参数矩阵,其作用于$x_l\in\mathbb R^{n\times d}$后的输出仍有$n$个通道。这样即可使得短接部分被扩展为了$n$个并行流,拓宽了其宽度。

  • $\mathcal H_l^{pre}\in\mathbb R^{1\times n}$为pre mapping,用于将输入$x_l\in\mathbb R^{n\times d}$的$n$个通道特征聚合为$1$个通道,然后这个$1\times d$的单通道向量即可输入网络$F_l(\cdot)$进行正常的网络变换,输出仍为$1\times d$的单通道向量

  • $\mathcal H_l^{post}\in\mathbb R^{1\times n}$为post mapping,用于将该层网络输出的$1\times d$单通道向量映射回$n$通道空间,从而和混合特征后的残差流相加后得到本层输出。

具体来说,这三个可学习变换矩阵均为由两部分组成,分别为输入依赖项(dynamic mappings)和全局项(static mappings):

\[\tilde x_l=\text{RMSNorm}(x_l)\] \[\mathcal H_l^{pre}=\alpha_l^{pre}\cdot\tanh(\theta^{pre}_l\tilde x_l^T)+b_l^{pre}\] \[\mathcal H_l^{post}=\alpha_l^{post}\cdot\tanh(\theta^{post}_l\tilde x_l^T)+b_l^{post}\] \[\mathcal H_l^{res}=\alpha_l^{res}\cdot\tanh(\theta^{res}_l\tilde x_l^T)+b_l^{res}\]

其中,$\alpha\in\mathbb R$为标量可学习门控因子,用于调控输入对于变换矩阵的影响度;$\theta$的维度和$\mathcal H$相同,也是可学习参数,其对于输入施加一个线性变换来构成动态映射部分;$b$为可学习偏置参数,为静态映射部分。可见,动态映射部分在推理时会随不同的输入而变化,静态映射部分则在训练结束后就确定了。

引入这些变换矩阵带来的额外计算开销很小,因为扩展因子$n$的取值一般远小于隐含维数$d$,如$n=4$。

然而,HC模式下多层传播后可能带来不稳定性的问题:

\[x_L=\textcolor{red}{\left(\prod_{i=1}^{l-1}\mathcal H_{L-1}^{res}\right)}x_l+\sum_{i=l}^{L-1}\left(\prod _{j=1}^{L-1-i}\mathcal H_{L-j}^{res}\right)\mathcal H_i^{post^T}F(\mathcal H_i^{pre}x_i,W_i)\]

可见,短接部分的变换多层叠加后的总变换$\prod_{i=1}^{l-1}\mathcal H_{L-1}^{res}$缺乏约束,无法确保维持特征的全局均值恒定,这种累乘随着层数加深会造成信号的爆炸或消失。反向传播时短接部分的梯度系数也不为1,容易造成梯度不稳定。另外,拓宽的残差流的也会造成额外的访存开销。这些因素使得原始的HC不适合在大规模训练中被使用。

为了解决这个问题,DeepSeek提出了mHC(manifold-Constrained Hyper-Connections,流形约束的超连接),从而约束变换矩阵$\mathcal H_l^{res}$的学习,使得其最终训练结果满足预期的约束条件,来避免HC中因它产生的训练不稳定问题。

具体而言,这里对该矩阵进行流形优化来避免其训练不受约束。流形优化指的是,希望最终得到的模型权重符合某限制条件,可以先让权重按照正常的目标函数进行训练优化,然后使用数学方法来将训好的权重“微调”成符合限制条件的样子,数学上相当于将其投影到某个流形上。

在mHC中,使用Sinkhorn-Knopp算法将$\mathcal H^{res}_l$变换为双随机矩阵(矩阵元非负,且所有行/列之和都是1),双随机矩阵的流形被称为Birkhoff polytope,它具有若干好的性质:

  • 范数保持:双随机矩阵的谱范数$||H_l^{res}||_2\leq 1$,这就可以确保得到的可学习映射是non-expansive的,避免梯度爆炸问题

  • 乘法封闭:两个双随机矩阵相乘结果仍为双随机矩阵,这就使得无论模型层数多少,无论变换矩阵累乘多少次,多层复合的总映射仍保持在流形上,确保了长距离信号流的稳定性

  • 从几何上看,该流形是permutation矩阵的凸包,也即$\mathcal H_l^{res}$变换相当于多个permutation的凸组合,施加于$x_l$的$n$个特征通道相当于对这些特征进行加权融合(也可以从双随机矩阵行列之和为1的特性出发来理解对于多个特征通道的加权融合),从而促使了多条残差特征流之间的信息交换,同时不改变信息总量。

Sinkhorn-Knopp算法是一种可以将任意矩阵投影到Birkhoff polytope的迭代算法。在实践中为了保证优化效率只迭代20轮,因此最终结果不一定会精确落到Birkhoff polytope上,但足够接近。最终会将$\mathcal H_l^{res}$约束到这个流形附近,从而达到避免其爆炸的目的。

总的来看,普通残差直接恒等连接虽然在深层模型中传播稳定,但对于短接传播的信息没有任何筛选,HC对于短接信息进行了复杂变换但无约束,容易造成爆炸。mHC则在保持HC的能力的前提下尽可能往恒等连接的特性上靠拢,同时吸取了二者的优势。

Kimi AttnRes

https://arxiv.org/abs/2603.15031

https://zhuanlan.zhihu.com/p/2016957666388387770

alt text

在传统的PreNorm范式下,残差连接可以为深层网络提供梯度高速公路,使得深层梯度可以绕开网络变换直达浅层,缓解梯度消失问题:

\[h_l=h_{l-1}+f_{l-1}(h_{l-1})\] \[h_l=h_1+\sum_{i=1}^{l-1}f_i(h_i)\]

标准残差连接本质上是一种固定的信息累积方式。每一层接收到的都是初始embedding($h_1$)加上之前所有层输出的均匀累加和($\sum_{i=1}^{l-1}f_i(h_i)$),这种无差别且无界的累积使得hidden state的幅值随深度而线性增长,逐渐稀释了每一层的相对贡献,那些较浅层的信息被埋没而无法被选择性地检索到。较深层的也会倾向于学习输出更大的值来在残差流中获得更大贡献比例,这可能导致训练不稳定。另外,不同种类的模块如attention和MLP都会接收到相同的聚合结果,忽视了不同类型模块对于历史信息的需求差异。

本质上,标准残差连接和RNN都由于固定的信息累积方式导致了信息瓶颈,它们分别是在空间维度(沿模型深度方向)和时间维度(沿序列长度方向)上将之前的所有位置的信息压缩为一个隐含状态,而无法让每个位置自主选择对于之前所有位置信息的获取。Transformer则通过attention机制使得各个位置可以用data-depentent地选择性取得前边各个位置的信息(利用和当前输入相关的attention score来加权过去所有位置的信息),在时间维度上实现了输入相关的选择性历史信息提取。类似地,可以在空间维度上也设置注意力机制,让每层根据当前输入来灵活选择对于之前各层信息的关注程度,这就是Attention Residuals

\[h_l=\textcolor{red}{\alpha_{0\to l}}\cdot h_1+\sum_{i=1}^{l-1}\textcolor{red}{\alpha_{i\to l}}\cdot f_i(h_i)\]

其中,$\alpha_{i\to l}$表示第$l$层对于第$i$层的注意力分数,且对于每层而言都满足其对之前所有层的注意力分数之和为1($\sum_{i=0}^{l-1}\alpha_{i\to l}=1$)。相比于可达数百万的上下文长度,模型层数通常小的多($L<1000$),因此沿模型深度方向的$O(L^2)$复杂度的注意力计算开销是可以接受的。

AttnRes的完整实现形式称为Full Attention Residuals,也即在深度方向上对之前所有层进行完整注意力计算。注意力分数使用softmax attention来计算,保证注意力分数的归一化:

\[\alpha_{i\to l}=\frac{\phi(q_l,k_i)}{\sum_{j=0}^{l-1}\phi(q_l,k_j)}\]

其中注意力计算为:

\[\phi(q,k)=\exp(q^T\text{RMSNorm}(k))\]

每层的query向量是一个和输入无关的可学习参数(不是依赖当前输入token的动态query,减少计算复杂度),key向量是则是第$i$层的变换输出$f_i(h_i)$,$f_i(h_i)$同时也作为value向量用于代表第$i$层的被加权信息:

\[q_l=w_l\] \[k_i=v_i=\begin{cases} h_1,& i=0 \\ f_i(h_i),& 1\leq i\leq l-1 \end{cases}\]

由此,第$l$层的输入就可以写成对于之前所有层的完整注意力加权形式:

\[h_l=\sum_{i=0}^{l-1}\alpha_{i\to l}\cdot v_i\]

计算注意力分数时key要先经过RMSNorm归一化再和query做计算,这是因为key来自于之前层的直接输出,这样可以避免量级较大的输出(如深层输出)在注意力分数的计算中占据主导,确保注意力权重合理分布。另外,query在训练开始时初始化为零向量,使得训练开始时所有前序层的注意力分数完全均等,等价于标准的均匀平均,和标准残差的行为类似,保证了训练初期的数值稳定性,避免随机初始化引发的非均匀注意力分布和损失爆炸。

对于每个token来说,Full AttnRes需要$O(Ld)$的内存来储存这些中间层输出$f_i(h_i)$,用于反向传播时计算梯度。而实际上由于这些中间层输出本身就要用于网络参数的梯度计算,因此在标准残差连接下,使用简单训练策略时它们本来就会被全部储存起来,用于后续反向传播计算梯度,AttnRes只不过相当于将它们又复用了一遍,无需额外的内存开销。然而,标准残差连接下训练和存在AttnRes时的训练的差异在于,前者反向传播计算每层的梯度时只需用到它上一层的输出activation,不需要直接用到再上边的层的activation,这就使得gradient checkpointing策略得以实施,可以不用在任何时刻都存储所有层的输出,也使得流水线并行下每个device只需要存储自己那部分层的输出即可,不需要关心更远的那些层;而后者在计算每层梯度时都要用到其上边所有层的输出activation,这就导致需要在任何时刻都保留模型所有层的输出,难以实施gradient checkpointing,使用流水线并行策略时每个device也要从其他device处获得它们持有的所有前序层的输出,带来额外通信开销。

为了缓解大规模训练使用gradient checkpointing和流水线并行时完整AttnRes带来的问题,进一步提出一种分块的变体:Block Attention Residuals,其本质上是标准残差连接和AttnRes的混合。

具体而言,其将总共$L$层切分为$N$个块,每个块内部的$L/N$层是连续的。块内部采用轻量高效的标准残差连接,块之间则引入AttnRes机制,每个块通过将内部所有层的输出加和来获得一个压缩的块级表征,用于提供给后续层用做key和value。(也即,对于位于第$b$个block内的第$l$层,其输入一方面来源于本block内的那些前序层的直接短接,另一方面来源于对于之前所有block的attention)

这样一来,块内的固定残差确保了短距离深度依赖的高效处理,块间的注意力则支持长距离深度依赖的灵活检索。在这种优化下,和AttnRes有关的需要为反向传播保留的东西由所有层的输出减少为了各block的压缩表征,内存开销由$O(Ld)$降低到$O(Nd)$。

推理优化:留坑