Lecture 4:注意力替代方案与混合专家模型

9925 字
50 分钟
Lecture 4:注意力替代方案与混合专家模型
文章摘要

Stanford CS336 Language Modeling from Scratch | Spring 2026 | Lecture 4: Attention Alternatives,时长 01:26:20。

上一讲讨论了标准 Transformer 的各种改进——从基础 Transformer 出发,对其各个组件进行微调以得到现代语言模型。本讲将深入两个更为激进的架构创新:其一是注意力替代方案,即如何通过修改注意力机制(attention block)来实现对序列长度的线性时间依赖,从而支撑百万级别的超长上下文;其二是混合专家模型(Mixture of Experts, MoE),即如何通过修改前馈网络(MLP block),在不增加推理计算量的前提下大幅扩展模型参数量。

注意力的成本问题与优化工具箱#

长上下文的迫切需求#

从工业界模型的发展趋势来看,各大 LLM 厂商正竞相提供越来越大的上下文窗口。如果我们在对数坐标轴上观察不同模型随时间的上下文窗口大小,会发现一条明确的增长曲线——从早期的几千 token 到如今的数百万 token。这背后的驱动力很直接:更大的上下文意味着模型能容纳更多的知识、能作为 agent 操作更多信息。

但问题在于计算成本。如果我们考察注意力操作与前馈操作随序列长度增长的计算占比,会发现一个关键的交叉点:在较短序列长度下,前馈层的计算占比更大(线性增长),但注意力是所有位置之间的全对全连接(quadratic),因此随着序列长度增长,注意力的计算成本会迅速超过前馈层。过去,对于大模型在较短序列长度下,前馈层是主要瓶颈;而当我们走向更长的序列时,注意力越来越成为瓶颈。

控制注意力成本的基本工具#

在进入更激进的替代方案之前,有几种已经被广泛验证的方法可以有效控制注意力成本:

局部注意力 + 混合架构。不必在每一层都使用全局注意力。例如,可以每 8 层才做一次全局注意力,其余层使用局部注意力(只关注附近的 token)。这种混合方式可以显著降低整体注意力成本,同时保留全局信息聚合能力。

Flash Attention(系统级常数因子优化)。这里有一个经常被低估的事实:常数因子极其重要。受过理论训练的人容易只关注大 O 复杂度,但 Flash Attention 恰恰是一个”不改变渐近复杂度,但通过重新组织计算来最小化内存传输开销”的系统优化。它的效果是戏剧性的:在较短序列长度下可以获得 2 倍以上的加速;在更长的序列长度下,标准注意力甚至因为需要物化巨大的注意力矩阵而无法放入内存,而 Flash Attention 因为不物化这些矩阵,至少可以继续运行(虽然速度较慢)。Flash Attention 不解决 quadratic 的问题,但常数因子的威力不容忽视——我们将在后续系统讲座和系统作业中详细讨论它。

然而,当我们的目标是 500 万甚至 1000 万 token 的超长上下文时,这些技巧可能不够。我们需要更根本的突破:是否存在对序列长度线性依赖的注意力替代方案?

线性注意力:从矩阵结合律出发#

核心思想:重排乘法顺序#

要理解近年来涌现的各种线性时间注意力方法,只需要理解一个核心思想——矩阵乘法的结合律(associativity of multiplication)。

标准注意力可以写成紧凑形式:

Attention(Q,K,V)=softmax(QK)V\text{Attention}(Q, K, V) = \text{softmax}(QK^\top) V

其中 Q,KRN×dkQ, K \in \mathbb{R}^{N \times d_k}VRN×dvV \in \mathbb{R}^{N \times d_v}QKQK^\top 这一步产生 N×NN \times N 的注意力矩阵,计算复杂度为 O(N2dk)O(N^2 d_k),当 NN(序列长度)达到数百万时,这个平方项将变得不可承受。

现在,暂时忽略 softmax(即令 softmax 为恒等映射)。如果没有 softmax,注意力变为:

Output=(QK)V\text{Output} = (QK^\top)V

根据矩阵乘法的结合律,我们可以改变括号的位置:

Output=Q(KV)\text{Output} = Q(K^\top V)

这个看似微小的改变,带来了复杂度的根本变化:

  • 原始顺序 (QK)V(QK^\top)V:先计算 QKQK^\top 得到 N×NN \times N 矩阵,复杂度 O(N2dk)O(N^2 d_k)
  • 重排后 Q(KV)Q(K^\top V):先计算 KVK^\top V 得到 dk×dvd_k \times d_v 矩阵,复杂度 O(Ndkdv)O(N d_k d_v)

