知识蒸馏 - 视觉模型知识蒸馏Vision-KD-Trainer KL散度公式变化

知识蒸馏 - 视觉模型知识蒸馏Vision-KD-Trainer KL散度公式变化

flyfish

代码抄自
https://github.com/shaoshengsong/KDTrainer
代码分析的是Vision-KD-Trainer/helloworld.py

Python 版本: 3.12.9
PyTorch: 2.6.0+cu124
Transformers: 4.55.0
PEFT: 0.15.2
PyYAML: 6.0.2

环境搭建好之后 执行 python helloworld.py 纯绿色版

最好先看完基础知识

知识蒸馏 - 蒸的什么

知识蒸馏 - 通过引入温度参数T调整 Softmax 的输出

知识蒸馏 - 对数函数的单调性

知识蒸馏 - 信息量的公式为什么是对数

知识蒸馏 - 根据真实事件的真实概率分布对其进行编码

知识蒸馏 - 信息熵中的平均为什么是按概率加权的平均

知识蒸馏 - 自信息量是单个事件的信息量，而平均自信息量（即信息熵）是所有事件自信息量以其概率为权重的加权平均值

知识蒸馏 - 最小化KL散度与最小化交叉熵是完全等价的

对于两个概率分布 $P$（真实分布）和 $Q$（模型预测分布），KL散度的定义是：
$$D_{KL}(P\Vert Q)=\sum _{x}P(x)\log \left(\frac{P(x)}{Q(x)}\right)$$

怎么就变成如下公式呢

$$D_{KL}(P\Vert Q)=\sum _{x}P(x)\log P(x)-\sum _{x}P(x)\log Q(x)$$

商的对数等于对数的差

“商的对数等于对数的差”就是：两个数相除的对数，等于这两个数分别取对数后再相减。

${\log }_c\left(\frac{a}{b}\right)={\log }_{c}a-{\log }_{c}b$

先搞懂“对数”到底是啥

对数其实是“指数的反过来”。比如“以c为底a的对数”（写成${\log }_{c}a$），本质是在问：“c的多少次方等于a？”
假设这个“多少次方”是x，那就能写成：
$x={\log }_{c}a$
反过来，也可以说：$c^x=a$（c的x次方等于a）。

给两个数的对数起个名字

现在有两个正数a和b（因为对数里的数必须是正数）。
按上面的逻辑：

• 设“c的多少次方等于a”是x，也就是$x={\log }_{c}a$，所以$a=c^x$；
• 设“c的多少次方等于b”是y，也就是$y={\log }_{c}b$，所以$b=c^y$。

看看a除以b等于啥

我们想算$a÷b$（也就是$\frac{a}{b}$），代入上面的$a=c^x$和$b=c^y$，得到：
$\frac{a}{b}=\frac{c^x}{c^y}$

这里用到一个指数的常识：同底数的幂相除，底数不变，指数相减。比如$2^5÷2^3=2^{5-3}=2^2$。
所以$\frac{c^x}{c^y}=c^{x-y}$，也就是：
$\frac{a}{b}=c^{x-y}$

把上面的结果转成对数

现在看$\frac{a}{b}=c^{x-y}$，这句话用对数的话来说就是：“c的（x - y）次方等于$\frac{a}{b}$”。
按对数的定义，“c的多少次方等于$\frac{a}{b}$”，这个“多少次方”就是${\log }_c\left(\frac{a}{b}\right)$。
所以：
${\log }_c\left(\frac{a}{b}\right)=x-y$

换回原来的对数

前面已经知道$x={\log }_{c}a$，$y={\log }_{c}b$，代入上面的式子，得到：
${\log }_c\left(\frac{a}{b}\right)={\log }_{c}a-{\log }_{c}b$

这样两个数相除的对数，等于这两个数的对数相减。

求和可以和加法、乘法（常数倍）交换顺序

求和的线性性质是求和可以和加法、乘法（常数倍）“交换顺序”。

规则

设 $f(x)$ 和 $g(x)$ 是任意两个关于 $x$ 的函数，$c$ 是任意常数（与 $x$ 无关），则：

1. 对加法的分配性：
   两个函数相加后的求和，等于两个函数分别求和后再相加。
   公式：
   $$\sum _x\left[f(x)+g(x)\right]=\sum _{x}f(x)+\sum _{x}g(x)$$

2. 对常数倍的分配性：
   常数乘以函数后的求和，等于常数乘以该函数的求和。
   公式：
   $$\sum _x\left[c\cdot f(x)\right]=c\cdot \sum _{x}f(x)$$

“乘法分配律”（如 $c\cdot (a+b)=c\cdot a+c\cdot b$），求和的线性性质本质是“求和符号对加法和常数倍的分配”。

比如：

• 若 $x$ 只能取1和2，$f(1)=2$，$f(2)=3$；$g(1)=1$，$g(2)=4$，则：
  左边 $\sum [f(x)+g(x)]=(2+1)+(3+4)=3+7=10$，
  右边 $\sum f(x)+\sum g(x)=(2+3)+(1+4)=5+5=10$，两边相等，验证了第一条规则。

• 若 $c=2$，则：
  左边 $\sum [2\cdot f(x)]=2\cdot 2+2\cdot 3=4+6=10$，
  右边 $2\cdot \sum f(x)=2\cdot (2+3)=2\cdot 5=10$，两边相等，验证了第二条规则。

具体到KL散度

对数性质，将KL散度写成：
$$D_{KL}(P\Vert Q)=\sum _{x}P(x)\cdot \left[\log P(x)-\log Q(x)\right]$$

现在要拆分成两个求和项，步骤如下：

步骤1：先展开括号里的乘法

根据乘法分配律（即$a\cdot (b-c)=a\cdot b-a\cdot c$），括号里的减法可以和外面的$P(x)$分别相乘：
$$P(x)\cdot \left[\log P(x)-\log Q(x)\right]=P(x)\cdot \log P(x)-P(x)\cdot \log Q(x)$$

步骤2：应用求和的线性性质

将上式代入KL散度的表达式，得到：
$$D_{KL}(P\Vert Q)=\sum _x\left[P(x)\log P(x)-P(x)\log Q(x)\right]$$

此时，求和符号后面是两个项的差（即$f(x)-g(x)$，
其中
$f(x)=P(x)\log P(x)$，
$g(x)=P(x)\log Q(x)$）。

根据求和的线性性质$\sum (f-g)=\sum f-\sum g$，可以直接拆分：

$$\sum _x\left[P(x)\log P(x)-P(x)\log Q(x)\right]=\sum _{x}P(x)\log P(x)-\sum _{x}P(x)\log Q(x)$$