CS336讲座4注意力替代方案笔记

Flash attention

线性注意力机制

考虑常规的注意力操作，

这是典型的关于线性注意力（Linear Attention）的推导起点。

常规运算：

第一步：计算 Q×KT

• 矩阵形状：(n×d)×(d×n)

• 结果矩阵大小：n×n。这个巨大的矩阵就是注意力分数矩阵。

• 计算量：生成这个 n×n 的矩阵，需要n2次乘加运算，n2个数，每个数要 2d-1 次乘加运算。

第二步：将结果乘以 V

• 矩阵形状：(n×n)×(n×d)

• 结果矩阵大小：n×d

• 计算量：这也是一个矩阵乘法，但前提是必须先把第一步算完。

总复杂度：在第一步生成 n×n矩阵的时候，就已经引入了O(n2)的复杂度。这就是为什么普通 Attention 处理长文本时会崩溃——序列长度 n 变大一点，计算量就会平方级爆炸。

线性运算：

利用了乘法的结合律，我们改变括号的位置

第一步：计算 KT×V

• 矩阵形状：(d×n)×(n×d)

• 结果矩阵大小：d×d。注意这个关键点！结果是一个很小的矩阵，它跟序列长度 n 没关系，只跟模型隐藏层维度 d 有关（d 是固定不变的，比如 4096）。

• 计算量：因为结果矩阵很小，所以计算量也是固定的，每次计算只需要做 d×n×d 实际上是 d×d×(2n-1) 次乘法，与 n 是一次方成正比，即 O(n)。

第二步：计算 Q×(上一步的结果)

• 矩阵形状：(n×d)×(d×d)

• 结果矩阵大小：n×d

• 计算量：同样，这里计算量也是O(n)，因为随着输入序列长度 n 变大，计算量只会线性增加。

这是线性的计算时间（太棒了），但更妙的是，这看起来就像一个 RNN（循环神经网络）。

这种“对偶性”使得我们既能高效地训练（使用并行的、二次复杂度的形式），又能高效地推理（使用串行的、线性复杂度的形式）。

• 这里的线性 Attention，把历史所有 Token 的信息压缩到了一个固定的矩阵 S 中（大小为 d×d）。每当新来一个 Token（kt​ 和 vt​），只需将它“累加”进 S 里即可。

• 这个状态 SS 意味着什么？意味着你不需要像传统 Transformer 那样，为了生成下一个词，必须把前面所有 10 万个 Token 的 Key 和 Value 都重新算一遍。你只需要维护这一个 SS 矩阵就行了。这种“有记忆状态”的递归形式，就是 RNN 的核心特征。

“线性 Attention”虽然快，但它脑子不太好使，记不住新东西，也甩不掉旧垃圾。给它装上“记忆阀门”，让它该记的记，该忘的忘，变得更聪明。

🟢 第一步：认清老问题

上一页提到的公式，你可以把它想象成一个“永远在往杯子里倒水”的过程。

• 杯子就是 St​（记忆状态）。

• 每来一个新的字，就相当于往杯子里倒一点水（加上𝑘 𝑡 𝑣 𝑡 𝑇​）。

老问题的缺点（致命伤）：水永远在加，永远不减少。如果模型读了 100 万字，杯子里装满了水，最后面的字和前面的字混杂在一起，模型就彻底糊涂了。这叫“记忆饱和”。

🔴 第二步：Mamba-2 的解决方案

公式是 Mamba-2 用的：

新手解释：加入“遗忘阀”

• γt 这个希腊字母，叫 Gamma）：这就是一个“遗忘阀”，取值在 0 到 1 之间。

• 原理：每次要加入新东西之前，先拿这个阀门把旧杯子里的水倒掉一点点（比如说倒掉 10%）。

• 作用：这样杯子就不会满出来了。模型可以记住最重要的信息，慢慢遗忘掉不重要的老信息。

🔵 第三步：Gated Delta Net 的更高级解法

这是 PPT 的核心，公式变复杂了：

新手解释：加入“擦除刷子”
这里引入了一个新东西叫 βt（Beta，0到1之间的值）。

我们换个比喻：想象你在写课堂笔记。

• 以前的方式（线性 Attention）：老师每讲一句话，你就直接写在笔记本最后面。时间久了，本子写满了，前面的内容和后面的内容混在一起，根本找不到重点。

• Gated Delta Net 的升级：这个新发明给了你一把“修正液/擦除刷子”。

  • 当你看完当前的字，模型会判断：“这个新信息重要吗？”

  • 如果新信息非常重要（βt≈1）：它不会直接写在最后，而是会先找到笔记本里旧有的、意思相近的那句话（也就是 ​指的方向），把它擦掉/覆盖掉，然后再把当前的关键信息写上去。

  • 如果新信息不重要（βt≈0）：它就不写进去，保持原样。

MOE混合专家

🟦 首先看左侧：Dense Model（稠密模型/传统模型）

左边是传统的大模型（比如早期的 GPT-3、Llama 2）。

• 架构逻辑：输入进入“FFN Layer”（前馈神经网络）时，必须经过同一个巨大的、唯一的“大脑”来计算。

• 缺点：FFN累死（计算量大），而且单个模型无法无限做大。如果把这个FFN变得特别巨大（几百亿参数），推理速度会变得极慢。

🟨 再来看右侧：Sparse Model（稀疏模型 / MoE）

这就是这张图的核心——MoE 架构。

• 架构逻辑：原来只有一个巨大的 FFN 层，现在被拆成了多个小的 FFN（FFN 1 ~ FFN 4）。中间多了一个“Selector Layer”（选择器/路由器）。

