第 4 讲：注意力的替代方案与混合专家（MoE）

对应：Stanford CS336 Spring 2026 Lecture 4（Attention alternatives and mixture of experts，Tatsunori Hashimoto 主讲）。Spring 2025 的 Lecture 4 只讲 MoE，有公开录像。 配套代码：code/lecture04/ 预计阅读时间：3–4 小时

0. 本讲导航

0.1 学完本讲，你应该能够

1. 说清标准 Transformer 的两个”天生成本”：注意力随上下文平方增长；每个 token 都要经过全部参数。
2. 用矩阵乘法的结合律推导线性注意力，并解释它为什么既能像 Transformer 一样并行训练、又能像 RNN 一样常数代价推理。
3. 把线性注意力、RetNet、S4/Mamba、Mamba-2、DeltaNet、Gated DeltaNet 放在同一个框架里理解：“状态 = 衰减/擦除 × 旧状态 + 写入”。
4. 用”联想记忆”的视角解释固定大小状态的局限，以及混合架构为什么流行。
5. 解释 MoE 为什么能”参数多、计算少”，计算总参数与激活参数。
6. 说出 MoE 的三个设计维度：路由函数、专家大小（细粒度、共享专家）、训练目标（负载均衡）。
7. 比较辅助损失与 DeepSeek-V3 的”无辅助损失”负载均衡；解释专家容量、token 丢弃和 MoE 的额外随机性。
8. 描述 DeepSeek MoE 从 V1 到 V3 的演进，以及 MLA、MTP 的基本思想。

0.2 先修知识

0.3 本讲地图

第一部分：注意力的替代方案（解决"随上下文平方增长"）
   第1节 问题与基本工具箱
   第2节 线性注意力：一个结合律带来的对偶性
   第3节 从线性注意力到现代线性模型：衰减、门控、SSM、Delta 规则
   第4节 混合架构与稀疏注意力：当前的实践

第二部分：混合专家（解决"每个 token 都用全部参数"）
   第5节 什么是 MoE，为什么流行
   第6节 设计空间：路由、专家大小、训练目标、系统、稳定性
   第7节 案例：DeepSeek MoE V1 → V2 → V3（以及 MLA、MTP）

两部分的共同主题是”稀疏”与”压缩”：不是每个 token 都需要看所有历史（第一部分），也不是每个 token 都需要所有参数（第二部分）。

第一部分：注意力的替代方案

1. 问题与基本工具箱

第 3 讲算过：GPT-2 XL 的上下文从 1K 增加到 16K 时，注意力核心（$QK^\top$ 和注意力 × $V$）从 7% 的 FLOPs 变成了 55%。而今天的模型动辄支持 128K 甚至上百万 token 的上下文。注意力的成本有两方面：

• 训练：每层要算 $n\times n$ 的分数矩阵，计算量 $\propto n^2$；
• 推理：KV 缓存随已处理的 token 数线性增长，每生成一个 token 都要把它完整读一遍。

“基本工具箱”是第 3 讲介绍过的：局部 + 全局注意力交替（滑动窗口），以及系统工程上的优化（FlashAttention 等，第 5 讲）。但如果想要更根本、可能更大的收益呢？

2. 线性注意力：一个”很傻但很重要”的结合律

2.1 结合律

回忆注意力 $\text{Attn}(Q,K,V) = \rho(QK^\top)V$，其中 $\rho$ 是 softmax。如果 $\rho$ 是恒等函数（去掉 softmax），矩阵乘法的结合律告诉我们：

左边的中间结果是 $n\times n$ 矩阵，右边是 $d\times d$ 矩阵。计算量从约 $n^2 d$ 降到约 $nd^2$。配套代码 01_linear_attention.py 演示 1（$n=4096$，$d=64$）：

CS336 课上的评价是：这个想法”非常傻，但出奇地重要”（Shen et al. 2018；Katharopoulos et al. 2020 的”核函数”版本用 $\phi(q)\cdot\phi(k)$ 近似 softmax 中的 $e^{q\cdot k}$）。

