CS336 2026 · Lecture 4已发布

Attention alternatives and mixture of experts

日期: Wed April 8
讲师: Tatsu
课程资料: 9 项
内容产物: 4 项

课程简介

本讲讨论两条扩大模型能力的路线：用线性注意力、状态空间模型等结构降低长序列计算成本，以及用 Mixture of Experts 在近似固定 FLOPs 下增加参数容量。课程会拆解专家路由、负载均衡、容量溢出和通信开销，说明稀疏模型的算法收益为何必须与系统实现共同考虑。

讲义 / Slideslecture 4.pdf打开链接↗论文Transformers are SSMs - Generalized Models and Efficient Algorithms Through Structured State Space DualityMamba-2↗论文Switch Transformers - Scaling to Trillion Parameter Models with Simple and Efficient SparsityMoE 基础↗论文DeepSeekMoE - Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models课程案例↗论文Mamba - Linear-Time Sequence Modeling with Selective State Spaces状态空间模型↗论文ST-MoE - Designing Stable and Transferable Sparse Expert ModelsMoE 稳定性↗论文DeepSeek-V3 Technical Report现代 MoE↗论文Outrageously Large Neural Networks - The Sparsely-Gated Mixture-of-Experts LayerMoE 经典论文↗视频课程视频课程视频↗

CS336 Lecture 4 中文讲义：注意力替代方案与混合专家（MoE）

学习目标

完成本讲学习后，你应该能够：

理解线性时间注意力的核心思想——利用矩阵乘法的结合律改变对序列长度的复杂度依赖。
掌握从线性注意力到 Mamba-2、Gated Delta Net 的演进脉络，理解门控机制的作用。
了解 DeepSeek Sparse Attention（DSA）的 Top-K 索引器方法及其与线性注意力在思路上的根本差异。
理解 MoE（Mixture of Experts）的基本动机：在不增加 FLOPs 的前提下扩大参数规模。
掌握 MoE 路由的常见设计（Top-K、共享专家、细粒度专家）与训练中的负载均衡问题。
了解 MoE 的并行化维度、通信成本、稳定性问题及 upcycling 思想。

课程主线

本讲分为两个大主题：

第一部分：注意力机制的替代方案——从标准 softmax attention 的二次方复杂度问题出发，介绍线性时间注意力（linear attention）及其后续发展，并讨论稀疏注意力（DSA）作为另一条降低成本的路径。
第二部分：混合专家（MoE） ——将 MoE 视为更高效的 MLP 替代品，讨论其路由机制、训练难点、系统层面的并行化，以及 DeepSeek 系列模型中的具体设计。

两条线共同的底层思想是：稀疏性 + 廉价选择机制。线性注意力用固定大小的状态压缩历史，DSA 用轻量索引器选出少量 token 做全注意力，MoE 用轻量路由器为每个 token 选择少数专家。这些方法都试图在保持模型表达能力的同时，打破对序列长度或模型规模的直接缩放成本。

一、注意力机制的替代方案

1.1 背景：长上下文的计算瓶颈

大语言模型的上下文窗口正快速增加，但标准注意力（full softmax attention）的计算量随序列长度 二次方增长。前馈网络（FFN）的计算量是线性的，因此在短序列下 FFN 占主导，但一旦序列变长，注意力很快成为主要成本。典型的控制手段包括：

局部注意力 + 稀疏全局注意力：只在少数层（如每 8 层一次）做全局注意力，其余层用局部窗口。
系统工程优化：Flash Attention 通过重新安排注意力计算的内存访问模式，实现了约 2 倍的实际速度提升（常数因子改进），并避免了显式物化完整的注意力矩阵。

课程原意：常数因子非常重要。Flash Attention 没有改变二次方复杂度，但它在实际中带来的增益极为显著，说明不应只关注大 O 复杂度。

1.2 线性注意力的核心思想：乘法结合律

标准注意力可以写成：

Attention(Q, K, V) = softmax(Q K^T) V

忽略 softmax 后（设恒等映射），利用乘法的结合律重排计算顺序：

(Q K^T) V = Q (K^T V)

左侧的复杂度为 O(n²·d)，右侧为 O(n·d²)。由于序列长度 n 远大于隐藏维度 d（n 是百万级，d 是千到万级），右侧的复杂度依赖更有利。

解释：这里的关键不是让二次方消失，而是把二次方从 seq_len 转移到 d 上。实际维度 d 远小于 n，因此计算可行性大幅改善。

进一步地，右侧的矩阵乘法 K^T V 可以写成 RNN 风格的增量更新：

S_t = S_{t-1} + k_t v_t^T
y_t = q_t^T S_t

这种形式在推理时的状态 S 大小固定，不需要缓存整个 K/V 序列，非常高效。这就是**线性注意力（linear attention）**的基本形式。它的核心优势在于训练时可用并行矩阵乘法（对偶形式），推理时可用循环形式——二者在数学上完全等价。

1.3 Mamba-2：门控线性注意力

Mamba-2 本质上是对线性注意力的一个简单细化：

S_t = gamma_t ⊙ S_{t-1} + v_t k_t^T   （gamma_t 仅依赖当前输入 x_t，不依赖状态）

gamma_t 是一个输入相关的门控，控制多少历史信息被传向未来。
只要更新项仅依赖输入（无状态依赖），就保留了并行/串行的二象性。
关键：gamma 与状态无关，所以计算代价极低。

课程原意：Mamba-2 虽然最初从状态空间模型出发，但其运算机制可以看作对线性注意力加了一个简单门控。这是一种视角转换，而非严格的推导。

1.4 Gated Delta Net：双门控与投影更新

Gated Delta Net 在 Mamba-2 上进一步加入第二个门控 beta_t，控制当前 token 的信息是否进入状态：

S_t = (1 - beta_t k_t k_t^T) S_{t-1} + beta_t v_t k_t^T

beta_t 类似 LSTM 的输入门，控制当前信息写入状态的程度。
投影项 (I - beta_t k_t k_t^T) 的直觉：写入新信息的同时，擦除状态中与当前 key 方向相关的内容。这类似于一种最小二乘或快速权重编程的更新方式。

解释：课程特别强调，这个投影更新在多个领域中被独立重新发明过（元学习、快速权重编程、测试时训练），说明它是一个鲁棒的、根本性的更新规则。

实际应用：Qwen 3.5 系列采用了 3:1 的 Gated Delta Net 与全注意力的混合架构。混合架构在推理解码吞吐量上显著优于纯注意力模型，且不损害下游任务性能。

1.5 混合注意力架构的受控研究

关于混合比例（RNN 层与全注意力层的比例）的影响，目前受控研究仍然有限。已知的规律包括：

低混合比例（如 25%-50% RNN 层）下，性能损失几乎为零。
超过某一点后，长上下文检索任务（如 single-key retrieval）出现明显下降。
纯 RNN（全无注意力）在所有评估中都有显著性能下降。
不同架构的下降曲线有明显差异，Gated Delta Net 在低比例时表现最好。

注意：从线性注意力到循环形式是精确等价的；性能下降的根本原因是第一步丢掉 softmax 行归一化所造成的信息损失。

1.6 DeepSeek Sparse Attention（DSA）

DSA 是一种完全不同的思路——不是改变注意力计算本身的复杂度，而是减少需要做注意力的 token 数量：

先用一个轻量索引器计算所有 token 对的 qk 内积并加 ReLU，得到每个位置的得分。
用 Top-K 选出得分最高的少数位置。
仅在选出的子集上执行标准全注意力。

关键特点：

索引器的复杂度仍是二次方（必须扫描所有 token），但可以做得非常轻量：低精度、低维度投影。
Top-K 使得真正执行注意力计算的序列长度可控且有界。
训练时可以分两阶段：先正常训练标准 Transformer，然后在长上下文扩展阶段引入索引器并让模型适应它。
DSA 已被 DeepSeek v3.2 和 GLM-5 等模型验证为有效的低成本方案。

课程原意：DSA 在理论上仍是二次方的，但通过让索引器非常便宜（常数因子小）并限制后续注意力的规模，实际取得了很好的效果。这再次说明常数因子的重要性。

二、混合专家（Mixture of Experts）

2.1 MoE 的基本动机

MoE 本质上是一种"更高效的 MLP"：将一个大的 FFN 替换为多个小 FFN（专家），通过路由器为每个 token 只激活其中 K 个（或 K+共享专家）。这样：

总参数量增加为原来的 N 倍。
每次前向/反向的计算成本只对应 K 个专家的计算量。

核心收益：在不增加 FLOPs 的前提下获得更多参数，而经验上这通常带来更好的模型质量。

课程原意：MoE 不是概念上的革命，而是一个很实用的参数效率杠杆——即使同一时间只有部分参数处于活跃状态，更多参数本身通常是有益的。

2.2 路由机制设计空间

路由器（router）是一个极其轻量的模块，通常只是一个线性投影——计算输入与每个专家向量之间的内积，再经过 softmax 得到路由概率。以下是主要的设计选择：

路由方式	描述	使用情况
token 选择 Top-K	每个 token 挑选 K 个得分最高的专家	几乎一切现代 MoE 的标准做法
expert 选择	每个专家挑选它最偏好的 token	有少量成功案例，但不是主流
哈希路由	直接对输入进行哈希映射	论文中的基线，未在部署中使用
强化学习路由	将路由器当作策略（policy）学习	早期尝试，因开销大而不常用
线性分配求解	全局求解 token-专家最优分配	尚未大规模验证

Top-K 路由器具体实现：

s = softmax(W_r · x)     # 每个专家的得分
top_k_indices = topk(s, K)
gating = s[top_k_indices] / sum(s[top_k_indices])   # 归一化后作为门控
output = sum_i gating_i · FFN_i(x)

2.3 共享专家与细粒度专家

DeepSeekMoE 开创了以下设计，随后成为现代 MoE 的标准：

细粒度专家（fine-grained experts）：将专家切得更小、更多，使得每个专家的功能更专注。
共享专家（shared experts）：少量专家始终被激活，不参与路由——它们负责处理通用的、高频的模式，让路由专家更专精。

消融结果表明：细粒度化总是有增益；共享专家在某些任务上带来显著改进（如 TriviaQA、NaturalQuestions），但 OlMoE 的受控实验认为共享专家的帮助不显著。无论如何，现代模型（Qwen 3.5、GLM、DeepSeek v3）普遍采用了这一组合设计。

2.4 训练 MoE 的挑战

训练 MoE 的主要困难源于训练期间的稀疏激活：

路由器不可微（Top-K 是离散选择）。
无法观察到"反事实"——未被选中的专家没有梯度信号。
朴素训练会导致"富者愈富"：好专家被选中更多次、获得更多梯度、变得更强，最终导致专家坍塌（expert collapse）或专家饥饿（expert starvation）。

解决方案是一组启发式方法，核心是负载均衡损失：

Switch Transformer 的 auxiliary loss 可以写成：

loss_aux = N · sum_i f_i · P_i

f_i：被分派到专家 i 的 token 比例。
P_i：路由器分配给专家 i 的累积概率质量。

对 P_i 求导后可以看到，这个损失本质上是在梯度空间里压低那些获得 token 过多的专家——按比例消耗其概率质量。这是对"平衡"的一种平滑近似。

DeepSeek 的扩展：

逐专家平衡损失 + 逐设备平衡损失（保证多 GPU 上的负载均衡，避免某些设备空闲）。
DeepSeek v3 引入了在线学习的偏置项来替代部分辅助损失，但并未完全移除辅助损失。

移除负载均衡损失的灾难性后果（OlMoE 消融实验）：

训练损失显著上升。
几乎所有 token 被路由到少数两三个专家上，其余专家的参数形同虚设。

课程原意：负载均衡在早期 MoE 研究和代码中常常被忽略——这意味着某些专家会悄悄被丢弃，并可能静默返回零结果。现代框架已基本解决此类问题，但理解该机制的历史和必要性仍然重要。

2.5 微调 MoE 的稳定性问题

MoE 的路由中引入 softmax，而 softmax 的指数运算和除法操作在数值上容易失稳（与上一讲的稳定性讨论一致）。
常用策略：在路由器上使用 float32 精度、添加 z-loss 来约束路由 softmax 的分布。
MoE 微调时容易过拟合：稀疏模型在微调阶段训练-验证差距显著大于稠密模型。缓解手段包括：只微调注意力层、使用更多微调数据，或直接重用预训练数据。

2.6 MoE 的系统层面考量

并行化：MoE 提供了额外的并行维度——专家并行（expert parallelism），将不同专家放在不同设备上并路由 token 至对应设备。但代价是通信开销。近期的做法（如 Nemotron-3）是先对残差流降维再进行通信，以减少传输量。

计算模式：多个专家 + Top-K 激活，可以通过结构化的稀疏矩阵乘法（如块对角或更通用的结构化稀疏）在 GPU 上高效实现。这属于硬件-架构协同设计，也是 MoE 能快速落地的重要原因。

Upcycling（向上循环）：将一个训练好的稠密模型复制其 MLP 多次，随机初始化路由器，然后继续训练得到 MoE。早期实验表明这比继续训练稠密模型效果更好。但如今已被"直接训练 MoE"取代。

2.7 DeepSeek 系列演进

DeepSeekMoE v1：提出共享专家 + 细粒度专家，使用标准 Top-K 路由和辅助损失平衡。
DeepSeekMoE v2：扩大规模，细化到 2 个共享专家 + 更多细粒度专家，增加逐设备与通信层面的平衡损失，强调系统级的约束。
DeepSeekMoE v3：保留共享与细粒度设计，引入无辅助损失的平衡策略，并将专家加权方式改为 sigmoid + softmax。此外还引入：
- MLA（Multi-head Latent Attention）：将 Q/K/V 压缩为低维潜在变量 c，显著减少 KV cache 开销。相对复杂之处在于与 RoPE 的兼容性：位置编码需要施加在非潜在维度上。
- MTP（Multi-Token Prediction）：同时预测多个未来 token（本质上是把推测解码器训练进权重）。

三、课程作业关联

本讲内容与课程的注意力实现作业（如 Flash Attention）无直接依赖，但相关思想（常数因子的重要性、内存带宽优化）会在系统讲座和作业中展开。
作业可能要求你实现 MoE 路由和负载均衡损失，并体会"去掉负载均衡后训练崩溃"的现象。
如果你需要处理长上下文的序列建模任务，线性注意力的对偶形式（并行训练 + 串行推理）可能比标准注意力+Flash Attention 更省显存。

常见误区

把"线性注意力"理解为"完全丢掉 softmax 就行"：实际上，去掉 softmax 的归一化会造成信息丢失，这是线性注意力性能瓶颈的根源，同时也是后续所有改进（门控、投影等）的出发点。
认为 MoE 只是"多个 FFN 拼在一起"：路由机制、负载均衡、通信调度才是真正决定 MoE 能否工作的关键。
误以为"专家"具有语义知识（如"金融专家"、"日语专家"）：路由器只是基于内积的轻量打分，没有显式的领域知识。某些 token 类型（如标点、特殊字符）可能会集中到某几个专家，但这不是语义分工。
忽略负载均衡损失的必要性：只关注主损失而完全不使用任何均衡机制会导致训练迅速退化。
认为 DSA 是线性时间注意力：DSA 的索引器仍是二次方的，只是常数因子极小且后续注意力规模有界。复杂度上它不属于线性时间方法。

复习清单

[ ] 能否从矩阵乘法的结合律推导出线性注意力的对偶形式（稠密 vs RNN）？
[ ] 能否解释为什么线性注意力在推理时高效、训练时可并行？
[ ] Mamba-2 与 Gated Delta Net 分别在线性注意力上加了什么改动？各自的作用是什么？
[ ] 为什么混合注意力（部分 RNN 层 + 部分全注意力层）在低比例时性能几乎无损？
[ ] DSA 的工作原理是什么？它的索引器为什么可能比全注意力便宜？
[ ] MoE 的核心收益是什么？"更多参数但更少 FLOPs" 为什么有用？
[ ] 请写出 Top-K 路由的概率计算与门控融合步骤。
[ ] 什么是共享专家？细粒度专家？它们各自解决什么问题？
[ ] 为什么负载均衡损失是有效的？写出 f_i 与 P_i 分别代表什么。
[ ] 为什么 MoE 在微调时容易过拟合？有哪些公开的缓解手段？
[ ] upcycling 与直接训练 MoE 的差异与优劣。
[ ] DeepSeek v1 → v2 → v3 各自的架构变化是什么？
[ ] MLA 如何减少 KV cache？为什么需要特殊处理 RoPE？

课堂原话引用（选择性保留）：“注意力中的全对全连接极其强大，但如果你的状态大小远小于上下文长度，压缩必然会有信息损失——这是没有免费午餐的。”

CS336 Lecture 4 深度解析：从线性注意力到混合专家——突破 Transformer 的双重瓶颈

导语

当我们谈论现代大型语言模型（LLM）的能力边界时，两个数字始终如影随形：上下文长度和参数规模。用户渴望把整本书、整个代码仓库甚至一生的聊天记录塞进上下文窗口；研究者则痴迷于用更多的参数换取更强的智能。然而，这两个追求都撞上了 Transformer 架构的硬墙——注意力机制随序列长度呈二次方增长的计算成本，以及密集模型（dense model）中参数与计算量不可解耦的绑定关系。

在本课程第四讲中，我们深入探讨了突破这两重瓶颈的架构创新。线性时间注意力（Linear-time Attention） 通过数学上的巧妙重排，直指序列长度这一维度的复杂度问题；混合专家（Mixture of Experts, MoE） 则在参数维度上实现了“用更少的计算换取更多的智能”。这两个看似独立的方向，实则共享一个底层思维：在深度学习的世界里，常数因子与巧妙的结构设计，往往比纯粹的渐近复杂度更具决定意义。

本文将从注意力机制的成本困境出发，系统梳理线性注意力的数学原理及其在现代模型（如 Mamba、Gated Delta Net）中的演进，再深入 MoE 的架构设计与训练挑战，最后串联起本讲在整个 CS336 课程体系中的位置——为后续系统层面的并行化与推理优化打下坚实的架构基础。

第一部分：注意力机制的效率革命

1.1 成本困境：当二次方复杂度成为瓶颈

让我们直面一个现实：所有人都想要更长的上下文。无论是智能体（agent）需要处理多轮工具调用，还是模型需要携带更多知识，上下文窗口的大小已经成为衡量模型能力的核心指标之一。各大 LLM 供应商在时间轴上竞相扩大上下文规模，这背后是对更复杂工作负载的渴望。

但代价是什么？请看序列长度与计算成本的关系图：随着序列增长，前馈（feed forward）部分与注意力（attention）部分的比例发生戏剧性反转。在短序列下，前馈网络是计算主力；但注意力机制的对数线性增长——严格来说是二次方——很快便反超并碾压前馈部分。原因在于注意力的本质是全对全连接：每个位置都要与所有其他位置交互。当序列长度从 1,000 增长到 10,000，注意力成本增长 100 倍，而前馈部分仅增长 10 倍。

应对这一困境的工具箱中，第一个是局部注意力（local attention）：仅在少数层（例如每 8 层）使用全局注意力，其余层采用受限的局部窗口，从而有效控制成本。第二个被严重低估的工具是系统工程——本课程反复强调的一个主题。Flash Attention 便是绝佳案例：它通过重新安排注意力计算的执行顺序，最小化内存传输开销，将 PyTorch 原生实现的约 30-40 TFLOPs 吞吐量提升了约 2 倍。这不是理论上的复杂度突破，而是工程细节的胜利——在某些规模下，原生实现甚至因内存不足而无法运行，Flash Attention 却能继续推进。然而，当上下文长度冲向 500 万甚至 1,000 万 token 时，常数因子的优化终究有限，我们需要更激进的复杂度跳跃。

1.2 核心思想：乘法的结合性

如何让注意力机制的时间复杂度从二次方降至线性？答案藏在一个我们在小学就学过、但从未在深度学习语境中深思过的性质里：乘法的结合律。

让我们以紧凑的矩阵记号审视标准注意力： $A t t e n t i o n (Q, K, V) = softmax (Q K^{T}) V$ 。这里的维度标注： $Q, K \in R^{n \times d_{k}}$ ， $V \in R^{n \times d_{v}}$ ，其中 $n$ 是序列长度， $d_{k}$ 与 $d_{v}$ 是特征维度。计算瓶颈在于 $Q K^{T}$ 这一步——一个 $n \times n$ 的矩阵乘法，复杂度为 $O (n^{2} d_{k})$ 。

关键洞察来了：暂时忽略 softmax。如果去掉这一非线性归一化，纯粹的计算就变为 $Q K^{T} V$ 。根据结合律，我们可以重新安排运算顺序：与其先算 $Q K^{T}$ （产生 $n \times n$ 矩阵），不如先算 $K^{T} V$ ——一个 $d_{k} \times d_{v}$ 的矩阵乘法！

这意味着什么？原来 $n^{2}$ 这样可怕的二次项消失了，取而代之的是 $d_{k} \times d_{v}$ 。 $n$ 是数百万量级的上下文长度，而 $d_{k}, d_{v}$ 通常只有数千或数万——没人会在隐藏维度里放上百万个坐标。复杂度依赖从序列长度转移到模型宽度，这是质的飞跃。

但这还不是全部。如果我们将这种重排形式应用到自注意力场景，会发现更深刻的联系：右边的矩阵乘法 $S = K^{T} V$ 看起来就像一个 RNN。我不必采用左上角那种稠密矩阵形式，完全可以增量式计算：从上下文左到右扫描，每遇到一组新的 $K, V$ ，就更新状态 $S_{t} = S_{t - 1} + K_{t}^{T} V_{t}$ ，最后用 $Q S$ 得到输出。这两种形式在数学上是严格等价的。

