PaperlayerView exact revision on arXiv

Source-linked AI summary

Future-KL Regularized GRPO: Process-Level Credit Assignment from $f$-Divergence Regularization

Jiarui Yao, Ruida Wang, Hao Bai, Tong Zhang

arXiv:2601.10201v2cs.LGcs.AIcs.CL

TL;DR

The paper addresses the missing autoregressive policy-gradient signal in GRPO’s local KL regularization and the distortion caused by fusing KL into non-linearly normalized rewards. It derives a general on-policy gradient with future regularization credit and proposes FRPO, which improves pass@n while limiting policy drift and entropy collapse in mathematical-reasoning experiments.

  • Problem

    GRPO’s local loss-side KL penalty omits future autoregressive regularization effects, while non-linear group normalization makes reward and KL non-fuseable before normalization.

  • Method

    The paper derives the on-policy gradient for token-wise f-divergence regularization and uses a decoupled future-KL correction after constructing the GRPO outcome advantage.

  • Results

    FRPO improves pass@n while limiting policy drift and entropy collapse in mathematical-reasoning experiments.

  • Takeaways & Limitations

    GRPO outcome ranking and token-level KL credit assignment should remain separate under non-linear group normalization.

  • Takeaways & Limitations

    Practical rollout staleness can make the gradient no longer exactly on-policy unless importance weighting or fresh log-probabilities are used.

Abstract

from arXiv · show

Group Relative Policy Optimization (GRPO) is widely used for critic-free Large Language Model (LLM) post-training, but its KL regularization is usually implemented as a local loss-side token penalty. We show that this misses the policy-gradient signal induced by autoregressive KL regularization. Unlike standard KL-regularized Reinforcement Learning (RL) objectives, GRPO's group normalization induces a non-linear prompt-level utility; for binary verifier rewards, this utility is $2\arcsin\sqrt p$. As a result, reward and KL cannot be fused before normalization without changing the implicit objective. We derive the on-policy gradient of GRPO-style objectives with token-wise $f$-divergence regularization. The reward term recovers the standardized GRPO advantage, while the regularizer term includes a causal future-regularization return-to-go omitted by local KL losses. For reverse KL, this yields a simple future KL correction: add a reverse cumulative sum of per-token log ratios after advantage construction. The resulting method, Future-KL Regularized Policy Optimization (FRPO), requires no critic or extra model passes. On mathematical reasoning tasks, FRPO improves pass@16 in our main large-model setting while maintaining higher entropy and lower policy drift than conventional loss-side KL baselines.

1 Introduction

GRPO uses group-normalized outcome rewards and token-level KL regularization, but autoregressive KL creates a future credit-assignment signal that local penalties omit. Because group normalization induces a non-linear utility, reward and KL must remain separate, motivating FRPO’s post-advantage future-KL correction.

  • Motivation: GRPO samples multiple responses, scores them with outcome verifiers, and normalizes rewards within each prompt group without a learned critic.This critic-free design supports large-scale reasoning training.
  • The regularization gap: Local loss-side KL penalties do not capture how an early token changes future prefixes and later regularization costs.The missing signal is a causal future-regularization return-to-go.
  • Non-linear GRPO utility: GRPO differs from linear-reward KL-regularized RL because group normalization induces a non-linear utility, ℓ(p) = 2 arcsin √p for binary verifier rewards.This makes reward and KL non-fuseable before normalization.
  • Derived gradient: The on-policy gradient decomposes into a group-relative outcome advantage and a future regularization correction for token-wise f-divergence regularization.The reward term recovers the standardized GRPO advantage, while the regularizer includes a causal return-to-go.
  • FRPO: FRPO preserves the original outcome advantage and adds a token-level future-KL correction after advantage construction.This decouples trajectory-level ranking from process-level KL credit assignment.

2 Preliminaries and Notations

The preliminaries define autoregressive response generation, the current, reference, and rollout policies, token likelihood ratios, and sampled KL estimators. The commonly used k3 estimator is unbiased and nonnegative for reverse KL under on-policy sampling.

  • Autoregressive notation: An autoregressive policy generates response tokens o = (o1, o2, …, oT) conditioned on prompt x.The model parameters are denoted by θ.
  • Policy notation: πref is the frozen reference policy, while πθold is the older policy that generated rollouts for the current policy πθ.These policies distinguish reference regularization from rollout sampling.
  • Token quantities: The token log ratio δt compares current and reference token probabilities, while ρt compares current and rollout-policy probabilities.These ratios support KL estimation and PPO-style correction.
  • KL estimators: The framework considers token-level likelihood-ratio quantities and sampled estimators for reverse and forward KL divergences.The notation also distinguishes generated prefixes from remaining future tokens.
  • KL estimators: The k3 estimator is unbiased in value and nonnegative for reverse KL under on-policy sampling.The broader derivation applies to general token-wise f-divergence regularizers.

