TML OPD 的 γ 在哪里

TML OPD 博客的 $\gamma$ 在哪里？他们的原文写的是：

「折扣因子」是什么？它不在公式里。作为初学者我很困惑 🤔。

$\gamma$ 在哪里？

Token-level OPD

如上图所示，TML 的逐词元 loss 已经定义为了：

$$\ell_t = D_{\mathrm{KL}} \left( \pi_\theta(\cdot\mid c_t) \;\|\; q(\cdot\mid c_t) \right) = \mathbb E_{y_t\sim\pi_\theta(\cdot\mid c_t)} \left[ \log\pi_\theta(y_t\mid c_t) - \log q(y_t\mid c_t) \right].$$

其中 $ c_t=(x,y_{<t})$。上式是 token-level 的 loss。我们先顺着这个熟悉的式子推一遍它的梯度。

$$\ell_t = \sum_{y_t} \pi_\theta(y_t\mid c_t) \Big[ \log \pi_\theta(y_t\mid c_t)-\log q(y_t\mid c_t) \Big].$$ $$\begin{aligned} \nabla_\theta \ell_t &= \sum_{y_t} \nabla_\theta \pi_\theta(y_t\mid c_t) \Big[ \log\pi_\theta(y_t\mid c_t)-\log q(y_t\mid c_t) \Big] \\ &\quad+ \sum_{y_t} \pi_\theta(y_t\mid c_t) \nabla_\theta\log\pi_\theta(y_t\mid c_t). \end{aligned}$$

常见 log-derivative trick：

$$\nabla_\theta \pi_\theta(y_t\mid c_t) = \pi_\theta(y_t\mid c_t) \nabla_\theta\log\pi_\theta(y_t\mid c_t),$$

所以有

$$\nabla_\theta \ell_t = \mathbb E_{y_t\sim\pi_\theta(\cdot\mid c_t)} \left[ \nabla_\theta \log\pi_\theta(y_t\mid c_t) \left( 1 + \log\pi_\theta(y_t\mid c_t) - \log q(y_t\mid c_t) \right) \right].$$

请注意 score-function identity：

$$\mathbb E_{y_t\sim\pi_\theta(\cdot\mid c_t)} \left[ \nabla_\theta\log\pi_\theta(y_t\mid c_t) \right] =0.$$

综上，有：

$$\nabla_\theta \ell_t = \mathbb E_{y_t\sim\pi_\theta(\cdot\mid c_t)} \left[ \nabla_\theta \log\pi_\theta(y_t\mid c_t) \left( \log\pi_\theta(y_t\mid c_t) - \log q(y_t\mid c_t) \right) \right].$$

定义：

$$s_t = \nabla_\theta\log\pi_\theta(y_t\mid c_t), \qquad r_t = \log\frac{\pi_\theta(y_t\mid c_t)} {q(y_t\mid c_t)}$$

其中 $s_t$ 是 score function，$r_t$ 是 log-ratio。整理梯度得：

$$\nabla_\theta \ell_t = \mathbb E_{y_t\sim\pi_\theta(\cdot\mid c_t)} \left[ s_t r_t \right].$$

所以，定义梯度的 Monte Carlo estimator：

$$\hat{g}_t = s_t r_t$$

它是无偏的（当把 $c_t$ 视为已经采样好的固定 context 时），如果 $y_t\sim\pi_\theta(\cdot|c_t)$，那么

$$\mathbb E[\hat{g}_t] = \nabla_\theta \ell_t.$$

对于采样一个 token $y_t$ 得到的单样本无偏 Monte Carlo 梯度估计量 $\hat{g}_t$，设计 surrogate loss，使得它 .backward() 之后产生 $s_t r_t$：

$$\tilde\ell_t = \operatorname{sg}(r_t) \log\pi_\theta(y_t|c_t)$$

等到真正训练的时候，因为 rollout 是由采样时的 policy $\pi_{\mathrm{samp}}$ 产生的，要再加一个 importance sampling ratio，

$$L_t^{\mathrm{IS}} = -\rho_t(\theta)A_t,$$

其中

$$\rho_t(\theta) = \frac{\pi_\theta(y_t|c_t)} {\pi_{\mathrm{samp}}(y_t|c_t)}, \quad A_t = -r_t^\text{samp}$$ $$r_t^{\mathrm{samp}} = \log\frac{ \pi_{\mathrm{samp}}(y_t|c_t) }{ q(y_t|c_t) }.$$

