cs.LGOct 16, 2026

On Reparameterizing the Score Function

Authors: Mohammad Hossein Amani, Leonardo Bocchieri, Maxime Peyrard, Clément Hongler, Robert West

Organizations: EPFL · CNRS, Univ. Grenoble Alpes · XENT Labs

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Policy gradients for autoregressive sequence models differentiate each token's log-probability while holding the conditioning history fixed. This discards information about how earlier decisions shape later conditional distributions, creating a credit-assignment bottleneck when rewards are sparse. We introduce History-Reparameterized (HR) policy gradients, which replace the frozen prefix in the score function with a differentiable reparameterized or relaxed history, allowing gradients of downstream log-probabilities to propagate backward through the generated trajectory. The resulting history-pathwise term is generally biased when added directly, so we derive Optimal HR (OHR), a centered control-variate estimator with a closed-form coefficient that minimizes the trace of the gradient variance at the population level. The variance reduction is governed by temporal structure in the policy dynamics: random or weakly structured policies provide little benefit, while policies with learned sequential dependencies yield informative history-pathwise corrections. For discrete tokens, the choice of relaxation is crucial: embedding-level relaxation yields substantial variance reduction that grows with horizon, whereas straight-through Gumbel-Softmax yields essentially none in our diagnostics. We validate the method on a Gaussian linear Markov model, a sparse-reward T-maze with a recurrent policy, and an autoregressive Transformer model for small Traveling Salesman Problem (TSP) instances.

Figures & tables

Figure 1: Embedding-based HR ( V=8V=8 , de=16d_{e}=16 ). Left: variance reduction ratio grows with horizon NN and temperature τ\tau . Center: PG variance (blue) grows with NN while OHR variance (green) remains flat. Right: α1∗≈1.1\alpha_{1}^{*}\approx 1.1 is stable across horizons, and the correction is anti-correlated with PG ( cos⁡≈−0.6\cos\approx-0.6 ).
Figure 2: T-maze training: PG (blue) vs. OHR (orange), 10 seeds per NN , τ=1.0\tau{=}1.0 . Top: reward curves with individual seed traces (thin) and mean ±\pm s.e. (thick). Bottom: estimator variance during training (log scale). As the corridor lengthens, PG increasingly fails to converge while OHR remains reliable.
Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 3: Gaussian linear Markov model ( d=5d=5 , ρ(W)=0.95\rho(W)=0.95 ). Top: linear reward. Bottom: L2L^{2} reward. Left: variance reduction ratio (OHR over baselined PG), growing monotonically with NN to 1.08×1.08\times (linear) and 1.18×1.18\times ( L2L^{2} ). Center: SNR of both estimators. Right: optimal coefficient α1∗\alpha_{1}^{*} .
Figure 4: Continuous Embedding training: PG (blue) vs. OHR (red), 10 seeds per NN , τ=3.0\tau{=}3.0 . Top: reward curves with mean ±\pm s.e. Bottom: estimator variance during training (log scale). OHR consistently achieves better training dynamics and lower variance.
Figure 5: T-maze with GRU policy and embedding-based action feedback (10 seeds). Left: variance reduction ratio grows monotonically with corridor length NN across all τ\tau . Center: PG variance (blue) grows steeply with NN while OHR variance (green) grows much slower (log scale). Right: α∗≈−2\alpha^{*}\approx-2 and cos⁡(gPG,C)≈0.3\cos(g_{\mathrm{PG}},C)\approx 0.3 are stable across NN .
Figure 6: TSP training: REINFORCE (blue) vs. OHR (red), M=16M{=}16 , τ=1.0\tau{=}1.0 , 4 seeds. Top: tour length converges to similar final values. Bottom: OHR gradient variance (red) is consistently lower than REINFORCE (blue), with VR annotations showing 1.151.15 – 1.25×1.25\times reduction.
Figure 7: T-maze composability ( τ=1.0\tau{=}1.0 , 10 seeds). At N=10N{=}10 , OHR (red) solves the task faster than REINFORCE (blue) regardless of baseline. Solid: mean baseline; dashed: learned critic.
Figure 8: TSP composability ( τ=1.0\tau{=}1.0 , M=16M{=}16 , 4 seeds). OHR + RLOO (dashed red) achieves the lowest variance at all problem sizes, with combined VR of 3.43.4 – 3.5×3.5\times at n=8n{=}8 – 1010 .
More in this topic