3 Classical GRPO and the Regularization Question

Classical GRPO assigns each response’s tokens a shared group-normalized outcome advantage and commonly adds a sampled token-level KL penalty. The paper argues that non-linear group utility makes KL-in-reward placement incorrect and that local KL losses miss future autoregressive credit assignment.

  • Classical GRPO: GRPO samples response groups, evaluates scalar outcome rewards, and optimizes a PPO-style surrogate with sampled token-level KL penalties.The ratio ρi,t compares current and rollout-policy token probabilities.
  • Group-normalized advantage: GRPO normalizes rewards within each prompt group and assigns every token in a response the same outcome advantage.This removes the need for a learned value model but leaves the outcome signal trajectory-level.
  • Implicit utility: For binary rewards, group normalization induces the non-linear utility ℓ(p) = 2 arcsin √p rather than the raw pass rate p.The standardized GRPO advantage is the score of this implicit utility.
  • KL placement: Because ℓ is non-linear, folding KL into rewards before normalization changes the implicit objective and is not equivalent to separately regularizing the policy.The discrepancy persists when the group has non-trivial KL spread.
  • Design implication: The proposed separation keeps outcome advantages tied to verifier rewards while assigning KL credit at the token level.This placement motivates the regularized objective and future-KL formulation.
  • Loss-side KL: A sampled k3 KL estimator can be unbiased in scalar value under on-policy sampling, but a local differentiable loss still omits future-prefix effects.A token changes later-state distributions and therefore later KL costs.

4 On-Policy f-Divergence Regularization

The on-policy gradient for GRPO-style objectives with token-wise f-divergence regularization combines the standardized reward advantage with local and future regularization terms. Because autoregressive actions affect later prefixes, local loss-side KL penalties omit a necessary causal return-to-go, and GRPO’s non-linear utility requires adding regularization after advantage construction.

  • General formulation: The on-policy gradient applies a differentiable convex f-generator to token-level likelihood ratios, with binary-reward GRPO using ℓ(p) = 2 arcsin √p.The framework covers token-wise f-divergences broadly, with KL as one special case.
  • Gradient decomposition: The gradient weight decomposes into a reward-side GRPO advantage and a regularization-side term containing local and future penalty contributions.The reward is scaled by the prompt-dependent ℓ′, while the regularizer is not.
  • Gradient decomposition: Changing an autoregressive action changes later prefix distributions, so every on-policy token-wise regularizer induces a causal future cost-to-go.The future contribution sums penalties from the current token onward; past costs vanish under the score-function identity.
  • Placement: The GRPO advantage and future-KL correction must remain separate because the non-linear ℓ′ scaling applies to reward but not genuine token-level regularization.The correction is inserted after advantage construction rather than fused into the normalized reward.
  • Loss-side omission: For reverse KL and k3, local differentiable losses capture only the direct token derivative, whereas the trajectory objective additionally includes future regularization terms.The missing term is therefore not specific to reverse KL, and k3’s value-estimator role does not remove its credit-assignment role.
  • Implementation scope: Importance weighting addresses rollout-policy mismatch but does not create the missing future regularization term, which remains necessary even on-policy.The off-policy caveat is implementation-level and orthogonal to the on-policy credit-assignment result.

5 Future-KL GRPO and Optimality View

FRPO implements the derived separation by adding a token-dependent future-KL correction to the original group-normalized verifier advantage. The correction is lightweight, preserves the reward normalization, and connects the policy-gradient derivation to a process-level interpretation of KL-regularized optimality.

  • The decoupled FRPO update: FRPO computes the group-normalized verifier advantage first, then adds a token-dependent future-KL correction.This decoupled update preserves GRPO’s outcome-reward normalization while restoring regularization credit assignment.
  • The decoupled FRPO update: The PPO update applies the same importance-ratio correction to the verifier and regularization signals after the future term is folded into the advantage.The finite-sample implementation uses rollout log-probabilities and the usual clipped PPO surrogate.
  • The decoupled FRPO update: FRPO requires about ten lines of code, one reverse cumulative sum over rollout log-ratios, no extra forward passes, and no learned value function.The structural requirement is placing the correction after advantage construction, not adding substantial computation.
  • The decoupled FRPO update: Alternative KL insertion strategies can alter group statistics or become fragile for small groups and variable-length responses.The proposed decoupled form keeps normalization on verifier rewards and adds future KL afterward.
  • Optimality view: Entropy-regularized optimality yields a reference-weighted exponential policy, and autoregressive likelihood-ratio decomposition converts this into a process-reward form.The optimal-policy view supplies a complementary interpretation of the future-KL structure.
  • Optimality view: The optimality interpretation identifies KL regularization as a process-level signal whose token weight is a future cost-to-go rather than an isolated penalty.FRPO uses this structure as a practical on-policy approximation with rollout log-probabilities and GRPO’s normalized advantage.

