On Reparameterizing the Score Function

AI summary

Overview

Research area: Reinforcement learning and gradient estimation — specifically policy-gradient methods for autoregressive sequence models, and variance reduction via control variates.

Technical level: Advanced. The paper is written for readers comfortable with REINFORCE, the reparameterization trick, and control variates; it is accessible to a determined intermediate reader willing to work through the derivations.

Scope (one sentence): The paper proposes History-Reparameterized (HR) policy gradients — which replace the frozen conditioning prefix inside the score function with a differentiable reparameterized or relaxed history — derives an unbiased, variance-optimal centered control-variate version (OHR), and tests it on a Gaussian linear Markov model, a sparse-reward T-maze with a GRU policy, and a small attention model for the Traveling Salesman Problem.

What This Paper Is About

Policy-gradient methods for autoregressive models differentiate each token's log-probability while holding the preceding tokens fixed, so a single scalar reward applied to a whole sequence is spread uniformly over all tokens and the model never sees how early decisions shaped later conditional distributions. The authors argue that in environment-less generation — unlike standard RL, where the environment is genuinely non-differentiable — this frozen-history treatment is a convention inherited from the MDP framework rather than a structural necessity. Their goal is to expose the gradient path through the model's own generated prefix, injecting credit-assignment information without requiring a differentiable reward.

Key Contributions

1. The History-Reparameterized (HR) estimator. The frozen history $x_{<t}$ in the score function is replaced by a differentiable reparameterized or relaxed history $\chi_{<t}(\theta)$. The chain rule splits the estimator into the usual score term $S_t$ and a new history-pathwise term $C_t$ that backpropagates gradients of downstream log-probabilities through the generated prefix. Because only the conditioning argument is modified, the construction applies to any algorithm built on the score function (REINFORCE, PPO, and descendants) and needs no differentiable reward.

2. An exact bias characterization (Proposition 1). The raw HR estimator is shown to be biased, with the bias written as the expectation of the reward-weighted difference between the reparameterized-history score and the frozen-history score.

3. Optimal HR (OHR), a principled control variate (Theorem 2). Centering the history-pathwise term yields a zero-mean control variate, and a closed-form scalar coefficient $\alpha_1^* = -\operatorname{Tr}(\operatorname{Cov}(g_{\mathrm{PG}}^{b}, \bar{H})) / \operatorname{Tr}(\operatorname{Var}(\bar{H}))$ minimizes the trace of the estimator variance. With exact centering the estimator is unbiased for any fixed $\alpha_1$ and the trace-variance can only decrease; Proposition 5 shows the fractional reduction equals a squared trace-correlation $\rho^2(g_{\mathrm{PG}}^{b}, \bar{H})$. A practical EMA-based centering scheme is given because centering with the current batch mean would cancel the correction entirely.

4. Evidence that the relaxation choice is decisive and that gains scale with horizon. Embedding-level relaxation produces large, horizon-growing variance reduction, while straight-through Gumbel-Softmax produces essentially none; the effect is demonstrated in three settings and shown to compose with reward-side methods.

Main Findings

• Embedding-level relaxation works; straight-through Gumbel-Softmax does not. In the minimal autoregressive model ($V=8$, $d_e=16$, embedding layer plus linear head), ST-Gumbel gives zero variance reduction at all temperatures and horizons (ratio $\le 1.01$), because hard one-hot tokens are used in the forward pass and the relaxed and true histories coincide. Continuous embedding reparameterization instead reaches $1.5\times$ at $N=3, \tau=1$ and $4.9\times$ at $N=15, \tau=3$ (the introduction rounds this to "approximately $5\times$").
• PG variance grows with horizon while OHR variance stays roughly flat. The optimal coefficient is stable at $\approx 1.1$–$1.3$ across horizons, and the correction is anti-correlated with the policy gradient ($\cos \approx -0.6$), confirming it behaves as an effective control variate.
• Closed-form validation on a Gaussian linear Markov model. For $N=1$ with a linear reward, the HR bias is exactly $W^{\top} a, x_{-1}^{\top}$ — precisely the history-pathwise component of $\nabla_W J$ — so the base estimator captures the credit-assignment signal but double-counts it, which the OHR correction resolves. Monte Carlo runs with $M=50{,}000$ show the OHR variance-reduction ratio growing monotonically with horizon, reaching $1.18\times$ at $N=20$ for the $L^2$ reward; in the scalar case the limiting reduction is $\rho^2 \to w^2/(1+w^2)$ as $N\to\infty$ for $|w|<1$.
• Benefits require learned temporal structure. On frozen pretrained T-maze GRU policies ($h=16$, pretrained to $\sim 0.78$ reward, $M=1{,}024$, 10 seeds), OHR reaches $3.4\times$ variance reduction at $N=30$ ($\tau=2.0$) and $3.1\times$ ($\tau=1.0$), with the ratio rising monotonically in corridor length. The authors tie this to their prediction that HR helps most in the fine-tuning regime, where sequential dependencies are already learned.
• Variance reduction translates into reliability on T-maze. With 10 seeds, $\tau=1.0$, $M=128$: at $N=10$ both methods solve the task; at $N=15$ PG fails on some seeds while OHR succeeds on all; at $N=20$ PG succeeds on half the seeds while OHR solves most (mean final reward $0.75$ vs. $0.90$). OHR keeps a $3.1\times$ variance ratio during active learning at $N=15$, and PG seeds stuck at $\sim 0.5$ reward illustrate the failure mode.
• Smaller but real gains on TSP. With a simplified Attention Model ($\sim 120$K parameters), both methods converge to similar final tour lengths, but OHR has $1.15$–$1.25\times$ lower gradient variance throughout training. The effect is smaller than in the embedding model or T-maze, attributed to the attention decoder's sensitivity to context perturbations; small batches matter ($M=16$ makes the gain visible, $M=128$ renders it negligible).
• Composability with reward-side methods is near-multiplicative on TSP. At $n=10$: RLOO alone gives $2.45\times$, OHR alone $1.27\times$, and OHR + RLOO $3.55\times$ versus a product prediction of $3.11\times$. On T-maze with a learned critic, OHR reduces variance at all corridor lengths tested, up to $1.70\times$ at $N=15$.
• The method does not help in every configuration. With a batch-mean baseline the PG/OHR variance ratio is $0.86\times$ for T-maze at $N=10$ and $0.94\times$ for TSP at $n=8$, i.e., OHR is slightly worse.
• Cost is modest. The history-pathwise term adds one backward pass through the reparameterized trajectory, $O(N)$ compute relative to REINFORCE, and the BPTT through history can be truncated to $K<N$ steps to trade bias for memory.