代价：结合律只在中间没有 softmax 时成立。softmax 是非线性的、按行归一化，不能拆开。放弃 softmax 会损失表达能力——后面所有的改进都在想办法把能力补回来。

2.2 因果版本就是一个 RNN

语言模型需要因果掩码：位置 $t$ 只能看 $1..t$。把 $y_t = \sum_{s\le t}(q_t\cdot k_s)v_s$ 改写一下：

这就是一个 RNN！状态 $S_t$ 是一个 $d_k\times d_v$ 的矩阵，大小固定，与已经处理了多少 token 无关。

2.3 对偶性：训练用并行、推理用递推

同一个计算有两种（实际上三种）算法，结果完全相同（演示 2，误差 $4\times10^{-7}$）：

推理时每一步只需要保存状态（演示 2，$d=128$、BF16、单头）：

实测运行时间（演示 3，CPU，$d=64$）：$n$ 从 1024 翻倍到 16384，softmax 注意力从 0.3 ms 增加到 26.5 ms（每翻一倍约 ×4），分块线性注意力从 0.4 ms 增加到 8.6 ms（每翻一倍约 ×2）。短序列时优化得很好的 softmax 注意力反而更快，序列越长，线性的优势越明显。

2.4 真实模型：MiniMax

MiniMax-Text-01 和 MiniMax-M1 使用 7:1 的混合：每 8 层中 7 层是线性注意力（他们称为 lightning attention），1 层是完整的 softmax 注意力。效果总体不错，而且在长上下文下成本近似线性增长。

3. 从线性注意力到现代线性模型

3.1 统一的框架

纯线性注意力的状态只进不出：$S_t = S_{t-1} + k_tv_t^\top$，所有历史信息都挤在一个固定大小的矩阵里，越来越”满”。之后几年的改进可以放在同一个框架里理解：

CS336 课上的一句总结：门控是好东西（gating is good）。关键在于这些系数都由当前 token 算出来（数据依赖），并且这些模型都保留了”并行形式”，训练仍然高效（演示 4：门控线性注意力的并行形式用衰减矩阵 $D$ 代替掩码 $M$，与递推形式的误差 $3\times10^{-6}$）。

3.2 另一条路：状态空间模型（SSM）

SSM 来自控制理论。最简单的离散形式：$h_t = Ah_{t-1} + Bx_t$，$y_t = Ch_t$。配套代码 03_state_space_models.py 讲了三件事：

(1) S4 等早期模型：递推与卷积的对偶。 如果 $A, B, C$ 不随时间变化（线性时不变，LTI），展开递推得到 $y_t = \sum_j K_j x_{t-j}$，其中 $K_j = CA^jB$ 是一个固定的卷积核。训练时对整段序列做一次卷积（并行，可用 FFT），推理时逐步递推（演示 1，两种算法误差 $10^{-7}$）。

(2) Mamba 的关键改动：选择性。 LTI 模型对每个输入一视同仁。演示 2 设计了一个任务：序列中第 10 个输入是重要信息，其余 89 个是噪声，要求在最后一步输出它。

• LTI 模型：$a$ 小（0.5）则很快忘掉；$a$ 大（0.99、1.0）则把后面的噪声也累加进来。不管 $a$ 取多少都做不到。
• 选择性 SSM：让衰减和写入强度 $\Delta_t$ 随输入变化——在重要位置大量写入，其余位置几乎不写也不忘。最后一步输出 0.347（目标 0.350）。

代价：参数随输入变化后，没有固定的卷积核了。Mamba 改用并行扫描：把每一步看成变换 $h\mapsto a_th + b_t$，这种变换的组合满足结合律，因此可以像求”前缀和”一样在 $\lceil\log_2 T\rceil$ 轮内并行算完（演示 3：$T=1000$ 只需 10 轮）。