这就是"二象性"（duality）的精髓：稠密形式支持并行训练（GPU 擅长的大矩阵乘法），循环形式支持高效推理（固定大小的状态 $S$ 只需一遍遍传递）。RNN 历来以推理友好、训练不友好著称，但线性注意力同时拥有两种形态，让我们能够两全其美。

1.3 从朴素到实用：门控的引入

朴素线性注意力的缺陷显而易见：它是线性（linear）的，意味着信息只能线性累积，无法有选择地遗忘或更新。这从直觉上就很脆弱。但修正方案同样直接——从 LSTM 的旧时代智慧中汲取灵感：门控（gating）。

Mamba-2 正是这一思路的代表。它的更新公式引入一个门控 $γ_{t}$ ，该门控仅依赖当前输入 $x_{t}$ （而非状态），因此计算极其轻量。它的作用如同一个可学习的遗忘开关：决定将多少当前状态携带到未来。关键的保真性质在于：只要门控项不依赖状态，二象性就依然成立——训练时可用并行形式，推理时可用循环形式。这一简洁的改进已在工业界落地验证：Nemotron-3 采用 Mamba-2 轻量层与标准 softmax 注意力层交替的组合，以管理推理成本与表达力之间的权衡。

Gated Delta Net 则更进一步，引入第二个门控 $β_{t}$ ，控制"当前信息进入状态的强度"。若 $β_{t} = 0$ ，意味着完全忽略当前 token 的输入——这是 LSTM 中输入门的流形变体。更为精妙的是其更新方向： $S_{t} = (I - β_{t} k_{t} k_{t}^{T}) S_{t - 1} + β_{t} k_{t}^{T} v_{t}$ 。蓝色项 $(I - β_{t} k_{t} k_{t}^{T})$ 本质上是沿着当前 key 方向投影掉旧信息。直觉上，当你要写入与当前 key 相关的信息时，你不仅想添加新内容，还想擦除任何先前写入的、与这个 key 重叠的旧内容。这是一个从"写入"思维到"写入同时清理"思维的转变。

值得玩味的是，这种更新机制已在多个看似独立的领域被重新发明：元学习（meta-learning）中的最小二乘问题、快速权重编程（fast weight programming）、测试时训练（test-time training）——不同的研究动机，收敛到同一数学形式。而 Qwen 3.5 和 Qwen Next 模型，作为最好的开源模型之一，采用 3:1 的 Gated Delta Net 注意力混合架构，在长上下文推理吞吐量上对 Qwen 3 展现出显著优势。

1.4 混合框架：从理论到实践

理解所有线性注意力变体的关键在于一个统一视角：在 RNN 形式中，只要所有依赖项都仅来自输入（而非状态），并行与串行的二象性就能存在。这一准则框定了设计空间，却也揭示了根本限制：当状态更新必须依赖过去的状态时，并行化就变得不可能。

因此，纯粹的线性注意力仍有局限。现实中的领先模型几乎都不是"纯线性"或"纯循环"的，而是混合架构：大部分层采用线性注意力（或状态空间模型），少数层保留完整 softmax 注意力。Minimax M1 采用 7:1 混合（7 层线性 + 1 层 softmax），在多个基准上逼近 O3 或 DeepSeek-R1 的水平；Qwen Next 采用 3:1 混合；Nemotron-3 同样交替使用两种层。

为什么需要保留少数 softmax 层？ 答案在于表达力与效率的权衡。当上下文越来越长时，压缩信息到固定大小的状态中必然造成损失。状态空间模型的本质限制是：状态大小与上下文长度的比值决定了信息密度。若状态等于上下文长度，则没有压缩损失，但代价是灾难性的高成本。没有免费午餐——想要小状态，就要接受压缩信息的损失；想要无损，就要付出大状态的代价。

课程展示的受控研究证实了这种权衡的残酷现实：随着非注意力层（RNN 层）比例的上升，性能呈现"平稳下降后加速恶化"的模式。在低比例（例如 25%-40%）时，混合架构的性能损失几乎可以忽略；但越接近纯 RNN，在长上下文检索等任务上的表现就越差。有趣的是，某些架构确实在各自专长任务上表现出色——每一位研究者都优化了自己最关心的指标。

1.5 另一个思路：稀疏注意力（DeepSeek Attention）

如果说线性注意力试图从理论上消除二次方，那么稀疏注意力则在常数因子上做文章。

DeepSeek 的 DSA（DeepSeek Attention）采用完全不同的哲学：先选后算。具体来说，训练一个轻量级索引器（indexer），从非常长的上下文中选出一个远小于完整序列的子集，然后仅在这个子集上执行标准的完整注意力运算。索引器的工作方式是：将 Q、K 通过降维投影，计算内积，经过 ReLU 激活，施加 Top-K 选择——每个位置对 k 个得分最高的位置施加关注。

这个机制的优雅之处在于系统工程视角的运用。有人会质疑：索引器不也是 $n^{2}$ 的吗？它当然必须查看所有内容才能决定选择什么，但其常数因子可以压缩到极小——通过更低的精度（如 FP4）运行索引计算，通过进一步投影降低索引器的维度。而底层真正昂贵的注意力部分，由于 Top-K 的限制，只需在很小规模的子集（ $n ≪$ 原始长度）上运行。这是"常数因子优先"哲学的又一次胜利——同一个问题，线性注意力通过改变复杂度依赖来解决，稀疏注意力通过昂贵的索引器与便宜的稠密部分分离来解决。

更令人惊讶的是，DSA 并不需要在预训练阶段就引入。DeepSeek 将索引器的学习放到独立的长上下文扩展（extension）阶段：先用标准 Transformer 完成预训练，此时模型学习基本的语言能力；然后在扩展阶段，注入轻量级索引器并让它适应，让模型学会"如何稀疏地看"。如果我们要做第二阶段（长上下文扩展），为什么不顺便把成本节约措施加进去呢？事实证明，这个看似吓人的不可微 Top-K 对象，在一个后续训练阶段竟然可以照常训练。

第二部分：混合专家（Mixture of Experts）

2.1 核心思想：用稀疏换规模

如果说注意力革命针对的是"能否处理更长上下文"，那么 MoE 革命针对的是另一个问题：能否在不增加计算成本的前提下拥有更多参数。

MoE 的概念本质并不高深：它就是一个更高效的 MLP。标准 Transformer 中，每个 token 的隐藏表示经过一个大 FFN 进行非线性变换。MoE 的替换方案是：将这个 FFN 扩展成 n 个独立的 FFN（称为"专家"），并通过一个路由器（router）为每个 token 挑选最合适的 k 个专家来执行计算（通常 $k$ 为 1 或 2）。

这个设计的收益立竿见影：4 个专家意味着 4 倍参数，但每次前向/反向传播只需 1 个专家的计算量。从参数中心的视角看，我们想要更多参数，因为更多参数往往意味着更强能力；但我们不愿为更多参数支付成比例的计算开销。MoE 成功地将参数规模与计算量解耦。

为什么 MoE 在近年如此流行？Fedus 等人的研究自 2022 年起揭示了一个引人注目的现象：保持计算量恒定，仅增加专家数量，语言建模损失持续下降。无论是训练速度还是最终性能，稀疏模型都优于同等计算量的密集模型。OlMoE 一年前的复现验证了这一结论：MoE 模型的训练损失、验证损失与下游任务表现全面领先，训练效率约为密集模型的 2 倍。

2.2 设计空间：路由、规模与训练

MoE 的设计可以沿三个维度变化：路由函数（如何选择专家）、专家规模（更多小专家 vs. 更少大专家）以及训练算法（如何处理不可微的路由决策）。

目前几乎所有部署模型都采用 Token 选择 Top-K 方式：每个 token 自行选择它最喜欢的 k 个专家。OlMoE 的消融实验清晰显示，Token 选择在验证损失和下游基准上均优于专家选择（Expert Choice）——后者由专家挑选它们偏好的 token。Token 选择更简单、更容易实现，且已被证实在实践中更可靠。

路由器的具体实现极为简单：一个线性投影（linear projection），为每个专家学习一个权重向量，然后计算输入与每个专家向量的内积，经过 softmax 后取 Top-K。这个模式——前馈层、得分计算、Top-K 选择——早已在 DSA 和 H-Nets 中反复出现，是一个重要的通用架构模式。

在专家规模维度上，DeepSeekMoE 开创了如今被广泛复制的设计：细粒度专家（Fine-grained Experts）与共享专家（Shared Experts）。细粒度专家将原始的大 FFN 切分为更多更小的专家，增加组合的灵活性；共享专家则始终处于开启状态，承担通用的信息处理任务，让路由专家更加专业化。这种设计优雅地解答了"专家到底应该有多专门"的问题——把通用的部分剥离出来固定，让路由部分专注于差异。

DeepSeek 的消融实验证实：细粒度切分与共享专家同时带来显著收益提升。虽然 OlMoE 对共享专家的效果略有分歧，但无疑认可了细粒度专家的价值。今天几乎所有的现代 MoE 模型（Qwen、GLM、DeepSeek v3 等）都遵循这一设计。

2.3 训练难题：不可微性与负载均衡

训练 MoE 的核心困难在于：我们在训练时就保持稀疏性——只激活一个或 k 个专家。这带来两个难题：其一，路由决策本身不可微（Top-K 是离散的）；其二，我们无法观察到"如果选择了另一个专家会怎样"的反事实结果。

理论上可行的方案包括：

强化学习：将路由器视为策略，用策略梯度优化。这确实可行，但梯度方差大、复杂度高，效果反而不如简单的启发式方法。
随机扰动（Noisy Top-K）：在路由得分上注入与输入规模成比例的噪声，允许近似打平的两个专家随机选择。这种探索机制在理论上合理，但后续消融显示，去除噪声实际上对模型稳定性与最终质量更有帮助。
启发式负载均衡：实践中真正使用的方法——一系列工程技巧的组合。

为何需要负载均衡？答案在于正反馈循环带来的坍缩：最初，某个专家可能因为随机初始化恰好擅长处理某些 token，从而获得更多梯度信号；于是它的权重被增强，变得更强；更强意味着被更频繁地选中；被更频繁选中又带来更多训练信号——最终这个专家"包揽一切"，其他专家被饿死。这就是"富者愈富"的专家坍缩或专家饥饿现象。

解决办法是在总损失上增加一个辅助平衡损失：不必期望它从第一性原理推导出来，而是通过理解梯度来把握本质。Switch Transformer 的平衡损失形式为损失 $\propto \sum_{i} f_{i} p_{i}$ ，其中 $f_{i}$ 是 token 被分派到专家 $i$ 的比例， $p_{i}$ 是分配给专家 $i$ 的概率质量。对 $p_{i}$ 求梯度，得到的是被分派到该专家的 token 比例。因此，这个损失在梯度空间中的效果是：某个专家获得的 token 越多，它获得的负梯度就越大——按比例压低受欢迎专家的概率质量。

DeepSeekMoE 将其扩展到设备级平衡：不仅专家要均匀，分配在每台设备上的负载也要均匀，以保证集群利用率。在 DeepSeek v3 中，他们用在线学习技巧与逐专家偏置项（bias）替代部分辅助损失，但依然需要保留少量负载均衡损失以确保极端不平衡不发生。OlMoE 的消融实验充分展示了负载均衡的关键价值：去除该损失后，几乎所有 token 都涌向两个专家，其余专家形同虚设——你购买了众多参数，却让它们闲置。

2.4 系统协同：MoE 的硬件优雅性与细微陷阱

MoE 在系统层面拥有独特的优势。当我们将多个专家放在同一 GPU 上时，稀疏矩阵乘法可以通过块对角矩阵或硬件原生的结构化稀疏形式高效执行——这是 MoE 与硬件协同设计的自然契合点。更重要的是，专家天然为模型并行提供了一个新的维度：每个专家就是一个自然的切分单元，可以分布在不同设备上，路由器负责将 token 路由到正确的设备——这就是"专家并行（Expert Parallelism）"的雏形。

然而，MoE 并非没有系统代价。通信瓶颈是专家并行的固有代价：token 必须从产生它的设备传输到存储专家的设备。Nemotron-3 提出一种优雅的缓解方案：在发送激活值前进行低维投影（降维），减少需要传输的数据量，虽然牺牲了部分表达能力，但显著节省通信开销。

另一个有趣且曾被忽视的问题是队列溢出与token丢弃。在早期 MoE 推理基础设施中，如果某个专家特别受欢迎，它的 token 队列会不断堆积，最终必须丢弃超额 token。更糟的是，多人共享同一专家队列时，用户请求可能因其他用户的查询而被悄然挤出——这是一种更隐晦的系统级随机性（nondeterminism）。现代框架（如 MegaBlocks）已经解决了这一问题，但这段历史提醒我们：稀疏性带来的随机性，远比理论模型所预想的更难以驾驭。

2.5 稳定性的额外考量

MoE 在架构中引入了额外的 softmax——路由器中的 softmax。而我们在上一讲已经强调：指数运算与除法是稳定性问题的危险区。MoE 的 softmax 操作因此成为架构中需要特殊对待的部分。

早期 Google 的 MoE 工作（Barret Zoph 等人）撰写了关于 MoE 稳定性的专项研究。常用的解决方案包括：在专家路由器上使用 Float32 精度，以及引入 z-loss——一种鼓励 logits 保持较小值的正则化项，防止 softmax 饱和与梯度消失。OlMoE 的消融实验证实 z-loss 对训练稳定性具有显著贡献，训练损失曲线从剧烈震荡变为平滑下降。

MoE 还存在微调泛化的特殊问题。由于参数数量巨大，稀疏模型在微调时极易过拟合——训练与验证的差距远大于密集模型。可行的策略包括：只微调注意力层（而非专家），或放弃微调整个 MoE，仅微调非 MoE 层。当然，"苦涩的教训"版本是最直接的建议：如果你有足够多的数据，直接在微调阶段重新训练 MoE——大的数据集可以完全消解稀疏模型的泛化差距。

2.6 从密集到稀疏：Upcycling 的智慧

Upcycling（向上循环）是一种独特的 MoE 初始化策略：与其从零开始训练 MoE，不如直接拿一个训练好的密集模型，复制其 MLP 作为多个专家，随机初始化路由器，然后继续训练。由于输入到 MLP 的随机性，这些复制出来的专家会逐渐专业化。

这一策略曾为早期 MoE 模型带来显著收益：MiniCPM 团队将 2.4B 密集模型 upcycle 为 13.4B MoE，Qwen 将 1.8B 模型 upcycle 为 2.7B MoE，均获得大量近乎免费的性能提升。为什么如今不再流行？因为既然最终目标是 MoE，不如一开始就用 MoE 训练——省去先训练密集模型再转换的额外步骤。但理解 upcycling 仍具价值，因为这揭示了参数初始化与稀疏性的深层互动。

2.7 案例研究：DeepSeekMoE 的演进

DeepSeek 系列是 MoE 设计的完美缩影。V1 提出细粒度专家与共享专家的原型架构，配合标准 Top-K 路由与辅助平衡损失；V2 扩展规模，引入设备级平衡与通信平衡损失，凸显"成功的训练不仅是深度学习，更是系统的哲学"；V3 则走向更精细的工程设计：保留共享与细粒度专家，但用 sigmoid + softmax 替代纯 softmax 加权，并尝试无辅助损失的平衡算法。

V3 还引入两个惊世设计：MLA（Multi-head Latent Attention） 与 MTP（Multi-token Prediction）。MLA 将 KV 缓存压缩到低维潜在变量 $c$ 中，避免缓存所有 $K, V$ ，仅存储潜在编码——唯一需要小心处理的是潜在维度与 RoPE 位置编码的兼容性。MTP 同时预测多个未来 token，既在统计学上可能带来更好的预测，又内置了推测解码（speculative decoding）的能力——这是推理加速的重要技巧，后续课程将会详细展开。

第三部分：本讲在 CS336 课程中的承上启下

回顾课程骨架：Lecture 3 从零构建了基础 Transformer，讨论了如何调整不同组件来构造现代语言模型；Lecture 4 则是在这个基础之上进行延伸与拓展——两个主题分别攻克模型的两个核心瓶颈。

理解线性注意力与 MoE 是迈向系统层面的必要桥梁。紧接着的课程中，我们将深入并行化（Parallelism）专题：数据并行、模型并行、张量并行与专家并行。MoE 在此扮演独特的角色——它为并行化提供了天然的维度。后续系统作业会要求我们处理网络拓扑、权衡通信与计算——这些不是抽象的理论推演，而是需要你理解"为何稀疏模型的通信瓶颈是不可避免的代价"。

推理课程同样需要本讲的知识：推测解码、KV 缓存优化、模型量化等主题，都与 MTP、MLA、稀疏计算等架构设计直接相关。例如，MTP 为推测解码提供了直接入口，MLA 大幅缩减 KV 缓存。更深层的视角是：本讲反复强调的"常数因子至关重要"与"系统工程决定成败"，将贯穿整个课程——Flash Attention 的 2 倍加速、共享专家带来的负载均衡、DSA 的轻量索引器，无不体现这一哲学。

结语：两个维度的突破，一种共同的智慧

回顾本讲的两条主线——注意力机制的效率革命与 MoE 的规模革命——我们能看到几个贯穿始终的思维模式：

第一，乘法的结合律不仅是数学工具，更是架构创新的透镜。 线性注意力通过重排运算顺序，将复杂度从序列长度转移到模型宽度；而 RNN 等价形式的发现，则让我们在并行与串行之间自由转换。

第二，现实世界的模型几乎都不是"纯"的。 混合架构（少量 softmax 注意力 + 大量线性注意力）已成为事实标准；MoE 中的共享专家 + 细粒度专家同样体现了"通用与专业并存"的设计哲学。最有效的解决方案往往是将多种策略融合，而非固执于某一种纯粹主义。

第三，工程细节与常数因子，往往比渐近复杂度更能决定成败。 从 Flash Attention 的 2 倍加速到 DSA 的轻量索引器，从 z-loss 到负载均衡损失，本讲反复印证：在语言模型的世界里，"大 O"固然重要，但"常数"同样决定命运。

第四，面对不可能，启发式方法往往比理论上更优雅的方案更有效。 稀疏不可微的路由问题，没有用强化学来解决，而是用简单的平衡损失和小技巧——现实世界的深度学习，常常是工程智慧与理论直觉的杂合体。

线性注意力与 MoE 并非互斥，而是从两个维度共同拓展了语言模型的边界：前者让我们可以处理更长的上下文，后者让我们在有限计算预算下获得更大的模型容量。正如课程最后所言："MoE 已经站稳脚跟，你应该理解它们是如何工作的。" 理解这些架构创新，不仅是知识上的积累，更是为未来设计更高效、更强大系统所做的必要准备。

我们正处在一个语言模型架构快速演进的黄金时代——从密集到稀疏，从二次方到线性。本讲所揭示的这些底层逻辑，将成为你理解下一代模型、甚至亲手设计新架构的基石。

CS336 2026 Lecture 4：完整中文逐字稿

视频：官方视频
来源：英文人工字幕提取后翻译为中文；保留段落起始时间戳，只修正明显的断句、重复和语序。
说明：这是逐字稿，不是讲义摘要；相邻且语义连续的片段已合并为阅读段落；术语保留必要英文。

正文

[00:05.200] 好，那我们开始吧。抱歉，今天这节课的讲义之前还没上传，不过现在已经传上去了。如果大家想随时在手边查看，应该可以通过课程网站获取。就算不行，GitHub 上也肯定已经有了。今天我们要讨论的，我认为是更进阶的架构思想。上一节课主要讲的是基础 transformer，以及我们如何调整基础 transformer 的不同部分来构建一个现代语言模型。现在，我想讨论的是在普通 transformer 之上更进阶、更复杂的一些发展。今天我要讲两件事——第一件是注意力机制的替代方案，通过架构上的改动来支持更长的上下文，这些改动大多能将时间复杂度从序列长度的二次方降到线性。

[01:06.558] 第二件我想讲的是专家混合（mixture of experts）的概念。所以第一部分是修改注意力模块。第二部分是修改 MLP 部分。专家混合能让我们在硬件利用效率上显著提升，也就是在相对于我们必须投入的计算量下，拥有更多参数。这就是今天的两大块内容和两个核心思想。那么我先从注意力开始讲。我来铺垫一下背景。现在很明显，大家都想要更长的上下文长度。你想把很多东西塞进上下文里，这样你的模型就能拥有更多知识。也许它是一个智能体（agent），需要处理很多东西。如果你看不同模型在时间轴上（这里是横轴）的上下文窗口大小，这是对数尺度，确实可以清楚看到很多顶级 LLM 供应商都在竞相提供越来越大的上下文规模，以支持这些更复杂的工作负载。

[02:05.530] 那么，我们如何控制这些成本呢？还不止这些。右边的图显示的是，随着序列长度增加，前馈（feed forward）部分和注意力（attention）部分之间计算成本的比例。你可以看到前馈部分一开始占比相当大，但注意力部分是线性增长的。它是一种全对全的连接，连接所有不同位置，是二次方的，所以随着序列长度增长，很快就超过了前馈部分。所以以前对于大模型来说，在较短的序列长度下，前馈是主要成本。但当你走向更长的序列长度时，注意力越来越成为问题。那么，我们控制这些的基本工具是什么？我们已经讨论过这些了。好吧，我们讨论过其中一个，我们还会讨论另一个。我们基本上可以用局部注意力（local attention）这类方法，然后用不同方式混合组合，来控制全局注意力的高昂成本。

