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

第 8 章 介绍的因果解码器包含两个稠密子层：注意力负责在不同位置之间混合信息，前馈网络（FFN）负责逐位置变换。本章沿两条相互独立的轴改变这两个子层。

混合专家（MoE）FFN 会存储多个专家网络，但每个词元只执行其中一部分。MoE 改变参数激活方式，并不会让注意力变稀疏。状态空间层或线性递归层改变序列混合方式，以及解码时保留的状态，并不会自行加入专家路由。混合架构可以只采用其中一种改动，也可以在同一堆叠中同时采用两种。把两条轴分开，才能避免混淆三个量：存储参数、单个词元实际执行的算术量，以及一条序列需要保留的状态。

两条相互独立的架构轴

在稠密 Transformer 块中，每个词元都会执行同一套 FFN 权重，每个注意力层也会为所有缓存位置保留键和值。这是一种具体设计，而不是语言模型的定义。替代架构分别回答两个问题：

替换稠密部件也会引入新成本。MoE 路由需要分发词元，并处理负载不均衡与通信。递归会把历史压缩进一个由架构固定宽度的状态，因此过去的位置不再以不断增长的键值缓存逐个保存。比较这些设计时，必须看完整工作负载，包括质量、训练算术量、内存、通信、批次形状和服务延迟。

MoE 让词元有条件地通过不同 FFN

稀疏专家模型早于 Transformer。Shazeer 等人 2017 年在循环语言模型中大规模使用了稀疏门控 MoE 层 (Shazeer et al. 2017)。GShard 随后把条件专家与 Transformer 和自动分片结合起来 (Lepikhin et al. 2020)，Switch Transformer 则研究了 top-1 路由和万亿参数模型 (Fedus et al. 2022)。这些工作的核心思想比规模纪录更具体：模型可以存储多于单个词元实际执行的专家权重。

路由是一项定义明确的数学运算

Mixtral 使用的一种常见 top-$k$ 形式，先计算路由分数，再从中选出专家 (Jiang et al. 2024)。这里的两步运算可以写成：

$$\begin{aligned} z_t &= W_rh_t, \\ \mathcal{S}_t &= \operatorname{TopK}(z_t,k). \end{aligned}$$

其中，每个入选专家的归一化门控权重为

$$g_{t,e}= \begin{cases} \dfrac{\exp(z_{t,e})} {\sum_{j\in\mathcal{S}_t}\exp(z_{t,j})}, & e\in\mathcal{S}_t,\\[6pt] 0, & e\notin\mathcal{S}_t. \end{cases}$$

MoE 输出是入选专家输出的加权和：

$$m_t=\sum_{e=1}^{E}g_{t,e}F_e(h_t).$$

其中，$h_t\in\mathbb{R}^{d}$ 是词元位置 $t$ 的隐藏向量，$E$ 是路由专家数，$W_r\in\mathbb{R}^{E\times d}$ 是路由矩阵，$z_{t,e}$ 是专家 $e$ 的路由 logit，$\mathcal{S}_t$ 是 $k$ 个入选专家的索引集合，$g_{t,e}$ 是归一化门控权重，$F_e:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$ 是专家网络，$m_t\in\mathbb{R}^{d}$ 是 MoE 输出。残差连接和任何共享专家都不包含在这条公式中。

这只是一个路由契约，并不是 MoE 的唯一定义。Switch 使用 top-1 路由。有些系统在选择前做 softmax，有些在选择后做，还有一些使用 sigmoid 亲和度。top-$k$ 之后的门控值也可能重归一化，也可能保持原值。这些选择会改变梯度和检查点语义，模型说明必须把它们写清楚。

图中的共享专家不经过路由门控。它是否存在、如何与路由专家组合，都属于架构契约。

存储参数、选中权重与运行时间不是一回事

假设每个专家都是无偏置的 SwiGLU FFN，模型宽度为 $d$，专家隐藏宽度为 $d_f$。根据 第 8 章，一个专家大约包含

$$P_e=3dd_f$$

个参数。这里 $P_e$ 包含门控投影、值投影和下投影。若有 $E$ 个路由专家且没有共享专家，该层存储和执行的参数量分别为

$$\begin{aligned} P_{\mathrm{stored}} &= EP_e+Ed, \\ P_{\mathrm{evaluated}}(t) &= kP_e+Ed. \end{aligned}$$

