Logit Gradient Correction: Deriving Score Centering from Logit Gradients

This post derives SCORE CENTERING from removing a bias term in the logit gradient before it is propagated through the logit Jacobian via chain rule, and then derives the corresponding surrogate objective.

Background: REINFORCE and the Logit Gradient

At a fixed prefix $𝒙$, a trainer with $𝐷$ parameters $𝜽\in {ℝ}^{𝐷}$ produces logits $𝒛\left(𝒙,𝜽\right)\in {ℝ}^{𝑉}$ over a vocabulary of size $𝑉$. Its probability vector $𝒑=softmax\left(𝒛\right)\in {ℝ}^{𝑉}$ defines the trainer distribution $𝑃$, with $𝑃\left(𝑎\right)={𝑝}_{𝑎}$. A sampler at the same prefix generates token $𝑎\sim 𝑄$, with probabilities $𝒒\in {ℝ}^{𝑉}$ and $𝑄\left(𝑎\right)={𝑞}_{𝑎}$. During backward, the recorded $𝒒$ and reward ${𝑅}_{𝑎}\in ℝ$ are held fixed. Scalars use plain letters, vectors bold lowercase, and matrices bold uppercase.

REINFORCE wants to increase the trainer’s expected reward ${𝐽}_{𝑃}$ under the trainer’s own distribution $𝑃$:

Holding the reward function fixed, differentiate ${𝐽}_{𝑃}$ with respect to $𝜽$ to obtain ${\nabla }_{𝜽}{𝐽}_{𝑃}\in {ℝ}^{𝐷}$:

The term ${𝑅}_{𝑎}{\nabla }_{𝜽}log{𝑝}_{𝑎}\in {ℝ}^{𝐷}$ is REINFORCE’s per-sample gradient contribution. With ${𝑅}_{𝑎}$ held fixed, it equals the negative gradient with respect to $𝜽$ of the reward-weighted cross-entropy loss $-{𝑅}_{𝑎}log{𝑝}_{𝑎}$.

We first differentiate $log{𝑝}_{𝑎}$ with respect to the logits $𝒛$. Given that $𝒑=softmax\left(𝒛\right)$, we have

Here, ${𝒆}_{𝑎}\in {ℝ}^{𝑉}$ is the sampled token’s one-hot vector. The gradient ${𝒆}_{𝑎}-𝒑$ contains the derivatives of $log{𝑝}_{𝑎}$ with respect to the logits of all vocabulary tokens. The mean logits contribution on the right therefore determines the mean parameter contribution.

Log-Prob Gradients Have Zero/Nonzero Bias when On/off-Policy

We consider a scenario that ${𝑅}_{𝑎}=1$ for every token, we have:

Thus the gradient of ${𝐽}_{𝑃}$ with respect to the model parameters $𝜽$ will be

The matrix $\frac{𝜕𝒛}{𝜕𝜽}\in {ℝ}^{𝑉\times 𝐷}$ is the trainer’s logit Jacobian: it describes how the logits change with the model parameters. Its transpose maps a $𝑉$-dimensional logits gradient to a $𝐷$-dimensional parameter gradient, applying the chain rule through backpropagation. At a fixed prefix and parameter value, this matrix is the same for every possible sampled token, so it can be moved outside the expectation. Score Centering supplies a corrected logits gradient to this same backward pass.

If the tokens are generated by the trainer distribution $𝑃$, the expectation ${𝔼}_{𝑎\sim 𝑃}\left[{𝒆}_{𝑎}\right]$ equals $𝒑$,

the logits gradient is unbiased and has zero mean when the reward is constant.

Logit Gradient Correction

If tokens instead come from the sampler distribution $𝑄$, then ${𝔼}_{𝑎\sim 𝑄}\left[{𝒆}_{𝑎}\right]=𝒒$. Still assuming ${𝑅}_{𝑎}=1$, the mean trainer log-probability gradient becomes:

Here $𝑄$ specifies where tokens come from; the differentiated log-probability still belongs to $𝑃$. This expression is the mean sampled contribution, not the gradient of the sampler’s expected reward.

Since every reward is $1$, the trainer objective ${𝐽}_{𝑃}=1$ has zero gradient. We want the corrected contributions to have zero mean even when tokens come from $𝑄$. Shift each one-hot contribution ${𝒆}_{𝑎}$ by $𝒑-𝒒$, giving ${𝔼}_{𝑎\sim 𝑄}\left[{𝒆}_{𝑎}-𝒒+𝒑\right]=𝒒-𝒒+𝒑=𝒑$.

Let $𝐽$ denote the per-sample surrogate objective that we will construct in the next section. We require its expected parameter gradient to equal the corrected mean contribution. With $𝑄$ held fixed during backward, this requirement is:

The essential of this correction is to replace each token’s logits gradient contribution ${𝒆}_{𝑎}-𝒑$ with ${𝒆}_{𝑎}-𝒒$. Averaging under the sampler distribution $𝑄$ then gives ${\sum }_{𝑎}{𝑞}_{𝑎}\left({𝒆}_{𝑎}-𝒒\right)=\mathrm{𝟎}$, restoring zero mean for constant rewards.

From Corrected Logit Gradient to a Surrogate Objective

We now construct the per-sample surrogate $𝐽$. For constant reward ${𝑅}_{𝑎}=1$, the required relation is:

For varying rewards, we weight each corrected contribution by ${𝑅}_{𝑎}$. The required gradient of the expected surrogate is

First, we seek a scalar surrogate objective $𝐽$ in terms of the logits $𝒛$, and whose gradient with respect to the logits $𝒛$ is:

The required gradient ${𝑅}_{𝑎}\left({𝒆}_{𝑎}-𝒒\right)$ is constant with respect to $𝒛$. An intuitive antiderivative is a linear function whose gradient equals its coefficient vector:

The last line uses ${𝒆}_{𝑎}^{\top }𝒛={𝑧}_{𝑎}$. The sum ${\sum }_{𝑣}sg\left[{𝑞}_{𝑣}\right]{𝑧}_{𝑣}$ is the sampler-weighted average of the trainer’s logits over the entire vocabulary. The operator $sg$ holds the probabilities fixed during differentiation while preserving gradients through the logits.

Second, we express the surrogate objective $𝐽$ in terms of log-probabilities. Softmax gives $log{𝑝}_{𝑣}={𝑧}_{𝑣}-log\left({\sum }_{𝑗}exp\left({𝑧}_{𝑗}\right)\right)$ for every token $𝑣$. Using ${\sum }_{𝑣}sg\left[{𝑞}_{𝑣}\right]=1$, we obtain:

We have the surrogate objective in terms of log-probability:

The surrogate loss replaces the per-sample REINFORCE objective ${𝑅}_{𝑎}log{𝑝}_{𝑎}$ with ${𝑅}_{𝑎}$ times the difference between the sampled token’s trainer log-probability and the $𝑄$-weighted average of the trainer’s log-probabilities over the vocabulary.

Averaging the per-sample surrogate over tokens drawn from $𝑄$ gives: