Best-of-Better-$N$: Generating Pre-Aligned Responses with In-Context Learning

2026-07-03

AI Summary

This work addresses the limitations of conventional inference-time alignment methods such as Best-of-N, which rely on high-quality responses generated by a reference model and struggle when such responses have extremely low probability. To overcome this, the authors propose the Best-of-Better-N framework, which retrieves high-reward examples and employs the reference model to rewrite them to match the target task’s style, thereby constructing high-quality in-context demonstrations that steer the generation distribution toward higher-reward regions. Integrating retrieval augmentation, in-context learning, and reward-guided rewriting, the method is theoretically shown to bias the output distribution of pretrained Transformers toward high-reward areas. Experiments demonstrate that Best-of-Better-N significantly outperforms baseline approaches on safety alignment and mathematical reasoning benchmarks, achieving superior performance with a fixed number of samples or matching baseline performance with fewer samples.

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.

Problem

• inference-time alignment
• response quality
• reward models
• large language models
• output distribution

Innovation

• in-context learning
• reward modeling
• response restyling
• distribution shift
• inference-time alignment