从残差到分位数平衡：拆解 Kimi K3 的 AttnRes 与 Stable LatentMoE

本文是 《从线性注意力到 KDA：拆解 Kimi K3 的核心注意力机制》 的续篇，主要基于 Moonshot AI 发表的技术报告 Kimi K3: Open Frontier Intelligence（arXiv 2607.24653）整理。上一篇讲了"序列维度"的 KDA + Gated MLA，这一篇补上另外两条轴：网络深度上的 Attention Residuals（AttnRes），和模型宽度上的 Stable LatentMoE。文中数字如无特别说明均引自官方技术报告。

Kimi K3 把架构改进归纳为三条轴：序列长度、网络深度、模型宽度。上一篇的 KDA、Gated MLA 解决的是"序列方向怎么高效检索历史 token"；这一篇的两个模块解决的是完全不同的两个问题——AttnRes 解决"深度方向怎么高效检索更早层的表示"，Stable LatentMoE 解决"专家数拉到 896 个之后怎么不炸"。有意思的是，AttnRes 用来解决深度瓶颈的思路，跟 KDA/MLA 解决序列瓶颈用的是同一套武器：attention。

一、标准残差连接：深度维度上的"RNN 瓶颈"

标准残差连接的公式很简单：

$h_l = h_{l-1} + f_l(h_{l-1})$

这个公式有个结构性问题——只会加，不会选。到第 50 层时，token 的表示是前面所有层输出等权重原封不动地累加在一起，模型没法说"第 3 层的特征对我现在这一层特别有用，第 30 层的没什么用"，只能被动接受"上一层给的这一份"。论文把这个类比得很精准：标准残差连接就像 RNN 沿深度方向把所有历史压进单一状态——跟朴素线性注意力沿序列方向的问题，是同一类瓶颈换了个维度。当年 Transformer 用 attention 解决了"沿序列方向不能选择性检索"的问题；Attention Residuals（AttnRes）就是把同一套方法论搬到深度方向：每一层不再无条件接收"上一层给的东西"，而是用 attention 主动去查前面所有层，自己决定要多少权重。

二、Full Attention Residuals：完整数学形式

对每一层 $l$，定义一个该层专属、可学习的伪 query（"伪"是因为它不依赖当前 token 内容，是每层自己固定学到的一个向量）：

$q_l = w_l \in \mathbb{R}^d$

Key 和 value 用的是同一个东西——之前每一层（以及最初的 embedding）的输出：

$k_i = v_i = \begin{cases} h_1 & i=0 \\ f_i(h_i) & 1\le i \le l-1 \end{cases}$

（embedding 记作第 0 号，永远作为一个可选的检索源。）注意力权重用一个软最大核：

$\phi(q,k) = \exp\big(q^\top \text{RMSNorm}(k)\big), \qquad \alpha_{i\to l} = \frac{\phi(q_l,k_i)}{\sum_j \phi(q_l,k_j)}, \qquad h_l = \sum_{i=0}^{l-1} \alpha_{i\to l} \cdot v_i$

两个细节值得注意：

1. key 要先过 RMSNorm 才拿去算相似度——如果某一层的输出恰好数值幅度特别大，它跟任何 query 的点积都会被这个"幅度"人为放大，看起来"权重很高"但其实跟内容相关性无关，纯粹是数值大小造成的假象。RMSNorm 把这个幅度差异抹掉，只留下方向上的相关性。
2. key 和 value 是同一个向量——不像正常 attention 里 K、V 是输入的两个不同线性投影，这里没有额外的投影，直接拿层输出本身既当 key 用来算相似度，也当 value 用来加权求和。

复杂度上，由于网络深度 $L<100$（远小于序列长度 $n$ 可以到 1M），这套完整版本的算术量 $O(L^2d)$ 完全不是问题；真正的开销是 $O(Ld)$ 的内存——必须把所有 $L$ 层的输出都留着不能释放，等着后面任何一层随时可能来查，流水线并行下还要跨阶段搬运这些留存的输出，这才是实际瓶颈。

