Best-of-Better-$N$: Generating Pre-Aligned Responses with In-Context Learning
2026-07-03 • Machine Learning
Machine LearningArtificial Intelligence
AI summaryⓘ
The authors looked at ways to improve how language models pick the best answer when given multiple choices, especially when the usual model struggles to find good answers. They proposed a new method called Best-of-Better-N (BoBN), which finds good example answers and rewrites them in a way that helps the model give better responses. This rewriting helps guide the model toward producing higher-quality answers more often. They tested BoBN on tasks like making answers safer and solving math problems, showing it works better with fewer tries than older methods.
Inference-time alignmentBest-of-NReward modelsIn-context learningLanguage modelsResponse generationSampling distributionSafety alignmentMathematical reasoningPretrained transformers
Authors
Eric Lei, Hsiang Hsu, Chun-Fu Chen
Abstract
Inference-time alignment methods, such as Best-of-$N$, offer a flexible alternative to training-based alignment by using reward models to select high-quality responses generated by a reference LLM. However, the efficacy of these methods is inherently limited by the response quality: if the reference LLM assigns negligible probability to high-reward responses, no selection strategy will succeed in finding aligned outputs. In this work, we propose Best-of-Better-$N$ (BoBN), an in context learning-based generation framework to address this challenge. Our method utilizes retrieval from high-reward examples relevant to the input query and task. Crucially, we introduce a restyling step where retrieved responses are rewritten by the reference LLM to align with the target task's format and style. These restyled examples are used in-context to shift the sampling distribution toward the high-reward region. We analytically characterize how in-context learning shifts the output distribution of pretrained transformers toward the high-reward region, resulting in provable benefits on the target task. We then evaluate BoBN on safety alignment and mathematical reasoning benchmarks across several reference LLMs. BoBN's higher-quality responses enable better performance to be achieved when the number of responses $N$ is fixed, and smaller $N$ required to achieve a target performance.