FKL 与 RKL 的梯度推导：为什么两者的优化行为不同？

关键是：FKL 中教师概率 q_j 是常量；RKL 中学生概率 s_j 本身也参与加权，因此求导时要同时对“权重”和“对数”求导。

下面只看一个预测位置，固定输入和前缀，教师不传播梯度，温度为 1，使用自然对数。

---

1. 符号与 Softmax 求导

设词表有 V 个 token：

• z_i：学生对第 i 个 token 的原始 logit。
• s_i = p_S(i)：学生概率。
• q_i = p_T(i)：教师概率，求导时视为常量。
• 学生和教师的概率分布均归一化。

$$\sum_j s_j=\sum_j q_j=1$$

学生概率由 Softmax 得到：

$$s_j=\frac{e^{z_j}}{\sum_{k=1}^{V}e^{z_k}}$$

令：

$$Z=\sum_k e^{z_k}$$

根据商的求导法则：

$$\begin{aligned} \frac{\partial s_j}{\partial z_i} &= \frac{\delta_{ji}e^{z_j}Z-e^{z_j}e^{z_i}}{Z^2}\\ &=s_j(\delta_{ji}-s_i) \end{aligned}$$

其中 δ_ji 是指示量：当 j=i 时为 1，否则为 0。

因此：

$$\frac{\partial s_j}{\partial z_i} = \begin{cases} s_i(1-s_i), & j=i\\ -s_js_i, & j\ne i \end{cases}$$

这说明：

修改一个 logit，不仅影响自己的概率，也影响其他所有 token 的概率。

后面还会用到：

$$\begin{aligned} \frac{\partial \log s_j}{\partial z_i} &= \frac{1}{s_j}\frac{\partial s_j}{\partial z_i}\\ &=\delta_{ji}-s_i \end{aligned}$$

---

2. FKL：为什么得到 s_i - q_i？

Forward KL 的定义：

$$L_{\mathrm{FKL}} = \mathrm{KL}(p_T\Vert p_S) = \sum_j q_j\log\frac{q_j}{s_j}$$

将其拆开：

$$L_{\mathrm{FKL}} = \underbrace{ \sum_j q_j\log q_j }_{\text{教师固定，与学生 logits 无关}} - \sum_j q_j\log s_j$$

第一项对 z_i 的导数为 0。

所以：

$$\begin{aligned} \frac{\partial L_{\mathrm{FKL}}}{\partial z_i} &= -\sum_j q_j \frac{\partial\log s_j}{\partial z_i}\\ &= -\sum_j q_j(\delta_{ji}-s_i)\\ &= -\sum_j q_j\delta_{ji} +s_i\sum_j q_j \end{aligned}$$

利用：

$$\sum_jq_j\delta_{ji}=q_i$$

以及：

$$\sum_jq_j=1$$

最终得到：

$$\boxed{ \frac{\partial L_{\mathrm{FKL}}}{\partial z_i} =s_i-q_i }$$

2.1 梯度的含义

• 当 s_i < q_i 时：梯度为负，梯度下降会提高 z_i。
• 当 s_i > q_i 时：梯度为正，梯度下降会降低 z_i。
• 当 s_i = q_i 时：该 logit 对应的梯度为 0。

因此，FKL 的 logit 梯度直接由学生和教师之间的概率差决定。

2.2 FKL 与标准 CE 的关系

如果教师是 One-Hot 标签，即：

• 正确 token 的 q_i = 1。
• 其他 token 的 q_i = 0。

那么：

$$L_{\mathrm{FKL}} = -\log s_{\mathrm{target}}$$

此时 FKL 就等价于标准监督交叉熵。

因此，标准 CE 可以看作教师分布为 One-Hot 时的 FKL 特例。

---

3. RKL：为什么需要减去整项 KL？

Reverse KL 的定义：

$$L_{\mathrm{RKL}} = \mathrm{KL}(p_S\Vert p_T) = \sum_j s_j\log\frac{s_j}{q_j}$$