三、Block Attention Residuals：K3 实际用的近似版本

为了压掉这个内存开销，K3 把 $L$ 层切成 $N$ 个 block，每个 block $S=L/N$ 层。block 内部退回最朴素的做法——直接求和：

$b_n = \sum_{j \in \mathcal{B}_n} f_j(h_j), \qquad b_0 := h_1$

block 之间才用完整的 attention，只对 $N$ 个 block 级别的汇总表示做检索。对 block $n$ 内的第 $i$ 层，参与检索的 value 矩阵是：

$V = \begin{cases} [b_0, b_1, \dots, b_{n-1}]^\top & i=1\ (\text{block 内第一层}) \\ [b_0, b_1, \dots, b_{n-1}, b_n^{i-1}]^\top & i\ge 2\ (\text{后续层}) \end{cases}$

其中 $b_n^{i-1}$ 是当前 block 内到上一层为止的部分和。也就是说：block 内第一层只能看"之前已经封口的 block"；block 内后面的层还能多看一眼"当前 block 到目前为止累积了多少"。

看到这个结构会发现一个熟悉的模式——“block 之间用聪明的 attention 精确检索，block 内部用便宜的直接求和”，这跟 KDA chunkwise/tile 并行反复出现的套路（粗粒度上用精确/聪明的机制，细粒度上退回简单/便宜的机制）是完全一样的设计哲学，只是换成了深度维度。

论文实验发现 $N\approx8$ 就能拿到大部分收益，K3 实际把 93 层切成 8 个 block，每 block 12 层（最后一个 block 不满），算上 embedding 作为 block 0，一共 9 个 block——任何时刻只需要留存 9 个 block 级汇总向量，内存直接降了一个量级：

	Full AttnRes	Block AttnRes
内存/通信开销	$O(Ld)$	$O(Nd)$
检索粒度	逐层精确检索	block 级检索 + block 内求和近似
推理时状态大小	随层数增长	被 block 数 $N$ 限死上界

推理（尤其逐 token 解码）时还有个工程细节：一次前向要合并"已经封口的 $N-1$ 个 block 做的那次 attention"（只依赖已完成的 block，可提前算好）和"当前 block 内部还在累积的部分和 $b_n^{i-1}$"（随层数变化）。论文用online softmax（跟 FlashAttention 里增量更新 softmax 归一化的技巧同源）把这两部分增量合并，不用每往下走一层就把整套 attention 重新算一遍。

四、AttnRes 到底接在模型的哪个节点上

这是最容易搞混的一点：AttnRes 跟 attention 内部（Q/K/V、softmax）完全没关系，它是层与层之间那根连接线，发生在一整层（token-mixing 子层 + MoE 子层都跑完）之后。论文原话说得很清楚：

“Each block contains three Kimi Delta Attention (KDA) layers followed by one Gated MLA layer, with each attention layer paired with a Stable LatentMoE feed-forward network.”

也就是说 $f_l(h_l)$ 指的不是"只做完 KDA/MLA 那一下"，而是**“KDA-或-MLA 子层 + 后面跟着的 Stable LatentMoE 子层，两个都跑完”**之后的输出：

标准架构里下一层直接拿上一层的最终结果往下传，$h_{l+1}=f_l(h_l)$；AttnRes 把这根"直接传递"的线，换成了"对所有前面层的最终产出做一次 attention 检索"。层内部（无论 KDA 层还是 Gated MLA 层）该怎么算 attention、该怎么算 MoE 都跟标准架构一样，层内部自己也有普通的小残差，这跟 AttnRes 是两个不同层级的东西。Block AttnRes 版本同理，只是检索单位从"层"变成了"block"，但检索的仍然是"整层跑完之后"的产出，不会更细到 attention 内部某个中间步骤。

