还真是！《score-centering的前世今生》与《OPENAI的神秘RL算法只是个STE的近似吗？解密OAI的RL 后训练算法》都是这么解释的。☺️

我学习了以上文章，做了阅读笔记。

先验知识

《BPO 与 Score Centering 笔记》

符号定义与前置

符号

在固定前缀（状态）$s$ 下，采样器按 $q(\cdot\mid s)$ 生成 token $a$，训练器给出 $p_\theta(\cdot\mid s)$。普通训练器 score 是

$s_p(a)=\nabla_\theta\log p_\theta(a\mid s).$

当 $a\sim p_\theta$ 时，$\mathbb E_p[s_p]=0$；当 $a\sim q$ 时，一般来说 $\mathbb E_q[s_p]\ne0$。

此非零均值使得 off-policy 更新对奖励零点、减 baseline 敏感。Score Centering（SC）减掉这个均值：

$\tilde s_p(a)=s_p(a)-\mathbb E_{v\sim q}[s_p(v)], \qquad \mathbb E_q[\tilde s_p]=0.$

若有完整的 $q$ 分布，构造 STE logits

$\tilde z=z_p+\operatorname{sg}(z_q-z_p)$

可得到与 SC 相同的单样本参数梯度。这个等价是梯度等价，不是前向 loss 数值相同，也不是说 SC 已纠正所有 off-policy 偏差。

固定一个前缀 $s$ 后，$z_p(s;\theta)\in\mathbb R^{|\mathcal V|}$ 是训练器给词表每个 token 的原始分数（logits），$z_q$ 是采样器的 logits。对应的概率向量为

$p_v=\frac{e^{z_{p,v}}}{\sum_u e^{z_{p,u}}},\qquad q_v=\frac{e^{z_{q,v}}}{\sum_u e^{z_{q,u}}}.$

$p,q,z_p,z_q$ 都是词表维的向量；$p,q$ 的分量非负、总和为 $1$。

本次抽到的 token $a$ 对应 one-hot 向量 $e_a$：第 $a$ 维是 $1$，其他维是 $0$。

Jacobian $J_\theta$ 是每个 logit 对每个模型参数的偏导数表，

$J_\theta=\frac{\partial z_p}{\partial\theta},\qquad (J_\theta)_{v,j}=\frac{\partial z_{p,v}}{\partial\theta_j}.$

前置：$\nabla_{z_p}\log p_a=e_a-p$

先取对数，再逐维求导：

$\log p_a=z_{p,a}-\log\sum_v e^{z_{p,v}},$ $\frac{\partial\log p_a}{\partial z_{p,j}} =\underbrace{\mathbf1[j=a]}_{\text{来自 }z_{p,a}} -\underbrace{\frac{e^{z_{p,j}}}{\sum_v e^{z_{p,v}}}}_{p_j} =\mathbf1[j=a]-p_j.$

把各维排起来，便是 $e_a-p$。

梯度再传回参数：

$s_p(a)=\nabla_\theta\log p_a=J_\theta^\top(e_a-p).$

前置：PG

固定状态下，PG 的上升方向为 $\mathbb E_{a\sim p}[A(a)s_p(a)]$。

若样本实际来自 $q$ 而未做分布校正，则用 $\mathbb E_{a\sim q}[A(a)s_p(a)]$。

单样本在 logit 空间的上升方向为 $A(a)(e_a-p)$。

前置：STE（Straight-Through Estimator）

$\operatorname{sg}(u)$（stop-gradient）前向值仍为 $u$，反向导数按 $0$ 处理。

STE 在反向时人为把硬阈值当成恒等函数，传递非零的替代梯度。这是估计，不是该函数的真实导数。

STE 是 Hinton 讲课时提出的。

构造 STE

构造 STE：

$\tilde z=z_p+\operatorname{sg}(z_q-z_p).$

前向：$\tilde z=z_p+(z_q-z_p)=z_q$，因此 $\operatorname{softmax}(\tilde z)=q$。

