FlashAttention 原理之 softmax 分块计算

标准 Softmax

Softmax 函数(也称为归一化指数函数)是一个将向量转换成概率分布的函数。对于输入向量 x，softmax 函数将其转换为一个概率分布向量，其中每个元素的值在 (0,1) 之间，且所有元素之和为 1。

$$softmax(x_i)=\frac{e^{x_i}}{{\sum }_j{e}^{x_j}}$$

归一化 Softmax

$$m(x):=\underset{i}{\max }{x}_i$$

这里，$x=[x_1,x_2,\dots ,x_B]$ 表示一个有 B 个分量的向量（例如，模型输出的对各类的打分）。$m(x)$ 则是向量 $x$ 中所有分量的最大值。

$$f(x):=[e^{x_1-m(x)},e^{x_2-m(x)},\dots ,e^{x_B-m(x)}]$$
这里做了一个“减最大值”的操作，即把每个 $x_i$ 都减去整个向量的最大分量 $m(x)$，然后取指数。这样做的好处是数值更稳定：当 $x_i$ 很大时，直接算 $e^{x_i}$ 容易导致溢出；但减去最大值以后，指数部分变为 $x_i-m(x)$（一个相对较小的或非正的数），从而避免数值爆炸。

$$\ell (x):=\sum _{i}f(x{)}_i$$
这里把向量 $f(x)$ 的各个分量加起来得到标量 $\ell (x)$。

$$\text{softmax}(x):=\frac{f(x)}{\ell (x)}$$
把每个分量 $f(x{)}_i$ 除以 $\ell (x)$ 后，就得到了标准的 Softmax 输出向量。

$$\text{softmax}(x{)}_i=\frac{e^{x_i-m(x)}}{{\sum }_j{e}^{x_j-m(x)}}.$$

由于每一项都经过指数函数且被总和归一化，它满足所有分量都非负且所有分量之和为 1，因此是一个有效的概率分布。

分块 Softmax

假设我们有两个同维度向量
${\mathbf{x}}^{(1)}和{\mathbf{x}}^{(2)}\in {\mathbb{R}}^B$，把它们拼接(concatenate)成
$$\mathbf{x}=[{\mathbf{x}}^{(1)},{\mathbf{x}}^{(2)}]\in {\mathbb{R}}^{2B}.$$

下面的公式说明，如何在不重复完整计算的情况下，用“各自的部分计算结果”组合成拼接后向量的 Softmax。先给出它的步骤，再解释其意义和好处：

最大值 m(x) 的分块计算

定义单个向量的最大值

对于 ${\mathbf{x}}^{(1)}\in {\mathbb{R}}^B$，我们先定义
$$m\left({\mathbf{x}}^{(1)}\right)=\underset{i}{\max }\left({\mathbf{x}}_i^{(1)}\right),m\left({\mathbf{x}}^{(2)}\right)=\underset{i}{\max }\left({\mathbf{x}}_i^{(2)}\right).$$

定义拼接向量的最大值

由于 $\mathbf{x}$ 是把 ${\mathbf{x}}^{(1)}$ 和 ${\mathbf{x}}^{(2)}$ 拼到一起，那么
$$m(\mathbf{x})=m([{\mathbf{x}}^{(1)},{\mathbf{x}}^{(2)}])=\max (m({\mathbf{x}}^{(1)}),m({\mathbf{x}}^{(2)})).$$

这样就不需要对拼接后的 $\mathbf{x}$ 再扫描一次去找最大值，而是只要比较两个子向量各自的最大值即可。

“指数向量” f(x) 的分块计算

Recall
$$f(\mathbf{x})=[e^{x_1-m(\mathbf{x})},e^{x_2-m(\mathbf{x})},\dots ,e^{x_{2B}-m(\mathbf{x})}].$$

由于 $\mathbf{x}$ 拆成了两块 ${\mathbf{x}}^{(1)}$、${\mathbf{x}}^{(2)}$，我们分别对每块计算其对应的“指数向量”：
$$f\left({\mathbf{x}}^{(1)}\right)\text{和}f\left({\mathbf{x}}^{(2)}\right).$$

然后拼起来即可。但要记住，每一块真正要减去的“中心化值”是整个 $\mathbf{x}$ 的最大值 $m(\mathbf{x})$
，因此它们之间会出现一个额外的“补偿系数”：

$$f(\mathbf{x})=[e^{m({\mathbf{x}}^{(1)})-m(\mathbf{x})}\cdot f({\mathbf{x}}^{(1)}),e^{m({\mathbf{x}}^{(2)})-m(\mathbf{x})}\cdot f({\mathbf{x}}^{(2)}\left)\right].$$

直观上看，如果某一块（比如 ${\mathbf{x}}^{(1)}$）的最大元素是整个拼接向量的最大元素，那么它带来的指数因子就会是 $e^{m({\mathbf{x}}^{(1)})-m(\mathbf{x})}=e^0=1$。而另一块若不是最大的，就会额外乘上一个小于 1 的因子。

归一化项 $\ell (x)$ 的分块计算

Softmax 要把向量的指数项归一化到和为 1，所以我们需要计算
$$\ell (\mathbf{x})=\sum _{i=1}^{2B}{e}^{x_i-m(\mathbf{x})}.$$

利用分块思想，可以分成两段求和，再用与上一步相同的补偿系数连接起来：

$$\ell (\mathbf{x})=\underbrace{e^{m({\mathbf{x}}^{(1)})-m(\mathbf{x})}\ell \left({\mathbf{x}}^{(1)}\right)}\text{第1块贡献}+\underbrace{e^{m({\mathbf{x}}^{(2)})-m(\mathbf{x})}\ell \left({\mathbf{x}}^{(2)}\right)}\text{第2块贡献},$$

同理，也只需要各块自己内部的和，再用一个相对的尺度因子即可。

最大值 $\mathrm{softmax}(x)$ 的分块形式

把上面得到的 $f(\mathbf{x})$ 和 $\ell (\mathbf{x})$ 带入到
$$\mathrm{softmax}(\mathbf{x})=\frac{f(\mathbf{x})}{\ell (\mathbf{x})},$$

就得到在分块后的 Softmax 形式：

$$\mathrm{softmax}(\mathbf{x})=\frac{\left[e^{m({\mathbf{x}}^{(1)})-m(\mathbf{x})}f\right({\mathbf{x}}^{(1)}),e^{m({\mathbf{x}}^{(2)})-m(\mathbf{x})}f({\mathbf{x}}^{(2)}\left)\right]}{e^{m({\mathbf{x}}^{(1)})-m(\mathbf{x})}\ell \left({\mathbf{x}}^{(1)}\right)+e^{m({\mathbf{x}}^{(2)})-m(\mathbf{x})}\ell \left({\mathbf{x}}^{(2)}\right)}.$$

为什么这样做？

1. 数值稳定性
   跟单向量计算 Softmax 类似，这里也要减去整段向量的最大值 $m(\mathbf{x})$，以避免 $e^z$ 里的 $z$ 太大或太小导致溢出/下溢。
2. 减少重复计算
   如果我们已经知道各块各自的 $\max$ 值和求和 $\ell ({\mathbf{x}}^{(k)})$，那就无需把 $\mathbf{x}$ 整体重新扫描、求最大值、求指数和，总结出公式即可快速拼成拼接后向量的软最大值。
3. 方便分布式或分块处理
   在实际系统里，${\mathbf{x}}^{(1)}$ 和 ${\mathbf{x}}^{(2)}$ 可能来自不同子网络或不同设备。这种分块计算可以让每一块先在本地完成自己的 Softmax 部分计算，最后再做一次简短的组合归一化即可。