[03:09.770] 如果你只在每八层中做一次全局注意力，而其他所有层都用这种非常局部的注意力，那么你就能很好地控制成本。另一件可以做的事情是系统工程。我认为一个非常被低估的事实是，常数因子真的非常非常重要。这门课的主题之一就是，你需要注意细节。对于接受过更经典的理论导向计算机科学训练的人来说，很容易觉得，没错，重要的是大 O 复杂度，是线性还是二次方。但我认为在上下文和注意力成本方面发生的最重要的事情之一就是 Flash Attention。Flash Attention，我们会在系统相关的讲座和系统作业中更详细地讨论。但 Flash Attention 本质上只是一种非常巧妙的方式，重新安排了注意力操作的执行方式，使其更符合系统的需求，从而最小化内存传输的开销。

[04:07.540] 通过这样做，你可以获得在注意力性能上非常非常显著的提升。所以基础的 PyTorch 是这条蓝色的线。你可以看一些较短的序列长度，你大概能达到每秒 30 到 40 万亿次浮点运算（teraflops）。然后你使用 Flash Attention，这实际上是一个常数因子上的系统工程改进。然后你就能获得非常大的提升，大约 2 倍左右的改进。而且在某些规模下，你甚至都无法将它放进内存了。但有了 Flash Attention，你不再需要具体化这些大的注意力矩阵，你可以继续运行，尽管速度会慢一些。这并不能解决任何二次方成本的问题，但常数因子是非常非常强大的。话虽如此，如果我们要处理 500 万、1000 万个 token，这些技巧可能就不够了。我们想要在处理长上下文的能力上有更激进、更大的提升。

[05:01.360] 这就引出了一个问题，有没有办法让我们对序列长度实现线性时间或线性依赖？那些方法会是什么样子呢？就是这样。事实证明，多年来有很多人尝试了很多不同的方法。我认为有一些是失败的起步。但在最近几年，我会说，大概过去两年左右，出现了一套在线性时间注意力方面相当有效的方法，这些方法现在已经在较大规模上经过了实战检验。这就是为什么今年是我第一次讲线性时间注意力，因为现在很明显，这些东西确实在大规模和生产环境中是有效的。要真正解释所有这些效果不错的不同方法，我认为你只需要从一个核心概念开始理解就行了，然后我们会在这个概念上进行阐述和构建。这个概念就是乘法的结合性。

[06:03.050] 我们来想想注意力是怎么运作的。我们可以用这种紧凑的记号来写，我们说，注意力操作将作用于 Q、K 和 V，上面这些是矩阵维度。我要把我的 Q 和 K 拿来，进行一个全对全的交互。所有 n 个位置都会通过矩阵乘法相互交互。然后我会通过 softmax 对它们归一化，就是这里的这一行。然后我会把它和我的值 V 相乘，得到注意力输出。这就是经典的自我注意力（self-attention）或者其他任何注意力操作。现在，我们先暂时忘记注意力里有 softmax。我直接去掉那一行。暂时忽略它。如果我们没有这一行，也许我们真的可以改变这个操作的复杂度。如果你看注意力，它是 n 平方乘以 dk，因为我们做的是 QK 转置的乘法。

[06:58.810] 但如果那个 Rho 是恒等函数，我就可以稍微移动一下括号。我可以说，好吧，与其做 QK^T V，实际上我要做的是 QK^T V。如果我们做这个重新排序，嗯，我们就改变了哪一部分是二次方的。与其有那些 n 平方项，它们是很糟糕的，因为这些 n 非常大。它们是数百万级别的。这是上下文长度。相反，我可以依赖 ndvdk。现在，我们本质上是在乘键和值的维度。这些原则上可能很大。我可以有一个很大的模型。但这些通常是在数千、数万的量级。没有人会在隐藏维度里有上百万的坐标。所以相较于我们做原始注意力操作时的这个 n 平方项，这是一个更有利的依赖项。

[07:53.660] 所以这是你们应该知道的第一件事，即乘法的结合律让我们能够改变对序列长度的依赖关系。现在，我们能做的另一件事是，如果我们考虑这种重排，即 QK 转置再乘 V，这里更妙的一点是——我认为这确实引领了接下来几张幻灯片中我将提到的许多研究——在于，在自注意力情形下，右边的这个矩阵乘法看起来就像一个 RNN。我不必像上面那样写成这种稠密形式，我可以增量式地写它。当我在上下文上从左向右扫描时，我可以乘上我的 K 和 V，并更新这个S 对象，它是一个状态，我可以把它加到之前的状态上。这就是在增量地乘累加我的 KV。然后我可以用它乘上我的 Q 来得到输出。而上面这个稠密运算和这个 RNN 形式是等价的。那么我们做了什么呢？

[08:58.280] 我们从左上角的线性注意力出发，然后我们以某种方式得到了一个看起来非常像 RNN 的东西。RNN 有什么好处呢？嗯，RNN 相当不错，因为在推理时，我有这个 S，即我的状态，它的大小是固定的。我只是一遍又一遍地把它向前传递。所以 RNN 在推理方面非常友好。它们在训练方面就不那么友好了。但线性注意力的好处在于，你可以用任何一种方式来使用它们。你可以用这种稠密形式，这对训练来说很棒，它是并行的。或者你可以用这种类似 RNN 的串行形式，这对推理来说很棒。这样你就可以两全其美。至少就所有这些对象的依赖性而言，这很棒，这非常妙。不幸的是，这是线性的，这不太好。

[09:46.600] 这确实是我们讨论的起点。到目前为止有问题吗？这确实是基础。有一件很酷的事情是，已经有人基本上将这种线性注意力用作其注意力机制的一部分。比如 Minimax M1，这是一个规模相当大、性能相当高的中国开源模型，它采用了 7 比 1 的混合方式，也就是七层线性注意力加上一层完整的 softmax 注意力。性能总体上是强的。你可以把它和许多强模型比较，比如 OpenAI 的 O3 或 DeepSeek-R1，它相当有竞争力。而且因为它是线性注意力，大部分对上下文长度的依赖是线性的。它不是完全线性的，因为里面还有这些 softmax，但依赖关系要温和得多。

[10:42.790] 所以这是一个例子，表明即使这种非常容易理解的基础线性注意力，只要你有一些完全的二次连接参与其中，也能在较大规模上相当有效地使用。到目前为止，还没有人真正在规模上证明完全线性时间的注意力机制。我接下来几张幻灯片要讲的所有内容都是混合式的。那么现在，我们已经讲完了线性注意力。但你可能会说，线性注意力对我来说太天真、太简单了。这不是我想要的。我想要一些更严肃、更复杂的神经网络东西发生。嗯，你现在可以做的是——哦，抱歉，我有点扯远了，看一下线性注意力的 RNN 形式。我把线性注意力方程放在这儿了，这样你就不用来回翻了。我们可以在线性注意力上加一个小的细化，以获得更具表达力的更新。

[11:46.770] 这就是你最终得到 Mamba-2 的方式。你们中有些人可能听说过 Mamba、Mamba-2，现在我想还有 Mamba-3。它们是由 Albert Gu、Tri Dao 及其合作者提出的一系列状态空间模型。它们最初源于状态空间理论和视角。但实际上，如果你看其背后的具体运算机制，你会发现它实际上可以看作是对线性注意力机制的一个非常简单的细化。所以思路如下。我们从线性注意力开始。你可能会说，线性注意力的主要问题在于我一直把我的状态向前传递。我们从 LSTM 的旧时代就知道，什么时候传递信息、什么时候不传递信息——即遗忘某些东西并把它们归零——是很重要的。所以我要做的，是受到这个启发，加上一个门控，gamma 下标 t。

[12:40.480] gamma 下标 t 不是有状态的。gamma 下标 t 只依赖于我当前的输入 x 下标 t。而 gamma 下标 t 将调节我携带多少当前状态进入未来。现在，如果你想一想，状态相关的项全都是这些 S。这个 gamma 完全不是状态相关的，所以这意味着它非常容易计算。而且这里仍然存在这种二象性。你可以把这个 Mamba-2 项计算为一个大稠密矩阵乘法，或者在推理时使用这种循环形式来获得很好的推理效率提升。所以这就是这个想法的核心构思或视角，即你可以像线性注意力一样，以并行或序列的方式使用它。而且人们现在已经在较大规模上使用了它。比如 Nemotron-3，我上一节课简单提到过，就使用了这种将 Mamba-2 作为其轻量层组合起来的思路。

[13:37.773] 当然，你时不时也会用上你的大 softmax 注意力。他们以各种方式交替使用这些层，以管理推理成本和表达力之间的权衡。与 Qwen3 thinking 相比，与 GPT OSS 相比，Nemotron-3 也取得了相当不错的性能。而且由于所有这些 Mamba-2 层，在相当大的上下文长度下仍能获得不错的吞吐量。所以你可以看到它是有效的。至少，这些是小型前沿模型，不算真正的大型前沿模型，但它们是开源的较小前沿模型。不错。那么现在，我们已经在这个门控思想上做了拓展，这很酷。不过你可能会说，我们能不能继续沿着这个思路推进，把这个循环机制做得越来越复杂？在某种意义上，你在这里所能够做的操作的准则是，只要你在所谓的 RNN 中得到的各项仅仅依赖于输入的项——

[14:44.310] 也就是没有状态依赖——那么你就仍然拥有这种相当不错的二象性，即可以用于训练的并行运算和用于推理的串行运算之间的二象性。那么 gated delta net 可能是现在使用最广泛的状态空间模型之一了，我想说。我觉得它已经在不少论文里被测试过了。而且对于这些模型，有一个规模扩展到 3.5 的非常好的大规模实践。那么让我讲讲 gated delta net 相对于 Mamba-2 的区别。Mamba-2 正如我上面所写的，我们取这个线性循环，这个状态更新。而我要做的是，用我的 gamma 来给它加门控。哦，我忘了解释这个 v 下标 t。抱歉，我为此道歉。我加它是为了完整性。Mamba-2 还有一个修改，他们在 y 下标 t 的更新上加了一个基本上是残差层的东西。

[15:43.530] 它对于状态更新或者他们论证为什么这是个好主意的论据来说，并不是核心内容，但我把它放进来是为了完整，确保我给出的是真正的 Mamba-2 更新公式，而不是一个只传达核心思想的精简版。但目前你可以忽略这个 v 下标 t。这实际上就像一个架构上的改进，并不是状态更新的核心。Gated delta net 在某种意义上沿用了这个想法，并加入了第二个门控，beta 下标 t。那么 beta 下标 t 是做什么的？我把 beta 下标 t 看作一个无输入的操作门控。如果 beta t 是 0，那基本上意味着，不要把我当前的任何信息加入进来，不要把它加进我的状态里。我觉得你们中大多数人可能都熟悉 LSTM。这非常让人想起 LSTM 里的思想。你有门控，控制是否要遗忘，以及是否要把当前时间步的信息放进来。

[16:39.960] 所以它的运行原理非常相似，尽管，我想，最初推导的方式非常不同。现在，关于门控 Delta Net，另一个有趣的地方——这个来自 Delta Net 而不是门控 Delta Net——是这个更新方向。所以不仅仅是做门控，你本可以把它做成 1 减去 beta t，比如说。他们做的是，如果你要做一次更新，把 beta 放进来，你要做的是，你要尝试在某种意义上投影掉我当前的 key 方向。所以这里的直觉是，我要从我当前的 key，也就是 k 的 t 时刻值，写入信息。而在这样做的时候，我不仅想放入新信息，我还想擦除任何之前进入过的、与当前 key 相关的旧信息。这是你思考更新的一种方式——不只是放东西进来，我想要清掉已经进去的当前 key。

[17:28.993] 这就是为什么有这个蓝色的项，这个单位矩阵减去 beta t kt kt 转置，它本质上起着投影器的作用，把 kt 维度上的东西投影掉。这不完全准确，因为你并没有做单位归一化这类事情，但你可以大致理解这个直觉。在继续讲门控 Delta Net 之前，我还要提一件事——哦，抱歉，我这里有个蓝色框，我继续往下。就是这已经被在多种场景中以某种方式重新发明过了。这个更新、这个投影器，在你尝试解决某些类型的元学习最小二乘问题时就会出现。它也在快速权重编程或测试时训练中以其他形式被重新发明过，那些领域的研究者基本上是通过非常不同的设计原则，最终得到了完全相同的解决方案。这里我要说的是，Qwen 3.5 和 Qwen Next 模型，也就是之前的那一代模型，

[18:28.190] 是当今可用的最好的开源模型之一。它们使用的正是这种架构。它们用的是 3 比 1 的门控 Delta Net 注意力混合架构。性能非常合理，推理特性也很强，你可以在右边看到。随着上下文长度不断增加，Qwen Next 的解码吞吐量相对 Qwen 3 高很多。中间的图是与各种闭源模型和上一代模型的性能对比。这种混合架构似乎完全没有损害性能。最后我要说的是，关于混合架构表现如何，还没有太多好的受控研究。可能有一个来自字节跳动 Seed 和 UC Santa Cruz 的人做的研究，他们比较了 Mamba-2、门控 Delta Net 以及其他架构，看混合层数量如何影响表现。有些结果，我得说，比较混乱，但我要指出的是，右边这个图里的这条虚线，这个家伙是全注意力性能。

[19:34.870] 然后你从左往右看，这是增加非全注意力层的数量。也就是那些 RNN 层，我们就这么叫它们，往右走时你增加它们，然后你会看到性能下降。但也许这里值得注意的一点是，对于一些你拥有的最好的架构——黄色、橙色、蓝色这些，这些像是门控 Delta Net 以及各种其他变体。你会看到在低比例时，基本上没有任何损失。然后一旦你过了某个点，在长上下文性能上就开始看到更明显的下降，在这个例子中是这样。基本上到最后，当你走向全 RNN 时，所有这些架构都会有非常明显的性能下降。在各种其他评估中也是如此。其中一些，我觉得，已经被这些架构本身优化过了。单键检索是一个我认为所有这些长上下文架构都明确优化的任务。但如果你看 QA 性能，我们能看到同样的故事，随着混合比例增加，性能会相当稳定且清晰地下降，最终到达这个纯 RNN 的点。

[20:36.710] 关于线性时间注意力这块，有什么问题吗？我还有最后一个高效注意力的事情要讲，但请说。［听不清］看起来像是注意力的并行形式对比线性形式。［听不清］它们在数学上不应该是等价的吗？是的。那为什么你会看到性能下降，如果是……？抱歉。我说的循环，是指一般的 RNN，任何循环形式。所以门控 Delta Net 就是一个循环。它和上面这个不是等价的。这是带行的完整 softmax 注意力。第一步是有损的，我们丢掉行，然后我们变成线性的。这是所有这些方法最开始的第一步。所以这部分会是有损的。然后在那之后，从线性形式到循环形式的这种等价，那是精确的。

[21:36.100] 好。是。［听不清］哦，对。是的。对，就是这个项，对吧？是的。所以在线性注意力形式中，我们实际上做的是，我们取这些 kv，用队列（queues？还是 queue）累积它们。在对偶形式中，这就是它的运作方式。但你实际上可能希望在你的输出中也有来自当前时间步的一些值。你可能会想，哦，我当前时间步的值，对我这个 token 来说，实际上对下游的各种事情是有用的。所以这东西有点像一个残差连接，我当前输入的值直接传到了注意力输出里。而 D 是另一个调制门控，你可以控制发生的直通量。好的。好。所以这就是线性时间注意力。有很多非常有趣的东西。同样有趣的是，很多被测试过的方法，最终确实有这种简洁、简单的线性注意力风格、循环形式。

[22:41.440] 有一些例外，但大多数情况下，这些方法最终都收敛到了几乎很像 LSTM 的对象。我想说什么来着？好。所以还有另一种替代方案。我讲的是线性注意力，它有一种非常好的复杂度理论味道。我想讲另一种注意力优化，我还没讲过，但它属于这一类重要的高效注意力方法。所以 DeepSeek，在 DeepSeek 3.2 中，我记得他们有 DSA，不是 Do Sparse Attention，那是 DeepSeek 注意力（DeepSeek Attention），他们做的是，不是计算注意力——这不是有意为之的部分——不是对所有 token 计算注意力，而是你会先有一个轻量级的索引器，用来选出一部分 token。所以你会看你的非常长的上下文。索引器会从中选出一个子集，比完整序列小得多。

[23:45.973] 然后你在这个较小的子集上做全注意力。所以这个机制相当简单，至少在正向传播中是这样，你做的就是你有正常的 Q 和 K，但你把它们通过这个索引器。索引器做的是，它取 qk 内积，加一个 ReLU。然后它会有这些由前面 token 推导出来的权重。然后这会给你每个位置的激活值。然后你把这个通过一个 top k，其中一定数量的这些 top 位置会被包含进你的注意力计算中。然后我会执行我通常的注意力计算，它会产生一个键值对。这样做的好处在于，它非常非常稀疏。如果这个索引器足够轻量，那么底下这个执行完整注意力的部分，就只需要在一个非常小的子集上完成即可。

[24:43.050] DSA 和 GLM 还共同表明了另一点：你其实不一定要在训练阶段就用上这个机制，因为用这个索引器来训练模型可能会相当麻烦且复杂。你的做法是，先正常训练一个标准的 Transformer 模型。然后，在进行长上下文扩展时，你再引入这个轻量级索引器，并且在一个与预训练分开的扩展阶段里，让模型学会适应这个索引器。因此，实际上，这样做本身的计算成本是非常低的。我想特别提一下 DeepSeek 的 v3.2 版本，就是他们提出 DSA 的那个模型，这个模型表现相当出色。

[25:22.400] v3.2 版本在当时的水平可以匹配其他所有前沿模型，比如 Claude 4.5 Sonnet 或 Gemini 3。而且他们展示了一张非常漂亮的图表，表明无论是处理输入时的预填充阶段，还是实际生成内容时的解码阶段，其扩展性相比前一代没有使用稀疏注意力的 DeepSeek 模型都有非常显著的优势，而且这很合理。如果你只对一小部分内容进行注意力计算，那么注意力机制的成本自然会大幅降低。我认为另一个很好的正交验证案例是去年发布的 GLM5。我觉得 GLM5 是目前可用的最好的开源模型之一，就这么简单。他们也采用了这种 DSA 方法。而且，如果你感兴趣的话，他们的论文里有相当不错的消融实验，他们比较了不同情况，比如带有些许预热步骤的 DSA、直接使用 DSA，或者完全不用 DSA，结果会怎样。

[26:26.602] 他们主要展示的是，如果进行完整的 DSA 训练，与完全注意力相比，性能损失非常小，即使在那些对循环神经网络架构来说相当困难的长上下文检索任务上也是如此。所以，这提供了一种完全不同的视角来降低注意力成本。注意，我退回上一张幻灯片，这个方法的复杂度并不是线性时间的。索引器仍然需要对所有的内容进行两两注意力计算，因为那些 qk 内积的计算成本不可能是线性的。但是，你可以让这个索引器变得便宜得多，例如，通过降低它的维度或者让这个 w 矩阵变得非常小。所有这些做法都能在很大程度上降低其成本。所以这是一个完全不同的思路，而且效果非常好。我不打算深入展开讲这个。但事实证明，这种 Top-K 选择的思想，将会是接下来讲座这部分内容的核心。在继续讲 MoE 之前，我想先停下来，看看大家对于高效注意力还有什么问题。

[27:29.750] 好的。那么索引器的时间复杂度是 n 平方吗？[听不清] 是的。我的意思是，为了知道要选择什么，它必须查看所有内容。这里没有什么巧妙的状态转换技巧，它真的就是蛮力计算内积。所以如果你要遍历所有这些内容，那么后续步骤也同样是二次方的复杂度。我们怎么能说时间上[听不清]有优势呢？我的意思是，这几乎就像系统 z 的技巧。你会让索引器变得非常轻量。有很多方法可以做到这一点。你可以在更低的精度下运行它，这是他们论文里提到的。你可以通过进一步投影 q 和 k 来降低维度，仅用于索引目的。所以，这里有很多你可以想象的技巧，但关键在于索引器会非常轻量。

[28:18.890] 所以常数因子是非常好的。然后对于第二步，即使它是二次方的，它也是作用于一个较短的上下文长度上的，因为通过 Top-K，我们可以控制 K 的大小。所以，现在这部分虽然代价高，但规模小。这再次说明了，有时候，不要太纠结于二次方还是非二次方。有时候，常数因子真的非常重要。还有别的问题吗？[听不清] 没错。或者它是持续预训练式流程的一部分，因为对于所有模型，你不会从头开始训练长上下文，出于计算和其他各种原因，你会先训练一个较短上下文的模型，然后才有这些长上下文扩展阶段。所以这里的一个巧妙想法是，我们反正都要做第二阶段。为什么不把长上下文的成本节约措施同时加到这个阶段呢？而且令人惊讶的是，这居然有效。说实话，他们能够把这个看起来相当吓人的、不可微的 Top-K 对象直接加进去。

[29:16.630] 但正如你将在下一部分看到的，也许你不应该对 Top-K 选择那么害怕。[听不清] 所以这也是和其他后训练步骤一起做的一部分吗？是的。据我所知，许多开源模型的常规流程通常是这样的：先进行短上下文预训练，然后是长上下文扩展，最后是后训练。是的。这里的 k 大小，和短的[上下文长度]类似吗？[听不清] 是的，没错。你会选择更接近短上下文性能的值。而且你会给它设定一个界限。所以无论输入大小如何，它都是有界的。还有问题吗？好的。移除 softmax 是否会导致[听不清]训练不稳定？训练稳定性问题。这是个好问题。我不这么认为。据我所知，没有文献记载线性 softmax 会导致很多注意力问题。如果有的话，通常 softmax 本身反而会因为各种原因带来更多问题。所以我认为这实际上会改善架构的稳定性。