关键在于:dkd_kdvd_v(key/value 的维度)通常在几千到几万的量级,而 NN(序列长度)可以是数百万。因此,从 N2dkN^2 d_kNdkdvN d_k d_v 的变化意味着从序列长度的平方依赖变为线性依赖

线性注意力的 RNN 对偶形式#

重排后的线性注意力 Q(KV)Q(K^\top V) 还有一个极其优美的性质:它可以等价地写成 RNN 的递推形式

我们定义状态矩阵 StRdk×dvS_t \in \mathbb{R}^{d_k \times d_v},随着从左到右扫描上下文,递推更新如下:

St=St1+ktvtS_t = S_{t-1} + k_t v_t^\top yt=qtSty_t = q_t^\top S_t

其中 kt,qtRdkk_t, q_t \in \mathbb{R}^{d_k}vtRdvv_t \in \mathbb{R}^{d_v}。这本质上是在逐步累积 KV 的外积到状态矩阵中,然后用当前的 query 去检索。

这个对偶性意味着线性注意力具有两种等价的计算形式

  1. 稠密矩阵形式(上面的 Q(KV)Q(K^\top V)):适合训练,可以充分并行化
  2. 递推 RNN 形式:适合推理,状态 StS_t 大小固定(dk×dvd_k \times d_v),不随上下文长度增长

这是两全其美的:训练时用并行的矩阵形式,推理时用高效的 RNN 形式。

线性注意力在工业规模的验证#

尽管线性注意力因丢弃 softmax 而损失了表达能力,但它已经在工业规模上得到了验证。MiniMax M1 是一个颇具代表性的大规模中国开源模型,它采用 7:1 的混合架构——每 7 个线性注意力层搭配 1 个完整的 softmax 注意力层。其性能与 OpenAI 的 o3、DeepSeek R1 等强模型相当,说明即使是最基本的线性注意力,只要与少量 softmax 注意力混合使用,就能在规模化场景下有效运作。

一个重要的观察是:迄今为止,没有人在规模化场景下成功验证了纯线性时间注意力机制。所有实际部署的方案都是混合架构——大部分层使用线性注意力/SSM,少部分层保留 softmax 注意力来维持全局信息聚合能力。

从线性注意力到 Mamba 2:引入遗忘门#

线性注意力的一个核心问题是:它总是无选择地将所有过去的信息向前传递。从 LSTM 的经验我们知道,“知道何时传递信息、何时遗忘信息”是至关重要的。因此,一个自然的改进是引入门控机制。

Mamba 2 正是这种改进的代表。Mamba 系列模型(由 Albert Gu 和 Tri Dao 等人提出,目前已发展到 Mamba 3)最初源自状态空间理论(state space models),但从机制上看,Mamba 2 实际上可以被理解为线性注意力的一个非常简单的扩展:

St=γtSt1+ktvtS_t = \gamma_t \cdot S_{t-1} + k_t v_t^\top yt=qtSt+dtvty_t = q_t^\top S_t + d_t \cdot v_t

其中 γt\gamma_t 是一个遗忘门(forget gate),仅依赖于当前输入 xtx_t(不依赖于状态 StS_t)。γt\gamma_t 控制了上一步状态 St1S_{t-1} 中有多少信息被保留。dtvtd_t \cdot v_t 是一个类似残差连接的项,让当前输入的 value 可以直接传递到输出。

关键约束在于:所有门控项(γt\gamma_t)只依赖于当前输入,不依赖于状态。这保证了 Mamba 2 仍然保持线性注意力的核心性质——训练时可以展开为稠密矩阵并行计算,推理时使用 RNN 递推形式。

NeMo Tron 3 是 Mamba 2 在工业规模应用的一个例子。它将 Mamba 2 作为轻量层与 softmax 注意力交替使用,在性能上与 Qwen 3 Thinking、GPT-OSS 等模型相当,同时在推理时对上下文长度的依赖更为温和。

Gated DeltaNet:最广泛使用的状态空间模型#

在 Mamba 2 的基础上,可以进一步增强递推更新的表达力。Gated DeltaNet 是目前最广泛使用的状态空间模型,已在多篇论文中经过测试验证。它的核心思想可以通过与 Mamba 2 的对比来理解。

双门控 + 投影更新#

Gated DeltaNet 在 Mamba 2 的基础上增加了两个改进:

第一,引入第二个门 βt\beta_t(输入门)βt\beta_t 控制”是否将当前信息写入状态”。当 βt=0\beta_t = 0 时,当前时间步的信息完全不会被写入状态。这与 LSTM 的输入门(input gate)和遗忘门(forget gate)的设计原则非常相似——通过独立的门来控制”遗忘什么”和”写入什么”。

