Sys4AI-Gradient

一、总论

虽然我已经在之前的文章中讨论过一遍张量是如何求导的了，并且附上了详细的数学推导。但是我觉得那次的讨论还是有些过于偏向于数学的严谨，而忽略了实际使用中的直观。

所以我打算再推一遍，省略一些数学细节，但是更加注重实际的使用，包括对于计算和内存开销的估算，矩阵的形状等。

---

二、规律

2.1 LOSS 是标量

无论 LLM Forward 最终是生成一个 token，还是一段 output 序列，还是一个 batch 的 output 序列组，最终的 loss 都是一个标量。

这点似乎还是有些反直觉的，这是因为在很多科普文章中，为了简化描述，往往只会用一个 token 来举例子，而实际上，训练往往是以 batch 的形式进行，并且与不是只生成一个 token，而是一串 token。

在实际生产中，大模型 forward 的结果其实一个形状为 $[B, S, V]$ 的三维张量，用于表示概率，其中有：

• $B$：Batch Size
• $S$：Output Sequence Length
• $V$：Vocabulary Size

对于每个 token 来说，它生成的是一个概率分布，形状是与词表大小相同的一维向量 ，它会与一个标量的 label（可以理解为 “正确答案”）去计算交叉熵，也就是看一下这个概率分布合理与否。最终会得到一个标量的 loss 。

那么此时，我们会得到一个 loss 矩阵 $L$ ，形状是 $[B, S]$ 。问题在于，我们会直接拿着这个二维张量去进行 backward 吗？并不会，我们会对 $L$ 进行 reduce（非常简单，就是把 $L$ 的所有分量求平均值），最终得到一个标量 $l$：

$$l = \frac{\sum^B_j \sum^S_i L_{i, j}}{BS}$$

至于为什么不直接用二维张量去 backward，我觉得是出于计算的简便性的考虑。我们都知道雅可比矩阵（Jacobian Matrix）是一个二维张量，而其实因变量和自变量都是一维张量。之所以会发生 “升维”，是因为我们要描述 “每个自变量的分量” 对于 “每个因变量的分量” 的影响，所以维度就升高了。

如果我们直接使用 $[B, S]$ 的 $L$ 进行 backward，显然梯度的形状应该会变成高维张量，这显然不利于我们进行梯度下降。而且最关键的是，就算我们能忍受高维张量，到最终这些高维张量还是要 reduce 后再更新模型的。

也就是即使我们维护了 “权重张量对于不同输出的影响” 这个张量，最后我们还是要把不同输出对应的影响全加起来，然后再更新权重，相当于没变化。

总结一下，我们将上面的标量 loss 记录为 $l$ 。

2.2 形状相同

梯度的形状是与原张量的形状是相同的。

这个规律其实和前面的规律息息相关。

然后我们需要理解一下在 backward 中提到的 “梯度下降法” 的具体含义。梯度最重要的，是确定 “因变量” 和 “自变量”。自变量很好理解，我们需要计算哪个张量的梯度，哪个张量就作为自变量。而因变量就比较迷惑了，先说结论，它一直是 $l$ 。

那么为什么说因变量比较迷惑呢？这是因为 backward 中充分应用了链式法则，在链式法则中，会引入很多的 “临时梯度”，这些临时梯度的因变量不再是 $l$ 了。当这些 “临时梯度” 与 “最终梯度” 一起出现在公式中的时候，就容易让人感到迷惑了。

那么考虑一个在 LLM 中出现的张量 $X$ ，无论 $X$ 在模型的哪个位置（是最后一层，还是第一层），发挥什么作用（是 FFN 的权重，还是 Attention 的权重，还是激活值，还是偏移值），形状如何（是一维向量，还是二维矩阵，还是考虑 batch 的高维张量），它其中的每个分量都会影响 $l$ 。

