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

arXiv:2603.06957v2 Announce Type: replace-cross Abstract: We study post-training linear autoregressive models with outcome and process rewards. Given a context oldsymbol{x}, the model must predict the response oldsymbol{y} in Y^N, a sequence of length N that satisfies a gamma 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 alpha, a variant of policy gradient (PG) can achieve likelihood 1 - arepsilon with an essentially minimax optimal number of reward queries ilde{O}((alpha^{-1} + arepsilon^{-1})/gamma^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.

Source: arXiv cs.AI | 2026-08-13