超越稠密 Transformer：MoE、状态空间模型与混合架构

第 8 章 把稠密块作为默认形态来讨论。这个默认形态让每个词元都走过每一个参数，也因此带来两项前沿构建者迟早会遇到的限制。第一项限制在前馈块里：容量和每词元计算量绑定在一起，想增加容量，就只能增加计算。第二项限制在注意力里：让每个词元与其他每个词元混合，代价是序列长度的平方，还会留下一个无界增长的 KV 缓存。混合专家（MoE）层把词元路由到众多专家中的少数几个，用来突破第一项限制。状态空间模型（SSM）用线性递归换掉注意力的二次方混合，用来缓解第二项限制。混合架构则把两种做法放进同一个堆叠。稀疏化这笔账只有在路由器做对时才划算：参数量要以很低的每词元计算成本增长，专家不能长期拿不到训练信号，负载不能集中到少数专家，通信也不能反过来抵消省下的计算。

稠密块的两项限制

稠密前馈块让每个词元都走过它的全部参数。这把前沿构建者本想分开的两件事耦合住了：模型的容量，即它能知道多少；以及每词元计算量，即每一次前向与反向步骤的代价。耦合之下，想增加知识就只能增加 FLOPs，而 FLOPs 正是 第 5 章 要分配的预算。一个想要的容量超出算力预算所能负担的模型，在稠密设计里没有直接办法。这就是容量限制。

第二项限制来自注意力。自注意力让每个词元与其他每个词元混合，代价随序列长度的平方增长，而它留下的 KV 缓存随上下文线性增长，主导着服务时的内存（第 31 章）。对长上下文来说，二次方成本就是主要限制。所以稠密 Transformer 同时背着两项低效：前馈块在每词元未必用得上的容量上消耗算力，注意力块在一种可能比任务所需更稠密的混合模式上消耗算力与内存。接下来依次处理这两项限制。MoE 处理容量，状态空间模型处理上下文成本，混合体则同时保留两种思路。

突破容量限制：混合专家

把参数与 FLOPs 解耦

混合专家的答案，是把参数与活跃 FLOPs 解耦。一个 MoE 层持有 $N$ 个专家前馈网络（FFN），却只把每个词元路由到其中 $k$ 个，且 $k \ll N$。关键的恒等关系在这里：总参数决定容量与知识，活跃参数，即被选中的这 $k$ 个专家加上模型里始终在线的部分，决定每词元的计算量与单步成本。相对于 FLOPs，容量变得廉价，这正是 MoE 之所以吸引前沿的全部理由。同样的算术也削减推断 FLOPs，不过服务一个 MoE 自有一套问题，不在此处讨论。图 9.1 把这种分离画了出来：随着专家数 $N$ 增长，总参数不断攀升，而决定每词元 FLOPs 的活跃参数量，因 $k$ 固定而保持平坦。

扫一遍专家数 $N$，看总容量攀升，而决定每词元 FLOPs 的活跃参数几乎保持平坦。

import numpy as np, matplotlib.pyplot as plt

# MoE 把总参数（容量）与激活参数（每词元 FLOPs）解耦。
d, ffn = 4096, 14336           # 隐藏维度与单个专家的 FFN 宽度
k = 2                          # 每个词元激活的专家数
per_expert = 2 * d * ffn       # 单个专家的上投影 + 下投影
N = np.arange(1, 65)           # 这一层的专家数
total  = N * per_expert                    # 全部专家都要存储
active = k * per_expert + N * d            # k 个激活专家 + 路由映射

plt.figure(figsize=(5, 3))
plt.plot(N, total / 1e9,  label="总参数（容量）")
plt.plot(N, active / 1e9, label="激活参数（每词元 FLOPs）")
plt.xlabel("专家数 N"); plt.ylabel("参数量（十亿）"); plt.legend()
plt.tight_layout(); plt.show()

print(f"N=64, k={k}: 总计 {total[-1]/1e9:.1f}B，但激活仅 {active[-1]/1e9:.2f}B")
print(f"容量/计算比：{total[-1]/active[-1]:.0f}x")

新增的学习对象是路由器，也叫门控。它是一个小型线性映射，把一个词元的隐藏状态映射成每个专家一个得分。softmax 或 sigmoid 再把得分转成权重，选出 top-$k$ 个专家，其中 $k$ 是每个词元实际调用的专家数；用重归一化后的门控值加权组合它们的输出。写成草图就是：