与 FKL 不同，这里求和项外面的 s_j 也会随着学生 logits 改变。

因此，不能简单地把 s_j 当作常量。

3.1 第一步：先对学生概率 s_j 求导

考虑单个 token 对应的项：

$$f(s_j)=s_j\log\frac{s_j}{q_j}$$

使用乘积求导法则：

$$\begin{aligned} \frac{\partial}{\partial s_j} \left( s_j\log\frac{s_j}{q_j} \right) &= \log\frac{s_j}{q_j} + s_j\frac{1}{s_j}\\ &= \log\frac{s_j}{q_j}+1 \end{aligned}$$

因此：

$$\frac{\partial L_{\mathrm{RKL}}}{\partial s_j} = \log\frac{s_j}{q_j}+1$$

注意：

这里多出来的 +1，正是对外部权重 s_j 求导产生的。

3.2 第二步：通过 Softmax 使用链式法则

根据链式法则：

$$\frac{\partial L_{\mathrm{RKL}}}{\partial z_i} = \sum_j \frac{\partial L_{\mathrm{RKL}}}{\partial s_j} \frac{\partial s_j}{\partial z_i}$$

代入前面的结果：

$$\frac{\partial L_{\mathrm{RKL}}}{\partial z_i} = \sum_j \left( \log\frac{s_j}{q_j}+1 \right) s_j(\delta_{ji}-s_i)$$

展开成两项：

$$\begin{aligned} \frac{\partial L_{\mathrm{RKL}}}{\partial z_i} ={}& \sum_j s_j \left( \log\frac{s_j}{q_j}+1 \right)\delta_{ji}\\ &- s_i\sum_j s_j \left( \log\frac{s_j}{q_j}+1 \right) \end{aligned}$$

3.3 第三步：分别化简两项

第一项只有 j=i 时保留：

$$s_i\left( \log\frac{s_i}{q_i}+1 \right)$$

第二项中的求和为：

$$\begin{aligned} \sum_j s_j \left( \log\frac{s_j}{q_j}+1 \right) &= \sum_j s_j\log\frac{s_j}{q_j} +\sum_j s_j\\ &= L_{\mathrm{RKL}}+1 \end{aligned}$$

因为：

$$\sum_j s_j=1$$

将两部分代回：

$$\begin{aligned} \frac{\partial L_{\mathrm{RKL}}}{\partial z_i} &= s_i\left( \log\frac{s_i}{q_i}+1 \right) -s_i(L_{\mathrm{RKL}}+1)\\ &= s_i\left( \log\frac{s_i}{q_i} -L_{\mathrm{RKL}} \right) \end{aligned}$$

最终得到：

$$\boxed{ \frac{\partial L_{\mathrm{RKL}}}{\partial z_i} = s_i \left[ \log\frac{s_i}{q_i} - \mathrm{KL}(p_S\Vert p_T) \right] }$$

减去整项 KL 不是额外设计的约束，而是 Softmax 概率相互耦合后，自然产生的结果。

---

4. 怎么直观理解 RKL 梯度？

定义每个 token 的对数概率比：

$$a_i=\log\frac{s_i}{q_i}$$

整项 Reverse KL 就是它在学生分布下的加权平均：

$$L_{\mathrm{RKL}} = \sum_j s_ja_j = \mathbb{E}_{j\sim p_S}[a_j]$$

于是 RKL 梯度可以重新写为：

$$\boxed{ \frac{\partial L_{\mathrm{RKL}}}{\partial z_i} = s_i \left( a_i-\mathbb{E}_{p_S}[a] \right) }$$

可以读成：

学生给这个 token 的概率 × 这个 token 的对数概率比相对平均值的偏差。

4.1 RKL 不是简单比较 s_i 和 q_i

FKL 直接看：

$$s_i-q_i$$

而 RKL 看：

$$\log\frac{s_i}{q_i} - \mathbb{E}_{p_S}[a]$$

所以在多 token 情况下：