第二,DeltaNet 式的投影更新。Gated DeltaNet 不是简单地做 1βt1 - \beta_t 的衰减,而是在写入新信息时,同时擦除状态中与当前 key 方向相关的旧信息。具体来说,更新规则中包含一个投影项:

St=γt(Iβtktkt)St1+βtktvtS_t = \gamma_t \cdot (I - \beta_t \cdot k_t k_t^\top) \cdot S_{t-1} + \beta_t \cdot k_t v_t^\top

其中 IβtktktI - \beta_t k_t k_t^\top 近似地起到投影算子的作用(严格来说并不是精确的投影,因为没有做单位归一化,但直觉上是正确的),它在写入新的 ktvtk_t v_t^\top 之前,先将状态中 ktk_t 方向上的旧信息投影掉。直觉上,当你要向状态中写入新的 key-value 对时,你不仅要”添加新信息”,还要”清除旧的同方向信息”,避免冲突。

值得注意的是,这种投影更新机制在不同的研究社区中被独立发现了多次:在快速权重编程(fast weight programming)和测试时训练(test-time training)等领域,研究者通过完全不同的设计原则,最终得到了几乎相同的数学形式。这种”殊途同归”暗示了某种深层的最优性。

Gated DeltaNet 是否真正有效?#

Qwen 3.5(以及其前身 Qwen Next)是目前最好的开源模型之一,它们使用的正是 3:1 Gated DeltaNet + 注意力混合架构——每 3 个 Gated DeltaNet 层搭配 1 个 softmax 注意力层。实验数据表明:

  • 性能方面:与闭源模型相比,这种混合架构几乎没有性能损失
  • 推理效率方面:Qwen Next 相比 Qwen 3(纯 Transformer)在长上下文下有显著更高的解码吞吐量

混合比例的实证研究#

ByteDance Seed 和 UC Santa Cruz 的研究者进行了一项较为系统的对照实验,比较了 Mamba 2、Gated DeltaNet 等架构在不同混合比例下的表现。主要发现:

  • 在低比例的 RNN 层替换下(如 3:1 或 7:1),最佳架构(如 Gated DeltaNet)几乎没有性能损失
  • 超过某个临界点后,性能开始显著下降
  • 当完全替换为纯 RNN(不保留任何 softmax 注意力层)时,所有架构都出现了明显的性能退化

这个结论在长上下文检索任务和其他评测基准上都是一致的。一个直觉性的解释是:纯 RNN 需要将所有信息压缩到固定大小的状态中,这在信息论意义上不可避免地会丢失信息。如果状态大小与上下文等长,信息就不会丢失,但那样就失去了效率优势。

关于训练复杂度的问题#

SSM 相比 Transformer 在训练上是否更复杂?Transformer 的 softmax 注意力有一个巨大的优势:它是一个纯粹的全对全矩阵乘法,既强大又易于训练(回想 LSTM 时代,Transformer 的硬件效率优势是其胜出的关键因素之一)。SSM 之所以能在今天获得关注(尽管与 LSTM 有相似的表达能力局限),是因为线性注意力的对偶性——从 RNN 形式到稠密矩阵形式的转换——使得 SSM 可以在训练时保持计算效率。因此,这个权衡已经被”打勾”了。但剩下的权衡是:有限状态带来的信息压缩损失是不可避免的

DeepSeek Attention(DSA):稀疏注意力的另一条路#

除了线性注意力/SSM 路线,还有一种完全不同的思路来降低注意力成本——稀疏注意力。DeepSeek 在其 V3.2 版本中引入的 DSA(DeepSeek Attention)就是这种方法的代表。

DSA 的工作原理#

DSA 的核心思路非常直接:不对所有 token 计算注意力,而是先用一个轻量级索引器(indexer)选出最相关的 token 子集,然后只对这个子集做完整注意力

具体流程如下:

  1. 索引阶段:计算 query 和所有 key 的内积 QKQK^\top,通过 ReLU 激活和一个权重矩阵 WW(来自前序 token 的信息)得到每个位置的激活分数
  2. 选择阶段:对激活分数做 Top-K 选择,挑出最相关的 K 个位置
  3. 注意力阶段:仅对这 K 个位置做标准的 softmax 注意力计算

需要注意的是,DSA 不是线性时间的。索引器仍然需要计算所有 token 之间的内积,本质上是 quadratic 的。但通过多种手段可以使索引器极其轻量:使用更低精度(如 FP4)、更低维度的投影等。而后续的完整注意力虽然也是 quadratic,但作用在 Top-K 选出的小子集上,K 的大小可以控制(通常与短上下文长度相当),因此实际计算成本大大降低。

“非端到端训练”的惊人有效性#

