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

📝 Paper Summary

Reinforcement Learning for LLMs Theoretical Analysis of RL Process vs. Outcome Rewards

Theoretical analysis proves that outcome-based RL is inefficient for generating responses outside the base model's support, a barrier that process rewards overcome by leveraging token-level feedback.

Core Problem

RL with outcome rewards (RLHF) typically fails to generate correct responses for prompts where the base model has negligible likelihood (off-support), merely sharpening the existing distribution.

Why it matters:

Current LLM post-training heavily relies on outcome rewards, but it is unclear if this can genuinely create new knowledge or just refine existing capabilities
Understanding the theoretical limits of RL post-training is crucial for designing algorithms that can generalize beyond pre-training data
Prior theory works on RL post-training often ignore the specific role of the base model's initial support in determining sample complexity

Concrete Example: Consider a sequence of length N where the correct response has exponentially small probability under the base model (e.g., a complex reasoning chain). Outcome-based PG requires exponentially many samples to learn this, effectively failing, while process-based PG can learn it with samples linear in N.

Key Novelty

The Base Model Barrier and Likelihood Quantile

Identifies a 'Likelihood Quantile' (LQ) property of the base model that governs the sample complexity of outcome-based post-training
Proves that while outcome rewards face an exponential barrier for off-support samples, process rewards (verifying each token) break this curse of dimensionality

Evaluation Highlights

Outcome-based PG achieves error ε with Õ(1/(αγ²ε)) samples, efficient only if base likelihood α is non-trivial (polynomial in N)
Online PG with uniform policy achieves minimax optimal mistake bound Õ(k^N/γ²), matching information-theoretic limits
Process rewards reduce the worst-case sample complexity dependence on sequence length N from exponential to linear

Breakthrough Assessment

9/10

Provides a rigorous theoretical foundation explaining the empirically observed limitations of RLHF (outcome rewards) and formally proving the necessity of process rewards for true generalization.

⚙️ Technical Details

Problem Definition

Setting: Contextual bandit problem with linear autoregressive models under a margin condition

Inputs: Context vector x in X

Outputs: Sequence of tokens y = (y_1, ..., y_N) in Y^N

Pipeline Flow

Input Context x
Linear Autoregressive Policy (generates y token-by-token)
Reward Mechanism (Outcome or Process)
Policy Update (PG)

System Modules

Linear Autoregressive Policy

Generate response sequence y conditioned on x

Model or implementation: Linear model on top of feature map φ(x, y_{1:i-1})

Outcome Reward Model

Evaluate full sequence correctness

Model or implementation: Binary indicator r(x, y)

Novel Architectural Elements

Theoretical abstraction of LLM generation as a linear autoregressive process with token-level margin separability

Modeling

Base Model: Linear autoregressive model (representing the final layer of a Transformer)

Training Method: Policy Gradient (REINFORCE variant) with optional importance weighting

Objective Functions:

Purpose: Update policy weights to maximize expected reward.

Formally: w_{t+1} = w_t + η * r_t * ∇ log q_t(y_t|x_t) (Standard PG)
Purpose: Pre-training via SGD with adaptive learning rate.

Formally: η_t = (4 + 2 ||∇ log p_w(y|x)||)^(-1)

Key Hyperparameters:

learning_rate: 1/(2N) (constant) or adaptive
initialization: Zero (w_0 = 0) for analysis simplicity

Compute: Not reported in the paper

Comparison to Prior Work

vs. Banditron: This work achieves optimal Õ(k/γ²) mistake bound for N=1, whereas Banditron achieves suboptimal Õ(k²/γ²)
vs. Prior PG Theory (e.g., Agarwal et al.): This work provides non-asymptotic bounds specifically for autoregressive sequence generation with 0-1 rewards, rather than general tabular settings
vs. Best-of-N Sampling [not cited in paper]: This work analyzes how to bake the Best-of-N performance into the model via training, rather than just sampling at inference time

Limitations

Analysis is restricted to linear autoregressive models, though applicable to deep networks with frozen features
Assumes a strict margin condition (separability) for the correct response
Focuses on 0-1 binary rewards, which may not capture nuanced human feedback
Worst-case analysis may be overly pessimistic compared to average-case practical performance

Reproducibility

Theoretical paper. Detailed proofs are provided in the appendices. No code or data artifacts are associated with this work.

📊 Experiments & Results

Evaluation Setup

Theoretical analysis of mistake bounds and sample complexity

Benchmarks:

Online Multiclass Linear Classification (N=1) (Contextual Bandits)
Autoregressive Sequence Generation (Sequence Learning)

Metrics:

Mistake Bound (Online Learning)
Sample Complexity (Number of reward queries)
Expected Error Rate
Statistical methodology: Not explicitly reported in the paper

Key Results

Benchmark	Metric	Baseline	This Paper	Δ
Theoretical bounds comparing the proposed Policy Gradient variant against classical baselines in online learning settings (N=1).
Online Linear Classification (N=1)	Mistake Bound order	k^2/γ^2	k/γ^2	Improved by factor of k
Comparison of sample complexity for pre-training SGD variants showing the necessity of adaptive learning rates.
Pre-training Convergence	Convergence Rate (Dependence on N)	N	1	Removed linear dependence on N

Main Takeaways

The 'Base Model Barrier' is a fundamental limit: Outcome-based RL cannot efficiently learn responses that are 'off-support' (low probability) in the base model
Process rewards (verifiers) are theoretically necessary to break the exponential dependence on sequence length N for learning new reasoning paths
Adaptive learning rates are crucial for both SGD pre-training and PG post-training to achieve length-independent convergence rates
A simple Policy Gradient algorithm with uniform exploration achieves minimax optimal mistake bounds for online contextual bandits, resolving an open problem for efficient algorithms

📚 Prerequisite Knowledge

Prerequisites

Contextual Bandits
Policy Gradient (REINFORCE)
Linear Autoregressive Models
Minimax Analysis

Key Terms

PG: Policy Gradient—an RL algorithm that optimizes a policy by following the gradient of expected reward

Outcome Reward: Feedback provided only at the very end of a generated sequence (e.g., correct/incorrect final answer)

Process Reward: Feedback provided at intermediate steps (e.g., per-token or per-step correctness)

Likelihood Quantile: A proposed theoretical property of the base model that determines the sample complexity of post-training; essentially how 'covered' the correct solution is by the base distribution

Margin condition: A geometric assumption stating that the correct class/token is separated from others by a gap of at least γ in the feature space

SGD: Stochastic Gradient Descent—the standard optimization algorithm used for pre-training the base model

Base Model Barrier: The theoretical finding that outcome-based RL cannot efficiently learn samples that have negligible probability under the base model

SFT: Supervised Fine-Tuning—training on labeled demonstrations

VC dimension: Vapnik–Chervonenkis dimension—a measure of the capacity or complexity of a space of functions