[30:23.960] 你能预测下一个[听不清]注意力吗？考虑到[听不清]注意力机制的需求，因为它计算成本低、延迟低。鉴于你的预测，从高层次来看，哪些类型的特征会[听不清]？是的。我不确定我对未来的注意力架构有什么很好的预测。我有点敷衍的回答是，我认为很大程度上我们会把这些技巧都整合进去。就像架构通过吸收所有成功的配方而变得非常复杂一样，我认为在更高层面上，基本上是利用后训练来让模型管理自己的上下文。压缩、检索，这些东西基本上都是加在顶层的一个额外层次。我认为很多工作会是这些不同层次的整合，因为我们看到，至少线性时间注意力方面的研究已经大量收敛到类似 LSTM 或线性注意力的架构上。而且我认为在不久的将来，这种情况不会有太大改变。好的。哦，是的。最后一个问题。你能谈谈低精度注意力吗？FP4 注意力是否可行？我认为索引器的想法部分就是受这类事情启发的。

[31:54.600] FP4 注意力，绝对可行。我觉得缺点在于，对于像 softmax 这样的操作，精度的损失可能成为一个问题，比如小数下溢、溢出这类情况真的会开始产生影响，有时会相当麻烦。所以，我认为这种动机，即“我们用低精度做选择，然后再用全精度进行计算”，是因为这样值向量可以更精细地加权和累加。这是一种有道理的方法。我不知道是否有真正好的方法能在 FP4 甚至更低精度下完成完整的注意力计算。好的。好吧，最后一个问题了，然后我就继续往下讲。那么与Transformer相比，与状态空间模型相比，训练更复杂的状态空间模型的缺点是什么？是的。这是个好问题。一定得有某些缺点，否则大家就都——嗯，我认为缺点在于表达能力。

[32:57.810] 软注意力中的全对全连接极其强大。它也非常容易训练。我认为，与比如LSTM相比，注意力曾经有一个非常显著的优势，就是硬件效率。你可以并行训练。我认为状态空间模型之所以流行起来，尽管它们与LSTM有相似之处，并且在表示能力方面也有类似的缺点，是因为现在，由于我之前提到的线性注意力，存在一种被充分理解的二元对偶关系，你可以从这种RNN形式转换到更稠密的矩阵乘法形式。而矩阵乘法形式可以实现计算效率。所以那个权衡已经被解决了。但你仍然面临这样一个权衡：如果你有一个有限的状态，并且你必须把所有信息都通过它传递，相对于直接携带所有信息，你一定会丢失一些信息。[听不清]在状态空间中将会是一个问题。

[33:50.690] 所以[听不清]上下文本身，[听不清]。嗯，在状态空间中，你基本上必须考虑相对于你状态大小的上下文长度。我的意思是，如果你的状态大小等于你的上下文长度，那没问题，但那样你就要付出非常大的代价。所以关键在于，没有免费的午餐——如果你想要一个非常小的状态，要把大上下文中的所有信息都压缩进去是非常困难的。也许有一天可以没有任何权衡，但我认为到目前为止，这就是权衡所在的地方。好的。好。那么现在，我要来讲讲混合专家（mixture of experts）。我认为混合专家是很值得讨论的。从概念上讲，它们并没有真正改变游戏规则。理解混合专家的一种方式是，它们只是一个更高效的MLP。你拿你的MLP，

[34:37.179] 然后不知怎么的，有人给了你一个更高效的MLP。这就是混合专家的本质。但我认为你需要理解什么是混合专家，我认为有两个原因——一是因为现在每个人都在做MoE，每个人都在发布MoE产品。如果你想要成为那些深入理解语言模型及其工作原理的人，你就应该理解它是什么。这是第一个原因。另一个原因是，MoE实际上有一些非常有趣的机械组件。如果你看看MoE是如何构建的，以及它们使用了哪些原语，你会很快在很多地方看到类似MoE的原语。然后你会意识到，哇，我可以用神经网络做更多有趣的事情。所以有一个更广泛的原因让我觉得MoE很有趣、很酷，你应该学习它们。那么什么是MoE？MoE在某些方面是非常简单的东西，它们只是你MLP的替代品。

[35:33.420] 左边是Transformer中你常用的MLP。你有注意力机制。你有某种归一化。然后紧接着就是一个前馈网络。这是大量密集的、大规模信息处理发生的地方。现在，我可以做的是，假设我拿我的大FFN，然后我把它们切成更小的FFN。所以我现在有四个前馈网络。实际上，假设我不是把它们切成更小的块。假设它们和我原来的FFN大小相同。然后通过某种魔法，我有一个系统告诉我为每个输入选择哪个FFN。现在我所做的是，我有了4倍的参数，或者至少对我的FFN来说，我有四个FFN的参数。但在任何前向或反向传播中，我只需要支付一个FFN的计算成本。这很棒。所以从某种意义上说，MoE最初的动机，如果你读原始论文，有一种非常以参数为中心的观点。

[36:33.710] 如果你从这样的想法开始，嗯，假设你只是想要更多参数，因为你相信更多参数是好的。好。好的。如果我们想要更多参数，但又不想为更多参数付出计算成本，你就需要右边的东西。所以这是你对MoE应该有的心智模型。你希望在不影响FLOPs的情况下增加参数。那么，为什么它们如此受欢迎？现在你在Hugging Face或其他地方得到的每个模型，超过一定规模之后都是MoE。超过一定规模，你得到的一切似乎都是MoE。原因是因为，无论出于什么原因，如果你保持总计算量相同，但只是增加稀疏参数的数量，模型似乎在总体上会变得更好。所以这是证据，也许支持了这样一种观点：更多参数实际上通常是有益的，即使一次只有其中一部分是活跃的。

[37:31.170] 在本节课中，我会大量引用一些Google的原始论文。这篇来自Fedus等人2022年的工作，我认为他们在这篇论文中做的是switch transformer。他们展示了，哦，随着你增加专家数量——也就是说活跃参数保持不变，专家数量增加——你的语言建模测试损失就是在不断下降、下降、下降。如果你只关心前向传播的FLOPs，那这很棒。在右边，你可以看到类似的情况，训练速度也更快了。随着你在x轴上增加训练计算量，如果对于相同的训练计算量你有更多专家，你就能获得更好的性能。所以无论是训练还是推理，MoE都能给你带来免费的收益。这非常棒。这非常好。而且这确实是已经被多次验证的现象。左边的图还是之前Fedus论文里的。去年，我认为AIQ的人发表了那篇OlMoE论文，那是他们的开源混合专家分析和训练论文。

[38:36.160] 他们也看到了同样的现象。如果你看训练损失、验证损失或下游基准性能，你会发现这些MoE模型在训练时间内比它们对应的稠密模型表现好得多。比如训练一个MoE相对于一个稠密模型，大约快两倍之类的。所以你真的能看到稀疏性带来的好处。我认为事实胜于雄辩。但如果你看看所有已发布的MoE，它们在激活参数方面确实比稠密模型强得多——如果你关心推理FLOPs，这很重要；或者甚至笼统地说，就训练和服务这些系统的总成本而言，也就是全部成本。所以在DeepSeek v2以及更早的DeepSeekMoE发布时，相对于其他所有人都在训练的许多稠密模型，这确实是一个巨大的转变。你会看到，哇，我们有更少的活跃参数，但MoE性能跟其他人一样好，甚至更好。

[39:37.780] MoE正在带来一些非常好的成果。好的。好。那么希望大家已经被我说服了，MoE是一个有趣的东西，值得你花时间去思考和钻研。哦，太好了。然后我要说的关于为什么这很重要的最后一件事是，MoE为你提供了另一个并行化的维度。我们将在之后的三节课中讨论系统方面的并行化问题。但并行化实际上是LLM中非常重要的事情之一。LLM非常庞大。无论是训练还是推理，它们都无法放入单个设备中。所以你需要很多不同的方式来切分你的模型，以便高效地服务或训练它们。混合专家模型还有一个额外的好处，就是每个FFN内部的专家天然地分成不同的块。每个专家就是一个自然的块。你可以把它们放在不同的设备上，然后将激活值路由到不同的设备。

[40:31.540] 我们稍后会讲到专家并行，就是这个东西。但这给了你一个可以优化的额外系统层面上的东西，如果你想优化你的推理服务，这是非常有帮助的。在西方，我认为有几个开源发布都是大型、强力的MoE模型。在大模型方面，有Llama 4和OpenAI的GPT-OSS。不幸的是，我想，西方开源模型的发布基本上已经停滞了。所以现在发布的不太多了，但至少有两个很好的例子，而且它们本身都是顶级模型。但我认为MoE研究和训练的大部分活跃进展都发生在中国。Qwen、DeepSeek，还有其他一些，比如MiniCPM，在训练和推广MoE方面做了一些最早的工作。他们继续在MoE上做非常出色的工作。如果你看一些早期的Qwen MoE模型，比如Qwen 1.5 MoE。

[41:36.510] 他们展示了他们那个只有27亿参数的活跃模型，表现优于当时很多70亿参数的模型。我认为DeepSeek和Qwen的这些早期成果确实说服了开源社区里的其他人，让大家相信这是正确的方向。说到DeepSeek——有问题吗？好。我关于专家并行有一个问题，这些专家之间互相通信，难道不会产生通信瓶颈吗？会。所以从某种意义上说，你需要付出将激活值传输过去的通信成本。我不会详细解释，但我会展示一个减少通信成本的技巧。但总的来说，你是在做某种权衡。你获得了更多的总计算量，减少了内存使用。作为交换，你需要付出通信开销。所以这是否能带来净收益，高度取决于你的网络拓扑和其他各种因素。

[42:33.070] 好的。哦，是的。所以在训练过程中[听不清]。好的，这是关键部分。在训练过程中，它们也是稀疏的。这确实是让MoE变得困难的关键所在。我这是在提前剧透几张幻灯片之后的核心观点，但如果我们训练时激活所有专家，那就会很容易。因为那样我们就能看到，哦，哪个专家对我的输入是合适的？这样你就能学习如何路由。MoE的难点在于，训练时你只有一个或k个专家是激活的。所以你不知道其他专家发生了什么。你只知道那些被激活的专家。尽管如此，你必须以某种方式学会路由。所以这个问题带有强化学习赌博机的味道，但我们不会用强化学习或赌博机来解决它。我们会用启发式方法和深度学习魔法的力量来解决。好的。专家之间切换的粒度是多少？

[43:24.360] 粒度，也就是我们路由专家的尺度，是在token级别。我不知道是不是——好，好。这就是问题。是的，在token级别。所以每个token都会分配一个专家。而且路由器，为了让你有个正确的认知，路由器是非常简单的。它就是一个输入和什么东西之间的矩阵乘法。所以你不会做任何复杂的事情。你不知道这是不是一个医疗问题。你只是，哦，这个token看起来像是，我不知道，日语。让我们把它路由到专家7。是的。好。那么在训练或推理中，并行化这些MoE模型有上限吗？有的，一般来说，并行化是有上限的，因为当你把模型切分到越来越多的设备上时，通信成本会爆炸式增长。

[44:09.860] 嗯，如果我们的作业设计得好的话，你们会在作业中处理其中一些问题。我们会让你们去搞清楚网络拓扑之类的东西，这些都关系到切分是否合理，或者反过来。给定一个网络拓扑，你们需要搞清楚如何切分模型，诸如此类。是的，好。DeepSeek，现在以DeepSeek R1等著称。但我真正欣赏DeepSeek的地方——最初这门课第一版就是专门讲DeepSeek论文的——是他们确实在LLM科学和架构方面做了很长时间真正出色的工作。DeepSeek是最早一批系统性地展示MoE为什么好的团队之一。他们做了大量很好的消融实验，比如，如果使用稠密层会怎样？如果使用哈希路由呢？如果使用switch路由器呢？所有这些不同类型的设计决策，你都可以在他们早期的一些DeepSeekMoE论文中看到。

[45:04.203] 所以如果你对架构、设计这类问题感兴趣，我建议你读一下DeepSeek早期的论文。从这些论文里可以学到很多东西。当然，进展并没有停止。DeepSeek v3、v3.2，还有GLM等等，都构建了这些非常强大的MoE模型。很明显，这就是未来大型模型在可预见未来的发展方向。至少在未来几年，我想都会是MoE的天下。好。那么现在的问题是，为什么它们没有更流行呢？在很多方面，我觉得MoE的普及速度相当慢。你看一些论文，从2022年就有了。谷歌当时在研究并真正推动MoE。但MoE真正开始流行起来是2024年以后的事了。另外，如果你在做各种LLM研究，你可能主要使用的是稠密模型，既然MoE这么好，为什么不用MoE模型呢？

[46:04.710] 我们可以节省计算量。但MoE带来了很多复杂性。训练MoE、使用MoE或者做很多事情都不容易。基础设施非常复杂。我的意思是，刚才有几个关于并行化的问题，但大体上这是对的，即高效利用资源地并行化专家是很困难的。MoE有很多参数，很难在单个设备上放下。如果你在训练它们，MoE可能会给你带来很大的麻烦。稍后几张幻灯片你们就会看到为什么MoE真的很难处理，它们确实不是容易驾驭的东西。我的意思是，现在有了一些好的经验法则，但确实不容易。我会用一张幻灯片简要提一下，也有人对Transformer中的注意力块做混合专家。有少数几篇论文，但要少得多。而且我看过的结果是，相比于直接替换FFN或MLP层，它们没那么容易被驯服或正常工作。所以我只讲左边这个。这基本上就是所有大模型的做法。大家基本上不会把注意力头拆分成专家。好。这就是对MoE以及当前状况的广泛行业概述。

[47:30.260] 但现在，让我们谈谈如何设计一个MoE，如何看待MoE。为此，我要讲三个不同的变化维度。MoE有哪些方面可以变化呢？嗯，你可以改变路由函数。我们知道MoE的核心机制是路由器必须将token发送给专家。所以我们可以改变这一点。我们可以改变不同专家（expert）的规模。在给定预算下，你可以拥有更多参数更少的专家，或者更少参数更多的专家。最后，我们将讨论训练问题，因为训练将会非常困难。一旦你开始理解 MoE 的工作原理，你就会意识到，哇，训练一个能将 token 路由到不同位置的东西可不是件容易的事。好的。所以在所有情况下，我们都不希望激活所有专家。我们只希望有一部分专家子集处于激活状态。因此，我们总是会选择某个 TopK。

[48:28.800] 你可以做不同类型的选择。可以让 token 来选择专家。可以让专家来选择 token。也就是说，每个专家挑选它们最喜欢的 token，或者你可以让每个 token 挑选它们最喜欢的专家。或者你也可以通过某种方式，用一个非常复杂的路由器来全局决定，哦，是的，你，这个 token，应该去这个专家那里。而另一个 token 应该去另外那些专家那里，试图进行全局优化的分配。所以基本上，几乎所有 MoE 都采用 token 选择 TopK 的方式。也就是说，将由 token 来负责选择专家。并且会有 k 个不同的专家被选中进入这个池子。

[49:09.740] 你可以看到，这是来自 OlMoE 论文的内容，展示了 token 选择和专家选择之间的差异。你会发现与专家选择相比，token 选择能获得更低的验证损失、更高的下游基准测试得分。也有一些成功的专家选择模型。我的意思是，这种模型也能训练得很好。但总的来说，token 选择更容易实现并使其正常工作，而且已经成为我们今天看到的所有模型的标准做法。我认为那四个模型中，可能有一个是专家选择的。但我不认为这一定是对这种方法的强力背书。我觉得，恕我直言，那是未发布的 Llama 4 模型之一。好的。所以我认为有很多不同的方法可以进行路由选择。其中一种路由方法是，你实际上可以使用一个很小的线性投影（linear projection），它说，好的，我将获取我的输入。

[50:04.790] 我的路由器将计算一个内积（inner product）。每个专家都有一个向量方向。在内积空间中距离最近的将被选中。这是目前最常用的路由器。它被用于经典的 Switch Transformer、Gshard 中。这两篇是谷歌早期的论文。它也被用于 Grok、Mixtral、DBRX、DeepSeek 等模型中。你选择的 k 值取决于你看的是哪篇论文。但在大多数情况下，你只是在进行内积运算。有一件事一直让我感到有点神秘，那就是很多很多论文都表明，实际上你不需要任何学习式的路由。你可以直接拿着你的 X，对它们进行哈希处理，然后发送到不同的前馈网络。实际上，这样做也可以。这也能带来提升，虽然没有 TopK 那么多，但不知何故，基于哈希的专家路由也是人们会做的事情。这是论文中常见的基线方法，但由于种种原因并未在部署中使用。

[51:05.350] 你可以使用强化学习（RL）来学习路由。你基本上可以有一个内积路由器。你可以使用强化学习，试图将这个路由器视为一种策略（policy），然后学习它。如果你是一个偏向经典机器学习学习理论风格的人，我认为这应该是思考这个问题的自然方式。你会想，好的，我在做什么？这是一个老虎机问题（bandit problem）。我要从 k 个中选择一个，或者从 n 个中选择 k 个。而我无法观察到所有这些结果。这是一个老虎机问题。我会使用老虎机或强化学习算法。这被用于一些最早期的研究中，因为这是思考该问题的自然方式，比如 Bengio 在 2013 年的工作。

[51:45.460] 但这并不是一种常见的方法。我认为它不常见的原因在于，进行强化学习所需的方法会引入大量开销，无论是强化学习算法本身还是随机性方面。而且人们已经发现，有一套很好的启发式方法，这些技巧可以应用于这个非常简单的经典 TopK 路由方案，基本上就能让它很好地工作。因此，没有理由去做更复杂的事情。最后，我想提一个我认为非常酷的想法，作为一个喜欢有道理事物的人，我很喜欢这个想法。你可以解决一个线性分配问题（linear assignment problem），你这样说，好的，我将从全局角度计算一个得分，这个得分表示，哦，这个 token 去那个专家那里会有多好。所以你要计算所有这些成对的、关于分配好坏的得分。然后你可以使用线性分配来精确求解从专家到 token 的最优分配。在某些情况下，它已被证明能带来良好的效果，但完全还没有在大规模场景下得到验证。

[52:50.320] 我认为这是因为，让这种方法工作起来，相较于其他方法要昂贵得多。所以，正如我所说，TopK 路由基本上是共识性的路由机制。它是如何工作的呢？非常简单。如果你在 DSA（动态稀疏注意力）那页幻灯片上认真听了，你会发现这看起来很像 DSA。如果你读过一些其他论文，比如 H-Nets 或其他类似的东西，你也会注意到这一点。这会出现在很多不同的地方。所以这是一个很好的、值得识别的模式。那么你要做的是，接收你的输入。你将它通过一个前馈层。然后我计算某种得分。我要计算的得分将是在专家和我的输入 u 之间的 softmax 结果。然后这个 s 就是得分。我将从中选择 TopK。

[53:47.677] 那将是控制我每个专家的门控（gate）。所以这就像输入。这是残差项。这是 FFN（前馈网络）。我有门控。而这些门控是由 TopK 控制的。这就是你对 TopK 路由器所想象的那个样子对应的方程形式。而你在这里可能学到唯一细节是，如何学习我的门控？我的门控只是通过对每个专家的权重与我的输入进行内积来学习的。非常非常轻量级。所以在我如何选择门控方面，并没有什么复杂的东西。即便在这个简单的框架内，DeepSeek 及其相关团队也有一些创新和变体。DeepSeekMoE 开创了这个现在被广泛使用的想法，叫做共享专家（Shared Experts）。这是一个非常直观的想法，在经典的 TopK 路由器中，你会有 k 个不同的——或者抱歉，在这种情况下，是 n 个不同的专家。

[54:44.410] 你会在这 n 个专家中进行选择。但有时候你可能会说，嗯，有一些非常通用的处理我想要应用到所有 token 上。也许你不希望专家每次都不同。你希望一些专家总是被应用，然后另一些专家按条件被应用。所以 DeepSeekMoE 所做的是，他们采用了原始的专家设计。他们把专家切分成更小的、更细粒度的块。其中一些细粒度专家的子集被设定为共享专家，始终处于开启状态。这些专家绕过了路由器。它们总是处理我们的输入。然后将结果放回输出中。抱歉。这是一个好主意，因为事实证明，如果你使用这种经典的专家设计，你实际上是在大量复用这些权重来做通用建模。因此你可以将其卸载到一个共享专家中，让其他专家更加专业化。

[55:43.200] 正如我一直在说的，DeepSeek 非常好的一点是，他们在很多论文中做了非常仔细的消融实验。所以他们有这样一个非常好的消融对比：一边是零共享专家，然后你有一些大的专家，16个路由专家。这实际上是老式的 GShard Google 设计；对比之下，他们则采用了非常细粒度的专家切分，并且在这些专家中设置一个共享专家。当你把专家做得越来越小，并且有一个共享专家时，你会看到这两种干预都带来了显著的收益提升。在某些任务上，共享专家带来了非常显著的改进，比如 TriviaQA、NaturalQuestions。

[56:20.600] 这个从蓝色到黄色的差距，对吧，就是有没有共享专家的区别。所以更多的专家、共享专家，总体来看似乎都有帮助。如果你想要这些说法和想法的更多验证，你也能从 OlMoE 中看到。OlMoE，我认为是西方那边一个做得很好的、控制严谨的 MoE 研究。他们做了类似“共享专家是否有帮助？”这样的研究。他们得出的结论是共享专家帮助不大。而在专家数量方面，他们展示了做细粒度的、更多的专家是有帮助的。所以他们在共享专家这一点上意见有些分歧，但他们也做了非常漂亮、控制严格的消融实验。很多近期的 MoE 基本上都遵循了 DeepSeek 的设计。在那之前有少数几个早期的 MoE。前三个——我一直夸 Google 做了很多有趣的事情——前三个来自 Google，而 Mixtral、DBRX 和 Grok 是西方早期的一些 MoE 尝试。