其中，$Ed$ 是路由矩阵。每个词元都要计算全部 $E$ 个路由分数；$EP_e$ 包含所有已存储专家，$kP_e$ 只包含入选专家。这里省略了偏置和其他始终执行的块参数。文献中的“活跃参数”还可能包含嵌入、注意力、归一化和共享专家，因此报告时必须说明计数口径。

图 9.2 按 $d=4096$、$d_f=14336$ 和 $k=2$ 计算这些公式。随着 $E$ 增加，入选专家项不变，路由器项却会增长。图中的两条线都不能直接当作实际运行时间，因为运行时间还取决于词元分发、填充或非齐整内核、all-to-all 通信、内存流量，以及最慢的专家分片。

下面的可运行单元使用同一套公式，不依赖绘图库。可以修改 $E$、$k$ 或两个宽度，检查其他层配置。

model_width = 4096
expert_width = 14_336
selected_experts = 2

def moe_parameter_counts(experts):
    if not 1 <= selected_experts <= experts:
        raise ValueError("selected_experts 必须介于 1 和 experts 之间")
    per_expert = 3 * model_width * expert_width  # 无偏置 SwiGLU
    router = experts * model_width
    stored = experts * per_expert + router
    evaluated = selected_experts * per_expert + router
    return stored, evaluated

for experts in (8, 16, 32, 64):
    stored, evaluated = moe_parameter_counts(experts)
    print(
        f"E={experts:2d}, k={selected_experts}: "
        f"存储={stored / 1e9:.3f}B, "
        f"每词元执行={evaluated / 1e9:.3f}B, "
        f"比值={stored / evaluated:.2f}x"
    )

路由会产生分发问题

专家权重通常分布在一个专家并行组的不同设备上。一次前向传播因此包含四个阶段：计算路由分数；按目标专家整理词元副本；把副本发送给持有相应专家的设备；再把专家输出恢复到原来的词元顺序。常见实现会使用两次 all-to-all，一次用于分发，另一次用于返回。不过，当所有专家都在本地，或者采用其他放置方式时，并不一定需要 all-to-all。因此，all-to-all 并不是 MoE 的固有要求。

专家 FFN 省下的算术量，可能被网络时间或规模很小、利用率很低的矩阵乘法抵消。应测量的量包括分发字节数、各专家词元数、填充槽与有效槽、集合通信时间、通信与计算的重叠，以及尾延迟。这些指标都不能由“活跃参数”单独推出。

假设一个路由组包含 $T$ 个词元、$E$ 个专家，每个词元产生 $k$ 次专家指派。均匀负载是每个专家 $kT/E$ 次指派。固定容量的实现可以为每个专家预留

$$C_e=\left\lceil c\frac{kT}{E}\right\rceil$$

个槽位。这里 $c>0$ 是容量系数，$C_e$ 是专家 $e$ 能接收的指派数，$kT$ 是路由组中的指派总数。超出容量的指派可以跳过、改路由或延后。Switch 论文给出的公式是 $k=1$ 的情形 (Fedus et al. 2022)。另一些系统使用无丢弃的非齐整或块稀疏执行，不设置固定容量上限 (Gale et al. 2023)。

均衡损失改变优化过程，而不只是利用率

不加约束的路由可能把指派集中到少数专家，使其他专家训练不足，并让负载最重的专家设备成为慢节点。Shazeer 等人使用了相互独立的重要性目标和负载目标 (Shazeer et al. 2017)。Switch 为 top-1 路由提出了更简单的辅助损失：

$$\begin{aligned} \mathcal{L}_{\mathrm{bal}} &=\alpha E\sum_{e=1}^{E}f_eq_e, \\ f_e &= \frac{1}{T}\sum_{t=1}^{T} \mathbf{1}\!\left[\arg\max_j p_{t,j}=e\right], \\ q_e &= \frac{1}{T}\sum_{t=1}^{T}p_{t,e}. \end{aligned}$$

其中，$p_t=\operatorname{softmax}(z_t)$ 是词元 $t$ 的完整路由分布，$f_e$ 是实际送到专家 $e$ 的词元比例，$q_e$ 是该专家的平均路由概率，$T$ 是路由组中的词元数，$E$ 是专家数，$\alpha\ge0$ 是损失系数，$\mathbf{1}[\cdot]$ 是指示函数。乘积会惩罚同时拥有高流量和高概率质量的专家。Top-$k$ 系统会采用相关但并不完全相同的定义。