# 路由器：把 h 线性映射到 N 个专家的得分
scores = h @ W_gate                  # (词元数, N)
weights = softmax(scores)            # 或逐专家 sigmoid
top = topk(weights, k)               # 专家索引与门控值
y = sum(g_i * Expert_i(h) for i, g_i in top)

这个微小的映射是唯一真正新增的组件，也几乎是所有 MoE 毛病的源头。设计里其余的全部，都是为了约束路由器的行为。图 9.2 追踪一个词元穿过该层：门控为所有专家打分，选出 top-$k$ 并组合，而一个共享专家（一项后来的改进）在旁路上始终在线。

路由器及其毛病

路由器最先出问题的地方，是均衡。放任不管，门控会在一个赢者通吃的回路里集中到少数几个受偏爱的专家上：被选中的专家拿到更多梯度，变得更好，于是被选得更多，其余专家则得不到训练信号，沦为死参数。修复的办法有两类。第一类是辅助负载均衡损失，出自 Shazeer 等（2017）(Shazeer et al. 2017)，又在 GShard (Lepikhin et al. 2020) 和 Switch (Fedus et al. 2021) 里得到改进。它加一个惩罚项，惩罚「路由到每个专家的词元比例」与「落在其上的门控质量」的乘积，把分布推向均匀。这有效，可它注入了一个损失，损失的权重让均衡与任务质量之间互相拉扯。第二类更新，是 DeepSeek-V3 的无辅助损失（loss-free）、基于偏置的均衡 (DeepSeek-AI 2024)，直接把辅助损失去掉。它改为给每个专家维护一个路由偏置，两步之间对用得不够的上调、对用得过头的下调，不引入干扰梯度就把选择引导过去 (Wang et al. 2024)。这样就绕开了辅助损失带来的均衡与质量之争。

另外两项改进重新定义了「专家」是什么。细粒度专家把每个专家切成更小的，并按比例提高 $k$，于是相同的活跃 FLOPs 能换来多得多的路由组合，专业化也更锐利。共享专家则是少数几个每个词元都必经的专家，它们吸收常见的、与领域无关的通用计算，让被路由的专家可以专注于专业化，而不必各自把基础内容重学一遍。DeepSeekMoE 把两者结合起来 (Dai et al. 2024)，如今这一对已是大型开放 MoE 模型的常见默认配置。

有一项系统层面的事实，会打破这幅本来清楚的图景。路由是动态的，硬件却想要固定形状的缓冲区，于是每个专家被分配一个容量上限，大致是 capacity_factor * (tokens / experts)。路由到一个已满专家的词元会被丢弃，要么经残差直接传过去，要么改投次优选择。容量系数是一个直接的控制参数：调低省内存、省通信，但丢弃更多词元，损失算力并增加训练噪声；调高则在未满的专家上浪费填充。

这一层如何走到今天

稀疏专家作为一个泛泛的想法，早于 Transformer 时代，但此处要紧的这条线始于 Shazeer 等（2017）(Shazeer et al. 2017)。他们在一个 LSTM 语言模型内部引入了 top-$k$ 门控与辅助均衡损失，其中 $k$ 表示每个词元被送往多少个专家。GShard（Lepikhin 等，2020）(Lepikhin et al. 2020) 把这一层带进 Transformer，也带进了跨设备的自动分片。Switch Transformers（Fedus 等，2021）(Fedus et al. 2021) 把路由器推到最廉价的 $k=1$，并推到万亿参数规模，表明 top-1 路由配上恰当的技巧也能稳定训练。GLaM（Du 等，2022）(Du et al. 2022) 在规模上把效率论据立了起来，Mixtral（Jiang 等，2024）(Jiang et al. 2024) 则把一个强劲的开放 $k=2$ 模型带入了广泛使用。

随后有两条线继续打磨这一层。在路由上，专家选择（Zhou 等，2022）(Zhou et al. 2022) 把问题反了过来：不是每个词元挑自己的专家，而是每个专家挑自己的词元，这从构造上让每专家负载精确，消除了丢弃，代价是有些词元一个专家都分不到。BASE Layers（Lewis 等，2021）(Lewis et al. 2021) 干脆把路由表述成一个均衡指派问题。在专业化上，DeepSeekMoE（Dai 等，2024）(Dai et al. 2024) 引入了细粒度专家与共享专家，DeepSeek-V3（2024）(DeepSeek-AI 2024) 在前沿规模上用基于偏置的均衡换掉了辅助损失，而这一独立方法由 Wang 等（2024）(Wang et al. 2024) 单独做了记录。