[57:21.890] DeepSeek v1 提出了细粒度专家和共享专家的设计。然后这个设计流行起来，其他人都跟随 DeepSeek。我的意思是，我认为这是一种类似的现象，就像 Llama 的设计之于稠密 Transformer 一样，基本上成了标准；而 DeepSeekMoE 和 DeepSeek v3 则是人们在 MoE 方面竞相复制的标准设计。如果你去看 Qwen 3.5，或者 GLM，或者任何这些现代模型，你会持续看到共享专家和细粒度专家在这些模型中被广泛使用。所以这是一个经过实战检验的设计。有什么问题吗？哦，有的。好的。

[58:02.050] 是的。如果有一个共享专家或者少数几个共享专家，在这种情况下并行化是如何工作的呢？是的。所以在共享专家的情况下，你不会得到任何并行化的节省。每个激活都必须经过那些共享专家。你可以复制共享专家，基本上用内存开销来换取通信成本的降低。好的。没问题。那么这就是路由器的设计空间。现在，我想谈谈训练 MoE。这会涉及到一些非常令人惊讶的，而且我觉得非常深度学习化的事情。我想，最初我学习 MoE 的时候，我觉得我们不可能把这些东西训练得好或者训练得合理。但事实证明，有很多技巧，当它们组合在一起时，出于某种原因，就是非常有效且稳健。正如我们在之前的问题中提到的，我们不希望在训练期间激活所有专家，因为如果那样做，我们将为所有专家支付全部的 FLOPs 成本。

[59:05.050] 而我们有很多专家。所以我们在训练时需要稀疏性来提高效率。但如果我们有稀疏性，那么门控决策就不再是可微的了。而且我们也看不到所有我们本可以选择的其他反事实专家。所以这两个都是大问题。所以有不同的解决方案。我们可以使用强化学习来优化我们的门控策略。我之前提到过这个，不是特别流行。你可以做像随机扰动之类的事情。所以如果你熟悉探索、利用、赌博机那类东西，你可以做其中的一些。现在，第三个是，你可以有一大堆奇怪的启发式方法，来处理门控决策是稀疏的这个问题，从而平衡哪些专家去哪里。猜猜人们在实践中用哪个？我一直在暗示。是第三个。它是一系列真正有趣的启发式方法，最终效果很好。

[59:56.090] 我会先讲另外两个。它们没有被广泛使用。我先说明这一点。但我认为了解一下是有用的。所以有些人做过研究，这是来自 Clark 2020 的工作，研究诸如“你能用强化学习来学习路由器吗？”这样的问题。所以它确实有效。你可以用强化学习，就是这里的绿线。那是其中一种基准方法。它效果不错。但由于梯度的方差和复杂性，它实际上不是最优的做法。而且这篇论文中提出的甚至一些方法也能轻松击败使用强化梯度来做路由决策的方法。现在，有一个与强化学习想法非常接近，但开始更多地转向这种“写下目标然后尝试”的启发式世界的东西，就是随机近似这个想法。它来自 Noam、Shazier 等人最早的 MoE 论文，其做法是，你在训练期间并不真正做出一个硬性的决策。

[01:00:59.220] 你做的事情是做出随机决策。所以做法是，你仍然有这个路由设置。这个 H，是你参数和输入之间的内积。这是你标准的路由内积。但是，与其直接这样做，我要注入一些依赖于输入规模的噪声。然后那将是我的 H。现在，我会保留 TopK 和 softmax，就像我之前解释 TopK 那样。然后那将是我输入到 TopK 函数中的输出。现在，如果你看这个，它在做什么呢？嗯，它只是意味着如果你有两个专家非常接近，或者它们彼此非常接近，那么随机地，你会选择其中一个而不是另一个。当你反向传播时，你会以这种方式反向传播，使得有帮助的专家会获得高权重。所以这种扰动允许你打破平局并多探索一点。并且希望梯度结构是这样的：如果一个专家有帮助，它会更多地强调那个专家，

[01:02:04.300] 并且总体上它会获得越来越高的权重。所以这是你可以做的一件合理的事情。它会导致专家更加稳健。而 softmax 允许你学习如何对这 k 个专家进行排序。而不是仅仅对 TopK 做硬性选择，你是通过这些元素来加权的。还有一个非常密切相关的方案，是在 Fedus 等人 2022 年的工作中做的，它做了一个均匀的乘法扰动。他们这样做是为了“硬化”专家，让它们变得稳健。这是一件有用的事情，但在后来 Google 做路由的论文中被移除了。而且不太清楚这是否真的有必要。你会在后来的消融实验中看到，表明不做任何这些随机性或稳健性技巧，实际上对最终训练好的 MoE 的稳定性和整体质量都有帮助。所以这暗示，你可能不一定需要随机探索项。

[01:03:13.030] 所以这让我们来到了你实际会做的事情，那就是一系列启发式方法，最终会给我们一个可用的 MoE。其中很多本质上都与平衡有关。那么，如果你只做普通的梯度下降，忽略任何这些探索、利用的担忧，问题会是什么呢？嗯，会发生的事情是，你首先会路由你的 TopK 专家。然后那些最强的专家会获得更多的信号。你的反向传播会说，那个专家很好。你增加那个专家的权重，把你的参数往那个方向移动。于是你得到了这种“富者愈富”的效应，被选中的专家会获得非常强的权重。强权重意味着它们被更频繁地选中。然后它们就会失控，包揽一切。所以这种专家坍缩现象，或者说专家饥饿现象，是一个非常非常真实的问题。

[01:04:06.260] 从某种意义上说，这就是你在用启发式方法训练MoE时必须解决的核心问题。那么该怎么做呢？嗯，你要做的就是在总建模损失上增加一个启发式损失。这个损失的目标是平衡分配到不同专家的token。有一种方法，是在这里的Switch Transformer中实现的，你可以看看它是怎么做的。你有从1到n的不同专家，还有一批token。你想要做的是，确保这两个内积很小。好的。那么这里发生的是什么？F是被分派到任意单个专家的token比例。然后P I是更平滑的版本。这不仅仅是token是否被分派到某个专家，而是路由器分配的累积概率质量。

[01:04:59.140] 你把这两项相乘。至少对我来说，这不是一个你能从第一性原理推导出来的东西。我觉得理解这个东西在做什么的最简单方式是取这个目标的导数。所以如果你对P I求导——记住，P I是分配给一个专家的概率质量——嗯，导数是分配给该专家的token比例。所以从某种意义上说，你可以把这看作是在梯度空间中的一种惩罚，你获得的token越多，你得到的负梯度就越大。所以这是试图按比例压低非常受欢迎的专家上的概率质量，与其分得的比例成正比。这就是这个损失的作用。这是理解它的方式。所以方程4中的目标函数本身一开始可能不清楚，但一旦你推理一下梯度的作用，它想做什么就变得很清楚了，那就是压低受欢迎的专家。

[01:06:00.150] 然后我们实际上可以把这些想法拿过来，基本上重建出DeepSeek v1和v2的MoE系统。那么要构建DeepSeekMoE，你会怎么做呢？嗯，你会做你正常的损失。然后你直接穿过专家进行反向传播，忽略所有这些不可微性和探索方面的顾虑。但你会增加一个逐专家（per expert）的平衡损失，这和Switch Transformer里的完全一样，就是我刚才给你展示的那个。所以你会有一个损失，它乘以F，即比例，和P，即总概率质量。但是DeepSeek那帮人，他们对系统设计非常在行，所以他们不仅仅平衡专家，他们还想要按设备（device）来平衡。

[01:06:40.590] 所以如果你的专家被分配到某些设备上——比如说，我这台机器上有四个专家，那边那台机器上也有另外四个专家，你希望那两台机器是平衡的，这样两台都在满负荷运行。所以你不仅想让专家变得均匀，你还需要确保设备之间也是均匀的。所以这是一个次要目标。它的作用方式完全相同。如果你看这里的函数，不是用专家，而是直接对每个设备的比例进行操作，你会应用这个次要的设备平衡损失，这纯粹是为了利用率的原因。在DeepSeek v3中，他们使用了一个逐专家的偏置项。他们用了一种在线学习的技巧来平衡专家，去掉了一些这些辅助损失，但他们最终确实还是得加一些这种辅助损失，以确保极端不平衡不会发生。所以我想，我把这个作为一个补充说明加上，DeepSeek v3和其他一些工作已经开始去掉一些这些，我会说，比较丑陋的辅助损失，

[01:07:44.770] 但到目前为止，还没有任何解决方案能完全去掉它们。关于负载均衡损失我要讲的最后一点是，你可能会问，我需要处理这个吗？这全是一堆复杂的启发式方法。我真希望我能直接去掉负载均衡损失。但如果你去掉负载均衡损失，后果是非常灾难性的。这两张图都来自OlMoE论文，他们做了一个很好的消融实验，完全去掉了负载均衡损失。如果有负载均衡损失，你会得到顶部的这条粉色曲线，非常漂亮的正常曲线。去掉负载均衡损失，你的损失显著增加，而且你的训练损失表现也不太好。也许更能说明问题的是，如果你看看哪些专家在做什么，有负载均衡和没有负载均衡的情况下——就是这里的下面两个面板——你会看到一个非常非常鲜明的差异。

[01:08:38.870] 没有负载均衡，几乎所有的token都流向两个专家——你在这里看到的黄色专家和粉色专家。而如果有负载均衡损失，即使我们这里有这种启发式的东西在起作用，所有这些专家都在被跨token地使用。所以我们得到了这种漂亮、均匀的利用率。所以我们确实看到负载均衡损失做了我们期望它做的事情。而没有它，我们相当于扔掉了一大堆参数。那些专家在大部分训练时间里什么都没做。所以这就是MoE的训练方式。我想，回到前面来说，这对我来说是相当令人惊讶的——你有这么一个不可微的东西，你做这种TopK选择。老实说，它是个相当复杂的对象，而你所需要做的全部事情，实际上就是加上这个平衡损失来均匀化专家。然后剩下的，你就像可以随意把梯度泵入整个系统一样处理它，模型就训练得非常好了。

[01:09:34.380] 我认为很大一部分原因在于这种动态：如果一个专家是有用的，你就会强化它。这是一个正强化循环，通过均匀化很好地平衡掉了。这两种动态相互抵消。酷。哦，我最后要说的是，同样的技巧在DSA中也有使用。我们已经看到了TopK这种东西。它被用在比如H-Nets这样的东西里。如果你熟悉那个，那是一次试图移除tokenizer的尝试。所以我想，你将来会更多地看到这种技巧，这种TopK选择加上使用负载均衡或其他类型的辅助损失来应对不可微性的想法，将会成为未来架构设计的一个组成部分。是的。[听不清] 混合专家有点像一种系统优化。但如果你真的去探究，你能看看它们吗？它们是否真的在不同领域做不同的事情？是的。所以有些论文有可视化展示哪些token被哪些专家激活。

[01:10:37.080] 因为路由器太简单了，所以并不是说专家就像聪明的专家。它们不像医学专家、法律专家之类的。你确实会看到某些token被路由到不同的地方，比如标点符号可能会被路由到一个专家，或者其他符号可能会被路由到一个专家，非英语字符集可能会被路由到另一个专家。但它真的不是那种你看着就能说“哦，这是华尔街日报专家”之类的东西，完全不是那样。没有这种语义，也许可以这么说。是的。[听不清] 是的，逐专家损失，原则上如果它被完美地强制执行，当你的专家是均匀分配的时候，它也会强制执行逐设备的平衡。但我认为，至少我对此的理解是，你并不想把每个专家的损失调得过高，以至于达到真正的完全均匀分布，因为那会对训练动态产生不利影响。但每台设备的均衡又足够重要，以至于你愿意增加一点额外的损失，来鼓励设备间的均衡而非其他目标。

[01:11:37.550] 好的。那么现在，在结束之前，我想稍微讲一下系统方面的事情。训练MoE会引入一些有趣的额外系统动态。当我谈到并行时，我主要谈论的是标准的，我们姑且称之为，并行形式。你可以将数据分成许多小批次。你可以把整个数据集拆分到多台机器上。你可以做模型并行，也就是把你的模型拆分到不同的机器上。你还可以用各种方式将这两种方法结合起来，当然，这是可以做到的。但我认为我们意识到的一点是，每种并行技术都有其局限性。对于数据并行，你的批次大小有上限。一旦达到批次大小的限制，你就无法再通过数据维度进行并行化了。对于模型规模，存在自然的切分点，你可以在这些地方对模型进行切分。

[01:12:33.430] 一旦你穷尽了这些，就无法再并行化了。如果你将此与专家并行结合起来，你就获得了另一个很好的并行维度。当你这样做时，许多底层实现都能让你利用GPU中内置的稀疏矩阵乘法。另外我要提到的一点是，如果你在单个GPU上有多个专家，比如说，一种思考方式是，我要执行多个小矩阵乘法。但这并不好，因为你理想中希望有更大的矩阵乘法，这样可以重用缓存并进行各种优化。因此，相反，你可以巧妙地利用稀疏性。块对角矩阵当然是最基本的形式，但你可以使用硬件原生支持的更复杂的结构化稀疏形式，使你能够以非常非常干净、快速的方式将专家和输入相乘。

[01:13:34.880] 这是MoE好用的另一个原因。如果你考虑MoE的计算模式，它们几乎对应于这类结构化矩阵乘法，这类乘法在硬件中非常容易且高效地支持。所以MoE正在发生这种硬件架构协同设计。在并行方面，我要提到的一点是非常近期的发展。这来自Nemotron-3。我忘了是谁问的这个问题。当我们进行专家并行时，我们必须将激活值从设备间传输，以便说，哦，你属于那个专家，我需要把你传输过去。这会导致显著的通信开销。因此人们开始思考的一些事情是，我的共享专家不需要通信，这个可以是大维度的。但我的专家，我需要通信那些激活值，我希望这些是更小的低维向量。

[01:14:34.960] 所以你实际上可以先将残差流进行下投影，然后再决定执行发送这些激活值的集合通信调用。这将显著节省通信开销，同时不会完全承受更小隐藏维度带来的缺点。所以你可以实际使用这类投影技巧来很好地控制通信和MoE并行化之间的权衡。最后一个有趣的琐碎细节，我们姑且称之为MoE的细节，我认为它很有趣，但近年来已经得到解决，就是随机性问题。如果你天真地做MoE基础设施，你会发现，除非你的专家在输入分布上完美平衡，否则你会发现某些专家可能比其他专家更受欢迎。如果某个专家，比如说，在这个例子中，专家0就是一个非常受欢迎的专家，你实际上会开始遇到这样的情况：这个专家的token队列开始不断堆积、堆积、再堆积。最终你会到达一个节点，你说，我的队列太长了，我必须开始丢弃token。在早期一代的MoE推理基础设施和代码中，很多情况下会发生的事情是，你实际上会悄悄地丢弃这个专家。

[01:15:56.510] 然后你就像什么都没发生一样继续你的计算。你只是传回一个零，假装那没问题。这意味着你可能会有这些奇怪的随机性，如果其他用户发送的查询命中了你在使用的专家，那么你实际上可能得到更差的结果，因为他们会把你的请求挤出专家队列。这类非常奇怪的事情可能会发生。但现在，架构已经减少了这些问题的发生，许多这类问题已经被消除了。所以这不再是一个真正的问题。MegaBlocks和其他非常常见的开源MoE框架已经不再有这个问题了。关于MoE的设计，我要讲的最后一点是稳定性问题。上一讲我们已经讨论了架构设计背景下的稳定性问题。但如果你记得我当时说的，我说指数运算不好，除法不好，这意味着softmax是稳定性问题的危险区域。

[01:16:58.100] 那么我们对混合专家模型做了什么？嗯，我们在路由中又引入了一个softmax，因为我们要做TopK。我们要做softmax。这就是我们路由到专家的方式。长期以来人们已经注意到，MoE中的softmax操作可能非常非常危险。Barret Zoph和其他人，在Google MoE设计的早期，写了一整篇关于MoE稳定性的论文。这是他们注意到的其中一个问题。我将深入探讨一些细节。你会经常看到解决方案用于一些更敏感和棘手的部分。你可能最终会在专家路由器上使用Float 32之类的。你可能也会有z-loss。我们之前讨论过z-loss，用来控制我们看到的softmax稳定性问题。我认为z-loss实际上在早期就非常流行，用于实现MoE路由器的稳定性。

[01:17:58.200] 我一直广泛讨论的OlMoE，做了关于添加或移除z-loss的消融实验。从这些非常尖峰的训练损失曲线中可以很清楚地看到，路由器上的z-loss可能非常有帮助。MoE还可能遇到其他问题。而且我想你们可能最近在使用语言模型时都经历过这个问题，那就是MoE微调起来可能相当烦人。它们有太多参数了，如果你试图微调专家，实际上会遇到非常严重的过拟合问题。所以你可以看到，如果你有密集模型，训练和验证相当接近，它们不会疯狂过拟合。但如果你看稀疏模型，训练和验证的差距非常大，当你在微调这些模型时，或者在当前这种情况下，是下游任务。我想那可能是GLUE基准测试中的某个任务。

[01:18:51.700] 现在，你可能会有一种情况，你微调非MoE的前馈网络。你也可以只微调注意力层。这些是人们经常做的干预措施，尤其是注意力微调是我在最近的MoE工作中经常看到的做法。但如果你的模型包含非 MoE 层，那么，当然，你也可以直接微调这些层。而且通常还有其他方法，我们也倾向于推荐这些方法。当然，还有“苦涩的教训”版本的思路，那就是，也许你应该直接用 140 万个样本，而不是用你手头的那个数量。如果你有大量数据，你基本上可以在微调阶段完全重新训练 MoE。这样你就不会有那么大的泛化差距。好的。我要提的最后一个 MoE 特定技巧是“向上循环”（upcycling）这个想法。实际上，这个方法在过去一年里已经不那么流行了。

[01:19:49.790] 我觉得今年我还没见过一个使用 upcycling 的模型。但这是个很酷的想法。而且，早期一些非常不错的模型也用过这个技巧。所以我想提一下。Upcycling 的想法是，如果你想要一个 MoE，也许你可以直接拿一个你训练好的密集模型，然后基于这个密集模型来实例化一个 MoE。具体做法是，你复制所有东西，包括 MLP。你只需复制一大堆这些 MLP。你有一个路由器，它是随机初始化的。然后你直接训练这个模型。如果你这样做，那么由于输入到 MLP 的随机性，这些专家就会开始专门化。最终，你确实会得到一个 MoE。在一些最早提出 upcycling 的论文中，他们展示了 upcycling 方法可以在语言建模方面获得更好的性能，在这种情况下，也就是准确性、性能，将密集模型进行 upcycling 比继续训练同一个模型要好，这是一个很酷的成果。

[01:20:59.470] 人们已经在许多方面大规模地验证了这一点。所以 Many-S-CPM 是我另一个非常喜欢的中国模型，因为他们做了很多精心控制的消融实验。他们有一个 MoE，他们把他们的 miniCPM 2.4b 模型拿来，向上循环成一个 134 亿参数的模型。他们从这中间获得了许多几乎是无本万利的收益。Qwen 也是，当他们做他们的 Qwen MoE 时，从他们的 1.8B Qwen 模型初始化。他们把它向上循环成一个 Qwen 1.5.A，2.7B 的模型。最终他们得到了首批高性能模型之一。我认为这是首批较大规模的 upcycling 成功案例之一。当然，现在，我觉得没人真的在搞 upcycling 了，因为你不再需要先训练密集模型然后再转换。你还不如一开始就用 MoE 来训练你的主力模型。

[01:21:59.820] 所以我想我们现在不太看到这个了，但我想提一下，因为我认为这是混合专家模型行动空间中需要了解的重要一点。好的。最后我想讲的是，我想带大家梳理一下 DeepSeek 模型——v1、v2 和 v3。我认为从 DeepSeek 模型的演进中可以学到很多东西。而且，如果你对 MoE 或架构设计普遍感兴趣，我鼓励你去读读 DeepSeek 的论文。我认为它们写得非常好。我觉得它们有很多值得深入挖掘的细节。DeepSeekMoE 架构在 v1 版本实际上已经是许多不同现代 MoE 的原型了。它已经具备共享专家和细粒度专家结构。它使用标准的 TopK 路由和辅助损失平衡。我的意思是，从某些方面来说，我会说，这就像是 MoE 模型的原型理想形态。

[01:22:55.040] DeepSeekMoE v2 把这个模型规模扩大了。所以你有两个共享专家，更多细粒度专家。而且你有这个设备路由组件，以及一个通信平衡组件。这两个本质上都是添加辅助损失来优化系统。我认为重要的是要认识到，成功的语言模型训练不仅仅是深度学习的问题，它也非常关乎真正尊重你的系统。所以 DeepSeek v2 确实是这种哲学的一个体现。然后在 DeepSeekMoE v3 中，他们仍然保留了共享和细粒度专家的设计，但他们有了不同的平衡方式。他们有了无辅助损失的方式来平衡不同的专家。并且他们改用 sigmoid 加 softmax，这是一种不同的专家加权方式，但基本上大同小异。最后，既然我们已经讲到这里了，我们不妨继续看看 DeepSeekMoE V3 的其他部分。

[01:23:54.890] 他们还有很多其他很酷的东西。我非常欣赏他们思考系统的方式。他们有多头潜在注意力（Multi-head Latent Attention），不是直接计算 QKV，而是将它们表示为一个更低维的潜在变量 c。所以你有你的隐藏输入。不直接生成你的 Q、K 和 V，你先产生 c。然后你根据 c 来产生你的 Q、K 和 V。这让你能够显著节省资源。不需要缓存所有的 K 和 V，我只需要存储这些 c。而这些 c 希望是更低维的。这就是为什么它被称为 MLA，潜在激活，因为 c 是潜在的，它们是你需要保存的 KV 缓存数量。你可能需要担心的唯一复杂之处是，当你做 KV 缓存时，这会与 RoPE（旋转位置编码）产生冲突。

