Runguo Li

Runguo Li

M.S. student in Information Science @ UIUC, advised by Prof. Minjia Zhang. Working on ML systems for large-scale models: MoE inference and expert streaming, distributed training and rollout, GPU compiler correctness.

第 4 讲：Attention 替代方案与混合专家（MoE）

目录 · ← l3 · l5 →

第 4 讲：Attention 替代方案与混合专家（MoE）

日期：4 月 8 日（周三，Spring 2026） | 讲师：Tatsu Hashimoto | 材料：lecture_04.pdf

概览

Attention 的复杂度对序列长度是平方级的，在长上下文下会成为主导成本。本讲讲两类应对方案：(1) attention 替代方案——线性注意力、状态空间混合模型（Mamba-2、Gated DeltaNet）、稀疏注意力（DSA）；(2) 混合专家（MoE）——用一个路由器把大 FFN 换成许多专家 FFN，从而把”参数量”与”每 token 的 FLOPs”解耦。内容包括路由方法、训练目标（辅助均衡损失）、系统实现、upcycling，以及 DeepSeek MoE v1→v3 的演进。

核心概念与定义

  • Attention 的成本：Attn(Q,K,V) = softmax(QKᵀ)V 中，QKᵀ 需 O(n²·d_k)、attention·V 需 O(n²·d_v)。上下文窗口越大，这一项越主导。
    • 类比：full attention 像派对里每个人都和每个人交谈——n² 次对话；线性注意力则像每个人只对记录员说一句摘要、由记录员统一广播——约 2·n·d 次交流。
  • 线性注意力（linear attention）：若把 softmax 换成恒等核，则 QKᵀV = Q(KᵀV)，成本从 O(n²d) 降到 O(2·n·d_v·d_k)——对序列长度线性。
    • 循环形式：S_t = S_{t−1} + k_t v_tᵀ，y_t = q_tᵀ S_t——这看起来就是 RNN！这种对偶性（duality）让你用并行（二次）形式训练、用串行（线性）形式推理。把 S_{t−1} 乘以权重 γ 就得到 RetNet。
    • 类比：一个上课认真记笔记的学生（状态 S_t）随时能回答问题，而不必为每个新问题把之前所有讲义（全部历史）重读一遍。
  • Mamba-2：加入逐位置门控：S_t = γ_t S_{t−1} + k_t v_tᵀ，其中 γ_t = f(x_t)——门控让线性注意力表达力更强，同时保留对偶性。
  • Gated Delta Net（GDN）：再加输入门与选择性状态擦除：S_t = γ_t(I − β_t k_t k_tᵀ) S_{t−1} + β_t k_t v_tᵀ。β=0 表示”不写入”；擦除项 (I − β_t k_t k_tᵀ) 会忘掉与当前 key 同方向的内容。与 fast weight programmer / test-time training 关系密切。Qwen 3.5/Qwen Next 使用 3:1 的 GDN/attention 混合。
  • 混合架构（hybrids）：不做”全 attention”或”全线性”，而是交错：Minimax M1（7:1 线性:full）、Nemotron 3（约 3:1 Mamba:attention）、Qwen 3.5（3:1 GDN:attention）。受控消融显示小混合比例下损失几乎不退化，而推理收益巨大（常量大小的状态 vs O(n) 的 KV cache）。
  • 稀疏自适应（DSA，DeepSeek Sparse Attention）：不再关注所有 token，而是用一个轻量 indexer 选出 top-k 个 token 参与 attention；可在稠密短上下文预训练之后”事后适配”（DeepSeek v3.2、GLM-5）。
  • 混合专家（MoE）：用许多 FFN（”专家”）加一个路由器（router）取代单个大 FFN，每个 token 只走 top-k 个专家。总参数量随专家数增长，但 FLOPs 基本不变（每 token 只激活 k 个专家）。
    • 类比：不是让一位全科医生看所有病人，而是开一家有 256 位专科医生的诊所，由分诊护士（router）把每位病人分给最相关的 2–8 位专家。诊所人手很多（参数多），但每位病人只见少数几位（FLOPs 不变）。
    • MoE 为何流行：同等 FLOPs 下参数更多 → 损失更低；单位算力下训练更快；与稠密模型相比很有竞争力；天然可跨设备并行（每个专家可放在一张卡上）。实例：Mixtral（8 专家、top-2）、Grok、DBRX（16、top-4）、Qwen（60、top-4）、DeepSeek v3（256 专家、激活 8 + 1 共享）、Llama 4 Maverick（128 路由 + top-1 + 共享）。
  • 路由方法：token 选择专家的 top-k（主流：Switch k=1，GShard/Grok/Mixtral k=2，Qwen/DBRX k=4，DeepSeek k=7-8）、哈希路由（基线）、RL 学习路由（早期工作，Bengio 2013）、线性分配路由（Clark ‘22）。
    • 路由打分方式：逻辑回归门控（DeepSeek v1-2、Grok、Qwen）vs 在 top-k 之后做 softmax（Mixtral、DBRX、DeepSeek v3）。
    • 近期变体：细粒度专家（很多小专家）+ 少量常开的共享专家（DeepSeek、Qwen；最初来自 DeepSpeed MoE）。消融大体显示细粒度专家有帮助；共享专家的收益存在分歧（OlMoE 认为没有）。
  • MoE 训练——不可微问题：稀疏路由决策不可微。可选方案：(1) 用 RL 优化门控策略（可行但方差大，未被广泛采用）；(2) 随机扰动（Shazeer 2017 的高斯噪声、Switch Transformer 的乘性抖动）；(3) 启发式均衡损失（auxiliary balance loss）——实践标准：例如 Switch 的负载均衡损失会把过热的专家降权；DeepSeek v1-2 同时做逐专家与逐设备均衡；DeepSeek v3 用无辅助损失的均衡（aux-loss-free）——通过在线调整每个专家的偏置实现。
  • MoE 的系统侧：专家可跨设备并行（expert parallelism），每个 FFN 能放进一张卡；稀疏矩阵乘需要专门 kernel（如 MegaBlocks）；Nemotron 3 通过把激活降维来减少通信量。
  • MoE 的稳定性/微调问题：router logits 建议用 fp32 + z-loss；稀疏 MoE 在小数据上微调容易过拟合（Zoph 等：只微调非 MoE 的 MLP；DeepSeek：用大量数据，1.4M 条 SFT）。
  • Upcycling（升级复用）：用预训练的稠密模型初始化 MoE（拆分/复制专家权重）。实例：MiniCPM-MoE（来自 MiniCPM）、Qwen-MoE（来自 Qwen 1.8B）。
  • DeepSeek MoE 谱系：v1（16B/激活 2.8B：共享 2 + 细粒度 64，辅助损失均衡）、v2（236B/激活 21B：共享 2 + 细粒度 160、激活 6，通信均衡损失，top-M 设备路由）、v3（671B/激活 37B：共享 1 + 细粒度 256、激活 8，sigmoid+softmax top-k + top-M，无辅助损失 + 序列级辅助损失）。
  • DeepSeek v3 的配套技巧：MLA（多头隐层注意力，把 KV 压到低维隐向量 c，见第 10 讲）与 MTP（多 token 预测：用小的轻量头预测多步）。