Policy Gradient

Policy gradient methods improve a decision-making policy by estimating how its actions affect expected reward.

See all 69 papers

Explore similar work

CardsList
Oct 23, 2026

Direction-Aware Token Credit Assignment with Reward-Model Gradients

GRPO shares a response-level advantage across tokens, leaving unresolved how strongly each token should contribute to learning. Updates induced by the shared advantage reinforce or suppress the sampled token at each position, but these updates do not always point toward increasing reward at the token level. This motivates a work-based view of credit: just as work depends on the component of force along a displacement, token credit should reflect the local reward change along the policy update direction. We introduce Direction-Aware Token Credit Assignment (DATC), which maps frozen reward-model gradients into policy logit space and allocates weights by evaluating the alignment between the reward direction and GRPO's signed sampled-action update direction, while preserving the sign and mean of the response-level advantage. DATC outperforms GRPO in all four settings across two model families.

Yikuan Hu, Jifeng Zhu, Minle Su +6

Nanjing University · National University of Singapore · Tsinghua University +1

Token-Level Credit AssignmentPolicy Gradient MethodsGroup Relative Policy OptimizationGradient-Based Attribution
AskPDF
Mar 15, 2026

Delightful Policy Gradient

Standard policy gradients weight each sampled action by advantage alone, regardless of how likely that action was under the current policy. This creates two pathologies: within a single decision context (e.g. one image or prompt), a rare negative-advantage action can disproportionately distort the update direction; across many such contexts in a batch, the expected gradient over-allocates budget to contexts the policy already handles well. We introduce the Delightful Policy Gradient (DG), which gates each term with a sigmoid of delight, the product of advantage and action surprisal (negative log-probability). For KK-armed bandits, DG provably improves directional accuracy in a single context and, across multiple contexts, shifts the expected gradient strictly closer to the supervised cross-entropy oracle. This second effect is not variance reduction: it persists even with infinite samples. Empirically, DG outperforms REINFORCE, PPO, and advantage-weighted baselines across MNIST, transformer sequence modeling, and continuous control, with larger gains on harder tasks.

Ian Osband

Google DeepMind

Policy GradientPolicy Gradient Methods
AskPDF
Jun 15, 2026

A First-Principles Derivation of LLM Policy Optimization: From Expected Reward to GRPO and Its Structural Extensions

Policy gradient algorithms for language models optimize the same objective J(θ)=E∗τ∼p∗θ(τ)[R(τ)]J(θ) = \mathbb{E}*{τ\sim p*θ(τ)}[R(τ)], which has exactly two factors: the trajectory probability pθ(τ)p_θ(τ) and the reward R(τ)R(τ). Every method from REINFORCE to PPO to GRPO and their descendants modifies one or both factors to address a specific failure in the preceding formulation. Existing surveys organize these methods by domain or chronology, which obscures the rationale behind each design choice and the precise location of its intervention within the gradient estimator. This survey revisits the landscape of LLM policy optimization from J(θ)J(θ) on first principles and uses the trajectory side, induced by pθ(τ)p_θ(τ), and the reward side, induced by R(τ)R(τ), as the two axes along which methods are located. It covers the path from REINFORCE and PPO to GRPO, as well as post-GRPO variants, Agentic RL, and GRPO-OPD. The resulting framework is unified, diagnostic, and extensible: it analyzes methods from a shared objective, identifies which side each method modifies and why, and applies the same trajectory and reward axes across these settings. Across these settings, the framework also exposes compound failures that no single-side fix resolves and that therefore require joint design of the trajectory side and the reward side. The boundary cases and coupled failures identified by this map mark where existing solutions run out and provide a principled starting point for designing the next generation of LLM policy optimization algorithms.

Jianghan Shen, Siqi Luo, Yue Li +9

1Nanjing University · 2Shanghai Artificial Intelligence Laboratory · 3Shanghai Jiao Tong University +4

RL for Language ModelsGroup Relative Policy OptimizationPolicy Gradient MethodsProximal Policy OptimizationPolicy Learning
AskPDF