Methodology in Plain English

The starting point is the standard REINFORCE estimator, which multiplies a reward-minus-baseline scalar by the sum of per-token score terms. The authors' move is to break the assumption that the prefix feeding each token's probability is a constant. Instead, they re-draw the whole trajectory as a deterministic, differentiable function of fixed random noise — the same trick used to train variational autoencoders — so that each generated token depends on the parameters both directly and through every earlier token.

Applying the chain rule to this re-drawn trajectory produces two pieces: the ordinary score term, and an extra "history-pathwise" term that asks how the downstream log-probabilities would change if an earlier step had been slightly different. Crucially, the reward itself is never differentiated — it is treated as a fixed scalar weight — so binary or otherwise non-differentiable rewards remain usable.

That extra term turns out to be biased. The fix is standard control-variate practice: subtract the term's mean so it has zero expectation, then scale it. The authors solve for the scaling that minimizes the trace of the resulting gradient variance, obtaining a closed-form ratio of covariances to variances. In practice the mean must be tracked with an exponential moving average rather than estimated from the current batch, because subtracting the current batch mean would make the correction average to exactly zero.

For discrete tokens, the differentiable history has to be manufactured by relaxation. Two options are compared: straight-through Gumbel-Softmax, which uses hard one-hot tokens in the forward pass and soft gradients only in the backward pass, and embedding-level relaxation, which feeds a soft weighted average of embedding vectors into the next step so the forward pass itself differs. The paper argues the latter is what makes the correction non-trivial.

Empirically, the authors first use a Gaussian linear Markov policy where the gradient, bias, variance, and optimal coefficient all have closed forms, then move to a minimal discrete autoregressive model to isolate the relaxation question, a sparse-reward T-maze with a GRU to test credit assignment over long horizons, and a small Transformer-style attention model for TSP. Variance is measured with frozen-policy diagnostics over large batches of per-sample gradients, separately from the variance of the batch-averaged training update, and supplemented with actual training runs from identical initializations.

Why This Matters

Impact on research. The paper reframes a near-universal convention — detaching the conditioning history in the score function — as a choice inherited from the MDP framework rather than a requirement of environment-less generation. It opens a variance-reduction axis that is orthogonal to and composable with reward-side techniques (baselines, critics, RLOO, process reward models), and it connects sequence-model policy gradients to the stochastic-computation-graph view in which history edges can be upgraded from detached stochastic nodes to reparameterized deterministic ones. The negative result on straight-through Gumbel-Softmax is itself useful, since that is the default relaxation many practitioners would reach for.

Real-world applications (the paper's own motivating settings and the domains its benchmarks proxy — note that the authors do not evaluate on natural language):

• Fine-tuning reasoning models where the reward is a single binary correctness signal for a multi-step answer.
• Neural combinatorial optimization, where tour construction is inherently sequential and the reward is terminal — the paper's TSP experiments.
• Sequential decision tasks with sparse terminal rewards and long credit-assignment horizons, exemplified by the T-maze.
• Any pipeline built on the score function where the reward is a black box (unit tests, simulators, human judgments) and therefore cannot be differentiated.

Industry relevance. The method is a drop-in modification to the conditioning argument of an existing policy-gradient implementation, requires no differentiable reward and no new model, and adds only $O(N)$ compute that can be truncated. Its composability with existing baselines means it can be layered onto current training stacks rather than replacing them. The caveat is that the paper's evidence comes from small models, so industrial adoption at language-model scale is not yet supported by the reported results.

Future Directions

• Scaling to large models and natural language. The authors explicitly state they do not evaluate on natural-language tasks and that whether HR remains beneficial in high-dimensional models with large parameter counts remains to be tested.
• Combining with learned per-token baselines and PPO-style value functions in large autoregressive models, which the paper lists as untested.
• A formal measure of the temporal structure that governs $\rho^2$. The link between policy Jacobian structure and variance reduction is characterized informally; a general-architecture analogue of the scalar $w^2/(1+w^2)$ result is posed as an open question.
• Richer coefficient estimators. The scalar trace-optimal $\alpha_1^*$ is only a scalar summary and does not minimize variance component-wise; diagonal or matrix-valued coefficients are suggested as a possible further gain, and the MSE-minimized two-coefficient variant PHR is deferred to an appendix. Better estimation of the centering mean beyond the EMA is also left open.

Target Audience

Reinforcement learning researchers working on gradient estimation and variance reduction; LLM post-training practitioners interested in credit assignment for sparse, non-differentiable sequence-level rewards; and anyone building on stochastic computation graphs or discrete gradient estimators (REBAR, RELAX, DiCE). Readers should be comfortable with the score-function estimator and the reparameterization trick, since the paper's core argument is a chain-rule decomposition of the score function.

Authors’ 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.