还有一条平行的线，让 MoE 不必为一次从零训练付费就能用上。稀疏升级改造（Komatsuzaki 等，2022）(Komatsuzaki et al. 2022) 把每个专家初始化为一个已训练稠密前馈块的副本，加上一个全新的路由器，再继续训练；随着路由把专家逐渐专业化，它们也逐渐分化。它用从零训练的一小部分算力，换来强劲的质量，其失效模式是：路由压力太弱时，专家可能始终几乎相同，于是只剩 MoE 的成本，没有它的收益。

在实践中稳住路由器

路由器必须稳住，它的不稳定有一个具体的成因。top-$k$ 选择是一个 argmax，也就是挑出得分最高的那几个专家，而 argmax 不连续：得分的一个微小变化，就会翻转一个词元的专家，让损失地形变得崎岖、容易出现尖峰，低精度下更糟。标准的缓解手段就那么几种，值得逐一说清。来自 ST-MoE（Zoph 等，2022）(Zoph et al. 2022) 的路由器 z-loss 惩罚过大的路由器 logits，让门控在数值上保持规矩；注意它有别于 第 5 章 稳定性工具包里的输出 logit z-loss。在 bf16 或 fp8 的主体下让路由器保持 fp32（这些都是数值精度格式，位数越少越省内存，但运算也越粗糙，详见 第 10 章），代价很小，却消掉了一类尖峰。还有，早期的路由决策可能在专家尚未分化前就被锁死，所以一次训练的最初阶段值得盯着这一点。同规模下，MoE 的损失曲线比稠密曲线更容易出现尖峰；一个无防护的路由器，既没有 z-loss、又用 bf16 logits，可能出现尖峰，并在训练数天之后仍恢复不过来。

几种失效模式值得逐一说明，因为它们各在不同时刻发作。路由集中把大多数词元导向少数几个专家，让其余沦为死参数；这正是负载均衡要预防的事，一旦均衡权重失当或被关掉，它就回来。负载不均衡与词元丢弃还没到集中失效的程度，但会超过某些专家的容量，悄悄丢掉这些词元的计算并增加噪声，还让最慢专家所在的设备停滞，这是一项要在 第 10 章 付出的系统成本。专家训练不足让很少被选中的专家见到的词元太少，无法专业化，正是升级改造在路由压力太弱时的风险所在。

什么时候划算

• 稀疏对比稠密。 MoE 在每个训练 FLOP、每个活跃 FLOP 上获得更多容量，但要付出代价：内存，因为所有专家都要存储与分片；通信，因为专家并行需要一次 all-to-all（每台设备都要与其他每台设备交换数据，第 10 章）；脆弱性；以及更棘手的服务问题。小规模或内存紧张时，稠密胜出。随着总规模增长，交叉点转向有利于 MoE。
• top-1 对比 top-k。 这里 $k$ 表示每个词元实际调用的专家数。$k=1$ 的路由器最廉价、最简单，但更脆弱、表达力也更弱 (Fedus et al. 2021)；$k=2$ 是常见的最佳折中 (Jiang et al. 2024)；再高的 $k$ 逼近稠密成本，把优势侵蚀掉。
• 容量系数。 越低越便宜，但丢弃更多词元、增加噪声；越高则浪费填充。没有单一正确值，且常在评估与推断时调高，因为那里丢弃一个词元是不可接受的。

突破上下文限制：状态空间模型

用固定状态替代可寻址的过去

状态空间模型回答的是另一项限制，即注意力的二次方成本。一个 SSM 维护一个固定大小的隐藏状态，通过线性递归逐词元更新它，每一步的状态都喂给下一步，再从中读出一个输出：

$$h_t = A\, h_{t-1} + B\, x_t, \qquad y_t = C\, h_t$$

其中 $x_t$ 是当前词元的输入，$h_{t-1}$ 是此前全部上下文压成的状态，$h_t$ 是更新后的状态，$y_t$ 是读出的输出。矩阵 $A$ 决定旧状态保留多少，$B$ 决定当前输入写入多少，$C$ 决定从状态里读出什么。由于状态大小固定，代价随序列长度线性增长，也没有不断增长的 KV 缓存要服务，只有一个常数大小的状态。代价是，一个词元的全部过去都必须被概括进那个固定状态，而注意力让每一个过去词元都保持可寻址。设计上的问题是：一个学得的、依赖输入的递归，能不能把过去压缩得足够好，去与之竞争。Mamba 的贡献，是让递归具有选择性，让 $A$、$B$、$C$ 依赖输入，使模型能挑选保留什么、遗忘什么，在保持线性的同时，把基于内容的注意力所做的大部分工作恢复回来 (Gu and Dao 2023)。