(3) Mamba-2：殊途同归。 Mamba-2 把 $A$ 限制为”标量 × 单位矩阵”（每步一个衰减系数 $\gamma_t$）。此时 SSM 与门控线性注意力完全相同：把 $C_t$ 看成 query、$B_t$ 看成 key、$x_t$ 看成 value 即可（演示 4，差异为 0）。Mamba-2 论文称之为”状态空间对偶”（SSD）。来自 RNN、控制理论、注意力三个方向的研究，最终汇聚到了同一个形式。

3.3 联想记忆的视角

为什么需要衰减和擦除？把状态 $S$ 看成一个联想记忆：写入 $S\leftarrow S + kv^\top$ 存一条”键 → 值”的关联，读取 $\hat v = S^\top q$。配套代码 02_associative_memory.py 的实验（$d = 64$，随机单位向量作键，余弦相似度越接近 1 读得越准）：

容量：存入 $N$ 条关联再逐条读出

$d_k$ 维的键最多只能互不干扰地存约 $d_k$ 条关联；存多了，累加式记忆读出来的是”目标值 + 其他所有值的干扰”。softmax 注意力原样保存每一条，读取几乎完美——代价是内存随 $N$ 线性增长。这就是固定状态模型与注意力模型的根本权衡：记忆容量 vs 内存与计算成本。（注意 Delta 规则的平均准确度并不更高：它会让新写入覆盖旧关联。它的优势在下面两点。）

更新：同一个键先写旧值、再写新值

Delta 规则 $S\leftarrow S + \beta k(v - S^\top k)^\top$ 先读出这个键当前存的值，只写入”差值”——就像变量被重新赋值、话题发生转移时，旧信息应该被替换而不是叠加。它等价于 $S\leftarrow(I - \beta kk^\top)S + \beta kv^\top$，也可以看成对 $\|S^\top k - v\|^2$ 做一步梯度下降（与”快速权重”、“测试时训练”等思想有密切联系）。

记忆满了之后：依次写入 256 条（远超容量 64），按写入时间分组读取

累加式记忆所有内容都一样模糊；门控让旧记忆淡出、最近的记忆变清晰；Delta 规则只改与当前键相关的部分，最近的记忆清晰，中等距离的也保留得更多。固定的 $\gamma$ 看不出”门控 Delta”的额外好处——在真实模型中 $\gamma_t$、$\beta_t$ 由内容决定（新文档开始时大幅遗忘，无关 token 不写入），这种按内容决定的灵活性才是门控的价值所在。

4. 混合架构与稀疏注意力：当前的实践

4.1 为什么要混合

上一节的容量实验揭示了线性模型的根本弱点：固定大小的状态无法精确回忆任意长的上下文。而第 3 讲的实验表明，“从上文中精确地找到并抄写信息”（归纳头）正是注意力的强项。

解决办法：大部分层用线性模型（便宜），少数层用完整注意力（精确回忆）。

受控的消融实验不多，但已有证据显示：只需要很少比例的注意力层，混合模型的损失就能接近全注意力模型。

4.2 另一种思路：稀疏注意力的”后期改造”

不换层，而是让每个位置只关注一小部分被选中的 token。DeepSeek-V3.2 的 DeepSeek 稀疏注意力（DSA）用一个非常轻量的”索引器”（lightning indexer）为每个 query 快速挑出最相关的若干 token，完整注意力只在这些 token 上计算。它可以在”短上下文、稠密注意力”的预训练之后事后改造上去，大幅降低长上下文的成本。据 CS336 课件，GLM-5 也采用了这种设计。

第二部分：混合专家（MoE）

5. 什么是 MoE，为什么流行

5.1 核心想法

把一个大的前馈层换成许多个前馈层（“专家”），再加一个”路由器”为每个 token 选择其中几个。

           ┌─► 专家 1 ─┐
           ├─► 专家 2 ─┤
 x ─► 路由器 ─┼─► 专家 3 ─┼─► 按门控权重加权求和 ─► 输出
