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

TL;DR

This paper analyzes the effectiveness of post-training policy gradient methods for autoregressive models, identifying optimal conditions and barriers related to model support and proposing solutions using process rewards.

Contribution

It introduces a theoretical framework for post-training with policy gradients, revealing barriers beyond the base model support and proposing process rewards to overcome these limitations.

Findings

Policy gradient can achieve near-perfect likelihood with minimal reward queries on test samples.

A barrier exists for going beyond the base model support, requiring exponential queries.

Using process rewards, PG variants can avoid the curse of dimensionality in sequence length.

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 $\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 - \varepsilon$ with an essentially minimax optimal number of reward queries $\tilde{O}((\alpha^{-1} + \varepsilon^{-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.