[01:24:51.150] 所以你需要小心如何处理不同维度的旋转等等。因此技巧是，你有编码位置信息的非潜在维度，但我不打算深入更多细节。最后，一个我认为对 MoE 来说非常酷的想法，实际上是 DeepSeekMoE v3 中的，但还没有太流行起来，这个想法是 MTP，也就是多 token 预测（Multitoken Prediction）。不是预测一个未来 token，而是同时预测多个 token。这在统计学上有理由说明为什么这是个好主意。也许它能让你更好地预测未来。而且，还有一个很棒的系统层面的理由，你现在内置了一个推测解码器。

[01:25:30.840] Percy 稍后会在讲推理的时候谈到这个，但这只是一个可以用来加速模型解码的技巧。好的。我想我们时间差不多了。所以，把所有内容整合起来，MoE 是一个非常聪明的想法，利用稀疏性，这样你可以拥有更多的参数，而你只需要为其中一部分付费。所以你不需要为所有参数的计算成本付费，但你却能获得参数带来的好处。你可能最初会觉得这种路由问题非常困难，但事实证明，非常简单的方法就能工作得很好，即使在大型规模下也是如此。而到了这个阶段，很明显 MoE 已经站稳脚跟了，所以你应该理解它们是如何工作的，以及它们是什么。

[01:26:13.960] 谢谢。

CS336 2026 · Lecture 4 · Attention alternatives and mixture of experts

视频：官方视频
字幕：en-US · manual captions
处理：连续字幕合并为自然段；每段保留首条字幕的起始时间戳
状态：原语言字幕稿，待校对

正文

[00:05.200] All right. We'll get started. Sorry, the lecture slides for this lecture was not up earlier today, but it is now up. So if you want to have that at your fingertips to look at, you should be able to get it through the course website. But if not, it's definitely up there on the GitHub. Today, we're going to talk about what I think of as more advanced architecture ideas. So last lecture was really about the basic transformer and how we might tweak different parts of the basic transformer to get to a modern language model. Now, I want to talk about things that I think are much more advanced, much more complex developments on top of the usual transformer. And the two things I'll talk about today-- the first one is attention alternatives, ways of going to much longer contexts using architectural changes that, for the most part, generally allow for linear time dependence on the length of the sequence rather than quadratic.

[01:06.558] The second thing I want to talk about is the idea of mixture of experts. So the first part is modifying the attention block. The second part is going to be modifying the MLP part. A mixture of experts are going to give us significantly better utilization in terms of our hardware, in terms of having more parameters relative to the compute that we have to spend on these models. So those are the two parts and the two ideas for today. So I'm going to start with attention. And I'm going to set the stage. It's clear now that people want longer context length. You want to pack a lot of things into the context, so that your model has more knowledge. Maybe it's an agent and operates on lots of things. If you look at different models over time on the x-axis over here, and the context window size, this is in log scale, there's really a clear rush by a lot of the top LLM vendors that provide larger and larger context sizes to support these kinds of more complex workloads.

[02:05.530] Now, how do we control these costs? Not yet. The right plot also shows the ratio between how much compute cost is spent in the feed forward part of the network versus the attention part of the network as you increase the sequence length. And you see that the feed forward, fairly large at the start, but grows linearly attention. It's an all to all connection between all the different positions, quadratic, so quickly outpaces feed forward as the sequence length grows. So it used to be that for big models, that smaller sequence lengths, feed forward was the dominant cost. As you go to these longer sequence lengths, attention increasingly becomes more and more of a problem. Now, what's the basic toolkit by which we can control these? We've already talked about these. Well, we talked about one of these, and we will talk about another one of these. We can basically take things like local attention and combine them in different kinds of hybrid ways to control the very high cost of global attention.

[03:09.770] If you're only doing global attention once every eight layers, and all the other ones are these very local attentions, you've very much controlled the cost. The other thing you can do is systems engineering. I think a very underappreciated fact is, constant factors really, really matter. Part of the theme of this course is, you need to pay attention to the details. And it's very easy for people trained in a more classical theory-oriented computer science tradition to be like, yes, it's the big O that matters, linear or quadratic. But I think one of the biggest things that has happened to context and attention cost is flash attention. Flash attention, we'll talk about this in a lot more detail once we get to the systems lecture and as you do, the systems assignments. But flash attention is really just a very clever way of rearranging the attention operation into a much more systems-friendly way, to minimize memory transfer overhead.

[04:07.540] And by doing so, you can get truly, truly dramatic improvements in the performance of your attention. So base PyTorch is this blue. You can just look at some of the shorter sequence lengths. You're doing 30 to 40 teraflops per second. And then you do flash attention, which is really a constant factor systems improvement. And you're getting these big, big improvements, factors of 2 improvements. And at some sizes, you can't even fit this in memory anymore. But with flash attention, where you're not materializing these big attention matrices anymore, you can continue to run these, albeit somewhat slowly. This doesn't fix any of the quadratic cost issues, but constant factors are very, very powerful. Now, that said, if we're going to 5, 10 million tokens, these tricks might not be enough. We want much more radical, much larger gains in our ability to handle long context.

[05:01.360] So this leads us to the question, is there ways for us to have linear time or linear dependence on the length of the sequence? And what would those kinds of things look like? There you go. It turns out that a lot of people have tried a lot of different things over the years. And I think there are a few false starts. But in the recent years, I would say, the last two-ish years, there's been an emergence of a set of recipes that are fairly effective in linear time attention, that are now battle-tested at scale. And that's why this is finally the first year when I'm talking about linear time attention, because it's clear now that these things really work at scale, and in production, and so on. To really explain all of these different methods that work well, I think you need to really only understand one core idea to start with, and then we're going to elaborate and build upon that idea. And that idea to start with is the associativity of multiplication.

[06:03.050] So let's think about how attention works. We can write it in this compact notation, where we say, the attention operation is going to be of Q, Ks, and Vs, with these matrix dimensions up here. And I'm going to take my Qs and the Ks, and I'm going to have an all to all interaction. All the n positions are going to interact with each other by matrix multiplication. And then I'll normalize them through a softmax, this row over here. And then I will multiply it with my values V to get the attention outputs. This is a classic self-attention or whatever other attention operation. Now, for the moment, let's forget that attention has a softmax. And then I'm just going to drop the row. Forget it for the moment. Now, if we didn't have a row, maybe we can actually change the complexity of this operation. So if you look at attention, it's n squared dk, because we're taking this QK transpose multiplication.

[06:58.810] But if Rho was the identity, I can just move the parentheses around a little bit. And I can say, all right, instead of QK transpose V, really, what I'm going to have is QK transpose V. And if we do this reordering, well, we've changed which part is quadratic. Instead of having the n squareds, which were pretty terrible, because these n's are very big. They're millions. It's the context length. Instead, I can have dependence on ndvdk. So now, we're multiplying essentially, the key and the value dimensions. These could be big, in principle. I can have a big model. But these are usually on the order of thousands, tens of thousands. No one has a million coordinates in their hidden dimensions. And so this is a much more favorable term to be dependent on compared to this n squared term when we were doing this original attention operation.

[07:53.660] So this is the first thing that you should know, which is like, associativity of multiplication allows us to change the dependence on the sequence length. Now, the other thing that we can do now is, if we consider this reordering, this QK transpose V, something that's even nicer about this, and I think this has led to really, a lot of the research that I'm going to mention in the next couple of slides, is that this matrix multiply on the right, in the self-attention case, looks just like an RNN. Instead of writing it in this dense form up here, I can write it incrementally. As I sweep left to right on my context, I can multiply my Ks and Vs, and I can update this thing S, which is a state, and I can add it to my previous state. This is incrementally multiplying and accumulating my KVs. And I can multiply it with my Q to get my output. And this dense operation up here and this RNN form are equivalent. And so what have we done?

[08:58.280] Well, we've started from linear attention on the top left. And then we've somehow gotten something that looks very much like an RNN. And what's nice about RNNs? Well, RNNs are quite nice, because at inference time, I have this S, which is my state, it's fixed in size. And I'm just carrying that forward and forward. So RNNs are very nice for inference reasons. They are not nice for training reasons. But nice thing about linear attention is, you can use them in either way. You can have it in this dense form, which will be great for training. Its parallel. Or you can have it in this serial form, like an RNN, which is great for inference. So you can get the best of both worlds. Now, at least in terms of dependence of all these objects, this is great. This is lovely. Unfortunately, this is linear, which is not very good.

[09:46.600] This is really the starting point of our discussion. Any questions so far? This is really the basics. One cool thing is, there have been people that have basically used this linear attention as part of their attention mechanism. So Minimax M1, which is a fairly large scale, fairly high-performance Chinese open model, uses a 7 to 1 hybrid, so seven linear attention layers plus one full softmax attention. Performance, generally strong. You can compare it to a number of strong models like OpenAI's O3 or DeepSeek-R1, fairly competitive. And then since it's linear attention, most of the dependence on context length is linear. It's not fully linear, because you've got these softmaxes, but the dependence is much, much milder.

[10:42.790] So this is one example, where people have shown that even this basic linear attention, which is very, very simple to understand, can be used fairly effectively at scale, as long as you've got some full quadratic connections involved. No one has, thus far, really proven out fully linear time attention mechanisms at scale. Everything that I'm going to talk about in the next couple of slides is a hybrid. So now, we've done linear attention. But you might say, linear attention is too naive and simple for me. This is not what I want. I want some more serious complex neural network stuff to be happening. Well, what you can now do is-- oh, sorry, I went a little too far, look at the RNN form of linear attention. So I've put the linear attention equation right here, just so you don't have to keep going back and forth. And we can add a small elaboration to linear attention in order to get more expressive updates.

[11:46.770] And this is how you end up getting Mamba-2. Some of you may have heard of Mamba, Mamba-2, Mamba-3 now, I think. These are a family of state space models by Albert Gu, and Tri Dao, and friends. And they're derived, originally, from this state space theory and view. But actually, if you look at the mechanics of what is being done, you can actually see it as a very simple elaboration of the linear attention mechanism. And so the idea is the following. So we start with linear attention. And you might say, really, the main problem with linear attention is that I'm always passing my state forward. We know, from the olden days of LSTMs, that it's important to when to pass information forward and when to just not pass information forward, to forget things and send them to 0. So what I'm going to do is, I'm going to be inspired by that. And I'm going to add a gate, gamma of t.

[12:40.480] Gamma of t is not stateful. Gamma of t only depends on my current inputs x of t. And gamma of t is going to modulate how much of my state I'm going to carry forward into the future. Now, if you think about this, the state-dependent terms are all these Ss. This gamma is not state-dependent at all, so this means that this is very simple to compute. And there's still this duality. You can compute this Mamba-2 term, either as a big dense matrix multiply or at inference time, use it in this recurrent form to get nice inference improvements. So that's the core conceit or view of this idea that you can either use it in parallel or sequence, just like linear attention. And people have now used this at scale. So Nemotron-3, I mentioned briefly last lecture, uses this idea of combining Mamba-2 as their lightweight layer.

[13:37.773] And of course, you've got your big softmax attention every now and then. And they alternate this in various ways to manage their inference cost, expressiveness trade-off. Nemotron-3 achieves pretty good performance compared to Qwen3 thinking as compared to GPT OSS. And because of all these Mamba-2 layers, has pretty good throughput at fairly large context length. So you can see that it works. At least, these are small frontier models, not really big frontier models, but small frontier models that are open source. Cool. So now, we've built on this gating idea, which is cool. You might say, though, can we just keep pushing this idea further, and taking this recurrence, and making it increasingly more complex? And the rule of thumb, in some sense, for what you can do with your operations here is, as long as you're getting the various terms in your RNN, so to speak, with only input-dependent terms,

[14:44.310] so no state dependence, then you will still have this fairly nice duality between parallel operations, which you can use for training, and serial operations, which you are going to use for inference. So gated delta net is probably among the most widely used, I would say, state space model now. I think it's been tested in quite a few papers. And there's a very nice large scale scale up in 3.5 for these models. So let me talk you through what gated delta net is in comparison to Mamba-2. So Mamba-2 is, as I have at the top here, we take the linear recurrence, this state update. And what I'm going to do is I'm going to gate it with my gamma. Oh, I forgot to explain this vt. Sorry. I apologize for that. I added this for completeness. Mamba-2 also has a modification, where they add what is basically a residual layer to the yt update.

[15:43.530] It's not really core to the state update or their arguments for why this is a good idea, but I put it in for completeness, just to make sure I'm actually giving the true Mamba-2 updates rather than a more abridged version that only gets the core ideas across. But you can ignore the vt for now. This is really just like an architectural improvement that's not core to the state updates. Gated delta net takes this idea in some sense and adds a second gate, beta t. So what does beta t do? Beta t, I think of as a no input operation gate. If beta t is 0, that basically means, don't take any of my current information, don't add it into my state. And I think, most or all of you are probably familiar with LSTMs. This is very reminiscent of the ideas in LSTM. You've got gates that control whether to forget and whether to put in information from the current time step.

[16:39.960] So it operates in a very similar set of principles, although, of course, derived, I think, originally in a very different way. Now, the other thing that's interesting about gated delta net, and this comes from the delta net rather than the gated delta net, is this update direction. So instead of just gating things, you could have made this 1 minus beta t, for example. What they do is, if you're going to do an update, where you're putting in this beta, what you're going to do is, you're going to try to project out, in some sense, my current key direction. So the intuition here is, I'm going to be writing in information from my current key, k of t. And when doing that, not only do I want to put in new information, I also want to erase any previous keys that have gone into it. That's one way that you can think about updates-- not just put stuff in, I want to clear out the current key that has gone in.

[17:28.993] And that's why there's this blue term, this identity minus beta tktkt transpose, which is essentially acting as a projector that is projecting out things in the kt dimension. That's not exactly right, because you're not doing things like unit normalization, but you get roughly the intuition here. The other thing I'll mention before-- oops, sorry, I have this blue box here, I'll move on with beta delta net, is that this has been reinvented, in some ways, in various settings. This update, this projector appears if you try to solve certain kinds of meta learning least squares problems. It's been reinvented in other forms in fast weight programming or test time training, where essentially, through very different design principles, researchers from those areas have basically ended up with the exact same kinds of solutions. And the thing that I'll say here is, the Qwen 3.5 and the Qwen Next models that were predecessors to that,

[18:28.190] some of the best open source models available today. And they use exactly this architecture. They use a 3 to 1 gated delta net attention hybrid. Very reasonable performance, strong inference characteristics that you can see on the right there. Qwen Next has much higher decoding throughput relative to Qwen 3 as the context length goes up and up and up. The middle panel is a comparison of performance versus various closed source models and previous generation models. This hybrid architecture doesn't seem to hurt performance very much at all. The last thing that I will say is, there hasn't been that many great controlled studies of how hybrid architectures perform. There's maybe one by ByteDance Seed and UC Santa Cruz folks, where they compared things like Mamba-2, and gated delta net, and other architectures as a function of how many hybrid layers you have. Some of the results, I would say, are messy, but the things that I would point at is, this dashed line over here on this right panel, this guy is full attention performance.

[19:34.870] And as you go from left to right, this is increasing the number of non-full attention layers. So the RNN layers, let's call them, you increase them as you go to the right, and you see performance degradation. But maybe the thing to see here is that for some of the best architectures that you have-- the yellow, the orange, the blue, these are like the gated delta nets and various other variants. You see that at the low ratios, there's basically no hit. And then once you pass a certain point, you start to see more significant degradation in long context performance, in this case. Basically, at the end, as you go to full RNN, you have very noticeable performance degradation in all these architectures. The same is true at various different other kinds of evaluations. Some of these, I think, have been optimized by the architectures themselves. Single key retrieval is a task that, I think, all of these long context architectures explicitly optimize for. But if you look at QA performance, we see a same story that as we increase the amount of hybrid ratio, we see decreases

[20:36.710] in performance, fairly steadily and fairly clearly, ending up at this pure point. Any questions for the linear time attention stuff? I have one last efficient attention thing to talk about, but yes. [INAUDIBLE] look like the parallel form of attention versus the linear form. [INAUDIBLE] Shouldn't they be mathematically equivalent? Yes. Why do you see degradation in performance if you [INAUDIBLE]? Sorry. By recurrent, I meant RNN in general, any recurrent formulation. So gated delta net would be a recurrence. That is not equivalent to this top one. This is full softmax attention with the row. The first step is lossy, where we drop the row, and then we become linear. That's the very, very first step of any of these. So this part is going to be lossy. And then after that, this linear form to this recurrent form, this equivalence, that is exact.

[21:36.100] Good. Yes. [INAUDIBLE] Oh, yes. Yes. Yeah, this term, right? Yeah. So in the linear attention form, what we're really doing is, we're taking these kv's, and we're accumulating them with queues. In the dual form, that's how this operates. But you might actually want in your outputs to also have some value from your current time step. You might think that, oh, the value for my current time step for my token is actually useful for various things downstream. And so this acts a little bit like a residual connection, where the value of my current inputs directly pass through to the outputs of my attention. And D is another modulated gate, where you can control the amount of pass through that happens. Good. OK. So that's linear time attention. There's a lot of really interesting stuff. It's also interesting that a lot of the methods that have gotten tested out really have this nice, simple, linear attention style, recurrent form.

[22:41.440] There's a few exceptions, but for the most part, a lot of this has converged to almost very LSTM-like objects in the end. What do I want to say? OK. So there's another alternative. What I talked about was linear attention, which has this very nice complexity theoretic flavor to it. I want to talk about another attention optimization, which I did not talk about, but falls into this family of efficient attention things that are important. So DeepSeek in DeepSeek 3.2, I think, had DSA, not Do Sparse Attention, that's DeepSeek Attention, where what they do is, instead of computing attention-- that was not intentional part, instead of computing attention over all the tokens, what you're going to do is you're going to first have a lightweight indexer that's going to subset a bunch of tokens. So you're going to look at your very long context. And the indexer is going to pick out some subset of these, much smaller than the full sequence.

[23:45.973] And then you're going to do full attention on that smaller subset. So the mechanics of this are pretty simple, at least in the forward pass, where what you do is you have your normal Qs and Ks, but then you pass this through this indexer. And what the indexer does is, it takes the qk inner products, it takes a ReLU. And then it's going to have these weights that are derived from the preceding tokens. And then this is going to give you the activations for each of the positions. And then you're going to take this through a top k, where some number of these top positions get included into my attention computation. And I'm going to do my usual attention computation that emits a key value. The nice thing about this is it's very, very sparse. And if this indexer is lightweight, now, this bottom term, where you're doing full attention, can be done on a very small subset.

[24:43.050] Another thing that both DSA and I think, GLM show is that, you don't actually have to train the model with this, because training with this indexer thing is maybe pretty annoying and very complex. What you do is you just train a normal transformer. And when you do your long context extension, you actually then drop in this lightning indexer, and then you train the model to handle this indexer in a extension stage that's separate from pretraining. So actually, the computational cost of doing this is itself, pretty light. And I want to highlight in DeepSeek's v3.2, where they proposed DSA, this model is quite good.

[25:22.400] The v3.2 matches all the other frontier models at the time, like Claude 4.5 Sonnet or Gemini 3. And they also have this very nice plot, where they're showing, well, our prefill, which is the attending to the input, and also decoding, where you're actually generating, both of these are very favorable scaling compared to the previous generation of DeepSeek models that did not use sparse attention. And that makes sense. If you're only paying attention to a very small subset, you're going to dramatically be reducing your costs of doing attention. And I think another nice orthogonal validation of this is GLM5, which came out last year. I think GLM5 is one of the best open models that are available, like period. And they also adopted this DSA approach. And they have actually, fairly nice ablations in their paper here, if you're curious, where they end up comparing with, oh, what happens if we do DSA with some warm up or just do DSA or don't do DSA at all?

[26:26.602] The main thing that they're showing is, if they do full DSA training, you don't lose very much in performance relative to full attention, even at long context retrieval tasks that are fairly difficult to do with RNN style architectures. So this is a totally alternative view of how to reduce attention costs. Notice, I'll go back one slide, this is not linear time. The indexer still has to operate on full all to all attention, because you got these qk inner products that are not going to be linear in cost. But you can make this indexer much, much cheaper, for example, by making this lower dimensional or making this w very small. All of this can be made much cheaper in its cost. So that's a totally alternative view, works very well. I won't talk much about this. But it turns out that this idea of top k selection is going to be core to the next part of this lecture. I'll stop here for any efficient attention questions before I move on to MoEs.

[27:29.750] Yes. So the indexer, the time complexity for that, is it n squared? [INAUDIBLE] Yeah. I mean, in order to know what to select, it does have to look at everything. And there's no clever state transition stuff that happens here. It is really brute force inner products. So if you're going through all of it, then the step afterwards is also quadratic as well. How are we getting time like [INAUDIBLE]? I mean, it's like a system z trick, almost. You're going to do the indexer much more lightweight. And there are many ways to do that. You can do that in much lower precision, which they talk about. You can do it in ways that are lower dimensional by projecting the q's and k's further, just for the indexing. So there are lots of tricks that you can do here, that you can imagine, but the indexer is going to be much more lightweight.

[28:18.890] So the constant factors are very good. And then for the second one, even though it's quadratic, it's quadratic on a shorter context length, because it's top k, we can control k. So now, this is expensive but small. So this is another one of these things of, sometimes, don't get too stuck up on the quadratic versus not. Sometimes, the constant factors are really, really important. There's another question? Yeah. [INAUDIBLE] That's right. Or it's part of a continued pretraining style thing, because all the models, you don't train long contexts from scratch for compute and various other reasons, you train a short-ish context model first, and then you've got these long context extension stages. And so the nice idea here is that, we're going to do the second phase anyway. Why don't we bolt the long context cost savings on at the same stage? And it's surprising that it works. Honestly, they can bolt on this frankly scary-looking TopK non-differentiable object.