从 S4 到 Mamba

状态空间这条线作为一个认真的注意力竞争者，要年轻得多。结构化状态空间模型 S4（Gu 等，2021）(Gu et al. 2021) 表明，一个精心参数化的线性递归能处理很长的序列；Mamba（Gu 与 Dao，2023）(Gu and Dao 2023) 让递归具有选择性、又对硬件高效，把它带到了规模化语言建模够得着的地方。Mamba-2（Dao 与 Gu，2024）(Dao and Gu 2024) 借一种对偶，把状态空间模型与注意力重新联系起来，让它们得以复用注意力的硬件。

Mamba 不是唯一的线性路线。线性注意力从另一头收敛到同一形态：把注意力里的 softmax 去掉，过去就能折进一个固定大小的运行状态，像 SSM 一样逐词元更新。Gated DeltaNet 是当前的改良，在 Mamba-2 的递归上加了一条门控遗忘规则 (Yang et al. 2024)；MiniMax 的 lightning attention 则是同一家族在生产规模上的实现。多年来，对这一整类模型的现实反对意见都落在工具链上：递归要快，得有一个对硬件高效的扫描，而这条路远不如稠密 Transformer 那条成熟。让这条反对意见过时的，正是下文的混合系统。

召回是主要难点

注意力让每一个过去词元都精确可寻址，擅长精确召回，代价是二次方成本和不断增长的 KV 缓存。一个 SSM 在长度上是线性的，带一个常数大小的状态，但必须压缩过去，在需要精确长程查找的任务上往往落后于注意力。图 9.4 是这一取舍的成本一侧：随着上下文增长，两条曲线分道扬镳，正是这一点，让召回的代价在长上下文下值得付出。下一节要谈的混合体之所以存在，恰恰是因为两种纯形式都不占绝对上风。

二者皆不独行：混合堆叠

实践中，前沿并不必须二选一。混合堆叠把少数全注意力层与多数 SSM 或线性递归层交错排布，于是少数几个全局混合层保住精确召回，而线性层廉价地承载序列的主体代价。MoE 与 SSM 也是正交的：一个混合体可以同时把前馈块放在一个 MoE 路由器之后，把序列混合放在一个状态空间递归之后。Jamba（Lieber 等，2024）(Lieber et al. 2024) 交错排布了 Mamba 与注意力层，又在其上加了 MoE，演示出这三种思想能组进同一个堆叠。Jamba 演示过的组合，2025 年这一代已经规模化上线：MiniMax-01 在 2025 年 1 月以 456B 总参数交错排布 lightning attention 与 softmax 注意力 (MiniMax 2025)；Qwen3-Next 以大约三层线性对一层注意力的比例，把 Gated DeltaNet 与门控全注意力搭配起来 (Qwen Team 2025)；Nemotron-H 用 Mamba 层做了同样的事 (NVIDIA 2025)。图 9.6 展示了这三种思想的组合：序列混合的位置在 SSM 与注意力之间交替，而前馈的位置是一个 MoE 层。

混合体里那少数几个注意力层仍带着一个 KV 缓存，所以混合体并没有完全逃开注意力的代价；它只是把这些代价付在少数层上，而非全部层上。图 9.7 把本章的几条线索铺成一张演进图：一条路由线、一条专业化线、升级改造这条捷径，以及独立的状态空间线，全部汇聚于混合堆叠。

有两项取舍贯穿全章，也框定了本书其余部分对这些架构的看法。

• 注意力对比状态空间。 注意力让每一个过去词元都精确可寻址，擅长精确召回，代价是二次方成本与不断增长的 KV 缓存。一个 SSM 在长度上是线性的，带一个常数大小的状态，但必须压缩过去，在需要精确长程查找的任务上往往落后于注意力。混合体之所以存在，正是因为两种纯形式都不占绝对上风 (Lieber et al. 2024)。
• 训练与服务复杂度。 这里的每一项选择都增加运维表面：一个要调的路由器、要监控的均衡、要排布的专家放置，还有一份 第 10 章 必须吸收的内存与通信画像。质量上的收获是实打实的，但工程成本也同样具体。

延伸阅读