DSA 的另一个重要特点是:不需要从头训练。实际做法是先正常训练一个标准 Transformer,然后在长上下文扩展阶段(long context extension stage)才引入索引器,通过少量继续训练让模型学会使用索引器。这利用了大多数模型本身就有的”短上下文预训练 → 长上下文扩展”的两阶段训练流程,将 DSA 的引入嫁接到已有的第二阶段上。

实验结果令人印象深刻:

  • DeepSeek V3.2 在性能上与 Claude 4.5 Sonnet、Gemini 3 等同期前沿模型相当,同时其预填充(prefill)和解码(decoding)的成本随上下文长度的增长非常缓和
  • GLM 5 独立地采用了相同的 DSA 方法,其消融实验表明在长上下文检索任务上,DSA 相比完整注意力几乎没有性能损失

稀疏注意力 vs 线性注意力#

DSA 代表了一种与 SSM/线性注意力完全正交的降低注意力成本的思路:

  • SSM/线性注意力:通过改变计算结构来实现线性复杂度,代价是丢弃 softmax 带来的表达能力损失
  • DSA/稀疏注意力:保留完整的 softmax 注意力,但只对最相关的 token 子集执行,代价是需要一个额外的索引步骤

两种方案在实践中都被证明有效。

对未来注意力架构的展望#

未来的注意力架构很可能是将上述所有技巧——局部注意力、Flash Attention、SSM/线性注意力层、稀疏注意力——综合集成到一个复杂的混合架构中。此外,利用后训练(post-training)让模型学会管理自己的上下文(context management)也是一个更高层次的优化方向。线性时间注意力领域的方法已经高度收敛到类 LSTM 的架构形式上,这一趋势在近期不太可能发生根本变化。

一个值得注意的论点是:去掉 softmax 反而可能改善训练稳定性。Softmax 通常是注意力中不稳定性的主要来源(指数运算导致的上溢/下溢),因此线性注意力在稳定性方面可能反而优于标准注意力。这与后面 MoE 路由器中 softmax 引入不稳定性的讨论形成有趣的对比。

关于低精度注意力:目前主要的障碍在于 softmax 对精度非常敏感——小的下溢(underflow)和上溢(overflow)可能导致严重问题。DSA 的”低精度索引 + 高精度注意力”的分离式设计恰好规避了这个问题——索引阶段专门为 Q/K 做额外的降维投影(与实际注意力计算使用的 Q/K 维度可以不同),并在更低精度下执行,而后续的 value 聚合仍使用全精度。


混合专家模型(Mixture of Experts)#

MoE 的基本概念#

MoE 的概念非常直观:它是对 Transformer 中 MLP 层的替换。

在标准 Transformer 中,每个注意力层之后紧跟一个前馈网络(FFN)。MoE 的做法是:将一个大的 FFN 替换为多个小的 FFN(称为”专家”),对于每个输入 token,只激活其中的一个或少数几个专家

假设我们有 4 个与原始 FFN 同等大小的专家,并且有某种机制告诉我们每个输入应该选择哪个专家。那么:

  • 我们拥有 4 倍的参数量(4 个 FFN 的参数)
  • 但任何一次前向/反向传播只需要支付 1 个 FFN 的计算成本

这就是 MoE 的核心价值主张:在不增加推理计算量的前提下大幅扩展模型参数量

为什么 MoE 有效?#

原始 MoE 论文的动机可以概括为一个简单的信念:更多的参数是好事。如果我们想要更多参数但不想为此付出更多的计算成本,MoE 就是自然的选择。

实验证据强有力地支持了这一点。来自 Fedus et al. 2022(Switch Transformer)的实验表明:在保持活跃参数量不变的情况下,随着专家数量的增加,语言建模的测试损失持续下降。类似地,在相同的训练计算预算下,更多的专家意味着更好的性能。OLMo 团队(AI2)的独立研究也得出了相同的结论:无论是训练损失、验证损失还是下游 benchmark 表现,MoE 在训练时间上相比稠密模型能获得大约 2 倍的加速

从实际部署的角度来看,几乎所有超过一定规模的现代开源模型都是 MoE。DeepSeek V2 的出现是一个转折点——它展示了用远少于竞争对手的活跃参数量,就能达到同等甚至更好的 MMLU 表现。

MoE 的另一个好处:额外的并行化维度#

MoE 不仅在统计意义上有效,在系统层面也提供了价值。LLM 极其庞大,无法放入单个设备来训练或推理。标准的并行化方式(数据并行、模型并行)都有各自的上限——数据并行受限于最大 batch size,模型并行受限于模型的自然切分点。MoE 中的每个专家天然构成一个独立的计算单元,可以放置在不同的设备上,形成专家并行(expert parallelism),为并行化提供了额外的维度。当然,所有并行化都有上限——随着设备数量增加,通信成本会爆炸式增长。专家并行也不例外:需要将激活值跨设备传输到对应专家所在的设备,这引入了额外的通信开销。