• 工作机制：

  1. 当词 x1​ ("The") 进来时，Selector 会判断这个字应该交给哪个专家（比如箭头指向了 FFN 2）。

  2. 当词 x2 ("Dog") 进来时，Selector 判断它适合另一个专家（比如箭头指向了 FFN 1 ）。

MoE 可以让模型在“变大”的同时，不增加计算量

“通过 MoE，我把模型的‘知识容量’扩大了 10 倍，但是我每算一个字用的力气（FLOPs / 电费 / 推理时间）却和以前完全一样！”

在 1024 块显卡的训练场景下，这种“稀疏化”意味着：
你可以把模型塞到7000 亿参数（比如 GPT-4），但每次推理只激活一小部分专家，实际算力消耗相当于一个700 亿的模型。这就是大模型既能“懂得特别多”，又能“回答得特别快”的根本原因。

这张图告诉我们，MoE 不仅“理论上变聪明了（参数量变大，算力不变）”，而且在“工程上（分布式训练）”它是天然为大规模 GPU 集群设计的。

难点

1. 不是谁都能跑 MoE：训练 MoE 需要极昂贵的高速互联网络（把 1024 张卡串起来）。没有这个基础设施，MoE 跑不起来，或者跑得比普通模型还慢。

2. MoE 是“大厂特供”：只有当你有了成百上千张卡做数据并行时，你才有足够的算力余量去塞下那些“大模型参数”，MoE 的算力优势才能发挥出来。

3. 训练过程像“走钢丝”：想要让那个“路由器”不乱分配工作，需要很高的技巧，一旦平衡没调好，几百万美金的训练就会白费。

结合你的上下文：
你之前问了“用 1024 张 B100 训练 15T token”。要完成这个工程，你不仅要有 1024 张卡，你还必须拥有一个能够把卡片之间做到极高带宽（比如 NVLink + InfiniBand）的超级网络，以及一套能解决“路由坍塌”的复杂算法。这就是为什么全世界只有 OpenAI、Google、Meta、DeepSeek 等少数公司能做出万亿参数大模型的原因。

MoE 的核心命脉——路由（Routing）机制

1️⃣ 上半部分：Top-k 路由（当前主流方案）

这是目前几乎所有主流大模型都在用的动态路由。

2️⃣ 下半部分：哈希路由（简单粗暴的基准线）

这是早期或作为对照实验的最原始的静态路由。

稳定性问题

训练 MoE 时，如果路由算法出了点小毛病，整个模型的 Loss（损失值）就会突然暴涨（图里那个尖锐的反弹），导致模型瞬间变傻。在大规模集群上，一旦遇到这种情况，工程师就得强行停止训练、找回快照、修改超参数，再重新开机。每一次这样的折腾，都要浪费掉数百万的算力成本。

🔵 第一部分：为什么路由器会发疯？

罪魁祸首：低精度（如 bfloat16）下的舍入误差。

🟢 第二部分：底部解决方案（FP32 路由器 + Z-Loss）

既然低精度会导致这致命的 36% 偏差，各大厂是怎么补救的？

• 怎么运作（看中间的图）：
  输入一个词（比如 x1​ "The"），它会进入一个可学习的路由器（Router，通常是一个极小的神经网络）。路由器会算出它跟每个专家的“匹配度”（概率分数，柱状图里）。

• 什么是 Top-k？
  “Top-k”的意思是：只激活分数排名最高的前 k 个专家来处理这个词。

• DeepSeek (7)：每次激活 7 个专家。

  • 💡 重点解析 DeepSeek 的 k=7：通常 MoE 模型的专家总数可能高达 100 多个。DeepSeek 选择每次激活 7 个专家，是一个非常高的数值。这意味着，在保持总计算量（FLOPs）不变的情况下，它把更多专家拉进来“共同会诊”，让输出更接近普通稠密模型的效果，极大提升了模型的上限。

• ❌ 缺点：训练不稳定。如果路由器“发疯”，全把病人塞给同一个专家，就会导致路由坍塌，训练直接崩盘。工程师需要设计复杂的“负载均衡”辅助损失函数来修正它。

• 怎么运作（看下面的图）：
  没有神经网络，没有学习能力。输入一个词，直接扔进一个“哈希函数（Hash Function）”。这个函数就是一个固定的数学算式（比如直接拿这个词的拼音首字母，或者 ASCII 码做个换算）。

• 结果：
  计算出的哈希值直接决定了这个词去哪个专家（图里固定指向某一个专家）。

• 问题产生：路由器（Router）在计算分配概率时，需要计算一个Softmax（找出排名前 k 的专家）。路由器算出来的一组分数里，最大的一个是128.5。

• 毁灭性后果：为了数值稳定，Softmax 核心操作会统一减去最大值。本来正确的最大值是 128.5，由于舍入误差变成了 128。这导致原本有一点点优势的那个专家，瞬间失去了优势！

  • 本来应该分配给最相关专家的概率是14.2%。

  • 因为 0.5 的误差，概率硬生生被算成了9.1%。

  • 相差了整整 36%！

• 解法一（硬件层面的精准定位）：只用 FP32 来算路由器
  PPT 明确提出Solution: Use Float 32 just for the expert router。

  • 大白话翻译：其他几千亿的参数我们全用节省显存的bfloat16或者FP8，唯独这个非常关键的路由器（Router），我们要用到精确度极高的FP32（全精度单精度浮点数）。

• 解法二（数学层面的限制器）：Z-Loss 辅助损失
  底部的公式 Lz(x)就是大名鼎鼎的Z-Loss。

  • 大白话解释：路由器算出的分数不能太大。Z-Loss 就像给路由器的输出带上了一道“脚镣”，刻意惩罚它算出过于极端的数值（比如那个 128.5），逼迫路由器的 logits 保持在一个合理的范围内。