这个设计在推理侧还接上了另一个模块：K3 预训练自带一个 MTP（multi-token-prediction）层，微调成 EAGLE-3 风格的草稿模型时，草稿输入直接融合了目标模型第 1、第 4、最后一个 AttnRes block 的输出——低、中、高层特征的融合，正好利用了 AttnRes 天然把各层特征留存并可检索这个特性，省了专门再设计一套多层特征融合方案。

五、Stable LatentMoE：为什么 896 个专家会有稳定性问题

架构改进的第三条轴是模型宽度。传统 MoE 的问题：每个被选中的专家都要接收完整的 $d$ 维 token 表示，专家数一多，通信和专家权重的搬运量就跟着路由多样性线性涨。LatentMoE 的解法是把模型整体宽度 $d$ 和路由专家的宽度 $\ell$ 拆开——共享专家走全宽度 $d$，负责通用变换；被路由的专家在一个更窄的 latent 空间 $\ell$（$\ell \ll d$）里工作。这样才能把专家池扩到 896 个（16 个激活 + 2 个共享，稀疏度 56）还能负担得起：

$\bm{u} = \sum_{i\in\mathcal{T}_k(\bm{x})} p_i\, E_i^{\text{routed}}(\mathbf{W}^{\downarrow}\bm{x}), \qquad \bm{y} = \sum_{j=1}^{N_s} E_j^{\text{shared}}(\bm{x}) + \mathbf{W}^{\uparrow}\,\text{RMSNorm}(\bm{u})$

极端稀疏放大了两个失效模式：

失效模式 1：链路太深，数值爆炸。 路由这条路径是 $\mathbf{W}^{\downarrow}$（降维）→ 专家内部的门控 FFN（gate 分支 + up 分支）→ $\mathbf{W}^{\uparrow}$（升维），接近四次连续矩阵乘法、中间没有归一化兜底。每一次矩阵乘法都可能放大输入的数值尺度，四层叠加在 2.8T 参数规模下，路由分支内部激活爆炸的问题被指数级放大。

失效模式 2：专家太多，负载均衡的旧方法失灵。 K3 用 auxiliary-loss-free 路由——给每个专家的路由分数加一个偏置 $b_j$ 来控制它被选中的频率。原始方法用固定步长更新偏置：

$b_j^{(t+1)} = b_j^{(t)} + \gamma \cdot \text{sign}\big(\bar{\ell} - \ell_j^{(t)}\big)$

$\gamma$ 存在两难：步长大了容易震荡，步长小了收敛慢。专家数逼近 $10^3$ 量级时，这个两难被放大到现有方法不好用的程度。

Stable LatentMoE 用三个组件分别堵这两个坑：RMSNorm 抑制聚合后的数值爆炸，SiTU-GLU 抑制专家内部的数值爆炸，Quantile Balancing 重新设计负载均衡的偏置更新规则。

六、Normalized LatentMoE：RMSNorm 插在哪

原始 LatentMoE 直接把聚合结果 $\bm{u}$ 喂给 $\mathbf{W}^{\uparrow}$，但 $\bm{u}$ 的数值尺度会随着"当前 token 选中了哪些专家、这些专家各自的路由权重"剧烈变化。K3 在聚合之后、升维之前插一个 RMSNorm（W↑RMSNorm(u)），把这个尺度波动在传给下一步之前先磨平。论文提到这个改动不仅稳定了训练，还持续改善了验证损失和下游指标。

七、SiTU-GLU：给 GLU 装一个软天花板

这个针对专家内部（聚合之前）的激活爆炸。先看三代 GLU 门控设计对比：

结构	Gate 分支	Up 分支	是否有界
GLU	$\sigma(x)$	$x$	Gate 有界，Up 无界
SwiGLU	$x\cdot\sigma(x)$（Swish）	$x$	都无界，大 $x$ 时线性增长
SiTU-GLU	$\beta_1\tanh(x/\beta_1)\cdot\sigma(x)$	$\beta_2\tanh(x/\beta_2)$	两个分支都有硬上界