（给每个专家打分，├─► ……    ─┤        （只有被选中的 k 个专家参与）
  选出 top-k）  └─► 专家 E ─┘

• 参数量 ≈ E 个专家之和：很大；
• 每个 token 的计算量 ≈ k 个专家：很小。

配套代码 04_moe_layer.py 演示 2（$d=256$，每个专家是一个 $d_{ff}=704$ 的 SwiGLU）用 FLOP 计数器验证：

计算量由 $k$ 决定，参数量由 $E$ 决定。 另外，只有 1 个专家的 MoE 与普通 FFN 完全相同（演示 1，差异为 0）。

通常把 Transformer 中的前馈层替换成 MoE 层（也有少数工作把注意力头做成专家，如 ModuleFormer、JetMoE，但不常见）。

5.2 为什么越来越流行

1. 同样的 FLOPs，参数更多，效果更好：Switch Transformer（Fedus et al. 2022）显示，固定每 token 的计算量，增加专家数量会持续降低损失。
2. 训练更快：OLMoE 的实验中，MoE 达到同样损失所需的训练 token 远少于稠密模型。
3. 性能极具竞争力：如今最强的开放模型大多是 MoE——DeepSeek-V3、Qwen3 的旗舰版本、Llama 4、gpt-oss 等；据传 GPT-4 也是 MoE。国内团队（DeepSeek、Qwen）在较小规模上做了大量扎实的 MoE 研究和消融。
4. 天然适合多设备并行：每个专家可以放在不同的设备上（专家并行，第 7–8 讲）。

5.3 为什么之前没那么流行

• 基础设施复杂：路由、token 的分发与收集、专家并行，工程难度大；优势主要体现在多节点训练中。
• 训练目标偏启发式，而且有时不稳定（Zoph et al. 2022 ST-MoE 详细记录了训练崩溃的问题）。

5.4 一个重要的提醒：“FLOPs 少”不等于”便宜”

• 所有专家的参数都要放进显存。Mixtral 8x7B 每个 token 只用 129 亿参数，但要加载 467 亿参数（推理时尤其明显）。
• 路由带来的数据搬运和跨设备通信，在 FLOPs 里完全体现不出来。

6. MoE 的设计空间

MoE 之间的差别主要在三个方面：路由函数、专家大小、训练目标。

6.1 路由函数

几乎所有路由算法都可以归结为”选 top-k”，区别在于谁选谁：

$k$ 的取值：Switch Transformer 用 1；GShard、Grok、Mixtral 用 2；Qwen、DBRX 用 4；DeepSeek 用 6~8。

门控权重的两种算法（演示 4）：

• 先 softmax 再 top-k：对全部专家的分数做 softmax，取前 $k$ 个概率作为权重，权重之和小于 1。Switch、DeepSeek V1/V2、Grok、Qwen 采用。
• 先 top-k 再 softmax：只在选中的 $k$ 个分数上做 softmax，权重之和为 1。Mixtral、DBRX 采用；DeepSeek-V3 改用 sigmoid 打分，再对选中的归一化。

两种方式选中的专家相同，区别只在权重是否归一化。

6.2 专家大小：细粒度专家与共享专家

DeepSeek 和 Qwen 等团队（思路最早来自 DeepSpeed-MoE）的做法：

• 细粒度专家：把每个专家切成更小的几份，同时激活更多份。总参数和计算量不变，但可选的组合数暴增（演示 3：8 选 2 只有 28 种组合；切成 32 个细粒度专家选 8 个，有 1052 万种组合），每个 token 可以更灵活地”拼装”所需的知识。
• 共享专家：少数几个专家对所有 token 都激活，承载通用知识，让路由专家更专注于各自的领域。

真实模型的设置（来自 CS336 课件；“细粒度比例”= 每个专家相对标准 FFN 的大小）：

证据并不统一：DeepSeek 的消融显示更多（更细的）专家和共享专家都有帮助；OLMoE 的消融只发现细粒度专家有帮助，共享专家没有。