RKL 并不是只要 s_i > q_i，就一定降低该 token 的 logit。

而是要看：

$$\log\frac{s_i}{q_i} \quad \text{是否高于学生分布下的平均值}$$

具体来说：

• 如果 a_i 高于平均值，梯度为正，降低该 logit。
• 如果 a_i 低于平均值，梯度为负，提高该 logit。
• 如果 a_i 等于平均值，该 logit 的梯度为 0。

因此，RKL 的更新不仅与当前 token 有关，还与整个学生概率分布有关。

4.2 为什么 RKL 梯度之和为 0？

根据前面的公式：

$$\begin{aligned} \sum_i \frac{\partial L_{\mathrm{RKL}}}{\partial z_i} &= \sum_i s_i \left( a_i-\mathbb{E}_{p_S}[a] \right)\\ &= \sum_i s_ia_i - \mathbb{E}_{p_S}[a]\sum_i s_i\\ &= L_{\mathrm{RKL}}-L_{\mathrm{RKL}}\\ &=0 \end{aligned}$$

因此：

$$\boxed{ \sum_i \frac{\partial L_{\mathrm{RKL}}}{\partial z_i} =0 }$$

这符合 Softmax 的性质：

所有 logits 同时加上相同的常数，Softmax 概率不会变化，因此 KL 损失也不会变化。

---

5. 二候选情况下，RKL 梯度可以进一步化简

假设只有两个候选 token。

学生分布：

$$p_S=(s,1-s)$$

教师分布：

$$p_T=(q,1-q)$$

其中：

$$0<s<1,\qquad 0<q<1$$

5.1 FKL 梯度

对第一个 token：

$$\boxed{ \frac{\partial L_{\mathrm{FKL}}}{\partial z_1} =s-q }$$

这就是普通的概率差。

5.2 RKL 梯度

从一般公式开始：

$$\frac{\partial L_{\mathrm{RKL}}}{\partial z_1} = s \left[ \log\frac{s}{q} - L_{\mathrm{RKL}} \right]$$

其中：

$$L_{\mathrm{RKL}} = s\log\frac{s}{q} + (1-s)\log\frac{1-s}{1-q}$$

代入：

$$\begin{aligned} \frac{\partial L_{\mathrm{RKL}}}{\partial z_1} &= s\left[ \log\frac{s}{q} -s\log\frac{s}{q} -(1-s)\log\frac{1-s}{1-q} \right]\\ &= s(1-s) \left[ \log\frac{s}{q} - \log\frac{1-s}{1-q} \right]\\ &= s(1-s) \log\frac{s(1-q)}{q(1-s)} \end{aligned}$$

因此：

$$\boxed{ \frac{\partial L_{\mathrm{RKL}}}{\partial z_1} = s(1-s)\log\frac{s(1-q)}{q(1-s)} }$$

5.3 两个公式放在一起

Forward KL：

$$\boxed{ g_{\mathrm{FKL}}=s-q }$$

Reverse KL：

$$\boxed{ g_{\mathrm{RKL}} = s(1-s) \log\frac{s(1-q)}{q(1-s)} }$$

这就是两种梯度强度比较的来源。

• FKL：直接看学生与教师的概率差。
• RKL：看学生与教师的对数赔率差，同时乘上 s(1-s)。

需要注意，s(1-s) 是由 Softmax 导数引入的。

当教师概率 q 固定且严格位于 0 和 1 之间时：

$$\lim_{s\to 0}g_{\mathrm{RKL}}=0$$ $$\lim_{s\to 1}g_{\mathrm{RKL}}=0$$

这是因为：

$$\lim_{s\to 0}s\log s=0$$

并且：

$$\lim_{s\to 1}(1-s)\log(1-s)=0$$

因此：

在二候选情况下，当学生分布接近极端确定性时，RKL 的纠正梯度可能趋于很弱。

但是 FKL 不存在同样的饱和形式：