SwiGLU 效果好但两个分支都不封顶——gate 和 up 恰好同时遇到大数值时（低精度训练 / FP8 / FP4 量化下这种异常大值尤其危险，动态范围一旦被一个 outlier 撑爆，其它正常值的精度就被牺牲了），激活就会飙出去。GLU 的 sigmoid gate 虽然有界，但丢掉了 Swish 那种正半轴近似线性的好特性。

SiTU-GLU 的做法：对 Swish 内部的线性部分和 up 分支的线性部分，分别套一个软上限 $\text{softcap}(x,\beta)=\beta\tanh(x/\beta)$：

$\text{SiTU-GLU}(\bm{x}) = \Big[\beta_1\tanh\!\Big(\frac{\mathbf{W}_g\bm{x}}{\beta_1}\Big)\odot \sigma(\mathbf{W}_g\bm{x})\Big] \odot \Big[\beta_2\tanh\!\Big(\frac{\mathbf{W}_u\bm{x}}{\beta_2}\Big)\Big]$

K3 取 $\beta_1=4$（gate 分支），$\beta_2=25$（up 分支）。这个设计有两个漂亮的性质：

1. 原点附近几乎不变：$\tanh(x/\beta)\approx x/\beta$ 在 $x$ 小时成立，所以 $\beta\tanh(x/\beta)\approx x$——正常量级的激活跟 SwiGLU 几乎一样。
2. 大值处有硬边界：因为 $|\tanh(\cdot)|\le1$、$|\sigma(\cdot)|\le1$，整个输出满足 $|f(x)|\le \beta_1\beta_2=100$——这是数学上可证明的硬上限，不是"经验上通常不会太大"。

一句话：SiTU-GLU 只在异常大的尾部时才踩刹车，正常情况下的行为和 SwiGLU 无差别，专门堵住了 outlier 激活这一个漏洞。

八、gate 和 up 到底是什么：从通用 FFN 到专家内部

这里容易和上一节的外层 $\mathbf{W}^{\uparrow}$ 搞混，需要先说清楚 gate/up 是通用 FFN 概念，跟 MoE 没关系。现代 LLM（LLaMA、Qwen 等）的 FFN 权重永远是三个矩阵：gate_proj、up_proj、down_proj。最原始的 FFN 是两个矩阵夹一个激活函数：

$\text{FFN}(\bm{x}) = \mathbf{W}_2 \cdot \text{ReLU}(\mathbf{W}_1 \bm{x})$

GLU 系列的改进是把中间那一步拆成两条并行的线性分支，算完之后逐元素相乘，再投影回去：

gate_proj 和 up_proj 形状完全一样（都是 $d\to h$），是两个独立并行的线性层。直觉上：up 分支算的是"这一路要传递的实际内容"，跟没有门控的普通 FFN 里那一步性质一样；gate 分支算的是"这个内容的每一个维度应该放多少比例通过"，充当逐维度的软开关。gate ⊙ up 就是"用一套独立算出来的开关，逐维度地筛选另一套独立算出来的内容"，比"内容算完直接套一个固定激活函数"更灵活，因为开关本身也是数据依赖的。

放回 K3 语境，容易搞混的地方是：gate/up 分支发生在每一个专家（无论 shared 还是 routed）自己内部，而外层的 $\mathbf{W}^{\downarrow}/\mathbf{W}^{\uparrow}$ 是整个 MoE 层级别的降维/升维，两者是不同层级，刚好都用到"投影"这个操作：

外层 $\mathbf{W}^{\uparrow}$ 是"MoE 层为了省参数/省通信设计的宽度转换"，跟"content vs. gate"这个语义无关；内层的 gate/up 才是 GLU 那种"用开关筛内容"的语义，两者名字撞车但完全是两件事：外层是所有专家共用的一份矩阵，内层是每个专家各自独立拥有的一份小参数。

九、Quantile Balancing：一步求解负载均衡偏置

这个直接解决失效模式 2。核心思路的转变：不再"猜一个方向、挪一小步、观察、再挪"，而是从当前这批数据里直接算出应该是多少的偏置。