6 Experiments and Results

Experiments compare KL integrations, divergence directions, estimators, baseline subtraction, and advantage estimators across mathematical-reasoning settings. FRPO improves large-model performance while maintaining higher entropy and lower policy drift, whereas baseline subtraction alone does not improve performance.

  • Main Experiments: Nearly 5% absolute Pass@n gain over the strongest baseline is reported for FRPO on Qwen3-30B-A3B-Base.The gain is accompanied by higher final entropy and lower PPO-KL, indicating less policy drift from the reference policy.
  • Main Experiments: FRPO achieves better pass@n performance than direct loss-side KL on AIME24 and AIME25 for the large model.The Qwen3-1.7B-Base result is weaker, which the authors hypothesize may reflect proximity to its reasoning capacity under the current training budget and hyperparameters.
  • KL Direction and Estimator Ablations: Reverse KL works better than forward KL in the Qwen3-1.7B-Base setting.The authors suggest forward KL may keep entropy higher than is useful for mathematical reasoning.
  • KL Direction and Estimator Ablations: FRPO uses k1 by default because Figure 3 shows a gap between k1 and k3, especially for forward KL.Although k3 is commonly favored as a nonnegative, value-unbiased reverse-KL estimator, FRPO uses the future-KL term as a signed credit-assignment signal.
  • Subtracting a Baseline for k3 Estimators: Baseline subtraction alone does not improve performance for the nonnegative k3 estimator.The authors hypothesize that the estimated baselines are too coarse to approximate the true baselines accurately.
  • Different Advantage Estimators: Retaining GRPO’s outcome-level normalization is beneficial compared with integrating KL into rewards before estimating advantages.This comparison uses GAE and GRPO advantage estimation with future KL on Qwen3-30B-A3B-Base.

7 Conclusion and Discussion

The conclusion frames FRPO as an on-policy correction to local KL regularization in GRPO, extending the analysis to token-wise f-divergences. It separates group-normalized outcome advantages from future regularization credit and reports improved reasoning performance with limited policy drift and entropy collapse.

  • Conclusion: The policy-gradient formula applies to general token-wise f-divergence regularizers and reveals a future return-to-go term from autoregressive regularization.For reverse KL, this becomes the future-KL correction used by FRPO.
  • Conclusion: FRPO keeps GRPO’s group-normalized outcome advantage separate from token-level regularization credit.This separation follows from revisiting KL regularization through the on-policy objective rather than only as a local loss penalty.
  • Conclusion: In mathematical-reasoning experiments, the formulation improves pass@n while limiting policy drift and entropy collapse.The conclusion presents these as the reported empirical outcomes of the formulation.
  • Related Context: The paper situates its contribution within critic-free policy optimization and entropy-regularized reinforcement learning for LLM reasoning.The cited discussion references classical PPO and TRPO, entropy-regularized RL, and mathematical-reasoning applications.
  • Estimator Discussion: The estimator discussion covers reverse-KL estimators k1, k2, k3, and the exact full estimator, with k3 noted as common because of value unbiasedness and nonnegativity.The appendix states that k1 is unbiased but can be negative per token, k2 is generally biased as a value estimator, and the full estimator is exact.

B.1.2 Gradient Properties

The appendix compares KL estimators by their value and gradient properties under different sampling distributions. Reverse-KL sample estimators can become biased off-policy, while forward-KL estimators have distinct value and gradient guarantees.

  • Estimator-dependent gradients: Under πθold sampling, the expected gradient of a sampled KL estimator depends critically on the estimator choice.The reference policy is constant with respect to θ, but the rollout and current-policy distributions may differ.
  • Reverse KL: k1 has zero expected gradient under on-policy πθ sampling and is biased for reverse-KL gradients under πθold sampling.Its score-function gradient cancels under πθ, so it does not recover the desired reverse-KL gradient.
  • Reverse KL: k2 provides an unbiased reverse-KL gradient estimator on-policy and remains approximately unbiased when πθold is close to πθ.Under off-policy sampling, k2 also becomes biased, although its bias is generally smaller than k1 or k3.
  • Reverse KL: k3 combines an unbiased value estimate with a nonnegative penalty, but its expected gradient is generally biased for reverse KL.The straight-through k3+ construction instead combines k3’s value with k2’s gradient behavior.
  • Forward KL: For forward KL, k1 and k3 are value-unbiased under reference-policy sampling, whereas k2 is generally biased as a value and gradient estimator.k3 additionally remains nonnegative, while k1 may be negative for individual tokens.

C Implicit Arcsin Objective of Binary GRPO