反向：

$\frac{\partial\tilde z_i}{\partial z_{p,j}}=\mathbf1[i=j], \qquad \frac{\partial\tilde z}{\partial z_p}=I.$

故

$\nabla_\theta\log\operatorname{softmax}(\tilde z)_a =J_\theta^\top(e_a-q).$

SC 与 STE 的梯度相同

先算普通 score 在 $q$ 下的均值。因为 $\sum_vq_ve_v=q$、$\sum_vq_v=1$，所以：

$\mathbb E_{v\sim q}[s_p(v)] =\sum_vq_vJ_\theta^\top(e_v-p) =J_\theta^\top(q-p).$

所以中心化

$\begin{aligned} \tilde s_p(a) &=s_p(a)-\mathbb E_q[s_p]\,\\ &=J_\theta^\top(e_a-p)-J_\theta^\top(q-p)\\ &=J_\theta^\top[(e_a-p)-(q-p)]\\ &=J_\theta^\top(e_a-q). \end{aligned}$

这正是 STE 梯度。（条件：同一固定前缀、完整的 $q$ 分布、$q$ 与 advantage 在反向时 detach）

不过两种 surrogate 的前向 loss 不同：

$L_{\rm SC}=-A\left(\log p_a-\sum_vq_v\log p_v\right), \qquad L_{\rm STE}=-A\log\operatorname{softmax}(\tilde z)_a.$

SC 与 STE 作用相同

On-policy 的 score 均值为零：

$\mathbb E_{a\sim p}[s_p(a)] =\sum_a p_a\nabla_\theta\log p_a =\nabla_\theta\sum_a p_a=0.$

Off-policy 时 $a\sim q$，普通 logit score 的均值却是 $\mathbb E_q[e_a-p]=q-p$；参数空间为 $J_\theta^\top(q-p)$。因此状态 baseline $V(s)$ 不再自动消失：

$\mathbb E_q[Q(a)s_p(a)] =\mathbb E_q[(Q(a)-V)s_p(a)]+V\,\mathbb E_q[s_p(a)].$

若 $A^p=Q^p-V^p$，$V^p=\mathbb E_p[Q^p]$，则 $\mathbb E_p[A^p]=0$，但 $\mathbb E_q[A^p]=\mathbb E_q[Q^p]-\mathbb E_p[Q^p]$ 通常不为零。

SC 让 $\mathbb E_q[\tilde s_p]=0$，于是任意状态常数都可以从奖励中减去：

$\mathbb E_q[Q\tilde s_p] =\mathbb E_q[(Q-V)\tilde s_p] =\operatorname{Cov}_q(Q,s_p).$

它修复了 $q$ 下的平均 score 漂移 / baseline 不变性。

注意没有解决协方差仍然是在 $q$ 下计算的问题。

离策略时，SC 有 IS 后与 STE 相同

令 $q_{\rm old}$ 是旧 vLLM 采样器，$z_M(\theta)$ 是训练器 logits。构造固定的旧差值

$\Delta z=z_{\rm vLLM}^{\rm old}-z_M^{\rm old},\qquad \tilde z_\theta=z_M(\theta)+\operatorname{sg}(\Delta z),\qquad \hat p_\theta=\operatorname{softmax}(\tilde z_\theta).$

在 $\theta=\theta_{\rm old}$：$\tilde z_{\rm old}=z_{\rm vLLM}^{\rm old}$，故 $\hat p_{\rm old}=q_{\rm old}$，比率 $\hat p_{\rm old}(a)/q_{\rm old}(a)=1$。

更新参数后，$z_M(\theta)$ 改变，但旧差值 $\Delta z$ 固定，通常 $\hat p_\theta\ne q_{\rm old}$，比率不再是 $1$。

在整个轨迹 $x$ 下，重要性采样恒等式为