关键技巧——Top-$(k{+}1)$ 而不是 Top-$k$：正常路由只需要 Top-16，QB 多算一个：取 Top-17。这样对每个 token $i$，同时拿到两个信息——真正被选中的 16 个专家，以及第 17 名的分数 $\alpha_i^{(t)}$：这是"一个专家想要挤进这个 token 的 Top-16，必须超过的门槛"。

对每个专家 $j$，定义它跟每个 token 的 margin：$s_{i,j}-\alpha_i^{(t)}$。目标负载是 $q=mk/n$（$m$ 个 token，每个选 $k=16$ 个，$n=896$ 个专家平摊）。要让专家 $j$ 刚好拿到 $q$ 个 token，偏置就该取"margin 里第 $(q{+}1)$ 大的那个值"的负数——这正好是 margin 分布的 $(1-k/n)$ 分位数：

$\widehat{b}_j^{(t+1)} \leftarrow -\text{quantile}_{1-k/n}\big(\bm{s}_{:,j}-\bm{\alpha}^{(t)}\big), \qquad \bm{b}^{(t+1)} \leftarrow \widehat{\bm{b}}^{(t+1)} - \text{mean}(\widehat{\bm{b}}^{(t+1)})\cdot\mathbf{1}$

第二步减掉全局均值只是把偏置整体归零中心化，不影响 Top-$k$ 的相对排序。

这样一次前向就能精确算出让每个专家刚好达到目标负载 $q$ 的偏置，从"迭代逼近"变成"一步求解"，天然避免了旧方法在专家数暴涨后出现的震荡/收敛慢问题。论文给了个 $m=8,n=4,k=1$ 的小例子：4 个专家原始负载是 $(4,3,1,0)$（明显不均），套一次 QB 后变成完全均衡的 $(2,2,2,2)$。

规模化的工程细节——直方图估计：真实训练时一个 batch 的 margin 数量是百万级、分散在多个 GPU rank 上，直接收集所有 margin 算精确分位数不现实。做法是每个专家维护一个 margin 的直方图，各 rank 只需要 all-reduce 把直方图 bin 计数加起来（计数可加，跟 token 怎么分片无关），再从合并后的直方图近似恢复分位数——每个专家的通信量只有几百个 bin，跟 batch 大小无关。

因果性：新算出的偏置只在下一步生效，当前这批数据永远不会用"从它自己算出来的偏置"来路由自己，避免信息泄漏。推理时这个偏置直接冻结，不再更新。

十、小结

这一篇覆盖了 K3 架构三条轴里剩下的两条：

• AttnRes：把"层间残差连接"从"只传上一层"改成"对所有已完成层/block 做一次带权 attention 检索"，用伪 query + RMSNorm 软最大核实现，K3 落地为 8 block（每 block 12 层）的 Block 版本，把内存开销从 $O(Ld)$ 压到 $O(Nd)$，是深度维度上照搬"attention 替代 RNN 式累加"这套方法论的又一次应用。
• Stable LatentMoE：把 896 个路由专家的"极端稀疏"带来的两个失效模式（数值爆炸、负载失衡）拆开逐个解决——RMSNorm 和 SiTU-GLU 分别在聚合后和专家内部堵住数值爆炸，Quantile Balancing 用一步分位数求解替代传统的固定步长符号更新，解决负载均衡失灵。

三个模块背后有一条共同的设计哲学：遇到"简单规则在大规模下失灵"，K3 反复采用的策略是换成一个更精确、有理论保证的替代规则，而不是简单调参——AttnRes 用检索取代盲目累加，Quantile Balancing 用分位数求解取代符号步长，SiTU-GLU 用有硬上界的软天花板取代无界激活，都是同一条思路在不同模块的具体实现。

参考资料

• Kimi Team. Kimi K3: Open Frontier Intelligence. arXiv:2607.24653
• MoonshotAI/Kimi-K3 — 模型卡与权重
• Dai et al. DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models
• Vaswani et al. Attention Is All You Need — Full Attention Residuals 借鉴的方法论来源