Is Best-of-N the Best of Them? Coverage, Scaling, and Optimality in Inference-Time Alignment (2503.21878v2)

Abstract: Inference-time computation offers a powerful axis for scaling the performance of LLMs. However, naively increasing computation in techniques like Best-of-N sampling can lead to performance degradation due to reward hacking. Toward a theoretical understanding of how to best leverage additional computation, we focus on inference-time alignment, which we formalize as the problem of improving the quality of responses drawn from a pre-trained policy, given a prompt of interest and access to an imperfect reward model. We analyze the performance of inference-time alignment algorithms in terms of (i) response quality, and (ii) compute, and provide new results that highlight the importance of the pre-trained policy's coverage over high-quality responses for performance and compute scaling: 1. We show that Best-of-$N$ alignment with an ideal choice for $N$ can achieve optimal performance under stringent notions of coverage, but provably suffers from reward hacking when $N$ is large, and fails to achieve tight guarantees under more realistic coverage conditions. 2. We introduce $\texttt{InferenceTimePessimism}$, a new algorithm which mitigates reward hacking through deliberate use of inference-time compute, implementing the principle of pessimism in the face of uncertainty via rejection sampling; we prove that its performance is optimal and does not degrade with $N$, meaning it is scaling-monotonic. We complement our theoretical results with an experimental evaluation that demonstrate the benefits of $\texttt{InferenceTimePessimism}$ across a variety of tasks and models.

Optimality in Inference-Time Alignment: A Detailed Examination

The paper "Optimality in Inference-Time Alignment" presents an analytical framework and algorithmic advancements focused on optimizing inference-time computation in LLMs. The work fundamentally revolves around the challenges of inference-time alignment, where the aim is to refine a pre-trained LLM’s responses in real-time using a reward model despite its imperfections. This exploration is pivotal given how such computation can significantly influence the performance of LLMs in practice.

Overview of Key Contributions

The paper presents several notable contributions:

1. Theoretical Framework for Inference-Time Alignment: The researchers introduce a structured approach to understand the limits and potential of inference-time alignment algorithms. By treating inference-time operations as statistical problems, the authors formulate computational efficiency through query complexity—a method reminiscent of theoretical learning models.
2. Detailed Analysis of Best-of-N Alignment: Best-of-N sampling is widely adopted in practice for its simplicity, where multiple outputs are generated, and the response with the highest reward score is selected. However, it risks reward hacking, a form of overoptimization that leads to a decrease in genuine task performance with increased computation (N). The paper thoroughly examines these performance scalability issues, providing insights into its limitations due to imperfect reward models.
3. Proposal and Analysis of InferenceTimePessimism Algorithm: A cornerstone of this paper is the introduction of the InferenceTimePessimism algorithm, which aims to overcome weaknesses identified in Best-of-N alignment. This algorithm integrates principles from pessimistic reinforcement learning, introducing regularization to mitigate overoptimization. Through theoretical proofs, the authors demonstrate that the approach is optimal concerning regret, ensuring that performance does not degrade as the computation scale increases—a property denoted as scaling-monotonicity.
4. Empirical Evaluation: Complementing the theoretical assurances, extensive experimental validation across various tasks (e.g., GSM8K, MMLU, MATH) and models affirms the efficacy of InferenceTimePessimism. The empirical results confirm theoretical predictions, such as its scaling-monotonic property, thereby suggesting robust practical applicability.

Theoretical Implications and Future Directions

The implications of this research are manifold:

• Understanding Coverage: A significant aspect discussed is the role of coverage in determining algorithmic performance. In this context, 'coverage' reflects the extent to which a pre-trained LLM can generate high-quality responses that cover optimal outcomes. By formalizing this, the paper establishes foundational requirements needed to achieve optimal performance at inference time.
• Algorithm Design Innovations: Beyond surface-level results, this research promotes the idea of deliberate algorithm design incorporating regularization, notably manifest through the InferenceTimePessimism algorithm. The potential for applying similar principles to enhance other types of AI-based decision problems is substantial.
• Co-Design Strategies: This research invites further exploration in designing training and inference-time procedures that align optimally, ensuring maximal performance gain.

Speculation on Future AI Developments

While the current focus is on inference-time computation, these findings could pave the way for advancements in various AI applications. Areas of interest include:

• Inference-Aware Training Models: Refining model training methodologies to inherently optimize for inference-time alignment could significantly advance AI systems' efficiency and adaptability.
• Incorporation of Exploration Strategies: Exploring how inference-time algorithms might effectively utilize real-time exploration to refine outputs could mirror advancements in online reinforcement learning, seeking to leverage feedback loops more intelligently.

In summary, "Optimality in Inference-Time Alignment" provides an intricate understanding and robust solutions to scaling issues in LLMs, positing noteworthy theoretical frameworks and practical algorithms. Through its balanced approach, it sets a pivotal foundation for future explorations into both theoretical and empirical AI research realms.