MoE 的工业实践现状#

在西方,代表性的开源 MoE 模型包括 Llama 4GPT-OSS(OpenAI)。但 MoE 研究和训练的主要推动力来自中国:QwenDeepSeek 在早期就开始训练和推广 MoE 模型。Qwen 1.5 的早期 MoE 模型展示了一个 2.7B 活跃参数的 MoE 可以超过当时多数 7B 稠密模型。DeepSeek 系列的早期论文则提供了大量精心设计的消融实验——包括稠密层 vs MoE 层、hash 路由 vs switch 路由等各种设计决策的对比——非常值得深入阅读。

为什么 MoE 的普及速度较慢?#

尽管 MoE 的优势在 2022 年的 Google 论文中就已经被充分论证,但直到 2024 年之后才真正开始流行。原因在于 MoE 带来的工程复杂性:

  • 基础设施复杂:专家并行的高效实现需要精心设计的通信模式
  • 内存开销大:大量的专家参数需要被存储,即使大部分在任何时刻都不活跃
  • 训练不稳定:MoE 的训练容易出现”专家坍缩”(expert collapse)等问题

关于 MoE attention(将注意力头也做专家化):虽然有一些论文探索了这个方向,但目前不如 FFN 专家化成熟和易于训练。实际部署的 MoE 模型几乎都只对 FFN/MLP 层做专家化。

MoE 的设计空间#

MoE 的设计涉及三个关键维度:路由函数(如何决定 token 去哪个专家)、专家粒度(专家的数量和大小)、以及训练策略(如何稳定地训练这个本质上不可微的系统)。

路由选择:Token Choice vs Expert Choice#

路由的首要决策是:由谁来做选择?

  • Expert Choice:每个专家选择自己最喜欢的 token
  • Token Choice:每个 token 选择自己最喜欢的 K 个专家
  • Global Assignment:全局优化分配方案

Token Choice Top-K 是目前的共识方案,在几乎所有主流 MoE 模型中使用。OLMo 的对比实验表明,Token Choice 在验证损失和下游 benchmark 上都优于 Expert Choice。值得一提的是,Llama 4 的某个版本可能使用了 Expert Choice,但那是一个未发布的版本,且这并不构成对 Expert Choice 有效性的有力支持。

路由函数的设计#

标准 Top-K 路由是目前最广泛使用的方案:

s=softmax(xWrouter)s = \text{softmax}(x \cdot W_{\text{router}}) gates=TopK(s,K)\text{gates} = \text{TopK}(s, K) y=iTopKgatesiFFNi(x)+xy = \sum_{i \in \text{TopK}} \text{gates}_i \cdot \text{FFN}_i(x) + x

其中 WrouterRdmodel×NexpertsW_{\text{router}} \in \mathbb{R}^{d_{\text{model}} \times N_{\text{experts}}} 是路由器的参数——就是一个简单的线性投影。每个专家有一个向量方向,输入与这些方向做内积,内积最大的 K 个专家被选中。这个路由器被 Switch Transformer、GShard、Grok、Mistral、DBRX、DeepSeek 等几乎所有主流 MoE 模型使用。

路由器的极度简单性是值得注意的:它只是一个矩阵乘法,不做任何复杂的语义分析。路由不是在做”这是医学问题,所以路由到医学专家”之类的推理——它更像是在 token 级别做简单的模式匹配(例如”这个 token 看起来像日文,路由到专家 7”)。

有趣的是,一些研究表明甚至不需要学习路由——直接对输入做哈希然后分配到不同专家也能工作。这虽然不如 Top-K 路由效果好,但”随机分配也能获益”的事实暗示了 MoE 的有效性在很大程度上来自参数量的增加本身,而不仅仅是智能路由。

其他路由方案包括:RL 路由(最早由 Bengio 在 2013 年提出,是思考这个问题的最自然方式——将其视为多臂赌博机问题并使用 RL 算法求解);全局线性分配(global linear assignment,计算所有 token 到所有专家的分数矩阵,然后精确求解最优分配问题)。前者因为梯度方差过大而性能不佳,后者因为计算开销极大而无法在规模上应用。

共享专家(Shared Experts)#

DeepSeek MoE 开创了一个现在被广泛采用的设计:共享专家

核心思想是:有些处理是所有 token 都需要的通用处理,不应该交给路由器来分配。因此,将一部分专家设为”共享专家”——它们绕过路由器,对所有输入都执行处理。其他专家则按照正常的 Top-K 路由条件性地执行。

DeepSeek 的精心消融实验表明:

  • 从传统的 16 个大专家设计(GShard/Google 风格)转向更细粒度的小专家 + 共享专家,可以获得显著的性能提升
  • 共享专家在 TriviaQA、Natural Questions 等知识密集型任务上带来了特别大的改善