6.3 训练目标：负载均衡

核心难题：训练效率要求稀疏（每个 token 只算 $k$ 个专家），但”选 top-k”是离散的决策，不可导。可能的办法有三种：

1. 用强化学习优化路由策略——“正确”但方差大、复杂，不常用；
2. 随机扰动——Shazeer et al.（2017）在路由分数上加高斯噪声；Switch Transformer 用乘性的随机”抖动”，后来又被去掉了。这让专家更鲁棒，softmax 也让模型学会给专家排序；
3. 启发式的均衡损失——实践中大家用的就是这个。

为什么要均衡？ 训练 MoE 有”马太效应”：某个专家一开始碰巧好一点 → 被选得更多 → 训练得更好 → 被选得更多……最终少数专家包揽大部分 token，另一些专家”饿死”，参数白白浪费；在专家并行中，最忙的 GPU 决定整体速度。

Switch Transformer 的辅助损失：

$f_i$ 是实际分给专家 $i$ 的 token 比例（来自离散选择，不可导），$P_i$ 是路由器给专家 $i$ 的平均概率（可导）。对 $P_i$ 求导，梯度正比于 $f_i$：被选得越多的专家，它的概率被压得越狠。DeepSeek V1/V2 在此基础上增加了按设备的均衡项（专家分布在不同设备上时，保证每个设备的负载均衡），V2 还加了通信均衡损失。

DeepSeek-V3 的”无辅助损失”均衡：给每个专家一个偏置 $b_i$，只用于选择专家、不影响门控权重、也不参与梯度下降；每一步之后，负载高于平均的专家 $b_i$ 减小，低于平均的增大。这是一个简单的在线反馈控制，完全不改变训练的损失函数。（V3 仍保留了一个很小的”序列级”辅助损失作为补充，所以 CS336 课上说它”并不完全无辅助损失”。）

配套代码 05_load_balancing.py 在一个小回归任务上训练 8 个专家、top-1 路由的 MoE（数据来自 8 个簇，理想负载是每个专家 0.125），每种方法 4 个随机种子：

• 不做均衡时，多数种子出现了负载集中（最忙的专家承担约 1/4 的 token）和饿死的专家，任务损失也更差（浪费了专家容量）。
• 辅助损失需要调 $\alpha$：太小不起作用；调好后负载均匀。在大模型实践中，$\alpha$ 太大又会与任务目标争夺梯度。
• 偏置法同样实现了均衡，任务损失与调好的辅助损失相当，而且不改变损失函数。
• 这是一个很小的玩具实验，它说明”不均衡会发生、两种方法都能修正”，但分辨不出两者在大模型上的优劣。

去掉负载均衡会怎样？ CS336 课件引用的实验显示，去掉均衡损失后，专家的负载会迅速坍缩到少数几个专家上。

6.4 系统层面

• 专家并行：每个专家放在不同的设备上，token 通过 all-to-all 通信在设备间分发和收集（第 7–8 讲）。MoE 由此带来了一种新的并行维度。
• 稀疏矩阵乘法：MegaBlocks 等库（被许多开放 MoE 使用）用块稀疏矩阵乘法处理”每个专家分到的 token 数不同”的问题，避免丢弃 token 或浪费计算。
• 专家容量与 token 丢弃：传统实现为每个专家预留固定容量 $\lceil c\cdot Tk/E\rceil$（$c$ 为容量因子），超出的 token 被丢弃（这一层对它的输出为 0，它仍通过残差连接进入下一层）。演示 5：一个随机初始化的路由器，容量因子 1.0 时丢弃 10.4% 的 token，1.25 时 0.8%，2.0 时为 0。
• 额外的随机性：token 丢弃是在批次层面发生的——同一个 token，在批次 A 中正常处理，在批次 B 中因为排在 127 个同样想去这个专家的 token 后面而被丢弃（演示 5）。也就是说，别人的请求可能把你的 token 挤掉，这是 MoE 推理结果带有额外随机性的一个来源。
• 减少通信：Nemotron 3 等工作尝试在分发前先把激活值降维，以减少专家并行的通信量。