[29:16.630] But as you'll see in the next part, maybe you shouldn't be so scared of TopK selection. [INAUDIBLE] So this is also part of the post with other post-training steps? Yeah. The usual recipe that I am familiar with, when a lot of the open source models, is it usually goes like short context pre-training, long context extension, then post-training. Yes. Is the k here, similar in length [INAUDIBLE]? Yeah, that's right. You would pick values that are much closer to short context performance. And you would bound it. So regardless of the input size, they would be bounded. Any questions? OK. Does the removal of softmax for [INAUDIBLE] causing training instability? Training stability issues. That's a good question. I don't think so. I'm not aware of documented linear softmax causing lots of attention. If anything, softmaxes are usually more of the problems for various reasons. So I think this will actually improve the stability properties of the architecture.

[30:23.960] Can you predict the next [INAUDIBLE] attention. Consider the [INAUDIBLE] attention mechanism requests, because the low computing cost and the low latency. Given that you predict, I mean, in high level, what kind of features that will be [INAUDIBLE] the [INAUDIBLE]. Yeah. I don't know if I have good predictions for what the future attention architecture looks like. I mean, my somewhat cop out answer to this is, I think a lot of this will look like we throw all of these tricks in. In the same way as the architectures have gotten very complex by taking all of these successful recipes and throwing them in, I think also that there's an even broader layer or a higher level layer of basically using post-training to have the model manage its own context. Compaction, retrieval, all of these things are basically an added layer on top. I think a lot of the work that I imagine will be integration of these different layers, because we've seen that, at least the linear time attention stuff has converged a lot into LSTM-like or linear attention-like architectures. And I don't really see that changing too much in the near future. Cool. Oh, yes. One last question. Can you speak about lower precision attention? Is FP4 attention possible? I think the indexer is partially motivated by things like this.

[31:54.600] FP4 attention, definitely possible. I think the drawback is that the loss of precision for things like softmaxes, where small, underflow, overflow, these kinds of things can really start to matter, can, sometimes, be pretty problematic. And so I think this motivation of saying like, OK, let's do selection with low precision and then do a full precision computation afterwards, because then the value vectors can be more finely added together and weighted together. It's an approach that makes sense. I don't know if there's really good ways of doing full attention in FP4 or even lower precision of some kind. Yeah. OK. OK, one last question then, then I'll move on. So compared to transformers, compared to state space models, what's the downside of training more complicated state space models? Yeah. So that's a good question. There have to be some downside or otherwise everybody-- Well, the downside, I think, is expressive power.

[32:57.810] The all to all connection in softmax attention is incredibly powerful. It's also very easy to train. It used to be, I think, compared to say, an LSTM. Attention had the very strong advantage of hardware efficiency. You could train in parallel. I think the reason why the state space models have caught on, despite their similarity to LSTMs and similar drawbacks in terms of representational power, is that now, there's this kind of well-understood duality because of the linear attention thing that I talked about, where you can go from this RNN form to the more dense matrix multiply form. And the matrix multiply form allows for computational efficiency. So that trade-off has now been checked off. But you've still got the trade-off of, if you have a finite state and you have to carry everything through, you're going to be losing some information relative to just carrying everything. [INAUDIBLE] is going to be a problem in state space.

[33:50.690] So [INAUDIBLE] the context itself, [INAUDIBLE]. Well, in the state space, you have to basically think about the context length relative to the size of your state. I mean, if you have a state the size of your context, then you're good to go, but then you're paying these very large costs. So really, the point is, the free lunches, if you want a really tiny state, it's really hard to compress all the information in a big context. It might be possible one day to not have any trade-offs, but I think, thus far, that's where the trade-offs are showing up. Cool. OK. So now, I'm going to talk about mixture of experts. And I think mixture of experts are important to talk about. Conceptually, they don't really change the game. One way of thinking about mixture of experts is, they are just a more efficient MLP. You take your MLP.

[34:37.179] And somehow, someone gives you just a more efficient MLP. That is what mixture of experts is. But I think you need to understand what mixture of experts is, I think, for two reasons-- one is this slide of everyone's doing MoEs these days, everyone's shipping MoEs these days. You should understand what it is if you're going to want to be people who deeply understand language models and how they work. So that's one reason. The other reason is that MoEs actually has some really interesting mechanical components to them. If you look at how MoEs are built and what primitives they use, you will quickly see MoE like primitives in many places. And you'll realize, wow, I can do more interesting things with neural networks. So there's a broader reason why I think MoEs are interesting and cool, and you should learn about them. So what is an MoE? MoEs are very simple things in some ways, they're just replacements of your MLP.

[35:33.420] So on the left is your usual MLP for a transformer. You've got your attention. You've got some normalization. And then you've got a feed forward right after that. This is where a lot of the dense, big information processing happens. Now, what I can do instead is, let's say, I take my big FFN. And I'm going to cut them up into smaller FFNs. So I now have four feed forward networks. And actually, let's say, I don't cut them up into smaller pieces. Let's say, they're the same size as my original FFN. And somehow, by some magic, I have a system that tells me which FFN to pick for every input. Now, what I've done is, I have 4x the parameters, or at least for my FFN, I have four FFNs worth of parameters. But on any forward or backward pass, I'm only going to pay one FFN worth of cost. That's great. So in some sense, the original motivations for MoEs, if you read the original papers, have this very parameter-centric view.

[36:33.710] If you start something like, well, let's say, you just wanted more parameters, because you believe more parameters are good. OK. Good. If we want more parameters, and we don't want to pay the cost for more parameters, you need something on the right. So this is the mental model that you should have for MoEs. You want to increase parameters without affecting your flops. Now, why are they so popular? Every model that you get on Hugging Face or whatever these days are MoEs past a certain size. Past a certain size, everything that you get seems to be an MoE. And the reason is because it seems to be the case that for whatever reason, if you keep the total compute the same, but you just increase the number of sparse parameters, somehow, the models are generally getting better. So this is evidence, maybe in favor of the view that more parameters are, in fact, generally good, even if only a subset of them are active at a time.

[37:31.170] And I'm going to refer to some of the original Google papers extensively throughout this lecture. This one's from Fedus, et al, in 2022, where they were doing, I think, the switch transformer for this one. And they're showing, oh, as you increase the number of experts, so this is active parameter stays the same, number of experts increases, your test loss for language modeling is just decreasing, decreasing, decreasing. That's lovely if all you care about is your forward pass flops. On the right, you see similar things as you train faster. As you increase your training compute on the x-axis, if you have more experts for the same amount of training compute, you just get better performance. So for both training and inference, MoEs just give you a free win. This is pretty wonderful. It's pretty great. And this is really something that's been seen many times. This is the same plot from before on the left from the Fedus paper. Last year. I think, AIQ folks put out this OlMoE paper, which was their open source mixture of experts analysis and training paper.

[38:36.160] And they see the same thing as well. If you look at training loss, or validation loss, or downstream benchmark performance, you see that in training time, these models are doing much better as MoEs compared to their dense counterparts. So something like two times faster or something like that, training a MoE relative to a dense model. So you can really see the benefits of sparsity here. And I think the proof is in the pudding. But if you look at all the MoEs that got released, they are much stronger than dense models in terms of activated parameters, which is what matters if you care about inference flops, or really, even just in general, the cost to train and serve these systems, the all-in cost. So in DeepSeek v2 and the earlier DeepSeekMoEs came out, it was really a big shift from a lot of these dense models that everyone else was training. You saw that, wow, we have much fewer active parameters, just as good, if not better MoE performance compared to everyone else.

[39:37.780] This is some really good stuff happening with MoEs. Good. OK. So hopefully, I've convinced you that MoEs are an interesting thing that you should spend your time thinking about and studying. Oh, great. And then one last thing that I'll say for why this is important is, MoEs give you another axis of parallelization. We'll talk about parallelization in systems in three more lectures, I think. But parallelization is actually one of the really big important things in LLMs. LLMs are very big. They don't fit in single devices, either to train or inference. So you want many different ways of cutting up your model to efficiently serve or train them. Mixture of experts have this additional nice thing that you can do, where your experts within each FFN naturally come in different chunks. Each expert is a natural chunk. You can put them on different devices. And then you can route the activations to different devices.

[40:31.540] We'll talk about expert parallel later, which is what this is. But this gives you an additional systems thing that you can optimize, which is very helpful if you want to optimize your serving. In the West, I think there's been several open source releases that are big, strong MoEs. On the big model side, there was Llama 4 and GPT OSS from OpenAI. Unfortunately, I guess, open source model releases in the West have stalled, for the most part. So there aren't quite as many being released now, but there's at least two good examples of this, both top tier models in their own right. But I think a lot of the action for both MoE research and training has happened in China. Qwen, and DeepSeek, and some others, like mini CPM, did some of the earliest work training and popularizing MoEs. They continue to do really great work on MoEs. If you look at some of the early Qwen MoE models, this was Qwen 1.5 MoE.

[41:36.510] They showed that their like small 2.7 billion parameter active model was doing better than many of the 7B models at the time. And I think some of this early proof of work from DeepSeek and Qwen really convinced everyone else in the open source community that this was the right way to go. Speaking of DeepSeek-- was there a question? Yes. I guess one question I had about the expert parallel, isn't there like a communication bottleneck that's built up because these experts are communicating with each other? Yes. So in some sense, you pay the communication cost of shipping an activation over. And I will not explain, but I will show you one trick of how you reduce that. But in general, you're trading certain things off. You're getting more aggregate flops. You're reducing your memory use. And in exchange, you're going to pay for comms. So it's highly dependent on your topology and all these other things, whether this is going to be a net win.

[42:33.070] Cool. Oh, yeah. So during training [INAUDIBLE]. OK. This is the important part. During training, they are also sparse. This is really the key thing that makes MoEs hard. I mean, I'm spoilering the key idea from a few slides down, but it would be very easy if we activated all the experts during training. Because then we get to see, oh, which expert is good for my input? So you get to learn how to route. The hard thing about MoE is that during training, you only have one or k experts active. So you don't what happened with the rest of them. You only know the ones that activated. Despite this, you must somehow learn to route. So it's got this RL bandit flavor to the problem, but we're not going to solve it with either RL or bandits. We're going to solve it with the power of heuristics and deep learning magic. Yes. What is the granularity of switching between experts?

[43:24.360] So the granularity, as in the scale at which we route experts, is at token level. I don't know if-- OK, good. Good. That was the question. Yes, at the token level. So every token gets an expert. And the routers, just to get your mental model correct, the routers are super naive. They're a single matrix multiply between your input and your whatever. And so you're not going to do anything complicated. You don't know it's like a medical question or not. You're like, oh, this is a token that looks like it's in, I don't know, Japanese. Let's route it to expert 7. Yes. Yeah. So is there an upper limit to parallelizing these MoE models in training or inference? Yeah, in general, there's an upper limit to parallelization, because as you shard over more and more devices, the communication cost explodes.

[44:09.860] Well, if we get the assignment to design done correctly, you will have to deal with some of these in your assignment. We will ask you to figure out the networking topology and things like this that will make sharding make sense, or it's the reverse. Given a networking topology, you'll have to figure out how to shard models and things like this. Yeah. Good. DeepSeek, now known for DeepSeek R1 and so on. But I think the thing I really like about DeepSeek-- and originally, this lecture was, let's just go through a DeepSeek paper lecture, the very first iteration, is that they've done actually really good architecture in LLM science work for a long time. And DeepSeek were one of the earlier ones going through and showing why MoEs are good. And they have a ton of good ablation work on, oh, what happens if we use the dense layer? What if we use hash routing? What if we use a switch router? All these different kinds of design decisions, you can see in some of their earlier DeepSeekMoE papers.

[45:04.203] So if you're interested in a lot of these questions, like architecture, design and so on, I would suggest you read the earlier DeepSeek papers. There's a lot to learn from all of these. And of course, progress hasn't stopped. DeepSeek v3, and v3.2, and GLM, and so on, have built these wonderful, very strong MoE models. It's quite clear that this is where the future of these big models are going to be for the foreseeable future. At least the next couple of years, I imagine, will be MoEs. OK. So now, the question is, why haven't they been more popular? In many different ways, I feel like MoEs caught on pretty slowly. You look at some of these papers. They're like in 2022. Google was studying and really trying to push MoEs. But it's only in 2024 onwards that MoEs really started to catch on why. Also, if you're doing LLM research of various kinds, you're probably mostly working with dense models, like why not MoE models if they're so nice?

[46:04.710] We can save on compute. There's a lot of complexities that come from MoEs. It's not easy to train at MoE, or to use that MoE, or to do many things. The infrastructure is very complex. I mean, there were a couple questions about parallelization, but this is generally right, that it's hard to parallelize experts in ways that are efficient in utilization. MoEs have a lot of parameters. It's hard to fit them on a single device. And if you're training them, the MoEs can really blow up on you. And you'll see why in a few slides, why MoEs are just really hard to deal with, but they're just not easy objects to work with. I mean, there's good rules of thumb now, but it's not easy. I'll briefly mention with one slide that there are people that also mixture of experts, the attention block in transformers. There have been a couple of papers, but they are much less common. And I think the things that I have seen is that they haven't been quite as easy to tame or to get to work compared to just replacing the FFN or the MLP layer. So I'm only going to talk about the left. This is basically what all the big models do. People are not, for the most part, sharding or sorry, splitting up attention heads into experts. OK. So that was the broad industrial overview of MoE and what the state of things are.

[47:30.260] But now, let's talk about how you design an MoE, how you think about an MoE. So to do that, I'm going to talk about three different kind of axes of variation. So what kinds of things can change for MoEs? Well, you can change the routing function. We know that MoEs have this core thing of a router has to send tokens to experts. So we can change that. We can change the sizes of the different experts. For a given budget, you might have more experts with fewer parameters or fewer experts with more parameters. And then finally, we're going to talk about training, because training is going to be really difficult. Once you start to realize how MoEs work, you realize, wow, it is not easy to train something that routes tokens to different places. OK. So in all cases, we don't want to activate all the experts. We want to only have some subset of experts active. And so we are going to always be choosing some TopK.

[48:28.800] There is different kinds of choices that you can do. You can have the token choose the expert. You can have the expert choose the token. So each expert picks their favorite tokens, or you can have each token pick their favorite experts. Or you can, somehow, have this very complicated router that globally decides, oh, yes, you, token, should go to this expert. And this other token should go to that other experts to try to globally optimize assignments. So basically, almost all the MoE do token choice TopK. And so the token is going to be the one to choose the experts. And there's going to be k different experts that get selected to be in the pool.

[49:09.740] And you can see, this is from the OlMoE paper, differences between token choice and expert choice. You see that token choice is going to get lower validation losses, higher downstream benchmark scores compared to expert choice. There have been successful expert choice models. I mean, this also trains fine. But token choice has generally been much easier to get working and has been the standard for all of the models that we see today. I think one of the four models was maybe expert choice. But I don't think that's necessarily a strong vote of confidence. I think that was one of the unreleased Llama 4 models in my defense. OK. So I think there's many different ways that you can do routing. So one of the ways that you can do routing is, you can actually have a small, small linear projection that says, OK, I'm going to get my input.

[50:04.790] My router is going to compute an inner product. Each expert has a vector direction. And the closest in inner product space will get selected. This is, by far, the most common router. It's used in the classic switch transformer, the Gshard. Those two are the early Google papers. It's used in Grok, Mixtral plan, DBRX, DeepSeek, and others. The number of k that you choose varies, depending on which paper you're looking at. But for the most part, you're just taking inner products. One thing that has always been a little mysterious to me is that, many, many papers have shown that actually, you don't need to do any learned routing. You can just take your X's, you can hash them and send them to different feedforward networks. And actually, this is fine too. This gives gains, not as much as TopK, but somehow, hashing-based expert routing is also something that people do. Common baseline in papers, not used in deployment for many reasons.

[51:05.350] You can do RL to learn the routes. You can basically have a router that's an inner product router. And you can use reinforcement learning to try to treat this router as a kind of policy and then learn it. If you're a machine learning classic learning theory-type oriented person, I think this should be the natural way to think about the problem. You're like, OK, what am I doing? It's a bandit problem. I'm selecting one out of k or k out of n. And I don't observe all of them. This is a bandit problem. I will use a bandit or an RL algorithm. Used in some of the earliest work, because this is the natural way to think about the problem, like Bengio in 2013.

[51:45.460] But it is not a common approach. And I think the reason why it's not common is because the approaches needed to do RL introduce a lot of overhead in terms of both the RL algorithm and the stochasticity. And what people have figured out is, there's a good set of heuristics, these recipes that you can apply to this very simple classic TopK routing scheme that basically just has it working right. And so there's no reason to do something much more complicated. I will finally mention an idea that I think is really cool and that I love as a person that likes things that make sense. You can solve a linear assignment problem, where you say, OK, I'm going to globally compute the score that says like, oh, it would be this good for this token to go to that expert. So you compute this all pairwise scores of how good the assignments are. And you can use linear assignment to exactly solve for the optimal assignment from experts to tokens. It's been shown to do good things in some cases, hasn't been seen at scale at all.

[52:50.320] I would think, because it's extremely expensive relative to the others to get working. So TopK routing, as I said, is basically the consensus routing mechanism. How does this work? It's very, very simple. And if you were paying attention in the DSA slide, this looks a lot like DSA. You will also notice this if you were reading some other papers like H-Nets or these other kinds of things. This will appear in many different places. So this is a good kind of pattern to be able to recognize. So what you're doing is, you are going to take in your inputs. And you're going to take it through a feed forward. And then I'm going to compute some sort of scores. And the scores that I'm going to compute is going to be the softmax between the experts and my inputs u. And then this s is going to be the scores. And I'm going to select the TopK of those guys.

[53:47.677] And that's going to be the gate that's going to control each of my experts. So this is like input. This is the residual term. This is the FFN. I've got gates. And the gates are controlled TopK. This is the equation form of what you would imagine of a TopK router. And the only detail that you're probably learning here is, how do I learn my gates? My gates are just learned by taking inner product between a weight for each expert and the inputs that I have. Very, very lightweight. So there's nothing complicated about the ways that I'm going to select my gates. Even within this simple thing, there's been some innovation and variation from DeepSeek and friends. DeepSeekMoE pioneered this idea, which is now very widely used, called Shared Experts. This is a very intuitive idea, where in the classic TopK router, you would have k different-- or sorry, in this case, n different experts.

[54:44.410] And these experts, you would select amongst them. But sometimes, you might say, well, there's some very common processing that I want to apply to all tokens. Maybe you don't want experts to be different every time. You want some experts to always be applied and then some other experts to be applied conditionally. So what DeepSeekMoE did was, they took the original expert design. They cut up the experts into smaller, finer grained chunks. And some subset of those fine grained experts were as shared experts that are always on. So these bypass the router. They always process our inputs. And they go back on the outputs. Sorry. And this is a nice idea, because it turned out that if you had this classic expert design, you were just reusing a lot of these weights to do common modeling. And so you can offload this into the one shared expert, letting these guys specialize even more.

[55:43.200] And as I was kind of saying, DeepSeek is really nice in that they do careful ablations in a lot of their papers. And so they have this very nice ablation, where they look at zero shared experts and then you have these big experts, 16 rooted experts. And this is really the old GShard Google design, compared to, they have very fine grained segmentation of experts and one shared expert amongst these. And as you make the experts smaller and smaller, and you have one shared expert, you see significant gains from both of those interventions. And for some of these, the shared expert has very significant improvements, like TriviaQA, NaturalQuestions.

[56:20.600] This blue to yellow gap right is shared experts or not. So more experts, shared experts, all generally seem to help. And if you want more validation of these claims and ideas, you see this from OlMoE as well. OlMoE is, I think, the nice Western carefully controlled MoE study. They do things like, oh, do shared experts help? They conclude that shared experts don't help very much. And number of experts, they show, doing fine-grained, many experts is helpful. So they disagree a little bit on the shared expert point, but they also have very nice, carefully controlled ablations. Many of the recent MoE do follow essentially the DeepSeek design. There were a couple of early MoEs up until this point. The first three, I keep giving Google praise for doing lots of interesting things. The first three were from Google, Mixtral, DBRX, and Grok were some of the early MoE attempts in the West.

[57:21.890] DeepSeek v1 comes up with both the fine-grained and shared expert design. And then that catches on, and everyone else follows from DeepSeek. I mean, it's a similar, I would say, phenomenon to the Llama design for dense transformers being essentially, the standard DeepSeekMoE and DeepSeek v3 are the standard design that people have copied for the MoE side of things. And if you look at Qwen 3.5, or GLM, or any of these modern models, you continue to see both shared experts and fine-grained experts in widespread use amongst all of these people. And so that's been a very battle-tested design. Any questions? Oh, yes. Good.

[58:02.050] Yeah. If there is a shared expert or there are few shared experts, how would parallelization work in that case? Yeah. So in the shared expert case, you wouldn't get any parallelization savings. Every activation has to route through those. You can copy the shared expert to basically eat memory savings to reduce comms cost. Cool. All right. So that's the design space of routers. Now, I want to talk about training MoEs. And it gets into really surprising, and I would say, very deep learning things about this. I think, initially, when I was learning about MoEs, I thought, there's no way that we can train these things well or reasonably. But it turns out, there are lots of tricks that when combined together, just really work well and robustly for some reason. So as we'll raise in the previous question, we don't want to activate all the experts during training, because if we do that, we're going to pay the full flops cost of all of our experts.