系数本身就是优化选择。均衡压力太弱会造成利用率低下，太强又可能压过由任务驱动的路由。DeepSeek-V3 主要通过依据近期负载调整的逐专家偏置来引导 top-$k$ 选择，同时保留一个较小的逐序列辅助损失，以防出现极端不均衡 (DeepSeek-AI 2024; Wang et al. 2024)。偏置影响专家选择，不影响组合专家输出时使用的门控权重。因此，不能据此说模型完全移除了所有均衡损失。

ST-MoE 还提出了路由器 z-loss。这个量用来约束 log-sum-exp 项 (Zoph et al. 2022)：

$$\mathcal{L}_{z} =\frac{\beta}{T}\sum_{t=1}^{T} \left(\log\sum_{e=1}^{E}\exp z_{t,e}\right)^2.$$

其中，$z_{t,e}$ 是前文定义的路由 logit，括号中的 log-sum-exp 是词元 $t$ 的路由配分函数，$\beta\ge0$ 控制惩罚强度。Switch 还报告过在 bfloat16 模型中，只把路由计算保留为 float32 (Fedus et al. 2022)。这些是有实验依据的稳定手段，但不能证明硬 top-$k$ 或低精度是所有 MoE 故障的唯一原因。

不同改进解决不同的路由问题

以下设计是彼此不同的选择，并不构成通往同一个默认方案的单一路线：

• Top-1 与 top-$k$。 Switch 每个词元执行一个专家，Mixtral 执行两个。增大 $k$ 会提高专家算术量和通信量；只有当 $k$ 接近 $E$ 时，才逐渐接近执行所有专家的稠密形式。
• 专家选择与均衡指派。 Expert Choice 让每个专家从固定容量的词元桶中选择输入，因而不会发生专家侧溢出，但每个词元入选的专家数不固定，甚至可能为零 (Zhou et al. 2022)。BASE Layers 则把路由写成均衡指派问题 (Lewis et al. 2021)。
• 细粒度专家与共享专家。 DeepSeekMoE 把专家容量拆成更小的路由专家，并加入始终执行的共享专家 (Dai et al. 2024)。在该设计中，共享专家的输出不经过路由门控便直接相加。论文报告了更好的专业化结果，但“共享”这个名称并不能保证网络只学到了通用知识。
• 稀疏升级改造。 可以把一个稠密 FFN 复制成多个专家，再加入新路由器继续训练 (Komatsuzaki et al. 2023)。这种做法改变的是初始化成本，不会改变最终的分发和服务契约。

GLaM 和 Mixtral 展示了这片设计空间中的其他位置 (Du et al. 2022; Jiang et al. 2024)。它们的实验不能据此确定普遍最优的专家数、$k$、容量系数或路由损失。答案取决于训练预算、批次形状、网络拓扑、内存、内核和目标延迟。

状态空间层用递归替代不断增长的历史记录

第二条轴是序列混合。状态空间模型（SSM）的数学来源是动力系统和控制理论，远早于神经序列模型。S4 在 2022 年展示了结构化线性状态空间层对长序列的建模能力 (Gu et al. 2022)。Mamba 随后让递归中的关键项依赖输入，并提供了面向硬件的选择性扫描实现 (Gu and Dao 2024)。

从连续系统到离散层

线性时不变状态空间系统可以写成

$$\begin{aligned} \frac{dh(t)}{dt} &= Ah(t)+Bx(t), \\ y(t) &= Ch(t)+D_{\mathrm{skip}}x(t). \end{aligned}$$

其中，连续时间为 $t$；$x(t)\in\mathbb{R}^{d_{in}}$ 是输入，$h(t)\in\mathbb{R}^{n}$ 是状态，$y(t)\in\mathbb{R}^{d_{out}}$ 是输出；$A\in\mathbb{R}^{n\times n}$ 是状态转移矩阵，$B\in\mathbb{R}^{n\times d_{in}}$ 把输入写入状态，$C\in\mathbb{R}^{d_{out}\times n}$ 从状态读出输出，$D_{\mathrm{skip}}\in\mathbb{R}^{d_{out}\times d_{in}}$ 是输入直连路径。$d_{in}$、$n$ 和 $d_{out}$ 都由层配置固定。

按固定步长离散化后，同一层可以写成

$$\begin{aligned} h_t &= \bar A h_{t-1}+\bar Bx_t, \\ y_t &= Ch_t+D_{\mathrm{skip}}x_t. \end{aligned}$$

这里整数 $t$ 表示词元位置，$x_t$、$h_t$ 和 $y_t$ 是采样后的输入、状态和输出，$\bar A$ 与 $\bar B$ 是离散化后的转移矩阵和输入矩阵。即使实现从连续参数中推导它们，它们也不等同于连续矩阵 $A$ 和 $B$。