$\mathbb E_{x\sim\hat P_\theta}[R(x)] =\mathbb E_{x\sim P_{\rm old}^{\rm vLLM}} \left[\frac{\hat P_\theta(x)}{P_{\rm old}^{\rm vLLM}(x)}R(x)\right].$

因此 STE surrogate 目标的梯度可写为

$g_{\rm STE} =\mathbb E_{x\sim q_{\rm old}} \left[\frac{\hat p_\theta(x)}{q_{\rm old}(x)}R(x) \nabla_\theta\log\hat p_\theta(x)\right].$

若 SC 无 IS：

$g_{\rm SC} =\mathbb E_{x\sim q_{\rm old}} \left[R(x)\bigl(s_p(x)-\mathbb E_{v\sim q_{\rm old}}s_p(v)\bigr)\right].$

若有，则相同。事实上，SC 的论文中做了与 MIS/TIS 结合的实验。

IGO 与 PMD 形式相同，优化单位不同

（IGO 是希望实际策略向这个分布靠近。STE 和 Score Centering 的梯度等价推导并不需要 IGO。《OPENAI的神秘RL算法只是个STE的近似吗？解密OAI的RL 后训练算法》的作者是从 IGO 出发、并意识到 ratio 不为 1 而得到 STE 的。我这里做一个对比笔记。）

IGO 视角可写成在完整轨迹 $x$ 上的 KL 正则优化：

$\max_Q\;\mathbb E_Q[R(x)]-\beta D_{\rm KL}(Q\|P_{\rm old}) \quad\Longrightarrow\quad Q^*(x)=\frac{P_{\rm old}(x)e^{R(x)/\beta}} {\mathbb E_{x'\sim P_{\rm old}}e^{R(x')/\beta}}.$

BPO 论文的 PMD 则在每个前缀 $s$ 的下一 token上优化：

$\max_{\pi(\cdot|s)}\;\mathbb E_{a\sim\pi}[f(s,a)] -\frac1\eta D_{\rm KL}(\pi(\cdot|s)\|\mu(\cdot|s)) \quad\Longrightarrow\quad \pi^+(a|s)\propto\mu(a|s)e^{\eta f(s,a)}.$

其中 BPO 取 $f(s,a)=A^\mu(s,a)$。

再补充一个笔记，《OPENAI的神秘RL算法只是个STE的近似吗？解密OAI的RL 后训练算法》的作者如何从 IGO 出发的。

设 $b=P_{\rm old}^{\rm vLLM}$，IGO 给出的目标分布是

$Q^*(x)=\frac{b(x)e^{R(x)/\beta}}{Z}.$

其中 $Z$ 是归一化常数：

$Z=\sum_x P_{\rm old}(x)e^{R(x)/\beta} =\mathbb E_{x\sim P_{\rm old}}\!\left[e^{R(x)/\beta}\right].$

要最小化 $D_{\rm KL}(P|Q^*)$，代入并展开会得到

$\begin{aligned} \beta D_{\rm KL}(P\|Q^*) &=\beta\sum_xP(x)\log\frac{P(x)}{Q^*(x)}\\ &=\beta D_{\rm KL}(P\|b)-\mathbb E_{x\sim P}[R(x)] +\beta\log Z. \end{aligned}$

所以，

$\min_P D_{\rm KL}(P\|Q^*) \quad\Longleftrightarrow\quad \max_P\left\{\mathbb E_P[R]-\beta D_{\rm KL}(P\|b)\right\}.$

（没有 $Z$？因为 $Z$ 只由固定的旧分布 $b$、奖励 $R$ 和 $\beta$ 决定；优化变量是新分布 $P$。无论怎样改变 $P$，$\beta\log Z$ 都不变。）

其中的奖励项：

$\mathbb E_{x\sim b}\!\left[\frac{P(x)}{b(x)}R(x)\right] =\mathbb E_{x\sim P}[R(x)]$

接着把 $P^{\rm vLLM}$ 换成 $P^{\rm Megatron}$ 近似。

Palind's Blog

Score Centering 是 STE