Track papers across labs, topics, and trends.
Search papers, labs, and topics across Lattice.
This paper analyzes policy gradient estimation in off-policy reinforcement learning for large language models, particularly when importance sampling ratios are regularized. It shows that regularized importance sampling breaks reward shift invariance, which can be restored by using the covariance between rewards and importance-weighted scores as the estimand. The work clarifies the distinct roles of group-mean baselines and score centering in this covariance estimation, revealing their impact on estimator bias and variance depending on whether sequence-level or token-level importance sampling is employed.
Regularized importance sampling in LLM RL breaks reward shift invariance, but a covariance-based estimand recovers this property, with distinct centering strategies impacting bias and variance.
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.