On Reparameterizing the Score Function
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.
Related
- Optimal score function estimation via derivatives constraints
- Nonlinear Assimilation via Score-Based Sequential Langevin Sampling
- Notes on Uniform Policy Score Bounds and Lipschitz Continuity for Truncated Gaussian, Product Cauchy, Boltzmann, and Softmax Policies
- Keeping Score: Adaptive, Tuning-Free Loss Weighting for Score-Augmented Neural Ratio Estimation
- Covariance-Informed Ensemble Score Filtering: Embedding Localized Covariance Geometry into Score-Based Transport
BibTeX
@misc{amani2026reparameterizing,
title = {{On Reparameterizing the Score Function}},
author = {Mohammad Hossein Amani and Leonardo Bocchieri and Maxime Peyrard and Clément Hongler and Robert West},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.24188},
eprint = {2610.24188},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.24188},
note = {CC BY-NC-SA 4.0},
}
RIS
TY - UNPB
TI - On Reparameterizing the Score Function
AU - Mohammad Hossein Amani
AU - Leonardo Bocchieri
AU - Maxime Peyrard
AU - Clément Hongler
AU - Robert West
PY - 2026
DA - 2026/10/16
PB - arXiv
DO - 10.48550/arxiv.2610.24188
UR - https://arxiv.org/abs/2610.24188
AB - 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.
ER -