Post-Training with Policy Gradients: Optimality and the Base Model Barrier

Abstract: We study post-training linear autoregressive models with outcome and process rewards. Given a context $\boldsymbol{x}$, the model must predict the response $\boldsymbol{y} \in Y^N$, a sequence of length $N$ that satisfies a $γ$ margin condition, an extension of the standard separability to sequences. We prove that on test samples where the base model achieves a non-trivial likelihood $α$, a variant of policy gradient (PG) can achieve likelihood $1 - \varepsilon$ with an essentially minimax optimal number of reward queries $\tilde{O}((α^{-1} + \varepsilon^{-1})/γ^2)$. However, a barrier arises for going beyond the support of the base model. We prove that the overall expected error after post-training with outcome rewards is governed by a property of the base model called the Likelihood Quantile (LQ), and that variants of PG, while minimax optimal, may require a number of reward queries exponential in $N$ to go beyond this support, regardless of the pre-training algorithm. To overcome this barrier, we study post-training with a process reward model, and demonstrate how PG variants in this setting avoid the curse of dimensionality in $N$ via dependence on a token-level LQ. Along the way, we prove that under the margin condition, SGD with adaptive learning rate (LR) achieves a near optimal test error for statistical learning, and PG with adaptive LR achieves a near optimal number of mistakes for online learning while being computationally efficient whenever possible, both of which may be of independent interest.

• The paper establishes that outcome-based policy gradients cannot significantly improve off-support performance without an exponential number of reward queries.
• It provides minimax optimality guarantees for SGD and PG variants, detailing nonasymptotic convergence rates and tight complexity bounds.
• The study demonstrates that token-level process rewards break the base model barrier by reducing dependence on sequence length from exponential to linear.

Post-Training with Policy Gradients: Base Model Barriers and the Role of Process Rewards

Overview

The paper "Post-Training with Policy Gradients: Optimality and the Base Model Barrier" (2603.06957) rigorously analyses the theoretical limits and capabilities of post-training large autoregressive models—such as LLMs—using policy gradient (PG) based RL, focusing in detail on the crucial influence of the pre-trained "base model". It identifies a core phenomenon: if the base model's likelihood for producing correct sequences is not sufficiently large, outcome-reward RL fine-tuning cannot efficiently improve expected error below the level achieved in pre-training, unless the reward signal is provided at the process (token) level. The work provides nonasymptotic convergence and optimality guarantees for SGD and PG variants, formal lower bounds, and demonstrates consequences when using outcome- versus process-level rewards.

Problem Setting and Main Contributions

The authors consider sequence-level prediction $y$ of length $N$ from context $x$, with a token-level $\gamma$-margin assumption guaranteeing the existence of a separating hyperplane in feature space. The key questions are:

• How do PG convergence rates (in sample and reward query complexity) depend on the likelihood assigned by the base model to the ground-truth sequence?
• Can RL post-training achieve a substantially lower test error than the base model, and if so, under what conditions and at what cost (e.g., in the number of reward queries)?

Under outcome-reward models (ORM; reward after generating a full sequence), it is demonstrated that the number of reward queries required to boost the expected likelihood above a base model scales at least inversely with the Likelihood Quantile (LQ) of the base model. The LQ directly encodes the distribution of base model likelihoods on ground-truth pairs.

Strong result: For typical SGD pre-trained base models, LQ is exponentially small in $N$ for worst-case samples, so post-training cannot substantially beat the base model in expected error without an exponential number of reward queries as $N$ grows.

In contrast, with process-reward models (PRM; reward for correctness of each prefix), the dependence is only linear in $N$. Thus, PRMs can break the "base model barrier" and allow efficient exploration beyond pretraining support.

Additional contributions:

• Minimax lower bounds: Both reward query and sample complexity bounds for PG (on top of a base) are shown tight.
• SGD (with adaptive LR) achieves minimax rates for supervised pre-training.
• Efficient online PG algorithm with optimal mistake bounds for contextual bandit feedback in autoregressive settings.

Policy Gradient Post-Training: Outcome Rewards

With outcome rewards, post-training is formalized as a contextual bandit: the reward $r(x, y)$ is $1$ if the sequence is correct, $0$ otherwise. The main theoretical finding is that policy gradient can only efficiently improve test accuracy on samples for which the base model already achieves a non-negligible likelihood (on-support samples).