固定系数使递归与卷积等价

设 $h_0=0$，并且 $\bar A$、$\bar B$ 和 $C$ 固定，展开递归可得

$$\begin{aligned} K_i &= C\bar A^{i}\bar B, \\ y_t &= \sum_{j=1}^{t}K_{t-j}x_j +D_{\mathrm{skip}}x_t. \end{aligned}$$

其中，$i\ge0$ 是滞后量，$K_i\in\mathbb{R}^{d_{out}\times d_{in}}$ 是该滞后上的卷积核，$\bar A^i$ 是矩阵的第 $i$ 次幂，$j$ 是较早的词元位置。递归视角和卷积视角计算的是同一个线性算子。S4 利用 $A$ 的结构高效生成并应用长卷积核，但它仍有滤波器生成和 FFT 成本，不能笼统地把所有 SSM 实现都说成简单的 $O(S)$。

Mamba 让递归具有选择性

对一个简化的 Mamba 通道，选择机制从当前输入生成三个随词元变化的项：

$$\begin{aligned} (\Delta_t,B_t,C_t) &= s_\theta(x_t), \\ \bar A_t &= \exp(\Delta_tA), \\ \bar B_t &= \int_{0}^{\Delta_t} \exp\!\bigl((\Delta_t-\tau)A\bigr)B_t\,d\tau. \end{aligned}$$

这些项定义离散递归：

$$\begin{aligned} h_t &= \bar A_th_{t-1}+\bar B_tx_t, \\ y_t &= C_th_t+D_{\mathrm{skip}}x_t. \end{aligned}$$

其中，$s_\theta$ 是参数为 $\theta$ 的学习投影，$\Delta_t>0$ 是依赖输入的离散步长，$B_t$ 和 $C_t$ 是依赖输入的写入与读出映射，$A$ 是学习得到但不依赖输入的连续转移矩阵，$\exp$ 是矩阵指数，$\tau$ 是积分变量。也就是说，$\Delta_t$、$B_t$ 和 $C_t$ 随输入变化，连续转移矩阵 $A$ 不随输入变化 (Gu and Dao 2024)。

输入依赖性意味着上一节的单一平稳卷积核 $K$ 不再存在。Mamba 在训练时对仿射递归更新执行融合的结合扫描，从而恢复并行性；自回归解码时则使用递归形式。Mamba-2 把一类受限 SSM 转移与半可分混合矩阵联系起来，并使用分块矩阵乘法。论文报告其核心层在测试设置中比 Mamba 的选择性扫描快 2 到 8 倍 (Dao and Gu 2024)。这是特定实现和测试条件下的实测结果，不是与架构无关的速度比例。

线性注意力也是一条固定状态路线

核化因果注意力也可以重排计算，得到递归状态。一种简化的门控形式为

$$\begin{aligned} S_t &= \lambda_tS_{t-1}+\phi(k_t)v_t^{\top}, \\ z_t &= \lambda_tz_{t-1}+\phi(k_t), \\ y_t &= \frac{\phi(q_t)^{\top}S_t} {\phi(q_t)^{\top}z_t+\varepsilon}. \end{aligned}$$

其中，$q_t$、$k_t$ 和 $v_t$ 是查询、键和值向量；$\phi$ 是使核可分解的特征映射；$S_t$ 是固定形状的键值汇总矩阵；$z_t$ 是归一化状态；$\lambda_t\in[0,1]$ 是可选的学习遗忘门；$\varepsilon>0$ 防止除零。消除所有查询与键两两配对的关键是结合律分解，而不只是去掉 softmax。Gated DeltaNet 在这一家族中加入了有针对性的 delta 规则更新 (Yang et al. 2025)。

线性序列扩展不等于自动提速

固定模型维度时，完整序列上的稠密注意力会形成与 $S^2$ 成正比的查询键关系，而递归扫描执行与 $S$ 成正比的状态更新。图 9.4 只画出以 1,000 个词元归一化后的序列长度增长项。它不是 FLOPs 或实测运行时间曲线，投影宽度、状态大小、内核融合、内存流量和硬件共同决定常数项。

解码时，稠密注意力层会让一个新查询与每个已缓存位置的键比较。递归层则更新一个形状不随上下文长度变化的状态。“不随上下文长度变化”并不表示状态很小或没有成本。对原始 Mamba 实现，持久解码状态可以近似写成