代码示例：线性注意力的循环形式（概念）

代码（Python）：

# 并行（训练）形式：  Y = Q(K^T V)   -- O(n^2 d)，只有两次矩阵乘法
# Q: [n, dk], K: [n, dk], V: [n, dv]
KV = K.transpose(-2, -1) @ V          # [dk, dv]  -- 所有 key/value 外积之和
Y = Q @ KV                           # [n, dv]

# 循环（推理）形式：        -- 每 token O(n)，状态 O(1)
S = zeros(dk, dv)                    # 状态
for t in range(n):
    S = S + K[t][:, None] * V[t][None, :]   # S_t = S_{t-1} + k_t v_t^T
    y[t] = Q[t] @ S                        # y_t = q_t^T S_t

代码做了什么： 展示线性注意力的两种等价计算：训练用的并行矩阵形式，与推理用的顺序状态更新形式（永不回看旧 token）。

实现深挖：

  • 为什么对偶性重要：训练用并行形式以充分利用 GPU 矩阵乘法；推理用循环形式，每 token 显存 O(1)（KV cache 不增长）。这正是状态空间模型实用的”训练并行、推理串行”技巧。
  • 为什么加门控（Mamba-2/GDN）：最朴素的递推无法遗忘；逐位置门控 γ_t、输入门 β_t 与擦除项 (I − β k kᵀ) 赋予选择性记忆，这在语言任务上经验性收益很大。

与作业的联系：这些内容对作业 1–5 属于概念性背景（作业 2 你实现的是完整 attention 的 FlashAttention-2），但这里培养的算术强度推理（第 2、10 讲）正是 GQA/MLA 取舍与 FlashAttention tiling 的理论依据。