因此

$$\nabla_\theta L_t^{\mathrm{IS}} = r_t^{\mathrm{samp}} \rho_t(\theta) \nabla_\theta\log\pi_\theta(y_t|c_t).$$

如果正处于完全 on-policy 点，即当前 learner 就是刚刚用于 rollout 的 policy，那么

$$\pi_\theta=\pi_{\mathrm{samp}},$$

那么

$$\rho_t(\theta)=1,$$

那么

$$\nabla_\theta L_t^{\mathrm{IS}} = r_t s_t = \hat g_t.$$

以上就是 TML 的 token-level OPD。

上面一直看的是单个位置 $t$。对于一条长度为 $T$ 的 student rollout，

$$y=(y_1,\ldots,y_T),$$

每个位置都有一个局部梯度估计

$$\hat g_t=s_tr_t.$$

TML 的 token-level OPD 对整条 rollout 的梯度估计，就是把各位置的局部贡献直接相加。estimator 定义为：

$$\boxed{ \hat g_{\mathrm{tok}} = \sum_{t=1}^T\hat g_t = \sum_{t=1}^T s_tr_t. }$$

每个 token 只使用自己当前位置的 log-ratio $r_t$ 作为训练信号，并不会把未来位置的 $r_{t+1},r_{t+2},\ldots$ 回传给当前 token。

为什么 $\hat g_{\mathrm{tok}} = \sum_{t=1}^T\hat g_t$？这合适吗？如果目标是精确估计完整 trajectory reverse-KL 的梯度，它是有偏的。

Sequence-level OPD

参考了知乎文章《重探 On-Policy Distillation（OPD）：三类典型失败以及修复路径》

对于 prompt $x$，完整 response 为 $y=(y_1,\ldots,y_T)$。sequence-level OPD 的目标是

$$J_{\mathrm{OPD}}(\theta) = \mathbb E_{x\sim D} \left[ D_{\mathrm{KL}} \left( \pi_\theta(\cdot\mid x) \;\|\; q(\cdot\mid x) \right) \right],$$

用刚刚一样的推导梯度，

$$J_{\mathrm{OPD}}(\theta) = \mathbb E_{x\sim D} \left[ \sum_y \pi_\theta(y\mid x) \Big[ \log \pi_\theta(y\mid x) - \log q(y\mid x) \Big] \right].$$ $$\begin{aligned} \nabla_\theta J_{\mathrm{OPD}}(\theta) &= \mathbb E_{x\sim D} \Bigg[ \sum_y \nabla_\theta \pi_\theta(y\mid x) \Big[ \log\pi_\theta(y\mid x) - \log q(y\mid x) \Big] \\ &\qquad\qquad + \sum_y \pi_\theta(y\mid x) \nabla_\theta\log\pi_\theta(y\mid x) \Bigg]. \end{aligned}$$

利用常见的 log-derivative trick：

$$\nabla_\theta\pi_\theta(y\mid x) = \pi_\theta(y\mid x) \nabla_\theta\log\pi_\theta(y\mid x),$$

所以有

$$\begin{aligned} \nabla_\theta J_{\mathrm{OPD}}(\theta) = \mathbb E_{\substack{x\sim D\\ y\sim\pi_\theta(\cdot\mid x)}} \Big[ \nabla_\theta\log\pi_\theta(y\mid x) \big( 1 + \log\pi_\theta(y\mid x) - \log q(y\mid x) \big) \Big]. \end{aligned}$$

再利用 zero-mean score identity：

$$\mathbb E_{y\sim\pi_\theta(\cdot\mid x)} \left[ \nabla_\theta\log\pi_\theta(y\mid x) \right] = 0,$$

得到：

$$\nabla_\theta J_{\mathrm{OPD}}(\theta) = \mathbb E_{\substack{x\sim D\\ y\sim\pi_\theta(\cdot\mid x)}} \left[ \nabla_\theta\log\pi_\theta(y\mid x) \left( \log\pi_\theta(y\mid x) - \log q(y\mid x) \right) \right].$$

用刚刚一样的记号，记

$$s_t = \nabla_\theta\log\pi_\theta(y_t\mid c_t), \quad r_t = \log\pi_\theta(y_t\mid c_t) - \log q(y_t\mid c_t)$$