OLMo 的独立实验在共享专家的价值上存在一些分歧(他们认为共享专家帮助不大),但在”更多细粒度专家有帮助”这一点上达成了一致。

Qwen 3.5、GLM 等最新的顶级开源模型都同时使用了共享专家和细粒度专家设计,这已经成为一个经过充分验证的标准方案。DeepSeek MoE 和 DeepSeek V3 的设计在 MoE 领域的地位,类似于 LLaMA 设计在稠密 Transformer 领域的地位——它们是被广泛复制的标准。

MoE 的训练:优雅的启发式方法#

核心挑战:稀疏路由的非可微性#

MoE 训练面临一个根本性的挑战:为了效率,训练时也必须保持稀疏(不能激活所有专家),但稀疏的 Top-K 选择是不可微的,而且我们无法观察到未被选中专家的反事实结果。

这本质上是一个**多臂赌博机(multi-armed bandit)**问题:在每一步,我们选择 K 个专家,只观察被选中专家的效果,但必须学会做出好的选择。

理论上有三种解决方案:

  1. RL/Bandit 算法:正统的处理方式,但引入大量方差和额外开销
  2. 随机扰动:通过噪声注入来增强探索
  3. 启发式平衡:通过辅助损失函数来人为平衡专家的使用

令人惊讶的是,实践中使用的是第三种——一系列看似 ad-hoc 的启发式方法。而且效果出奇地好。

RL 路由:理论上优雅但实践中不佳#

Clark et al. 2020 的研究表明,使用 REINFORCE 算法来学习路由器是可行的,但性能不如更简单的启发式方法。原因是梯度方差(gradient variance)太大,导致训练不够稳定和高效。

随机扰动:可选但非必需#

最早的 MoE 论文(Noam Shazeer et al.)提出了一种随机扰动方案:在路由器的 logits 上注入依赖于输入尺度的噪声:

H(x)=xWrouter+SoftPlus(xWnoise)ϵ,ϵN(0,1)H(x) = x \cdot W_{\text{router}} + \text{SoftPlus}(x \cdot W_{\text{noise}}) \cdot \epsilon, \quad \epsilon \sim \mathcal{N}(0, 1)

这使得当两个专家的路由分数接近时,会随机选择其中一个,增加了探索(exploration)。直觉上,如果某个专家确实对某类输入有帮助,通过反复的随机探索和梯度更新,这个专家的权重会逐渐增大。

Fedus et al. 2022(Switch Transformer)也提出了类似的均匀乘性扰动,但目的不同——他们的目的是硬化专家选择(hardening the experts),即让专家做出更确定性的选择承诺。但在后续的 Google 论文和其他研究中,这些随机扰动被发现不是必需的——去掉它们有时反而改善了稳定性和最终质量。

负载均衡损失:MoE 训练的核心#

如果不加任何干预就直接用梯度下降训练 MoE,会发生什么?专家坍缩(expert collapse / expert starvation):

  1. 初始时,某些专家碰巧被更多 token 选中
  2. 被选中的专家获得更多梯度更新信号,权重变得更强
  3. 更强的权重导致更多 token 被路由到这些专家
  4. 形成”富者愈富”的正反馈,最终只有极少数专家被使用,其余完全闲置

Switch Transformer 的负载均衡损失 是解决这个问题的核心方法。对于每个专家 ii,定义两个量:

  • fif_i:分配到专家 ii 的 token 数量比例(硬指标)
  • PiP_i:路由器分配给专家 ii 的总概率质量(软指标)

负载均衡损失为:

Lbalance=αNi=1NfiPi\mathcal{L}_{\text{balance}} = \alpha \cdot N \sum_{i=1}^{N} f_i \cdot P_i

这个形式乍看不直观,但分析其梯度可以揭示其作用:LbalancePifi\frac{\partial \mathcal{L}_{\text{balance}}}{\partial P_i} \propto f_i。也就是说,一个专家获得越多的 token(fif_i 越大),它的概率质量 PiP_i 所受的负梯度就越大,从而被推低。这有效地抑制了”赢家通吃”的正反馈。

DeepSeek 的双层平衡设计#

DeepSeek V1/V2 在 Switch Transformer 的基础上增加了设备级平衡损失:不仅要确保各个专家之间的负载均衡,还要确保各个设备之间的负载均衡。如果专家分布在不同设备上,即使专家级别的负载均衡了,设备级别仍可能不均衡(因为被频繁选中的专家可能恰好集中在某些设备上)。设备级平衡损失的形式完全类似,只是将”每个专家”替换为”每个设备上所有专家的总和”。