6.5 稳定性与微调

• 稳定性：路由器的 softmax 是不稳定的来源。常见做法是路由器单独用 FP32 计算（配套代码的 TopKRouter 就这样做），有时再加上 router z-loss（第 3 讲 9.1 节的 z-loss 用在路由器的 logits 上）。ST-MoE 的实验显示去掉 z-loss 后训练会出现不稳定。
• 微调时过拟合：稀疏 MoE 在小的微调数据集上容易过拟合。ST-MoE 的办法是只微调非 MoE 的层；DeepSeek 的办法是用大量的微调数据（140 万条 SFT 数据）。

6.6 从稠密模型”升级”：Upcycling

能不能用一个训练好的稠密模型来初始化 MoE？做法是把稠密模型的 FFN 复制成多个专家，再继续训练。

• MiniCPM-MoE：8 个专家选 2 个，约 40 亿激活参数，在基础模型上继续训练约 5200 亿 token 后有明显提升。
• Qwen1.5-MoE：从 Qwen-1.8B 初始化，60 个专家 + 4 个共享专家，top-4。与 DeepSeekMoE 的结构类似，是最早被确认成功的 upcycling 之一。

7. 案例：DeepSeek MoE 的演进

做出 DeepSeek-V3 还需要两项与 MoE 无关的技术：

MLA（多头潜在注意力，V2 引入）：把 Q、K、V 表示成一个低维”潜在向量” $c^{KV}$ 的函数。推理时 KV 缓存只需存 $c^{KV}$，比完整的 K、V 小得多；K 的上投影矩阵 $W^{UK}$ 可以合并进 Q 的投影，计算时甚至不需要显式恢复 K（训练时他们还压缩了 Q，以节省激活显存）。麻烦在于 RoPE 与这种缓存方式冲突：

解决办法：保留少量不压缩、专门做旋转的 key 维度。

MTP（多 token 预测）：用一个很轻量的模块，在预测下一个 token 的同时额外预测后面的 token，提供更密集的训练信号（V3 实际只多预测一个 token）。推理时这个模块还可以用于推测解码（第 10 讲）。

8. 本讲小结

注意力的替代方案

1. 去掉 softmax 后，$(QK^\top)V = Q(K^\top V)$，计算量从 $n^2d$ 降到 $nd^2$；因果版本是一个状态大小固定的 RNN。
2. 对偶性：训练用（分块的）并行形式，推理用递推形式。
3. 统一框架：$S_t = (\text{衰减/擦除})S_{t-1} + (\text{写入})$。从线性注意力 → RetNet（衰减）→ Mamba/Mamba-2（数据依赖的门控，SSD 对偶）→ DeltaNet（擦除后写入）→ Gated DeltaNet（两者结合）。
4. 固定大小的状态是一个容量有限的联想记忆，无法精确回忆任意长的上下文 → 混合架构（少量注意力层 + 大量线性层）成为主流；另一条路是稀疏注意力的事后改造（DSA）。

混合专家 5. MoE：参数多、计算少。计算量由 $k$ 决定，参数量由 $E$ 决定；但所有参数都要占显存，路由还带来通信开销。 6. 路由几乎都是 token 选择 top-k；专家越来越多、越来越细，常配共享专家。 7. 路由不可导，实践中用启发式的负载均衡：Switch 辅助损失，或 DeepSeek-V3 的偏置法。不均衡会导致专家饿死、GPU 负载不均。 8. 容量限制会导致 token 丢弃，并让结果依赖于同批次的其他请求；路由器用 FP32、z-loss 保证稳定。 9. DeepSeek V1→V3：更细的专家、更好的均衡方法，再加上 MLA 和 MTP。

9. 与作业的关系