所以因此我们得到的梯度 $\frac{\partial l}{\partial X}$ 的形状就一定和 $X$ 保持相同，因为它就表示了 $X$ 中的每个分量对于 $l$ 的影响。

为了强调这个规律，我们引入了新的标记：

$$X_{G} = \frac{\partial l}{\partial X}$$

其中 $X_G$ 与 $X$ 的形状完全相同。

当然梯度下降法也就很好表示了，因为形状相同（这么看形状也必须相同），说白了就是：

$$X_{new} = X_{old} - \eta X_{G}$$

2.3 矩阵乘法很简单

对于 Forward 过程：

$$Y = A \cdot B$$

无论 $A, B$ 的形状是什么，有 backward 过程：

$$A_{G} = Y_{G} \cdot B^T$$

$$B_{G} = Y_{G}^T \cdot A$$

也就是说，虽然理解并推导链式法则 backward 是一件很困难的事情，但是最终的结论是非常简单且 general 的。当我们拿到一个因变量的梯度的时候，自变量的梯度计算会变得很容易记忆。

如果有两个自变量，那么某个自变量的梯度就是因变量梯度与另一个自变量的乘积（当然可能还需要转置）。

那难道 LLM 中就都是 $Y = AB$ 这种简单形式吗？难道 Attention 和 FFN 这种复杂的网络结构，也能用矩阵乘法表示吗？还真能，这是因为：

• 我们并不限制 $A, B$ 的形状，也就是无论他们是一维的，还是二维的，上面的式子都成立。
• 虽然 Attention 这种结构乍一看很复杂，但是它都可以被拆成很多个矩阵乘法，比如说 $QK$ 和 $PV$ 等。
• 确实激活函数或者 norm 函数（包括 softmax）无法用这个模式硬套，但是用这种 vec2vec 的梯度下降，本身也不是计算的大头。

---

三、应用

在这一章里面，我们用上面介绍的规律来实战一下，其实主要就是矩阵乘法梯度的规律。

3.1 FFN

可以被理解成两个线性映射层：

$$Y = WX$$

我们已经拿到了 $Y_G$ ，带入上面的规律可知：

$$W_G = Y_G \cdot X^T$$

$$X_G = Y_G^T \cdot W$$

在计算 $W_G$ 的时候，需要使用到激活值 $X$，这就要求在 forward 的时候，要将这个临时的激活值进行保留。

我们计算 $W_G$ 是出于更新 $W$ 的目的，那么我们为什么要计算 $X_G$，这是因为 $X$ 此时是自变量，而它同时也是因变量，所以需要计算它保证链式传播。

3.2 Attention

Attention 看起来很复杂，但是实际上经过拆解，并不难。在 forward 过程中，有：

$$S = Q \cdot K^T$$

$$P = Softmax(S)$$

$$O = P \cdot V$$

经过这么一拆解，就算不算，也知道 backward 可以表示成一个很简单的形式了。

我们已经拿到了 $O_G$ ，然后有：

$$V_G = O_G \cdot P^T$$

同时有：

$$P_G = O_G^T \cdot V$$

这两个计算的开销是非常大的，因为我们要保留形状为 $[n, n]$ 的 $P, P_G$ （$n$ 是序列长度），这都是非常高昂的开销。

我们不考虑 $softmax$ 的梯度，但是总之 $S_G$ 的形状也是 $[n, n]$ 。

当我们有了 $S_G$ 后，我们就可以推算 $Q_G, K_G$ 了，有：

$$Q_G = S_G \cdot K$$

$$K_G = Q^T \cdot S_G$$

FlashAttention 之所以这么厉害，不止是因为它在 forward 的卓越贡献，能够避免 $[n, n]$ 矩阵，在 backward 中，它利用保存在 SRAM 里的 Block 形式的 $Q, K, V$ 重新算一遍 Forward，当场算出局部的 $P$，随即算出梯度，直接加和，同样不需要保存 $[n, n]$ 矩阵。

---