$$\begin{aligned} M_{\mathrm{rec}} &=B_{\mathrm{batch}}L_{\mathrm{rec}}d_{\mathrm{inner}} \\ &\quad\times \bigl(n_{\mathrm{state}}+d_{\mathrm{conv}}\bigr)b \quad\text{字节}. \end{aligned}$$

其中，$B_{\mathrm{batch}}$ 是驻留序列数，$L_{\mathrm{rec}}$ 是递归层数，$d_{\mathrm{inner}}$ 是递归通道数，$n_{\mathrm{state}}$ 是每个通道的 SSM 状态宽度，$d_{\mathrm{conv}}$ 是短卷积缓冲区长度，$b$ 是每个存储元素的字节数。精确布局因实现而异。状态会随批大小、层数、通道宽度、状态宽度、卷积宽度和精度增长，但不会随已缓存的上下文长度增长。

代价体现在表示方式上。注意力为每个过去位置保留独立的缓存表示，并允许按内容直接访问。递归层反复把先前信息压缩进固定形状的状态，信息可能相互干扰或丢失。两种机制都不保证召回。必须在目标分布、目标长度和匹配的训练条件下测量检索质量。

混合架构选择层级排布

MoE 与递归序列混合可以组合，因为它们占据块中的不同位置。一个混合架构可以在部分层使用注意力，在其他层使用递归；每个混合器之后又可以接稠密 FFN 或专家 FFN。不存在固定比例。

已经发布的模型展示了不同的排布和参数报告方式：

这些例子证明各部件能够在大规模下训练和实现，但不能证明混合架构在每一种质量指标和服务工作负载上都优于完整注意力。MiniMax 后来的 M2 改用完整注意力，并报告混合架构在更大规模的复杂多跳任务上暴露出弱点；低精度状态、前缀缓存、推测解码、内核成熟度和评测也仍是实际问题 (Sun 2025)。这是一个团队的第一方证据，并不是反对递归或混合模型的普遍结论。

把稀疏与递归架构记录成一份契约

一份可复现的架构说明至少应写明：

• 哪些层使用注意力、SSM、线性注意力或其他混合器；
• 每种注意力的掩码、头布局、位置规则和 KV 缓存数据类型；
• 递归状态形状、离散化或门控参数化、卷积宽度、重置语义和状态精度；
• 专家数、共享专家数、入选路由专家数、专家 FFN 形状，以及哪些层使用稀疏 FFN；
• 路由打分、归一化、top-$k$ 平局处理、门控重归一化和输出组合规则；
• 均衡损失和 z-loss 的方程与系数，以及任何偏置更新规则；
• 路由组大小、容量系数、溢出或无丢弃策略，以及组合后的词元顺序；
• 专家放置、专家并行组、分发集合通信、路由器精度，以及总参数和激活参数的报告口径。

“MoE”“Mamba 混合架构”或“80B-A3B”这类名称并不能确定上述大部分字段。检查点张量、训练代码、推理内核和模型卡必须遵守同一份契约。

在扩大规模前分别验证两条轴

1. 核对参数量。 把存储参数、入选参数、共享参数、路由参数和始终执行的参数与序列化检查点逐项核对。
2. 精确测试路由。 用小批次比较分片后的分发与组合和单设备参考实现，覆盖平局、溢出、填充和零词元专家。
3. 测量负载分布。 在训练数据和评测数据上报告每个专家的指派数、不均衡程度、丢弃或改路由的指派、填充、熵，以及最慢的专家设备。
4. 测量通信。 在真实批大小下记录分发与组合字节数、集合通信时间、计算重叠和端到端每秒词元数。
5. 比较扫描与递归。 确认并行训练路径和逐词元递归路径在声明的数值容差内一致。
6. 测试状态边界。 在文档和序列边界重置状态，并测试打包样本、前缀复用、状态复制和推测解码回滚。
7. 压力测试检索与状态跟踪。 在部署长度范围内评测直接查找、顺序、重复键、干扰项和多跳使用，而不是只依赖一次针尖检索测试。
8. 对完整工作负载做基准测试。 在预算匹配的条件下，对照稠密模型和纯注意力模型比较质量、训练吞吐量、预填充、解码、内存、批次容量和故障恢复。

MoE 和递归层分别移除了一个稠密假设，但都没有取消资源核算。真正有用的抽象是一份契约：一个词元执行哪些权重，这些权重位于哪里，序列状态如何演化，以及哪些成本转移到了通信或服务环节。

延伸阅读