$$\lim_{s\to 0}g_{\mathrm{FKL}}=-q$$ $$\lim_{s\to 1}g_{\mathrm{FKL}}=1-q$$

这说明在教师仍赋予目标非零概率时，FKL 可以继续提供明确的纠正梯度。

---

6. FKL 与 RKL 的最终对比

对比维度	Forward KL（FKL）	Reverse KL（RKL）
KL 方向	KL(Teacher ∥ Student)	KL(Student ∥ Teacher)
加权概率	教师概率 q_j	学生概率 s_j
权重是否随学生变化	否	是
对 Logit 的梯度	s_i - q_i	s_i × (log(s_i/q_i) - RKL)
梯度依赖	当前 token 的概率差	当前 token 与整体分布的相对关系
低学生概率区域	教师概率非零时仍有直接纠正梯度	可能出现较弱的 logit 梯度
常见分布倾向	更倾向覆盖教师支持的模式	更倾向集中于部分高概率模式
主要风险	可能分散概率质量	可能忽视教师的低概率模式

需要注意：

Mode-Covering 和 Mode-Seeking 是 KL 方向的典型优化倾向，不代表任意模型、参数化方式和训练数据下都必然出现相同结果。

---

7. 温度参数对梯度的影响

以上推导都假设 Softmax 温度为 1。

如果使用温度 τ：

$$s_i = \frac{\exp(z_i/\tau)} {\sum_k\exp(z_k/\tau)}$$

此时：

$$\frac{\partial s_j}{\partial z_i} = \frac{1}{\tau} s_j(\delta_{ji}-s_i)$$

因此，如果 KL 损失没有额外乘以温度平方：

FKL

$$\boxed{ \frac{\partial L_{\mathrm{FKL}}}{\partial z_i} = \frac{1}{\tau}(s_i-q_i) }$$

RKL

$$\boxed{ \frac{\partial L_{\mathrm{RKL}}}{\partial z_i} = \frac{s_i}{\tau} \left[ \log\frac{s_i}{q_i} - L_{\mathrm{RKL}} \right] }$$

注意，这里的 s_i 和 q_i 均应理解为对应温度下的概率分布。

如果实现中额外乘以：

$$\tau^2$$

则上面的梯度还需要整体乘以该因子，最终显式系数由 1/τ 变成 τ。

但由于概率分布本身也随温度变化，不能仅根据显式系数判断实际梯度强度。

此外，上面推导的都是对学生 logits 的梯度。

实际参数更新还要经过模型参数的链式求导：

$$\frac{\partial L}{\partial\theta} = \sum_i \frac{\partial L}{\partial z_i} \frac{\partial z_i}{\partial\theta}$$

因此，logit 梯度大小不等于最终参数更新幅度。

最终更新还受到学习率、优化器、梯度裁剪、参数共享等因素影响。

---

8. 最终总结

FKL 的梯度：

$$\boxed{ \frac{\partial L_{\mathrm{FKL}}}{\partial z_i} =s_i-q_i }$$

RKL 的梯度：

$$\boxed{ \frac{\partial L_{\mathrm{RKL}}}{\partial z_i} = s_i \left[ \log\frac{s_i}{q_i} - \mathrm{KL}(p_S\Vert p_T) \right] }$$

两者最核心的区别：

FKL：

教师概率是固定权重，因此 logit 梯度直接表现为学生概率与教师概率的差值。

RKL：

学生概率既是优化对象，也是 KL 求和中的权重，因此必须同时考虑权重变化与对数概率比变化。

最终得到的梯度包含：

1. 当前 token 的学生概率。
2. 当前 token 的对数概率比。
3. 整个学生分布下的平均对数概率比，即 RKL 本身。

一句话总结：FKL 根据绝对概率差纠正学生；RKL 根据学生加权的相对对数概率比调整 logits。

对于 OPD、ASR 蒸馏等任务，这意味着两种 KL 即使在相同教师、相同学生前缀上训练，也可能产生不同的梯度分布与优化行为。