OPD_SNR

OPD中的梯度信号

OPD中，设学生分布为\(p\)，教师分布为\(q\)，其Loss，即RKL，表示为 \[ \mathbb{E}_{\hat{y}_{t}\sim p}[\log p(\hat{y}_{t}) - \log q(\hat{y}_{t})] \] 对应的总体梯度信号为 \[ \overline{g}_{OPD}(\theta) = \nabla_{\theta}\mathbb{E}_{\hat{y}_{t}\sim p}[\log p(\hat{y}_{t}) - \log q(\hat{y}_{t})] = \mathbb{E}_{\hat{y}_{t}\sim p}[u_{t}\nabla_{\theta}\log p(\hat{y}_{t})] \] 即Policy Gradient的梯度形式，每一个梯度项的权重为 \[ u_{t} = \log p(\hat{y}_{t}) - \log q(\hat{y}_{t}) \] 当然也可以把OPD的整体梯度信号表示成每一个token处对应的梯度信号的期望，即 \[ \begin{flalign} \overline{g}_{OPD}(\theta) &= \mathbb{E}_{\hat{y}_{t}\sim p}[g_{OPD}(\theta;\hat{y}_{t})]\\ \text{where} \ g_{OPD}(\theta;\hat{y}_{t}) &= u_{t}\nabla_{\theta}\log p(\hat{y}_{t}) \end{flalign} \]

关于权重项\(u_{t}\)

不难发现，\(u_{t}\)的期望为 \[ \overline{u} = \mathbb{E}_{\hat{y}_{t}\sim p}[\log p(\hat{y}_{t}) - \log q(\hat{y}_{t})] = \mathcal{D}_{KL}(p||q) \] 其方差为 \[ \mathrm{Var}_{p}(u_{t}) = \mathbb{E}_{p}[u_{t}^2] - \mathbb{E}_{p}[u_{t}]^2 = \mathbb{E}_{p}[u_{t}^2] - \overline{u}^2 \] 当两个分布\(p,q\)足够接近的时候，定义 \[ \epsilon(y) = \frac{q(y)-p(y)}{p(y)}, \text{hence} \ q(y) = p(y)(1+\epsilon(y)) \] 注意到\(p,q\)都是分布，因此有 \[ \mathbb{E}_{y\sim p}[\epsilon(y)] = \sum_{y}p(y)\epsilon(y) = \sum_{y}q(y)-p(y) = 1-1=0 \] 从而 \[ \begin{flalign} u_{t} &= \log p(\hat{y}_{t}) - \log q(\hat{y}_{t}) = -\log(1+\epsilon(\hat{y}_{t}))\\ &= - \epsilon(\hat{y}_{t}) + \frac{\epsilon(\hat{y}_{t})^2}{2} + O(\epsilon(\hat{y}_{t})^3)\\ \end{flalign} \] 因此\(u_{t}\)是误差项\(\epsilon\)的一阶小量。同时也有 \[ \overline{u} = \mathbb{E}_{p}(u_{t}) = \underbrace{\mathbb{E}_{p}[- \epsilon(\hat{y}_{t})]}_{=0} + \mathbb{E}_{p}[\frac{\epsilon(\hat{y}_{t})^2}{2}] = \mathbb{E}_{p}[\frac{\epsilon(\hat{y}_{t})^2}{2}] \] 在对\(u_{t}\)取期望之后，\(\epsilon\)的期望也变成0了，因此\(u_{t}\)的期望，也就是\(\overline{u}\)反而是\(\epsilon\)的一个二阶小量 将\(\overline{u}\)与误差项\(\epsilon(\hat{y}_{t})\)的关系代入\(\mathrm{Var}_{p}(u_{t})\)的定义式，有 \[ \begin{flalign} \mathrm{Var}_{p}(u_{t}) &= \mathbb{E}_{p}[u_{t}^2] - \overline{u}^2\\ & \approx \mathbb{E}_{p}[\epsilon^2] - (\frac{1}{2}\mathbb{E}_{p}[\epsilon^2])^2\\ & \approx \mathbb{E}_{p}[\epsilon^2]\\ &\approx 2\overline{u} \end{flalign} \] 最终我们发现，