代码示例：top-k 路由（概念 MoE router）

代码（Python）：

import torch
import torch.nn.functional as F

def topk_route(x, router_weight, num_experts, k):
    # x: [num_tokens, d_model]
    logits = x @ router_weight                # [num_tokens, num_experts]
    topk_logits, topk_idx = torch.topk(logits, k, dim=-1)   # 选中哪些专家
    # 只在被选中的专家上做 softmax（Mixtral 风格）
    probs = F.softmax(topk_logits, dim=-1)
    return topk_idx, probs                    # 把 token 分发到专家

代码做了什么： 计算每个 token 对各专家的 logits，取 top-k，并只在被选中的专家上归一化（Mixtral/DBRX/DeepSeek-v3 的做法）。

实现深挖：

  • 为什么用 top-k 而不是 k=1 的 argmax：k>1 给路由器更细的粒度并平滑梯度；Switch（k=1）更简单但脆弱。DeepSeek v3 在 256 个专家中激活 8 个。
  • 为什么要有共享专家：常开的共享专家负责捕获通用模式，让路由专家更专注；DeepSeek/Qwen 配 1–4 个共享专家。
  • 为什么”top-k 之后才 softmax”：只对选中的专家归一化，使路由权重与未选中专家的 logits 无关——这是与 DeepSeek v1-2 逻辑门控路由器有意的设计差异。

与作业的联系：必修作业不要求实现 MoE（第 8 讲会从概念上讲 expert parallelism），但作业 2 的分布式训练（all-to-all token 分发）正是 MoE 需要的系统基础；而”在预算下最小化损失”的排行榜思维与 DeepSeek 靠消融做选择的方式一致。

关键要点

  1. Attention 的 O(n²) 成本可用线性/状态空间替代方案（线性注意力、Mamba-2、Gated DeltaNet）来攻击——实践中多以混合形式使用（如 3:1），用少许精度换取长上下文推理的巨大成本下降。
  2. 线性注意力具有”训练并行/推理递推”的对偶性；门控是让递推具备表达力的关键。
  3. MoE 把参数量与 FLOPs 解耦：同等算力下容量更大，代价是复杂的路由、均衡损失与系统开销。
  4. top-k “token 选专家”路由 + 启发式均衡损失是实践共识；RL 路由理论优雅但方差太大。
  5. 现代 MoE 采用细粒度专家 + 共享专家、无辅助损失均衡、fp32/z-loss 稳定路由器；从稠密 checkpoint 做 upcycling 是获得 MoE 的廉价途径。

常见陷阱

  • 负载不均衡：没有均衡损失时少数专家吃掉所有 token——容量浪费，更糟的是设备闲置。
  • 路由器不稳定：fp16 的 router logits 会爆炸；用 fp32 + z-loss（Zoph 等 2022）。
  • 批次级 token 丢弃：路由在 batch 级别丢 token，意味着别人的 query 可能把你的 token 丢掉——这是额外的随机性来源。
  • 微调时 MoE 过拟合：稀疏模型在小 SFT 集上容易过拟合；要么加数据，要么冻结/降权路由专家。
  • MLA/RoPE 的 KV cache 盲区：RoPE 与 MLA 缓存冲突；需要保留非旋转的 key 维度（DeepSeek 的 64 维技巧）。

复习题

  1. 问： 为什么 Q(KᵀV) 能改变 attention 的成本？代价是什么？
    • 答： 矩阵乘法满足结合律：(QKᵀ)V = Q(KᵀV)。左边 O(n²·d)，右边 O(n·d²)。代价是：只有当 attention 核为恒等（没有 softmax）时才成立，这正是线性注意力的假设——softmax attention 无法这样因式分解。
  2. 问： 线性注意力的”对偶性”是什么？为什么在工程上重要？
    • 答： 同一个计算既能写成并行形式（Q(KᵀV)，适合 GPU 训练），也能写成循环形式（S_t = S_{t−1} + k_t v_tᵀ，推理时状态 O(1)）。训练用一种、推理用另一种。
  3. 问： 一个 256 专家、激活 8 专家的 MoE 为何能”参数更多但 FLOPs 不变”？
    • 答： 每个 token 只经过 k=8 个专家，因此每 token 的矩阵乘 FLOPs 大致相当于”激活规模”的稠密模型；其余 248 个专家的参数仍占显存（提供容量/知识），但对该 token 不产生计算。