作业 1 实现的是标准的稠密 Transformer，没有要求实现 MoE 或线性注意力。但作业 1 的排行榜允许你修改架构，本讲的内容是很自然的探索方向：

• 把部分层换成（门控）线性注意力，看看在固定的训练时间内能否处理更多 token；
• 把 FFN 换成 MoE（注意在单卡上，MoE 的优势是”同样 FLOPs 下更多参数”，但路由开销和显存开销会抵消一部分收益）。

配套代码中的 MoE 类可以直接替换第 3 讲 TransformerBlock 中的 SwiGLU（注意返回值形状一致）。

10. 练习

练习 1（结合律的收益）：$n = 32768$、$d = 128$ 时，$(QK^\top)V$ 和 $Q(K^\top V)$ 的 FLOPs 各是多少？比值是多少？$n$ 小于多少时，前者反而更省？

$(QK^\top)V$：$2n^2d + 2n^2d = 4n^2d \approx 5.5\times10^{11}$；$Q(K^\top V)$：$2nd^2 + 2nd^2 = 4nd^2 \approx 2.1\times10^{9}$；比值 $n/d = 256$。当 $n < d$ 时前者更省。

练习 2（递推形式）：写出因果线性注意力 $y_t = \sum_{s\le t}(q_t\cdot k_s)v_s$ 的递推形式，并说明推理时每生成一个 token 的计算量和内存与已生成的长度有什么关系。

$S_t = S_{t-1} + k_tv_t^\top$，$y_t = q_t^\top S_t$。每一步 $O(d_kd_v)$ 的计算和 $O(d_kd_v)$ 的内存，都与已生成的长度无关；softmax 注意力则需要 $O(td)$ 的计算和 KV 缓存。

练习 3（衰减矩阵）★：对 $S_t = \gamma_tS_{t-1} + k_tv_t^\top$，推导它的并行形式 $Y = ((QK^\top)\odot D)V$ 中 $D$ 的表达式。为什么实现时要用对数的累加和？

展开得 $S_t = \sum_{s\le t}\left(\prod_{r=s+1}^{t}\gamma_r\right)k_sv_s^\top$，所以 $D_{ts} = \prod_{r=s+1}^t\gamma_r$（$s\le t$），否则为 0。连乘很多个小于 1 的数会下溢，用 $\log D_{ts} = \sum_{r\le t}\log\gamma_r - \sum_{r\le s}\log\gamma_r$ 计算更稳定。

练习 4（Delta 规则）：设 $k$ 是单位向量，$\beta = 1$。证明经过 $S' = S + k(v - S^\top k)^\top$ 之后，$S'^\top k = v$。

$S'^\top k = S^\top k + (v - S^\top k)(k^\top k) = S^\top k + v - S^\top k = v$（因为 $k^\top k = 1$）。所以用 $k$ 读取时恰好得到新值 $v$，旧值被覆盖。

练习 5（并行扫描）：验证组合运算 $(a_1, b_1)\circ(a_2, b_2) = (a_1a_2,\ a_2b_1 + b_2)$ 满足结合律。为什么结合律是并行计算的关键？

$((a_1,b_1)\circ(a_2,b_2))\circ(a_3,b_3) = (a_1a_2a_3,\ a_3(a_2b_1+b_2)+b_3)$；$(a_1,b_1)\circ((a_2,b_2)\circ(a_3,b_3)) = (a_1a_2a_3,\ a_2a_3b_1 + a_3b_2 + b_3)$，两者相同。结合律意味着可以任意分组计算（比如先两两组合，再组合结果），不必严格从左到右，从而可以并行。

练习 6（MoE 参数）：一个 MoE 模型有 32 层，$d = 4096$，每层 64 个路由专家（每个是 $d_{ff} = 1408$ 的 SwiGLU）、2 个共享专家（同样大小），每个 token 激活 6 个路由专家。只考虑 MoE 层，总参数和激活参数各是多少？

