Source-linked AI summary
Baselines and Score Centering in Policy Gradient Estimation
Shuo Han, Ziqing Hu
TL;DR
Off-policy post-training can arise from differences between training and inference policies, while regularizing importance ratios breaks reward-shift invariance. The paper uses a covariance view of policy gradients and develops it for sequence- and token-level importance sampling, recovering that invariance; the covariance estimators also offer distinct variance trade-offs.
Problem
Method
The paper treats the policy gradient as a covariance between reward and weighted score and applies covariance estimation to sequence- and token-level importance sampling.
Source: §3 · p. 8, §4 · p. 16Results
The covariance form restores invariance to constant reward shifts under ratio regularization and reduces to the policy gradient without regularization.
Source: §4 · p. 16Takeaways & Limitations
The covariance view yields two distinct unbiased estimators whose variance comparison guides the choice of estimation algorithm.
Source: §3 · p. 8Sequence-level ratios can become numerically unstable and produce enormous estimator variance, while token-level optimization guarantees improvement over the inference policy, not necessarily the current training policy.
Source: §2.3 · pp. 7–8
Abstract
from arXiv · showhide
Reinforcement learning has become a widely used approach to post-training large language models. In a typical implementation, learning is off-policy: The policy generating the responses differs from the policy being trained because of training-inference mismatch or policy staleness in asynchronous training. This discrepancy is commonly addressed by importance sampling, whose ratio is usually regularized in practice through truncation, masking, or clipping. Under regularized importance sampling, the importance-weighted policy gradient is no longer invariant to constant shifts in the reward. This invariance is recovered by using the covariance between the reward and the importance-weighted score as the estimand, which coincides with the policy gradient when the ratio is not regularized. This note develops this covariance perspective for both sequence-level and token-level importance sampling. The covariance perspective reveals the roles of the group-mean baseline and score centering in policy gradient estimation. Both perform centering in covariance estimation, but they differ in the mean being subtracted: The group-mean baseline uses a sample mean, whereas score centering uses a population mean. The two lead to distinct unbiased covariance estimators with different variances. Whether both forms of centering should be used simultaneously depends on whether importance sampling is applied at the sequence level or at the token level. With sequence-level importance sampling, score centering adds no benefit once the group-mean baseline is used; with token-level importance sampling, using the group-mean baseline without score centering generally introduces bias when the importance ratio is regularized.
1 Background
The background frames language-model post-training as maximizing reward over responses to prompts, with the model represented as a policy. It then models token generation as an MDP with reward assigned only at the terminal step.
1.1 Post-training of language models
- Post-training seeks a policy that generates high-reward responses to prompts, with response quality measured by a reward function.The objective is defined over prompts and complete responses.Source: §1.1 · pp. 1–2
1.2 Token generation as a Markov decision process
- The objective is well-defined under the assumption that generation terminates almost surely, enforced in practice by a maximum response length.The paper also introduces policy value and advantage functions for this MDP.Source: §1.1 · p. 2, §1.2 · p. 2
- The token-generation process is an undiscounted indefinite-horizon MDP, with states holding the prompt and generated prefix and actions representing the next token.Transitions append each token; EOS moves the process to an absorbing state.Source: §1.2 · p. 2
- Reward is sparse: it is zero for nonterminal tokens and assigns the completed response’s reward when EOS is generated.The response length is the decoding step containing EOS.Source: §1.2 · p. 2
2 Importance Sampling in Policy Optimization
Importance sampling lets post-training policy gradients use responses from a mismatched inference engine, but sequence- and token-level ratios target different objectives. The section develops unbiased estimators and explains their variance and policy-improvement trade-offs.
2 Importance Sampling in Policy Optimization
- Importance sampling permits estimating gradients from inference-policy samples when exact objective expectations are intractable and training and inference engines produce different token probabilities.Source: §2 · p. 3
2.1.1 REINFORCE
- The sequence-level REINFORCE estimator using inference-policy samples is unbiased for ∇J(θ) when prompts are sampled from the prompt distribution.Source: §2.1.1 · p. 4
2.1.2 Policy gradient with a baseline
- Conditionally independent baselines preserve unbiasedness for ∇J(θ) in multi-sample estimation, while Dr. GRPO’s group-mean estimator is practically equivalent to RLOO up to (G −1)/G.Source: §2.1.2 · pp. 4–5
2.2.1 Trust-region policy optimization and proximal policy optimization
- ˆJ is affine in policy space, and improving it gives an ascent direction near the inference policy; TRPO constrains divergence, while PPO clips token ratios heuristically.Source: §2.2.1 · p. 6
2.2.2 Token-level GRPO
- Token-level GRPO substitutes reward for prefix advantage and yields an unbiased estimator of ∇ˆJ; clipping its ratio trades variance reduction for bias.Source: §2.2.2 · p. 7
2.3 Comparison between sequence-level and token-level importance sampling
- Sequence-level ratios can destabilize estimates through enormous variance, while token-level ratios behave better but optimize a surrogate that may not improve the current training policy under mismatch.Source: §2.3 · pp. 7–8
3 Policy Gradient as a Covariance
The section recasts policy-gradient estimation as covariance estimation, recovering reward-shift invariance under regularized importance sampling. It distinguishes sequence- and token-level centering: a group-mean baseline makes score centering redundant at the sequence level, but not generally at the token level.
3.2 The sequence-level estimand
- With regularized importance sampling, the reward–weighted-score covariance preserves invariance to constant reward shifts and recovers the policy gradient when the ratio is unregularized.The section presents this covariance as the sequence-level estimand instead of the regularized importance-weighted expectation.Source: §3 · p. 8, §3.2 · p. 11
3 Policy Gradient as a Covariance
- For the ordinary policy gradient, the reward–score expectation is itself a covariance because the policy score has mean zero.The section’s broader covariance framing applies to both the exact and approximate policy gradients.Source: §3 · p. 8, §3.2 · p. 11
3.1 Unbiased covariance estimation
- The covariance estimators are unbiased, but their variances differ: the known-mean estimator can be more efficient with a suitable baseline and much less efficient when that baseline is far from optimal.With the population-mean baseline, its variance is smaller by an O(n^-2) term; with the minimum-variance baseline, it is further reduced by an O(n^-1) term.Source: §3.1 · pp. 9–10
3.3 Sequence-level covariance estimation
- RLOO and Dr. GRPO correspond to sample-mean covariance estimators, for which score centering is unnecessary; the mean-reward baseline may nevertheless be far from variance-optimal.The variance-optimal sequence-level baseline depends on the covariance between reward and score norm, which the passage notes may be non-negligible when response length correlates with reward.Source: §3.3 · pp. 11–12
3.4 The token-level estimand
- At the token level, the estimand is a covariance between reward and D, and the mean weighted score can be computed per prefix by a finite sum over the vocabulary.The stated regularity conditions assume finite expected termination and a bounded conditional mean for the weighted score.Source: §3.4 · pp. 12–14
3.5 Token-level covariance estimation
- Unlike at the sequence level, a group-mean reward baseline does not generally make token-level score centering optional; omitting it yields a different estimator because prefix mean scores vary across rollouts.The optimal baseline is prompt-dependent, so estimating it per prompt still requires multiple rollouts for that prompt.Source: §3.5 · p. 15
4 Summary
Under regularized importance sampling, covariance estimation restores reward-shift invariance and clarifies how group-mean baselines and score centering affect policy-gradient estimation. Which centering to use depends on whether importance sampling is sequence-level or token-level.
4 Summary
- Covariance between reward and weighted score restores reward-shift invariance and equals the policy gradient when the importance ratio is unregularized.Without regularization, sequence-level and token-level importance sampling are unbiased for different objectives: J and the surrogate ˆJ, respectively.Source: §4 · pp. 15–16
- The group-mean baseline centers with a sample mean, whereas score centering uses a population mean, yielding distinct unbiased covariance estimators.The sample-mean estimator coincides with RLOO and, up to a multiplicative factor, Dr. GRPO; the known-mean estimator is unbiased for any constant baseline and is defined for a single sample.Source: §4 · p. 16
- The known-mean estimator has lower variance only when its baseline is sufficiently close to the variance-optimal baseline b⋆, and can become arbitrarily worse farther away.For token-level importance sampling, that baseline is prompt-dependent and must be estimated from multiple rollouts.Source: §4 · p. 16
- With sequence-level importance sampling, score centering is redundant once the group-mean baseline is used; with regularized token-level sampling, omitting it generally biases the group-mean estimator.At the token level, the estimator's expectation contains an extra term that vanishes under standard importance sampling but is generally nonzero when the ratio is regularized.Source: §4 · p. 16
- Use the group-mean baseline alone for sequence-level sampling; for regularized token-level sampling, apply score centering with either that baseline or a constant baseline.Both estimator families degrade as training–inference mismatch grows, so reducing mismatch at its source remains useful.Source: §4 · p. 16