对于这个梯度的权重项\(u_{t}\)，当教师和学生分布非常接近时，其方差与期望是同阶的 记\(\delta=\overline{u}=\mathcal{D}_{KL}(p||q)\)，有\(\mathrm{Var}_{p}(u_{t})=Θ(\delta)\) 此时不管是正向KL、反向KL、对称KL、还是对数比的方差,都会以同一个速度趋于零,只是常数系数不同

OPD梯度的信噪比

在此基础上，我们继续分析一下OPD的整体梯度\(\overline{g}_{OPD}(\theta)= \mathbb{E}_{\hat{y}_{t}\sim p}[u_{t}\nabla_{\theta}\log p(\hat{y}_{t})]\) 对于score function \(\nabla_{\theta}\log p(\hat{y}_{t})\)，其在自己的分布\(p\)下期望为0（显然），从而对\(\overline{g}_{OPD}(\theta)\)的权重项引入一个常数baseline并不会改变其期望（这也是REINFORCE with basline的动机来源），即 \[ \begin{flalign} \overline{g}_{OPD}(\theta) &= \mathbb{E}_{\hat{y}_{t}\sim p}[u_{t}\nabla_{\theta}\log p(\hat{y}_{t})]\\ &= \mathbb{E}_{\hat{y}_{t}\sim p}[u_{t}\nabla_{\theta}\log p(\hat{y}_{t})] - \overline{u} \underbrace{\mathbb{E}_{\hat{y}_{t}\sim p}[\nabla_{\theta}\log p(\hat{y}_{t})]}_{=0}\\ &= \mathbb{E}_{\hat{y}_{t}\sim p}[(u_{t}-\overline{u})\nabla_{\theta}\log p(\hat{y}_{t})] \end{flalign} \] 再次利用\(\mathbb{E}_{\hat{y}_{t}\sim p}[\nabla_{\theta}\log p(\hat{y}_{t})]=0\)，上式可以写成 \[ \begin{flalign} \overline{g}_{OPD}(\theta) &= \mathbb{E}_{\hat{y}_{t}\sim p}[(u_{t}-\mathbb{E}_{p}[u_{t}])(\nabla_{\theta}\log p(\hat{y}_{t})-\mathbb{E}_{p}[\nabla_{\theta}\log p(\hat{y}_{t})])]\\ &=\mathrm{Cov}_{p}(u_{t}, \nabla_{\theta}\log p(\hat{y}_{t})) \end{flalign} \] 利用Cauchy-Schwarz Inequality， 有 \[ ||\overline{g}_{OPD}(\theta)||_{2}=||\mathrm{Cov}_{p}(u_{t}, \nabla_{\theta}\log p(\hat{y}_{t}))||_{2}\leq \sqrt{ \mathrm{Var}_{p}(u_{t}) } \cdot \sqrt{ \mathrm{Var}_{p}(\nabla_{\theta}\log p(\hat{y}_{t})) } \] 注意到 \[ \mathrm{Var}_{p}(\nabla_{\theta}\log p(\hat{y}_{t})) = \mathbb{E}_{p}[(\nabla_{\theta}\log p(\hat{y}_{t}))^2] = Tr(\mathcal{F}(\theta)) \] 其中\(Tr(\mathcal{F}(\theta)) = \sum_{i}\lambda_{i}(\mathcal{F}(\theta))\geq dim(\theta) \cdot \lambda_{min}(\mathcal{F}(\theta))\)，可以看成关于\(\delta\)是\(O(1)\)的。而前面已经分析过\(\mathrm{Var}_{p}(u_{t})\)是\(O(\delta)\)的，因此\(||\overline{g}_{OPD}(\theta)||_{2}^2\)是\(O(\delta)\)的

