CS336 Lecture 04 - 注意力替代方案与混合专家
注:本文整理自 Stanford CS336 的 2026 版 Lecture 04 讲义;文末课程主页链接仍指向公开可访问的既有课程站点,用于补充背景资料。
Stanford CS336: Language Models From Scratch(2026 春季)第四讲,围绕两个核心问题展开:当上下文越来越长时,标准 attention 的代价怎么控制;当算力预算固定时,模型容量还能怎样继续扩大。
本讲在回答什么问题?
到了 Lecture 04,问题已经从“如何训练一个 Transformer”进一步推进到“当标准 Transformer 的成本开始失控时,还有什么替代路径?”
这一讲主要讨论两条今天非常现实的路线:
- 注意力替代方案:解决长上下文下 attention 二次复杂度的问题
- Mixture of Experts(MoE):在不同比例增加 FLOPs 的前提下扩展参数规模
前者关注的是 sequence length 扩展,后者关注的是 model capacity 扩展。两者共同指向一个主题:现代大模型的发展,越来越像是在和系统成本赛跑。
为什么标准 attention 会成为瓶颈?
标准 self-attention 的核心计算可以写成:
Attn(Q, K, V) = ρ(QK^T)V
其中最昂贵的部分是 QK^T。如果序列长度是 n,那么这一项的复杂度大致是二次的。当 context window 不断增大时,attention 的时间和内存成本都会迅速上升。
这就是 Lecture 04 的出发点:上下文变长以后,标准 attention 不是不能用,而是越来越贵。
因此,业界常见的应对方式大致分三类:
- 做更好的系统工程优化
- 采用局部 + 全局混合注意力
- 直接寻找 attention 的替代机制
而本讲重点讲的是第三类。
线性注意力:最直接的降复杂度思路
Lecture 04 从一个非常漂亮的代数重排开始:
QK^T V = Q(K^T V)
如果忽略某些非线性归一化细节,这个重排意味着我们可以把原来与 n^2 相关的计算,改写成更接近线性的形式。
为什么它重要?
因为一旦这个重排成立,计算量就不再需要显式构造完整的 n × n attention matrix。对于长序列来说,这会显著降低代价。
Lecture 04 的重点不是说“线性注意力已经全面取代标准 attention”,而是指出:只要愿意放弃一部分原始注意力形式,就能换来巨大的复杂度收益。
这类方法的核心吸引力在于:
- 上下文越长,收益越明显
- 推理时更容易维持线性扩展
- 对系统实现非常友好
线性注意力为什么又像 RNN?
Lecture 04 进一步指出,线性注意力还有一个更深刻的特性:它可以写成一种递推状态更新形式。
讲义给出的形式大致是:
S_t = S_{t-1} + k_t v_t^Ty_t = q_t^T S_t
这意味着什么?
意味着模型不一定非要把所有历史 token 两两做注意力对比,而是可以把历史压缩进一个不断更新的状态 S_t 里。这个视角会让线性注意力和传统 RNN / 状态空间模型之间出现明显的“家族相似性”。
这种 duality 有什么价值?
Lecture 04 很强调这点:
- 训练时,仍然可以使用更适合并行计算的形式
- 推理时,则可以切换到递推更新的线性形式
这是一种很强的工程性质,因为它允许模型在训练和推理阶段采用不同但等价的计算视角,从而同时兼顾吞吐与长上下文推理效率。
从线性注意力到 Mamba-2:增加表达力
纯线性注意力虽然便宜,但表达能力有限。Lecture 04 接着展示了一个重要方向:通过引入门控和位置相关权重,让线性状态更新变得更灵活。
Mamba-2 可以粗略理解为:
- 保留“状态递推”的优点
- 让状态更新不再是简单相加
- 对当前输入施加内容相关的 gating
这样做的结果是,模型仍然具备某种线性 / 递推结构,但表达能力更强,不再只是最朴素的线性注意力。
为什么 gating 这么关键?
因为一旦有了 gate,模型就不只是“把新信息加进状态”,而是可以进一步决定:
- 保留多少旧状态
- 让哪些新信息进入
- 在当前位置上如何调整更新强度
Lecture 04 传递出的核心判断是:很多 attention 替代结构之所以逐渐变强,不是因为抛弃了记忆机制,而是因为它们学会了更细致地控制记忆机制。
Gated Delta Net:进一步加入“擦除”和选择性写入
在 Mamba-2 的基础上,Lecture 04 又介绍了更进一步的 Gated Delta Net(GDN)思路。
它不仅能够决定写入多少新信息,还能够:
- 选择“不输入”
- 选择性擦除旧状态中的某些方向
这件事的重要性在于:模型的状态不再只是一个不断累加、越来越臃肿的缓存,而是一个可以被主动整理和重写的动态记忆结构。
从工程直觉上看,GDN 代表的是一类更普遍的方向:把“记忆”从静态缓存改造成受控状态机。
这也是为什么它和 fast weights、test-time training 等思路有亲缘关系——它们都在研究怎样让模型在推理过程中更高效地维护与更新内部状态。
混合架构:不是“全替代”,而是按比例混搭
Lecture 04 给出了一个很现实的观察:真正落地到大模型时,很多团队并不会彻底抛弃 full attention,而是采用 hybrid 路线。
例如:
- 多层线性注意力 + 少量 full attention
- Mamba / GDN 与 attention 交替堆叠
- 在长上下文场景保留部分稠密注意力以维持表达能力
为什么这样做?因为标准 attention 虽然贵,但它依然很强。完全替代往往有风险,而混合方案能在两个世界之间做折中:
- 既拿到更好的长上下文扩展性
- 又保留一部分传统 attention 的建模能力
Lecture 04 展示的一些新模型,本质上都在探索这个混合比例该怎么选。
另一条路线:Sparse Attention
如果不想完全改写 attention 机制,另一条常见思路是:不要让每个 token 都看见所有 token。
这就是 sparse attention 的核心逻辑。
Lecture 04 提到的一个方向是:
- 通过较轻量的索引器先选出值得关注的位置
- 再只在这些位置上做更昂贵的计算
这种方法的优点是:
- 可以部分继承标准 attention 的表达方式
- 复杂度更可控
- 有些方法甚至可以在稠密短上下文预训练之后再做后置适配
这反映出长上下文优化的现实图景:并不是所有人都在赌同一种新架构,很多团队是在“稠密 attention、线性替代、稀疏化”之间寻找各自最顺手的折中点。
Mixture of Experts:另一种扩容方法
Lecture 04 的后半部分转向了另一个今天极其重要的话题:Mixture of Experts(MoE)。
如果说 attention alternatives 主要在解决“序列太长”的问题,那么 MoE 更像是在解决“模型还想继续变大,但 FLOPs 不能按参数量线性一起涨”的问题。
MoE 的基本想法
讲义给出的直观定义是:
- 把一个大的前馈层替换成很多个专家网络
- 再用一个 selector / router 为每个 token 选择少数几个专家
于是就出现一个非常诱人的性质:
- 总参数量可以很大
- 每次实际激活的参数只是一小部分
这意味着,模型容量可以上升得比每 token FLOPs 更快。
为什么 MoE 会越来越流行?
Lecture 04 总结了 MoE 流行的几个核心原因。
1. 相同 FLOPs 下,更多参数通常更强
这几乎是 MoE 最有吸引力的一点。若路由设计得当,MoE 可以在接近固定计算预算下,提供更大的总模型容量。
2. 训练速度可能更有优势
当系统并行化做得足够好时,MoE 往往可以比同级别 dense 模型更快地逼近目标性能。
3. 天生适合多设备扩展
不同专家可以天然分散到多个设备上。对于大规模分布式训练来说,这种结构和现代硬件拓扑往往非常契合。
4. 实证结果已经非常有说服力
Lecture 04 提到,不只是少数实验模型,很多今天最强的开源模型已经广泛采用 MoE。也就是说,它不再只是“研究方向”,而是已经进入主流工业实践。
那为什么 MoE 没有更早统治一切?
Lecture 04 也很诚实地讨论了 MoE 的代价。
基础设施复杂
MoE 的难点从来不只是写一个 router,而是整套系统工程:
- token 如何分发到不同专家
- 跨设备通信怎么做
- 负载如何均衡
- 某些专家过载怎么办
这些问题会把工程复杂度迅速抬高。
训练目标往往带有启发式色彩
为了避免路由塌缩、专家不均衡、训练不稳定,MoE 往往需要额外的辅助损失或路由约束。这些设计通常相当经验化,不像标准 dense Transformer 那样“配方稳定”。
因此,MoE 的问题从来不是“有没有理论吸引力”,而是:你有没有足够成熟的系统和训练 recipe 去把它稳定跑起来。
把本讲连起来看:两类扩展,两个方向
Lecture 04 其实在讲两种不同维度上的扩展:
1. 面向更长上下文的扩展
代表方法:
- linear attention
- recurrent / state-space 风格更新
- Mamba-2
- Gated Delta Net
- sparse attention
- hybrid attention
这条线关注的是:上下文越来越长时,如何降低每一步的序列相关成本。
2. 面向更大模型容量的扩展
代表方法:
- Mixture of Experts
这条线关注的是:在单位 FLOPs 有限的情况下,怎样让模型拥有更多参数与更高容量。
这两条线并不冲突,未来的模型很可能同时结合两者:既有更高效的长上下文机制,也有更高容量的稀疏专家结构。
总结
Lecture 04 给人的最大启发是:Transformer 并不是一个已经定型、不可再动的终点。随着上下文长度和模型规模不断上升,标准 attention 与 dense FFN 的成本都会越来越突出,因此行业自然会往两个方向持续探索:
- 用更高效的序列机制替代或部分替代标准 attention
- 用稀疏激活的专家结构替代纯 dense 扩容
这节课最值得记住的不是某个单独模型名字,而是背后的系统逻辑:
- 长上下文问题本质上是在和二次复杂度搏斗
- MoE 问题本质上是在争取“更多参数但不同比例增加 FLOPs”
- 真正能落地的方案,往往既不是最纯粹的理论形式,也不是最激进的全替代,而是系统上最可控、收益最稳定的折中设计
如果说第三讲是在回答“现代 Transformer 为什么长这样”,那么第四讲就在进一步回答:“当这个范式继续扩张时,它可能往哪里变。”