每个专家 $3\times4096\times1408 \approx 1730$ 万参数。每层总计 $(64+2)\times1730$ 万 ≈ 11.4 亿，32 层约 365 亿；每层激活 $(6+2)\times1730$ 万 ≈ 1.38 亿，32 层约 44 亿。（这与 DeepSeekMoE 的专家设置类似。）

练习 7（辅助损失）：证明当所有 $f_i = P_i = 1/E$ 时，$E\sum_i f_iP_i = 1$；再举一个不均衡的例子，计算它的值。

均衡时 $E\cdot E\cdot\frac1E\cdot\frac1E = 1$。若 $E=2$，$f = (0.9, 0.1)$，$P = (0.8, 0.2)$，则 $2\times(0.72 + 0.02) = 1.48 > 1$。损失越大说明越不均衡。

练习 8（思考）：为什么说 MoE 的 token 丢弃会让同一个输入在不同时刻得到不同的输出？在”推理服务”场景下这会带来什么问题？有什么办法避免？

练习 9（动手）： (a) 在 01_linear_attention.py 中实现 Katharopoulos 的核函数版本：用 $\phi(x) = \text{elu}(x)+1$ 代替恒等函数，并加上归一化项 $y_t = \phi(q_t)^\top S_t / \phi(q_t)^\top z_t$，其中 $z_t = \sum_{s\le t}\phi(k_s)$。 (b) 在 02_associative_memory.py 的演示 1 中，把随机键换成正交的键（例如单位矩阵的行），$N \le d$ 时累加式记忆的读取准确度是多少？为什么？ (c) 在 05_load_balancing.py 中把 top-1 改成 top-2，负载不均衡的问题是更严重还是更轻？ (d) 把 04_moe_layer.py 的 MoE 接入第 3 讲的 TransformerBlock，在第 3 讲的名字复现实验上训练，和稠密模型比较。

11. 延伸阅读

• CS336 Spring 2026 Lecture 4 课件（GitHub stanford-cs336/lectures 中的 lecture_04.pdf）。
• Katharopoulos et al.，Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention（2020）。
• Sun et al.，Retentive Network（RetNet，2023）。
• Gu, Goel, Ré，Efficiently Modeling Long Sequences with Structured State Spaces（S4，2021）；Gu & Dao，Mamba（2023）；Dao & Gu，Transformers are SSMs（Mamba-2，2024）。
• Yang et al.，Parallelizing Linear Transformers with the Delta Rule over Sequence Length（DeltaNet，2024）；Gated Delta Networks（2024）。
• MiniMax，MiniMax-01 技术报告（2025）：7:1 混合线性注意力。
• Shazeer et al.，Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer（2017）。
• Fedus, Zoph, Shazeer，Switch Transformers（2022）；Zoph et al.，ST-MoE（2022）：训练稳定性。
• Dai et al.，DeepSeekMoE（2024）；DeepSeek-V2、DeepSeek-V3 技术报告。
• Muennighoff et al.，OLMoE（2024）：完全开源的 MoE，包含大量消融实验。

12. 下一讲预告

到这里，我们已经知道了模型长什么样、要多少计算。但同样的 FLOPs，在 GPU 上可能跑出完全不同的速度。下一讲（GPU 与 TPU）打开硬件的黑盒：

• GPU 由什么组成？SM、线程、warp、共享内存、HBM 分别是什么？
• 为什么”计算增长得比内存快”让大多数运算都受内存限制？什么是 Roofline 模型？
• 让 GPU 跑得快的六个技巧：避免分支发散、低精度、算子融合、重计算、合并访存、分块。
• 为什么矩阵边长从 1792 变成 1793，速度会突然下降？（波次量化）
• FlashAttention 如何用分块 + 在线 softmax，在不存储 $n\times n$ 矩阵的情况下精确计算注意力？
• TPU 与 GPU 有什么不同？

下一讲之前，建议你：跑通 code/lecture04/ 下的 5 个脚本，完成练习 3、4、6。