ABoN: Adaptive Best-of-N Alignment

2025-05-17

AI Summary

To address computational resource waste in test-time alignment of large language models (LLMs), this paper proposes a two-stage, fine-tuning-free method that adaptively allocates Best-of-N sampling budgets based on prompt difficulty. First, lightweight exploration estimates the reward distribution for each prompt; then, sampling budgets are dynamically allocated per prompt to achieve prompt-level adaptive optimization. This plug-and-play framework balances efficiency and alignment quality. Experiments across 12 LM/RM combinations and 50 prompt batches on AlpacaEval show: (i) significant superiority over uniform budget allocation under identical computational budgets; (ii) achievement of the performance of uniform allocation with +20% budget using only 80% of that budget; and (iii) consistently increasing gains as batch size grows. The core contribution is the first prompt-level adaptive budget allocation mechanism, which uniquely couples reward distribution estimation with sampling resource scheduling within the test-time alignment pipeline.

Technology Category

• Planning, Routing, and Scheduling: Planning with Language Models
• Search and Optimization: Sampling/Simulation-based Search
• Machine Learning: Large Multimodal Models (LMMs)

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

Recent advances in test-time alignment methods, such as Best-of-N sampling, offer a simple and effective way to steer language models (LMs) toward preferred behaviors using reward models (RM). However, these approaches can be computationally expensive, especially when applied uniformly across prompts without accounting for differences in alignment difficulty. In this work, we propose a prompt-adaptive strategy for Best-of-N alignment that allocates inference-time compute more efficiently. Motivated by latency concerns, we develop a two-stage algorithm: an initial exploratory phase estimates the reward distribution for each prompt using a small exploration budget, and a second stage adaptively allocates the remaining budget using these estimates. Our method is simple, practical, and compatible with any LM/RM combination. Empirical results on the AlpacaEval dataset for 12 LM/RM pairs and 50 different batches of prompts show that our adaptive strategy consistently outperforms the uniform allocation with the same inference budget. Moreover, our experiments show that our adaptive strategy remains competitive against uniform allocations with 20% larger inference budgets and even improves in performance as the batch size grows.

Problem

• Reduces computational cost of Best-of-N alignment
• Adapts compute allocation based on prompt difficulty
• Improves performance with same inference budget

Innovation

• Adaptive Best-of-N alignment strategy
• Two-stage algorithm for efficiency
• Prompt-specific compute allocation