On such samples, the likelihood can be boosted to $1 - \varepsilon$ with $$\tilde{\mathcal{O}}((\alpha^{-1}+\varepsilon^{-1})/\gamma^2)$$ reward queries ($\alpha$ is the initial likelihood). However, off-support samples—those with exponentially small likelihood under the base—cannot be improved efficiently.

Figure 1: PG with ORM cannot increase likelihood on off-support samples, whereas with PRM, the average likelihood improves; expected test error plateaus with ORM but decreases with PRM; sample-wise trajectories show ORM gets stuck on initially low-likelihood samples.

The global expected error after post-training is governed by the base model's likelihood quantile function $Q_q(\epsilon)$—the minimum value such that a fraction $\epsilon$ of the base model's test likelihoods fall below. The sample complexity for expected error $\epsilon$ is lower-bounded (and matched) by $$\tilde{\mathcal{O}}(Q_q(\epsilon)^{-1}/(\gamma^2\epsilon))$$.

Base Model Barrier: For pre-trained base models from SGD, $Q_q(\epsilon)$ is often $O(k^{-N})$. This makes sub-supervised error unattainable by outcome-reward RL unless exponential computation is permitted.

Figure 2: CDF and quantile evolution for base model likelihoods as SGD pre-training length increases; quantile improves towards 1 for fixed $\varepsilon$ as training improves.

Sample/Reward Query Complexity

• Best-of-$m$ exploration: If allowed to sample multiple times (from base model or uniform), the number of reward queries still scales inversely with LQ; that is, exponentially in $N$ in the worst case, matching the model coverage profile.
• Lower bounds: Any post-training RL algorithm with access to a base model cannot outperform these scaling laws.

Breaking the Barrier: Process Rewards

When intermediate rewards are available for correct token prefixes, the landscape changes fundamentally. The key technical object is the Token-level Likelihood Quantile (TLQ), which considers the worst likelihood at any token step. For a uniform base, TLQ is $k^{-1}$, not $k^{-N}$ as with LQ.

Under PRM, the number of reward queries to reach error $\epsilon$ is only $\tilde{\mathcal{O}}(Nk+\epsilon^{-1})/\gamma^2$ even if the base model fails to model most sequences well, and only grows linearly with sequence length.

Implication: PRMs make it feasible for RL to induce sequences outside the base model's support with polynomial effort, given access to token-level evaluators.

Statistical and Online Optimality

The paper precisely characterizes the minimax sample and reward query complexity for both online and statistical learning settings. It proves that:

• SGD with adaptive LR is minimax optimal for pre-training under the $\gamma$-margin assumption.
• PG with outcome/process rewards (plus variants with appropriate exploration strategies) is also minimax optimal for post-training, given the base model's coverage profile.
• There exist fundamental computational limits: for reasonable $k$ and $N$, expected error below the base model's is unattainable without exponentially many reward queries in the absence of process feedback.

Experimental Validation

Experiments on synthetic data verify and concretize the theoretical predictions:

• Outcome-reward PG (ORM) cannot increase likelihood on off-support samples; test error plateaus at the base model level (see Figure 1).
• Process-reward PG (PRM) improves likelihood across the support, outperforming ORM and reducing expected error further as sequence length increases.
• Quantile plots of base model likelihoods empirically match the predicted exponential scaling in $N$ for outcome rewards (Figure 2).

Implications and Future Directions

This analysis clarifies and strengthens theoretical understanding of RL post-training for large autoregressive models. It draws a sharp line between what can and cannot be achieved with outcome-based reward models, and rigorously motivates efforts on designing (and learning) process-based reward signals in RLHF and related settings.

Practical implication: In domains (e.g., code, math) where it is possible to define and generate process rewards efficiently, RL post-training can scale to generate new behaviors not seen or likely under the pre-trained base model. For most text generation tasks, unless such rewards can be synthesized, models are limited by their pre-training coverage.

Speculation: Future research directions include developing methods to learn process reward models efficiently, extending results to non-separable or noisy regimes, and bridging coverage gaps through reward modeling or more structured interventions during post-training.

Conclusion

The work provides a detailed, quantitative, and minimax-optimal theoretical framework for understanding the limits of RL-based post-training of autoregressive models. It proves that outcome reward-based RL is strictly limited by base model support, and only process reward models allow RL to efficiently generate and learn truly novel, low-probability behaviors outside pre-trained distributions.