Score Centering 的局部分布族解释与 STE 解释是一致哒。🫠
一言以蔽之，「STE 指定的代理梯度，可以由一个普通分布族在起点处的真实导数实现」。
符号约定
固定一个生成前缀，下文省略固定前缀条件。
- $a$：候选 token。
- $\theta_0$：训练器当前参数，固定为起点。
- $\theta$：可以变化的训练器参数。
- $\delta=\theta-\theta_0$：参数改变量。
- $z_\theta$：训练器 logits 向量。
- $p_\theta=\operatorname{softmax}(z_\theta)$：训练器概率分布。
- $z_q$：固定的采样器 logits。
- $q=\operatorname{softmax}(z_q)$：采样器概率分布。
- $\operatorname{sg}$：stop-gradient，前向保留数值，反向停止求导。
定义，训练器的 score 是
定义，SC 操作是
其中，$b$ 遍历词表所有候选 token，$\mathbb E_q$ 表示按采样器概率求平均。
回顾：STE
构造
构造
前向：
反向，括号内不传导数。变化只通过外面的 $z_\theta$。
因此，前向用采样器的概率，反向借用训练器的变化。
STE 的 score
记：
- $e_a$：token $a$ 的 one-hot 向量。
- $J_\theta=\partial z_\theta/\partial\theta$：训练器 logits 对参数的 Jacobian。
- $J_\theta^\top$：它的转置，将 logits 梯度传回参数空间。
普通 score 是
STE 的自动微分结果是
SC 恰好也是
因此，STE 与 SC 梯度等价。
回顾：局部分布族解释
现在构造一个数学上普通的、随参数变化的分布：
其中归一化因子为
这个分布，有三个特点：
- 从采样分布 $q$ 出发。
- 乘上训练器更新前后的相对概率变化。
- 重新归一化。
对于起点，也就是当 $\delta=0$ 时，$\nu_0=q$。所以，对数概率的普通导数在起点处为
这自然得到 SC。被减掉的平均 score 来源于归一化项产生的导数。
而这个构造局部分布族解释，对于 SC，有一个很好的洞察：SC 的期望更新可以理解为：在这个以采样分布为起点的分布族上，做局部梯度上升。 具体来说，先定义：在单步问题中，令 $R(a)$ 为固定的动作奖励，定义平均奖励
推一波 PG，
注意到 $\widetilde s_{\theta_0}$ 是中心化后的 score，因此，SC 的期望更新是：「在这个以采样分布为起点的分布族上，对奖励做局部梯度上升」。
Score Centering 的局部分布族解释与 STE 解释一致
把分布族写成 logits
把这个分布族写成 logits，就有：
怎么写出来的？之所以这样，是因为 softmax 概率比值中的归一化常数与 token 无关，会在再次归一化时消掉，具体来说，过程如下，先展开 softmax。我们先记三个归一化因子为
其中，$b$ 遍历词表。于是有
代入分布族的构造。对于刚刚 $\nu_\delta(a)$ 的定义式，分子是：
其前面的因子 $C_0/(C_qC_\delta)$ 对所有 token 都一样，因此，在「重新归一化」这一步，会从分子、分母中抵消，我们得到
其实，这是 softmax：
所以：概率比值构造与固定 logits 偏移构造，定义的是同一个分布族。它们是同一个分布族的两种写法。什么是固定 logits 偏移构造？定义固定偏移：
于是
注意，旧的差值被固定了，旧差值是 $\Delta z$，其中的减数项为 $z_{\theta_0}$。
Score Centering 的局部分布族解释与 STE 解释一致
比较两种 logits 构造：
| 当前差值 STE | 固定差值分布族 | |
|---|---|---|
| 公式 | $z_\theta+\operatorname{sg}(z_q-z_\theta)$ | $z_\theta+(z_q-z_{\theta_0})$ |
| 起点值 | $z_q$ | $z_q$ |
| 起点处反向变化 | 通过 $z_\theta$ | 通过 $z_\theta$ |
| 离开起点后的前向值 | 仍然是 $z_q$ | 随 $z_\theta$ 改变 |
| 导数 | 指定的代理导数 | 普通函数的真实导数 |
因此，在起点处，也就是 $\delta=0$ 处：
综上所述
STE 指定的代理梯度，是这个局部分布族，在起点处的真实梯度。