不应该把专家级平衡损失的系数调得过高以实现完美均匀分布——那会对训练动态产生有害影响。因此,设备级平衡作为额外的软约束存在是合理的:即使专家之间不是完美均匀的,至少要保证设备之间的利用率基本均等。

DeepSeek V3 进一步引入了逐专家偏置项(per-expert bias)作为一种更轻量的平衡方式,试图减少对辅助损失的依赖,但最终仍然需要某种形式的辅助损失来防止极端不均衡。目前没有任何方案能完全摆脱辅助损失。

负载均衡的实证验证#

OLMo 提供了一个非常直观的消融实验:完全移除负载均衡损失后,几乎所有 token 都被路由到仅仅 2 个专家(在可视化中表现为黄色和粉色专家独占所有流量),导致训练损失和验证损失都显著恶化。与之对比,有负载均衡损失时,所有专家被均匀使用。这直接验证了:没有负载均衡,绝大部分参数完全浪费了。

整个 MoE 训练方案的工作原理,从理论角度看,是令人惊讶的。一个本质上不可微的、基于 Top-K 选择的复杂系统,仅靠”直接反向传播忽略不可微性 + 负载均衡损失”就能稳定训练并产出高质量模型。其核心动态是:反向传播自然地强化有用的专家(正反馈),而负载均衡损失防止任何专家独占所有流量(负反馈),两种力量的平衡产生了稳定的训练。

Top-K 选择作为通用模式#

这种”Top-K 选择 + 辅助损失来处理不可微性”的模式不仅限于 MoE。它同样出现在 DSA(稀疏注意力)中、H-Net(一种去除 tokenizer 的尝试)中,并且预计会在未来的架构设计中越来越多地出现。这是一种值得内化的通用设计模式。

专家是否真的在”专精”?#

由于路由器极其简单(只是一个线性投影),专家并不是语义层面的”专家”——不存在”医学专家”或”法律专家”。但确实可以观察到某种模式:例如数字可能倾向于被路由到某个专家,某些符号被路由到另一个专家。但这种分化更接近于浅层的 token 级模式匹配,而非深层的语义专业化。

MoE 的系统层面考量#

稀疏矩阵乘法与硬件协同设计#

当多个专家位于同一个 GPU 上时,朴素的实现方式是执行多个小的矩阵乘法。但这在硬件上效率不高——GPU 更擅长处理大的矩阵乘法以充分利用缓存和并行计算单元。解决方案是利用结构化稀疏矩阵乘法:MoE 的计算模式天然对应于块对角矩阵乘法和其他形式的结构化稀疏矩阵操作,而这些操作在现代 GPU 上有原生的硬件支持。这种架构与硬件的协同设计是 MoE 高效实现的重要基础。

专家并行的通信优化#

专家并行要求将激活值从一个设备发送到另一个设备(“你的 token 属于那个专家,需要把它发过去”),这带来了显著的通信开销。

NeMo Tron 3 提出了一种精妙的优化:共享专家不需要跨设备通信(它在每个设备上都有副本),可以保持高维度;而需要跨设备路由的专家,可以先将残差流降维后再发送。具体来说:

  1. 将输入从 dmodeld_{\text{model}} 维投影到更低的维度 drouted_{\text{route}}
  2. 在低维空间中完成跨设备的集合通信(collective communication)
  3. 在目标设备上的专家接收低维输入进行处理

这样显著减少了通信量,同时因为共享专家承担了大部分通用处理,对整体性能的影响很小。

Token 丢弃问题#

在早期的 MoE 推理基础设施中,如果某个专家的 token 队列过长,系统会静默地丢弃(drop)超出的 token——直接返回零值假装什么都没发生。这导致了令人困惑的随机性:同一个请求在不同时刻可能得到不同的结果,因为其他用户的请求可能占满了你需要的专家队列。

现代框架(如 MegaBlocks)采用无丢弃架构(dropless architecture),通过更灵活的调度和内存管理来处理负载不均衡,不再静默丢弃 token。这个问题在当前已基本解决。

训练稳定性#

MoE 引入了额外的 softmax(在路由器中),而 softmax 中的指数运算和除法运算是训练不稳定的常见来源。解决方案包括:

  • 对路由器使用 float32 精度:即使模型其他部分使用更低精度,路由器的 softmax 也使用全精度来避免数值问题
  • Z-loss:控制 softmax 输入的 logits 不要变得过大。Z-loss 最初由 Barrett Zoph 等人在 Google 早期的 MoE 论文(ST-MoE)中提出,专门用于解决 MoE 路由器的稳定性问题。OLMo 的消融实验也明确验证了其重要性:移除 Z-loss 后训练损失曲线出现严重的 spike