OPD优化过程中，梯度的信噪比SNR定义为 \[ SNR(g) = \frac{||\overline{g}||_{2}^2}{Tr(Cov[g])} \] 分子部分是期望梯度的范数大小，也就是优化信号的强弱，而分母是估计梯度信号中的方差大小，这个式子衡量了估计过程中的信号强弱与噪声的比值。SNR趋于0时，代表训练中基本全是噪声信号，难以进行有效训练；而如果SNR趋于无穷，则代表梯度信号占主导。 \[ \begin{flalign} Tr(Cov[g]) &= Tr(\mathbb{E}_{p}[(g-\overline{g})(g-\overline{g})^T]) = \mathbb{E}_{p}||g-\overline{g}||_{2}^2 = \mathrm{Var}_{p}(g) \\ \end{flalign} \] 注意到 \[ \begin{flalign} \mathrm{Var}_{p}(g_{OPD}(\theta)) &= \mathbb{E}_{\hat{y}_{t}\sim p}[u_{t}^2||\nabla_{\theta}\log p(\hat{y}_{t})||^2_{2}] - \mathbb{E}_{\hat{y}_{t}\sim p}[u_{t}\nabla_{\theta}\log p(\hat{y}_{t})]^2\\ &=\mathbb{E}_{\hat{y}_{t}\sim p}[u_{t}^2||\nabla_{\theta}\log p(\hat{y}_{t})||^2_{2}] - ||\overline{g}_{OPD}(\theta)||_{2}^2 \end{flalign} \] 对于前半部分\(\mathbb{E}_{\hat{y}_{t}\sim p}[u_{t}^2||\nabla_{\theta}\log p(\hat{y}_{t})||^2_{2}]\)，通过引入\(\overline{u}=\mathbb{E}_{\hat{y}_{t}\sim p}[u_{t}]\)，将\(u_{t}\)分解为\(u_{t} = \overline{u} + (u_{t}- \overline{u})\)，有 \[ \begin{flalign} \mathbb{E}_{\hat{y}_{t}\sim p}[u_{t}^2||\nabla_{\theta}\log p(\hat{y}_{t})||^2_{2}] &= \overline{u}^2 \mathbb{E}_{\hat{y}_{t}\sim p}[||\nabla_{\theta}\log p(\hat{y}_{t})||^2_{2}] + 2\overline{u} \mathbb{E}_{\hat{y}_{t}\sim p}[(u_{t}-\overline{u})||\nabla_{\theta}\log p(\hat{y}_{t})||^2_{2}]\\ & + \mathbb{E}_{\hat{y}_{t}\sim p}[(u_{t}-\overline{u})^2||\nabla_{\theta}\log p(\hat{y}_{t})||^2_{2}]\\ \end{flalign} \] 为了说明信噪比是收敛的，我们需要给方差项找一个与\(\delta\)相关的下界。当分布\(p,q\)接近的时候，\(\mathrm{Var}_{p}(g_{OPD}(\theta))\)中\(||\overline{g}_{OPD}(\theta)||_{2}^2\)是\(O(\delta^2)\)的，\(\mathbb{E}_{\hat{y}_{t}\sim p}[u_{t}^2||\nabla_{\theta}\log p(\hat{y}_{t})||^2_{2}]\)的展开中，第一项和第二项分别是\(O(\delta^2)\)以及\(O(\delta)\)的。接下来尝试分析交叉项。 \[ \begin{flalign} \mathbb{E}_{\hat{y}_{t}\sim p}[(u_{t}-\overline{u})^2||\nabla_{\theta}\log p(\hat{y}_{t})||^2_{2}]&=\mathbb{E}_{\hat{y}_{t}\sim p}[(u_{t}-\overline{u})^2]\mathbb{E}_{\hat{y}_{t}\sim p}[ff^T], f=\nabla_{\theta}\log p(\hat{y}_{t})\\ &= \mathrm{Var}_{p}(u_{t})\cdot Tr(\mathcal{F}(\theta))\\ & \geq \mathrm{Var}_{p}(u_{t}) \cdot \mathcal{F}_{min}(\theta) \end{flalign} \] 这里\(\mathrm{Var}_{p}(u_{t})\)前面已经证明过是\(O(\delta)\)的，而\(\mathcal{F}_{min}\)可以看成\(O(1)\)的。最终，可以发现OPD梯度信噪比SNR的分母项\(Tr(Cov[g])\)与分子项都是\(O(\delta)\)的。因此，当训练后期两个分布接近时，SNR趋向于收敛到一个常数值，即噪声始终在训练过程中存在并影响模型训练。 最终对OPD的训练过程进行总结：

训练早期梯度范数\(||\overline{g}||_{2}^2\)较大时，梯度信号强烈，模型可以快速收敛；但是训练进行到后期，师生分布接近时，信噪比始终存在，干扰模型训练，导致性能始终上不去甚至下降。

这一现象，究其根本，是因为SNR的分母项，也就是梯度采样估计的方差，由于on policy采样的原因，始终存在。