Is a Good Foundation Necessary for Efficient Reinforcement Learning? The Computational Role of the Base Model in Exploration
Dylan J. Foster, Zakaria Mhammedi, Dhruv Rohatgi
TL;DR
This work formalizes a sampling-oracle framework for language-model alignment and analyzes the computational-statistical tradeoffs of exploration under a linear softmax policy. It identifies coverage of near-optimal responses by the base model as a key factor limiting computational efficiency, and introduces SpannerSampling to achieve near-optimal data efficiency with polynomial-time computation by leveraging inference-time exploration and a spanner of informative samples. The paper further shows that training-time interventions cannot, in general, be both data- and computationally-efficient, under ETH. It also proposes multi-turn exploration (MTSS) to exploit autoregressive representations, achieving improved runtime by shifting from sequence-level to token-level coverage under suitable assumptions. Overall, the results provide a foundational perspective on efficient exploration with powerful pre-trained models and highlight avenues for broader representation-based improvements.
Abstract
Language model alignment (or, reinforcement learning) techniques that leverage active exploration -- deliberately encouraging the model to produce diverse, informative responses -- offer the promise of super-human capabilities. However, current understanding of algorithm design primitives for computationally efficient exploration with language models is limited. To better understand how to leverage access to powerful pre-trained generative models to improve the efficiency of exploration, we introduce a new computational framework for RL with language models, in which the learner interacts with the model through a sampling oracle. Focusing on the linear softmax model parameterization, we provide new results that reveal the computational-statistical tradeoffs of efficient exploration: 1. Necessity of coverage: Coverage refers to the extent to which the pre-trained model covers near-optimal responses -- a form of hidden knowledge. We show that coverage, while not necessary for data efficiency, lower bounds the runtime of any algorithm in our framework. 2. Inference-time exploration: We introduce a new algorithm, SpannerSampling, which obtains optimal data efficiency and is computationally efficient whenever the pre-trained model enjoys sufficient coverage, matching our lower bound. SpannerSampling leverages inference-time computation with the pre-trained model to reduce the effective search space for exploration. 3. Insufficiency of training-time interventions: We contrast the result above by showing that training-time interventions that produce proper policies cannot achieve similar guarantees in polynomial time. 4. Computational benefits of multi-turn exploration: Finally, we show that under additional representational assumptions, one can achieve improved runtime (replacing sequence-level coverage with token-level coverage) through multi-turn exploration.