[59:05.050] And we have a lot of experts. So we need sparsity for efficiency at training time. But if we have sparsity, then gating decisions are no longer differentiable. And we also don't see all of the different counterfactual experts we could have picked. So these are both big problems. So there are different solutions. We could use RL to optimize our gating policies. I already mentioned this, not super popular. You could do things like stochastic perturbations. So if you're familiar with explore, exploit, bandit style stuff, you could do some of that. Now, the third one is, you could have a whole bunch of weird heuristics that deal with the fact that the gating decisions are sparse to balance out which experts go where. Guess which one people use in practice? I've been hinting at this. It's number 3. It's a collection of really interesting heuristics that end up working.

[59:56.090] I'll talk about the other two first. They're not widely used. I would open with that. But I think it's useful to know. So some people have done, this is from Clark 2020, research on things like, can you use RL to learn routers? So it does work. You can use RL, this is the green line over here. That's one of the baseline approaches. It works fine. But because of the gradient variance and complexity, it's not actually the optimal thing to do. And even some of the approaches that are proposed in this paper handily beat using reinforced gradients for routing decisions. Now, something that's very close to the idea of RL but starts to move more towards this heuristic world of write down objectives and try things is this idea of a stochastic approximation. And it comes from the earliest MoE paper from Noam, Shazier, and others, where what you do is, you don't really make a hard decision during training.

[01:00:59.220] What you do is you make stochastic decisions. So what you do is, you still have this routing setup. So this H is, you have the inner product between your parameters and your inputs. This is your standard routing inner product. But then, instead of just doing this, I'm going to inject some noise that's dependent on the scale of the input. And then that's going to be my H. And now, I'm going to keep the TopK and softmax, just like I was explaining with TopK. And then that's going to be the output that goes in to my TopK function. Now, if you look at this, what is this doing? Well, it just means that if you have two experts that are closely tied or it's very close to each other, then stochastically, you'll pick one of them and not the other. And as you backprop, you're going to backprop in a way that the experts that are helpful will get high weights. So this perturbation allows you to tie break and explore a little bit more. And hopefully, the gradient structure is such that, if an expert is helpful, it's going to emphasize that expert more,

[01:02:04.300] and it will get higher and higher weights in general. So this is one reasonable thing that you can do. It leads to experts that are a bit more robust. And the softmax allows you to learn how to rank all these k experts. Rather than just doing a hard selection of TopK, you're weighting by these elements. There's another scheme that's very closely related that's done in Fedus, et al, in 2022, which does a uniform, multiplicative perturbation. They do this for the reason of hardening the experts, making them robust. This was a useful thing to do, but this was later removed in later Google papers that do routing. And it's not really clear that this is necessary at all. You see later ablation, showing that not doing any of these stochastic or robustness tricks actually helps with both stability and overall quality of the final trained MoE. So this suggests, you don't need necessarily the stochastic exploration terms.

[01:03:13.030] So this takes us to what you actually do, which is a series of heuristics that, in the end, is going to give us a working MoE. And a lot of these have to do with essentially, balancing. So what's the problem if you just do normal gradient descent, ignoring any of these exploration, exploitation concerns? Well, what's going to happen is, you're going to first route your TopK experts. And then the strongest of those experts will get more signal. Your backprop is going to say, that expert was good. You increase the weight of that expert, move that parameter closer to that. So you get this rich gets richer effect, where experts that are chosen get very strong weights. Strong weights means that they're selected more often. And they run away, taking on everything. And so this expert collapse phenomenon or expert starvation phenomenon is a very, very real problem.

[01:04:06.260] And in some sense, this is the core issue that you have to solve with heuristic training of MoEs. And so what do you do? Well, what you do is, you add a heuristic loss to the total modeling loss. And this loss is supposed to balance out the tokens that go to different experts. And one approach, which is done in the switch transformer over here, you can look at what this does. You have different experts going from 1 to n And you have a batch of tokens. And what you want to do is, you want to make sure that these two inner products are small. OK. So what is happening here? F is the fraction of tokens that are dispatched to any single expert. And then P of I is the softer version. This is not just whether the token is dispatched to an expert, this is the total probability mass of the router that's allocated.

[01:04:59.140] And you multiply these two terms. And this is, at least to me, not a thing, where you would derive it from first principles. I think the easiest way to see what this thing is doing is to take the derivative of this object. So if you take the derivative of this with respect to P of I-- and P of I, remember is the probability mass allocated to an expert, well, the derivative is going to be the fraction of tokens that are allocated. So in some sense, you can think of this as a penalty in the gradient space, where the more tokens you get, the more negative gradient you get. So this is trying to push down the probability mass on very popular experts proportional to their fraction. And that's the action of this loss. That's the way to think about it. So the objective itself in equation 4 might not be clear to start with, but once you reason about the action of the gradient, it becomes clear what this is trying to do, which is to push down popular experts.

[01:06:00.150] So then we can actually take these ideas and essentially reconstruct DeepSeek v1 and 2s MoE system. So what do you do to build DeepSeekMoE? Well, you're going to do your normal losses. And you're just going to backprop straight through the experts, ignoring all this nondifferentiability and exploration concerns. But you're going to add per expert balancing that's identical to the switch transformer, what I just showed you. So you're going to have this loss that multiplies F, which is the fraction, and P, which is the total mass. But DeepSeek folks, they're very savvy with their systems design, so they don't just balance the experts, they also want to balance by device.

[01:06:40.590] So if your experts are allocated to some devices-- so let's say, my machine over here has four experts, this machine over here has another four experts, you want those two machines to be balanced, so both of them are running at full utilization. So you don't want to just make the experts even, you also want to ensure that the devices are even. So this is a secondary objective. It acts in exactly the same way. If you look at the function here instead of experts, you're just operating on each device's fraction, you apply the secondary device balancing loss, just for utilization reasons. In DeepSeek v3, they have a per expert bias term that they use. And they use an online learning trick to balance out experts, getting rid of some of these auxiliary losses, but they do, in the end, have to add some of these auxiliary losses in order to ensure extreme imbalances don't happen. So I think, I add this as a side note that DeepSeek v3 and others have started to get rid of some of these, I would say, uglier auxiliary losses,

[01:07:44.770] but there's no solution that fully gets rid of them so far. The last thing I'll talk about with load balancing losses is, you might ask, do I need to deal with this? It's all a bunch of complicated heuristics. I wish I could just get rid of the load balancing loss. But if you remove the load balancing loss, the effects are pretty catastrophic. These Are both from the OlMoE paper, where they do this nice ablation of removing the load balancing loss entirely. If you have the load balancing loss, you get this pink line at the top, very nice normal curves. Remove the load balancing loss, your losses significantly increase, and your training loss also is not doing so well. Maybe more tellingly, if you look at which experts are doing what, with load balancing and without load balancing, that's these bottom panels over here, you see a really, really stark difference.

[01:08:38.870] Without load balancing, almost all the tokens are going to two experts-- the yellow expert and the pinkish expert that you see here. Whereas if you have the load balancing loss, even though we had this kind of heuristic thing going on, all of these experts are being utilized across the token. So we have this nice, even utilization. So we do see that the load balancing loss does what we expect. And without it, we've thrown away a ton of parameters. Those experts are doing absolutely nothing for most of training. So this is how MoEs are trained. And I think, just to go back, this is pretty surprising to me, that you've got this thing that's nondifferentiable, you're doing this TopK selection. And it is, honestly, a pretty complicated object, and all you really need to do is to just add this balancing loss to even out the experts. And then the rest of it, you treat it as if you can just pump gradients through the system, and the model trains very nicely.

[01:09:34.380] And I think a big part of this is, this dynamic that if an expert is useful, you're going to reinforce it. And that's a positive reinforcement cycle that's nicely balanced out by evening things out. Those two dynamics cancel each other out. Cool. Oh, and the final thing I'll say is, that same trick is used in DSA. We already saw the TopK thing. It's used in things like H-Nets. If you're familiar with that, that was a temp to remove tokenizers. And so I think you'll see this trick more often in the future, where this idea of TopK selection and using load balancing or other kinds of auxiliary losses to enable that nondifferentiability will be an ingredient of future architecture design. Yes. [INAUDIBLE] mixture of experts is like a systems optimization. But if you actually exploit, can you look at them? Are they actually doing something in different domains? Yeah. So some papers have a visualization of which tokens are activated by which experts.

[01:10:37.080] And because the routers are so simple, it's not like, experts are like smart experts. They're not like medical experts, legal experts, or whatever. You do see that certain tokens are routed to different things, like punctuations might be routed to one expert or other symbols might be routed to one expert, non-English language character sets might be routed to another expert. But it's not really something where you can look at it and be like, oh, this is the Wall Street Journal expert or something, nothing like that. There's no semantics of maybe one way of putting it. Yeah. [INAUDIBLE] Yeah, the per expert loss, in principle, if it's perfectly enforced, will also enforce a per device balancing if your experts are evenly split. But I think, at least my interpretation of that is, you don't want to crank up the per expert loss so high that you get true full uniformity, because that has deleterious effects on training dynamics. But per device is important enough that you're willing to add on a little bit of extra loss to encourage per device balancing over others.

[01:11:37.550] Cool. OK. So now, I'm going to talk a little bit about the system side of things before we wrap up. So training MoEs has additional fun systems dynamics that are introduced. When I talk about parallelism, I'm mostly going to talk about the standard, let's call it, forms of parallelism. You can split up your data over many small chunks. You can take your one data set and split it up over many machines. You can do model parallelism, where you can take your model, and maybe you can split up your model over different machines. And you can combine those two in various ways, you can, of course, do that. But I think the thing that we realize is that each of those parallelism techniques has some limit. For data parallelism, you're maxed out your batch size. Once you hit your batch size, you can't parallelize by data anymore. For model size, there's natural cut points, where you can cut your model.

[01:12:33.430] Once you exhaust those, you can't parallelize anymore. If you combine this with expert parallelism, you get another additional nice access to parallelize on. And when you do this, a lot of the underlying implementations for this allow you to take advantage of sparse matrix multiply that has been built into GPUs. So another thing that I'll mention is, if you have multiple experts, let's say, on one GPU, one way of thinking about it is to say, I'm going to have multiple small matrix multiplies. But this is not nice, because you ideally want these bigger matrix multiplies, where you can reuse caches and do all sorts of things. So instead, what you can do is, you can nicely leverage sparsity. And block diagonal, of course, is the most basic form, but you can use much more of complicated forms of structured sparsity that are natively supported in hardware to allow you to essentially multiply experts and inputs in very, very clean ways, very fast.

[01:13:34.880] So this is another reason why MoEs are nice. If you think about the computation patterns of MoEs, they almost correspond to these kinds of structured matrix multiplications that are very easy and efficient to support in hardware. So there's this hardware architecture codesign that's happening with MoEs. And in terms of parallelism, one thing I'll mention, this is a very recent development. This is from Nemotron-3. I forgot who asked this question. When we do this expert parallelism, we have to be shipping activations from device to device in order to say like, oh, you belong in that expert, I need to ship you over. That can result in significant communications overhead. So some of the things that people have started to think about is to say, well, my shared expert, which I'm not going to communicate, this can be in big dimensions. But my experts, I need to communicate those activations, I want those to be in smaller lower dimensional vectors.

[01:14:34.960] So you might actually take your residual stream and downproject it first, and then decide to do the collective communication call of sending this activation out. And that will significantly save on communication without fully having the drawbacks of having a smaller hidden dimension size. So you can actually do these kind of projection tricks to nicely control the trade-off of communication and having your MoEs parallelized. One final trivia, let's call it detail of MoEs that I think is fun, but it has been solved in recent years, is stochasticity. If you do MoE infrastructure naively, what you'll find is that, unless your experts are perfectly balanced on your input distribution, you'll find that certain experts might be much more popular than others. And if some expert like, let's say, in this case, expert 0 is just a really popular expert, you're going to actually start running into situations where this expert queue of tokens just starts building up, and building up, and building up. And eventually, you'll get to a point, where you say, my queue is so long, I have to start dropping tokens. And the earlier generation of MoE inference infrastructure and code, a lot of what happened was, you would just actually silently drop this expert.

[01:15:56.510] And you just proceed on with your computation as if nothing had happened. You just send a zero back and pretend that that was fine. And this means that you can have these weird stochasticities, where if other users sent in queries that hit the experts that you're using, then you could actually get a worse result, because they would bump you out of the expert queue. These kinds of very strange things can happen. But these days, they drop less architectures that have gotten rid of a lot of these issues. So this is no longer really an issue. So mega blocks and other very common open source MoE frameworks don't have this issue anymore. So the last thing that I'll talk about in terms of the design of MoEs is this issue of stability. We've already talked about stability issues in the context of architecture design last lecture. But if you remember what I said then, I said exponentials are bad, divisions are bad, which means softmaxes are danger zone for stability issues.

[01:16:58.100] Now, what we done with mixture of experts? Well, we have introduced yet another softmax in the routing, because we're going to do TopK. We're going to do softmax. And that's going to be how we route to experts. Now, people have long noticed that the softmax operation in MoEs are potentially very, very dangerous. Barret Zoph and others, in the early days of Google MoE design, had an entire paper on MoE stability. And this was one of the things that they noticed. And I'm going to go into some of the details. You'll see the solution fairly often for some of the more sensitive and tricky parts. You might end up using something like Float 32 for just the expert router. You might have a z-loss as well here. We've talked about z-loss before in order to control the softmax stability issues that we saw. Z-loss, I think, was quite popular, actually, to do MoE router stability even in the early days.

[01:17:58.200] OlMoE, which I've been talking about extensively, has done ablations on adding or removing z-loss. It's quite clear, from these very spiky training loss curves, that z-loss on the router can be quite helpful. There are other issues with MoEs that you can end up with. And I think you probably all have experienced this maybe recently if you work with language models, which is that MoEs can be pretty annoying to fine tune. They have so many parameters that if you're trying to fine tune the experts, you actually end up with very serious overfitting issues. So you can see that if you have your dense models, train and val are fairly close, they don't overfit crazily. But if you look at the sparse, train and val gap is extremely large, as you fine tune these models, or, in this case, downstream task. I think it was one of the glue benchmark tasks.

[01:18:51.700] Now, you could have a case, where you fine tune non-MoE feed forwards. You could also fine tune just attention. These are pretty common interventions that people do, especially the attention fine-tuning is something that I see quite often with MoE works recently. But if your model has non-MoE layers, then, of course, you can also just fine tune those. And there is often others, we're kind of arguing for those. And of course, there's also the bitter lesson version of this, which is, maybe you should just use 1.4 million examples instead of whatever number you have. If you have a lot of data, you can retrain, basically, the MoE entirely during your fine-tuning stage. And you won't have quite as much of a generalization gap. OK. The last MoE-specific trick that I'll mention is this idea of upcycling. This has actually become a lot less popular in the last year.

[01:19:49.790] I don't think I've seen a single upcycled model this year. But it's a cool idea. And also, there were some very nice models earlier on using this trick. So I want to mention this. So upcycling is this idea that if you want a MoE, maybe you can just take a dense model that you've trained, and then you can instantiate a MoE based on that dense model. And the idea is, you just copy everything, including the MLPs. You just make a whole bunch of copies from these MLPs. You have a router, which have a random initialization. And then you just train this model. And if you do this, then because of the stochasticity of which inputs go to where and the MLP, you're going to get these experts that start to specialize. And in the end, you'll actually end up with MoE. And in some of the earliest papers that proposed upcycling, they show that an upcycled method could get much better language modeling, in this case, accuracy, performance, upcycling a dense model rather than continuing to train the same model, which is a cool thing that you can show.

[01:20:59.470] People have validated this at scale in many ways. So many CPM is another Chinese model that I quite like, because they do a lot of carefully controlled ablations. They had a MoE, where they took their miniCPM 2.4b model. They upcycled it to a 13.4 billion parameter model. And they got a bunch of almost free wins from doing that. Qwen also, when they did their Qwen MoE, initialize from there 1.8B Qwen model. They upcycled that to a Qwen 1.5.A, 2.7B model. And they ended up with one of the first and one of the very high performance models. It's, I think, one of the first larger scale upcycling successes. Of course, these days, I don't think anyone's really upcycling anymore, because you don't really train dense models and then convert them. You might as well just train your big hero run on an MoE to start with.

[01:21:59.820] So I don't think we see this quite as much, but I wanted to mention this, because I think it's an important thing to know in the action space of mixture of experts models. OK. So the last thing I want to end with is, I want to just walk us through the DeepSeek models-- v1, v2, and v3. I think there's a lot to learn from the evolution of DeepSeek models. And also, if you're interested, just in general about MoEs or architecture design, I would encourage that you read the DeepSeek papers. I think they're very well written. I think they have a lot of details that are worth digging into. So the DeepSeekMoE architecture at v1 is already, actually, the prototype of a lot of different modern MoE. It's got the shared and fine-grained expert structure. It uses standard TopK routing with the auxiliary loss balancing. I mean, in some ways, I would say, this is like the prototypical platonic ideal of MoE model.

[01:22:55.040] DeepSeekMoE v2 scales this guy up. So you've got two shared experts, many more fine-grained experts. And you've got this device routing component, and you've got a communication balancing component. Both of these are essentially adding auxiliary losses to optimize systems. I think it's important to realize that successful language model training is not just about deep learning, it's also about really respecting your systems. So DeepSeek v2 is really a reflection, I think, of this philosophy. And then in DeepSeekMoE v3, they've still got the shared and fine-grained experts design, but they've got different ways of doing balancing. They've gotten aux-loss-free ways of balancing the different experts. And they've switched to using sigmoid plus softmax, which is a different way of weighting their experts, but really, mostly similar. Finally, since we're already here, we might as well just walk through the rest of DeepSeekMoE V3.

[01:23:54.890] They've got lots of other cool things. I'm a big fan of how they think about systems. They've got multihead latent attention, where instead of doing QKV directly, what you're going to do is you're going to represent them as this lower dimensional latent c. So you have your hidden input. Instead of directly producing your Qs, Ks, and Vs, you're going to first produce the c. And then you're going to produce your Q, Ks, and Vs as a function of that. And that allows you to get significantly improved savings. Instead of KB cashing all your Ks and Vs, I only need to store these Cs. And the Cs are hopefully, lower dimensional. And that's why it's called MLA, latent activation, because the Cs, which are latent, are the KB cache quantities that you need to save. The only complexity that you might need to worry about is that this is going to conflict with rope when you do KB caching.

[01:24:51.150] So you have to be a little bit careful about how you rotate different dimensions and so on. And so the trick is, you have nonlatent dimensions that encode positions, but I'm not going to go into too much more details about that. Finally, something that I thought was a very cool idea for MoE, actually, DeepSeekMoE v3, but hasn't caught on very much, is the idea of MTP or multitoken prediction. Instead of predicting one future token, you predict multiple tokens all at once. That has a statistical arguments for why that's a good idea. Maybe it lets you predict the future a little bit better. But also, there's this nice systems argument that you now have a speculative decoder built in.

[01:25:30.840] Percy will talk about that when he talks about inference later, but it's just a trick that you can use to speed up your decoding from your model. OK. I think we're roughly at time. So to put everything together right, MoEs are this very clever idea to take advantage of sparsity, so that you can have more parameters than you're paying for. So you don't pay for the compute cost of all of your parameters, but you are getting the parameter benefits. You might initially think that this kind of routing problem is very hard, but it turns out that very simple things work well, even at scale. And at this point, it's clear that MoEs are here to stay, so you should understand how they work and what they are.

[01:26:13.960] Thanks.

Attention alternatives and mixture of experts

课程简介 ​

CS336 Lecture 4 中文讲义：注意力替代方案与混合专家（MoE） ​

学习目标 ​

课程主线 ​

一、注意力机制的替代方案 ​

1.1 背景：长上下文的计算瓶颈 ​

1.2 线性注意力的核心思想：乘法结合律 ​

1.3 Mamba-2：门控线性注意力 ​

1.4 Gated Delta Net：双门控与投影更新 ​

1.5 混合注意力架构的受控研究 ​

1.6 DeepSeek Sparse Attention（DSA） ​

二、混合专家（Mixture of Experts） ​

2.1 MoE 的基本动机 ​

2.2 路由机制设计空间 ​

2.3 共享专家与细粒度专家 ​

2.4 训练 MoE 的挑战 ​

2.5 微调 MoE 的稳定性问题 ​

2.6 MoE 的系统层面考量 ​

2.7 DeepSeek 系列演进 ​

三、课程作业关联 ​

常见误区 ​

复习清单 ​

CS336 Lecture 4 深度解析：从线性注意力到混合专家——突破 Transformer 的双重瓶颈 ​

导语 ​

第一部分：注意力机制的效率革命 ​

1.1 成本困境：当二次方复杂度成为瓶颈 ​

1.2 核心思想：乘法的结合性 ​

1.3 从朴素到实用：门控的引入 ​

1.4 混合框架：从理论到实践 ​

1.5 另一个思路：稀疏注意力（DeepSeek Attention） ​

第二部分：混合专家（Mixture of Experts） ​

2.1 核心思想：用稀疏换规模 ​

2.2 设计空间：路由、规模与训练 ​

2.3 训练难题：不可微性与负载均衡 ​

2.4 系统协同：MoE 的硬件优雅性与细微陷阱 ​

2.5 稳定性的额外考量 ​

2.6 从密集到稀疏：Upcycling 的智慧 ​

2.7 案例研究：DeepSeekMoE 的演进 ​

第三部分：本讲在 CS336 课程中的承上启下 ​

结语：两个维度的突破，一种共同的智慧 ​

CS336 2026 Lecture 4：完整中文逐字稿 ​

正文 ​

CS336 2026 · Lecture 4 · Attention alternatives and mixture of experts ​