For binary verifier rewards, population GRPO group normalization induces the variance-stabilizing utility 2 arcsin √p rather than the raw success probability. The standardized GRPO advantage is the score of this nonlinear utility, with finite groups introducing estimation noise and degenerate cases.

  • Implicit utility: For binary rewards, the implicit population GRPO utility is ℓ(p) = 2 arcsin √p.This transform is the variance-stabilizing transform for Bernoulli means.
  • Finite groups and degeneracy: Finite groups replace population statistics with empirical estimates, creating finite-sample noise and degenerate groups when all sampled rewards match.When σG = 0, implementations typically mask or skip the prompt.
  • Population gradient: Group normalization removes the direct pθ(x)(1 − pθ(x)) reward-variance factor from the population gradient.The resulting transform stabilizes the Bernoulli-mean variance dependence.
  • Population gradient: The standardized GRPO advantage equals ℓ′(pθ(x))(r(x, a) − pθ(x)).It therefore represents the score of the arcsin utility rather than an untransformed reward score.

D.1 Proof Sketch of Proposition 3.2

The proof shows that KL-shaped rewards do not preserve GRPO’s intended nonlinear utility, while autoregressive token regularization generates a causal future cost-to-go. FRPO therefore keeps outcome normalization separate and adds future KL on the regularizer side.

  • Non-commutativity: The reward gradient carries a prompt-dependent ℓ′ gain, but the regularizer gradient does not.This asymmetry prevents folding KL into the reward before group normalization without changing the implicit objective.
  • Non-commutativity: KL-in-reward changes the group statistics and generally does not target ℓ(E[r]) − βE[DKL] under nonlinear GRPO utility.That equivalence holds only for affine ℓ, whereas GRPO uses ℓ(p) = 2 arcsin √p.
  • Causal future regularization: In autoregressive generation, changing an earlier token changes the distribution of later prefixes and their KL costs.The resulting regularization signal is a causal future regularization return-to-go, not merely a local token penalty.
  • FRPO construction: FRPO computes the original GRPO outcome advantage first and then adds a token-wise future-KL correction without modifying group reward statistics.The algorithm uses the resulting corrected weights in PPO updates and requires no additional model passes in the supplied procedure.

F.1 Why Not Put KL into the Reward Before Normalization?

The section rejects coupling KL with rewards before or during group normalization because it changes GRPO’s nonlinear objective or yields unstable process-level statistics. FRPO instead decouples trajectory ranking from token-level KL credit assignment.

  • Coupled trajectory-level strategy: Trajectory-level KL shaping makes every token in a response share one advantage and changes the group mean, standard deviation, scale, and ordering.This discards the fine-grained future-KL credit assignment that process rewards are intended to provide.
  • Coupled step-level strategy: Step-level KL normalization can have high variance because later time steps contain fewer contributing trajectories, especially with variable response lengths.The resulting normalization constants can produce erratic advantages late in responses.
  • Decoupled FRPO strategy: FRPO preserves the implicit 2 arcsin(√p) objective by normalizing the original outcome rewards before adding future KL.Trajectory ranking uses the stable GRPO advantage, while token differentiation uses an within-trajectory future KL term without cross-trajectory normalization.
  • Experimental setting: The experiments use Qwen3 base models for mathematical reasoning, with training dynamics for Qwen3-4B-Base shown in Figure 5.The passage states that the trend is similar for Qwen3-1.7B-Base and Qwen3-30B-A3B-Base.

G.2 Other Evaluation Metrics

FRPO improves mean@n performance alongside pass@n on the evaluated reasoning benchmarks. Its future-KL formulation is especially helpful for stabilizing training on the Qwen3-30B-A3B-Base model, while directly placing reverse KL in the reward leads to model collapse.

  • Mean@n performance: FRPO improves pass@n performance and achieves comparable or better mean@n accuracy than methods using loss-side or reward KL integration.The comparison is reported through training dynamics for Qwen3-1.7B-Base, Qwen3-4B-Base, and Qwen3-30B-A3B-Base.
  • Mean@n performance: On Qwen3-30B-A3B-Base, future KL helps stabilize training and produces a smoother training process.The paper attributes this setting to more pronounced off-policy issues associated with the model’s MoE structure.
  • Evaluation setup: The evaluated benchmarks include MATH500, AIME24, AIME25, AMC23, OlympiadBench, and MinervaMath.Tables 1 and 9 summarize pass@16 and mean@16 results across models and KL integrations.
  • KL integrations: Using reverse KL directly in the reward leads to model collapse.This comparison is noted in the evaluation summary covering the different KL integrations.
  • Qualitative examples: Models trained with future KL typically produce longer responses while maintaining the reference-policy constraint.The paper presents this as evidence that FRPO can elicit long chain-of-thought behavior from the base model.
Loading 2601.10201v2…