MoE 的微调挑战#

MoE 模型因为拥有大量参数,在微调(fine-tuning)时容易出现严重的过拟合。对比实验表明:稠密模型的训练损失和验证损失差距较小,而 MoE 模型在微调过程中两者的差距极其悬殊。

应对策略包括:

  • 只微调非 MoE 层(注意力层是常见的选择)
  • 只微调共享专家或非 MoE 的前馈层
  • 使用足够大的微调数据集(“苦涩教训”版本的解决方案)——如果有足够多的数据(如 140 万样本),可以直接微调整个 MoE 而不严重过拟合

Upcycling:从稠密模型到 MoE#

Upcycling 是一种将已训练好的稠密模型转换为 MoE 模型的技术:

  1. 取一个已训练的稠密模型
  2. 将其 MLP 层复制多份,形成多个专家(每个专家的初始权重相同)
  3. 随机初始化路由器
  4. 继续训练——由于路由的随机性,各专家会逐渐专业化

早期的 upcycling 实验表明,相比继续训练稠密模型,upcycling 后的 MoE 在相同计算预算下能获得更好的性能。MiniCPM(从 2.4B 稠密模型 upcycle 为 MoE)和 QwenMoE(从 1.8B Qwen 模型 upcycle)都验证了这种方法的有效性。

但值得注意的是,upcycling 在近年来已不太流行。原因是:既然 MoE 的优势已经明确,不如从一开始就直接训练 MoE(hero run),而不是先训练稠密模型再转换。这是一种有趣的历史过渡技术。

DeepSeek V1 → V2 → V3 的架构演进#

DeepSeek 系列的架构演进是理解现代 MoE 设计的一条非常有价值的线索:

DeepSeek MoE V1 已经是现代 MoE 的原型:共享专家 + 细粒度专家 + 标准 Top-K 路由 + 负载均衡辅助损失。可以说它是 MoE 模型的”柏拉图式理想形态”。

DeepSeek MoE V2 在 V1 基础上进行了规模化扩展:2 个共享专家、更多的细粒度路由专家,并增加了设备级路由平衡通信平衡组件——这两者本质上都是为系统优化添加的辅助损失。这反映了一个重要的认识:成功的语言模型训练不仅仅关乎深度学习本身,还需要深度尊重底层硬件系统的约束。

DeepSeek MoE V3 的主要改进包括:

  • 辅助损失无关的平衡方式(aux-loss-free balancing):通过逐专家偏置项来替代部分辅助损失
  • Sigmoid + Softmax 路由:用不同的方式组合路由器的输出来加权专家
  • 多头潜在注意力(Multi-Head Latent Attention, MLA):不直接产生 Q/K/V,而是先产生一个低维的潜在表示 CC,再从 CC 导出 Q/K/V。这使得 KV 缓存只需要存储低维的 CC,显著减少了推理时的内存需求。但需要注意与 RoPE 的兼容性问题——需要保留部分非潜在维度来编码位置信息
  • 多 Token 预测(Multi-Token Prediction, MTP):同时预测多个未来 token。这既有统计上的论据(可能帮助模型更好地建模长程依赖),也有系统上的好处(内建了投机解码器 speculative decoder,可以加速推理时的 token 生成)。不过值得注意的是,MTP 目前尚未广泛流行——虽然 DeepSeek V3 使用了它,但在其他模型中并未被广泛采纳

核心要点总结#

注意力替代方案方面:线性注意力通过矩阵乘法结合律将 O(N2)O(N^2) 降为 O(N)O(N),且可以等价转换为 RNN 形式(训练并行、推理高效)。Mamba 2 和 Gated DeltaNet 通过引入门控机制增强了线性注意力的表达力,后者已被 Qwen 3.5 等顶级模型采用。DSA 则代表了一条完全不同的路径——通过轻量级索引器选出关键 token 子集来做稀疏注意力。所有方案在实际部署中都采用混合架构,纯线性/纯稀疏方案尚未在规模上被验证。

混合专家模型方面:MoE 通过将 FFN 替换为多个条件性激活的专家来实现”更多参数、同等计算”。Top-K 路由 + 负载均衡辅助损失是经过充分验证的标准方案。DeepSeek 的共享专家 + 细粒度专家设计已成为行业标准。MoE 的训练虽然涉及不可微的 Top-K 选择,但通过启发式平衡就能稳定工作,这一事实本身就是深度学习领域的一个有趣现象。

Lecture 4:注意力替代方案与混合专家模型
https://www.xwysyy.cn/posts/cs336/lec04/
作者
xwysyy
发布于
2026-05-17
许可协议
CC BY-NC-SA 4.0
© 2026 xwysyy. All Rights Reserved.
Powered by Astro & Firefly

文章目录