语言模型是自回归的，

$$\log \frac{\pi_\theta(y\mid x)} {q(y\mid x)} = \sum_{t'=1}^{T}r_{t'},$$

还有

$$\nabla_\theta \log \pi_\theta(y\mid x) = \sum_{t=1}^T s_t$$

因此，目标的梯度整理为：

$$\nabla_\theta J_{\mathrm{OPD}}(\theta) = \mathbb E_{\substack{x\sim D\\ y\sim\pi_\theta(\cdot\mid x)}} \left[ \sum_{t=1}^T s_t \sum_{t'=1}^T r_{t'} \right].$$

因此，sequence-level 的 estimator 定义为：

$$\hat{g}_{\text{seq}}(\theta) = \sum_{t=1}^T s_t \sum_{t'=1}^T r_{t'}$$

当 $t’ < t$ 时，因为

$$\begin{aligned} \mathbb E[s_t\mid c_t] &= \sum_{y_t} \pi_\theta(y_t\mid c_t) \nabla_\theta \log\pi_\theta(y_t\mid c_t)\\ &= \sum_{y_t} \nabla_\theta \pi_\theta(y_t\mid c_t)\\ &= \nabla_\theta \sum_{y_t} \pi_\theta(y_t\mid c_t)\\ &= \nabla_\theta 1\\ &= 0. \end{aligned}$$

所以

$$\begin{aligned} \mathbb E[r_{t'}s_t] &= \mathbb E\left[ r_{t'}\, \mathbb E[s_t\mid c_t] \right]\\ &= \mathbb E[r_{t'}\cdot 0]\\ &= 0, \qquad t'<t. \end{aligned}$$

这里全期望公式是因为固定 $c_t$ 时，$r_{t’}$ 不随当前的 $y_t$ 变化，我补一下解释：

$$\mathbb E[r_{t'}s_t] = \sum_{c_t}P(c_t) \sum_{y_t}P(y_t\mid c_t)\, r_{t'}\,s_t$$ $$= \sum_{c_t}P(c_t)\, r_{t'} \left[ \sum_{y_t}P(y_t\mid c_t)s_t \right]$$ $$= \sum_{c_t}P(c_t)\, r_{t'}\, \mathbb E[s_t\mid c_t].$$

因此可以把 gradient 改写成 causal return-to-go：

$$\mathbb E[\hat g_{\mathrm{seq}}] = \mathbb E \left[ \sum_{t=1}^{T} \left( \sum_{t'=t}^{T}r_{t'} \right)s_t \right].$$

回忆一下，token-level 的 estimator 是：

$$\hat g_{\mathrm{tok}} = \sum_{t=1}^T s_tr_t$$

由此可见，TML 的 token-level 的 estimator 是有偏的。

γ 在哪里？

参考了知乎文章《重探 On-Policy Distillation（OPD）：三类典型失败以及修复路径》

把 immediate token-level 和完整 sequence-level 写进同一个 estimator，这样它们就是这个折扣 return-to-go estimator 的特例：

$$\boxed{ \hat g_\gamma = \sum_{t=1}^{T} \left( \sum_{t'=t}^{T} \gamma^{\,t'-t}r_{t'} \right)s_t, \qquad \gamma\in[0,1]. }$$

$\gamma$ 就是在这里。😉

References

TML 让 $\gamma = 0$。这是有偏的。但是方差会更小。证明在Yuqian Fu 等人写的文章和他们的论文里。

他们还发现了 sampled-token OPD 的其他问题。而且，他们还研究了介于单 token 和全词表之间的「teacher top-K 局部匹配」。

我主要阅读了他们分享的知乎内容。

另外 TML 的思路首先是：

$$\boxed{ A_t = -\left( \log \pi_{\rm samp}(y_t\mid c_t) - \log q(y_t\mid c_t) \right) }$$

然后训练，只有当

$$\gamma=\texttt{kl\_discount\_factor}>0$$

时，它才额外调用：

$$A_t \leftarrow A_t+\gamma A_{t+1}+\gamma^2A_{t+2}+\cdots.$$

所以另一个视角就是，所谓 $\gamma=0$，从 RL 语言里几乎立即得到：

$$G_t = A_t = -r_t.$$

没有必要先证明

$$\nabla D_{\mathrm{KL}} = \mathbb E[s_tr_t].$$

欸，那我这篇博客不就是白写了。