Math-06.矩阵微积分-20.注意力与LoRA梯度

本页从矩阵微积分角度补充 Math-03/20 Attention 的梯度流。

段末注释：LoRA 训练时仅 $A,B$ 的 requires_grad=True；$W_0$ 冻结则 $\bar{W}_0=\mathbf{0}$，梯度只更新低秩分支。

系列入口：00.系列规划 | 前置：06 层梯度手册

---

1. Attention backward 要点（D3）

缩放点积 Attention（单头，略 batch）：

$$
S = \frac{QK^\top}{\sqrt{d_k}}, \quad P = \mathrm{softmax}(S), \quad O = PV
$$

已知 $\bar{O}$，核心步骤：

$$
\bar{V} = P^\top \bar{O}, \quad \bar{P} = \bar{O} V^\top
$$

$\bar{S}$ 来自 softmax backward（同 CE 结构，行和为 0 约束）：

$$
\bar{S}{ij} = P{ij}\left(\bar{P}{ij} - \sum_k P{ik}\bar{P}_{ik}\right) / \sqrt{d_k}
$$

$$
\bar{Q} = \bar{S} K / \sqrt{d_k}, \quad \bar{K} = \bar{S}^\top Q / \sqrt{d_k}
$$

再经 $Q=XW_Q$ 等 Linear backward（06 手册）。

---

2. LoRA 梯度（D3–D7）

$$
\mathbf{y} = W_0 \mathbf{x} + B A \mathbf{x}, \quad A \in \mathbb{R}^{r \times d_{\mathrm{in}}},\ B \in \mathbb{R}^{d_{\mathrm{out}} \times r}
$$

设 $\mathbf{z} = A\mathbf{x}$，$\Delta \mathbf{y} = B\mathbf{z}$。已知 $\bar{\mathbf{y}}$：

$$
\bar{B} = \bar{\mathbf{y}} \mathbf{z}^\top, \quad \bar{\mathbf{z}} = B^\top \bar{\mathbf{y}}, \quad \bar{A} = \bar{\mathbf{z}} \mathbf{x}^\top, \quad \bar{\mathbf{x}}_{\mathrm{lora}} = A^\top \bar{\mathbf{z}}
$$

$W_0$ 冻结：不计算 $\bar{W}_0$ 或不更新。$\bar{\mathbf{x}}$ 总梯度 = Linear 主路径 + LoRA 路径之和。

参数	典型 shape 梯度
$A$	$r \times d_{\mathrm{in}}$
$B$	$d_{\mathrm{out}} \times r$

Math-04/20 LoRA lr。

---

3. 实现注意（D8）

for p in model.base.parameters():
    p.requires_grad = False
# 仅 lora_A, lora_B 可训

问题	说明
合并权重推理	$W = W_0 + BA$，训练时分开
rank 小	$\bar{A},\bar{B}$ 低秩外积和
FlashAttention	backward 融合 CUDA，手推一致

---

4. 局限（D8）

陷阱	说明
$S$ 的 $L^2$ 内存	长序列 OOM
softmax 行	必须整行 backward
忘记 scale $\sqrt{d_k}$	梯度差常数因子

---

5. 小结

Math-06 矩阵微积分 全系列：从 $\partial L/\partial \mathbf{x}$ 到 Attention/LoRA 的 $\bar{Q},\bar{K},\bar{A},\bar{B}$。

系列导航：00 规划 | Math-03/20 Attention