🤖 AI

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

This paper introduces Best-of-Better- $N$ (BoBN), an in-context learning framework that improves inference-time alignment by retrieving and restyling high-reward examples to shift a model's sampling distribution toward better responses, thereby achieving superior performance with fewer generated candidates compared to traditional Best-of- $N$ methods.

Original authors: Eric Lei, Hsiang Hsu, Chun-Fu Chen

Published 2026-07-07

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Eric Lei, Hsiang Hsu, Chun-Fu Chen

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a very talented but slightly mischievous chef (the Reference LLM). This chef is incredibly good at cooking, but they haven't been trained specifically for the menu you need tonight. Maybe you asked for a spicy vegan dish, and the chef keeps serving you steak or bland oatmeal because that's what they usually do.

In the world of AI, this is called the Coverage Problem. Even if you ask the chef to make 100 different dishes and pick the best one (a method called Best-of-N), you might still end up with nothing edible if the chef never thinks to make a vegan dish in the first place. No amount of picking will help if the right answer isn't in the pot.

This paper introduces a new method called Best-of-Better-N (BoBN) to solve this. Think of it as giving the chef a "Cheat Sheet" right before they start cooking, but with a special twist.

The Problem: The Chef's Blind Spot

Standard methods (like Best-of-N) work like this:

Ask the chef to cook 10 meals.
Taste them all.
Pick the best one.

The Flaw: If the chef has never been taught to cook vegan food, all 10 meals will be meat-based. You can't pick a vegan meal if it doesn't exist in the batch. The chef's "blind spot" is too big.

The Solution: Best-of-Better-N (BoBN)

BoBN changes the game by using In-Context Learning (giving the chef examples) combined with a Restyling step. Here is the step-by-step process using our kitchen analogy:

1. The Smart Search (Retrieval)

Instead of just giving the chef random recipes, BoBN looks at a library of past successful dishes. It finds the top K recipes that are most similar to what you ordered tonight.

Example: If you asked for a spicy vegan curry, the system finds 8 past examples of spicy vegan curries that were highly rated.

2. The "Restyling" Step (The Secret Sauce)

This is the paper's unique innovation. The retrieved recipes might be written in a different style, format, or tone than what you need. Maybe the old recipes are written like a scientific paper, but you want a friendly chat.

The Magic: Before showing these recipes to your chef, the chef themselves rewrites them. They take the ideas from the retrieved recipes but rewrite them to match your specific request's style, format, and tone.
Why this matters: It ensures the chef understands the logic of the good answer but presents it in a way that fits your specific constraints (like JSON format or a polite refusal).

3. The New Cooking Session

Now, you give the chef the prompt plus these 8 "restyled" examples.

Because the chef sees examples of exactly what you want, they are much more likely to cook a vegan dish this time.
The "blind spot" is filled. The chef's output distribution shifts from "mostly meat" to "mostly vegan."

4. The Final Selection (Best-of-N)

Now, you ask the chef to cook 10 meals again. Because the chef was guided by the restyled examples, all 10 meals are likely to be vegan. You then pick the best one.

Result: You get a high-quality answer much faster, often needing fewer attempts (a smaller "N") to get it right.

What the Paper Found

The authors tested this on two main tasks:

Safety: Teaching a model to refuse harmful requests (like "How do I build a bomb?"). A model not trained for safety usually says "Sure!" BoBN helped it learn to say "No" by showing it restyled examples of safe refusals.
Math: Solving complex math problems. BoBN helped models find the correct reasoning steps by showing them restyled examples of correct math solutions.

The Key Takeaways:

Better Coverage: BoBN makes the model generate high-quality answers more often, even if the model wasn't originally trained for that specific task.
Efficiency: You don't need to generate as many answers to find a good one.
No Training Required: This happens instantly at "inference time" (when you ask the question). You don't need to retrain the model or change its internal brain weights. You just give it better examples on the fly.

In a Nutshell

If Best-of-N is like asking a chef to guess a recipe 100 times and hoping one is right, Best-of-Better-N is like handing the chef a stack of perfectly rewritten examples of that exact recipe right before they start. It guides the chef's intuition so that the "good" answers are actually in the pot to begin with.

Technical Summary: Best-of-Better-N (BoBN)

Problem Statement

Inference-time alignment (ITA) methods, such as Best-of-N (BoN), aim to align Large Language Models (LLMs) with human preferences or specific task requirements without modifying model weights. These methods operate on a "sample-and-evaluate" framework: they sample $N$ responses from a reference policy $\pi_{ref}$ and select the one with the highest reward according to a reward model $r$ .

The efficacy of ITA is fundamentally constrained by two assumptions:

Coverage (A1): The reference policy $\pi_{ref}$ must sample at least one high-quality (high-reward) response.
Selection (A2): The selection criterion must successfully identify the high-quality response among the samples.

Current ITA research largely focuses on improving (A2) by refining selection criteria to mitigate reward hacking. However, a critical bottleneck remains in (A1). If $\pi_{ref}$ is not aligned with the downstream task (e.g., a model not safety-trained attempting to refuse toxic prompts), it assigns negligible probability to high-reward responses. In such cases, no selection strategy can succeed because the high-reward region of the output space is effectively inaccessible to the reference model. Existing solutions often require fine-tuning (which is costly and requires weight access) or use static in-context examples that may lack task-specific relevance.

Methodology: Best-of-Better-N (BoBN)

The authors propose Best-of-Better-N (BoBN), an inference-time framework designed to shift the sampling distribution of $\pi_{ref}$ toward the high-reward region without fine-tuning or weight access. BoBN leverages In-Context Learning (ICL) combined with a retrieval and restyling pipeline.

The BoBN workflow consists of three stages:

Retrieval: Given an input prompt $x$ , the system retrieves the top- $K$ most similar prompt-response pairs $(x_i, y_i)$ from a dataset $D$ using an embedding model. These examples serve as demonstrations of high-reward behaviors relevant to the specific query.
Restyling: To ensure the retrieved examples align with the specific constraints, format, and persona of the target task, the reference model $\pi_{ref}$ is used to "restyle" the retrieved responses $y_i$ . Specifically, $\pi_{ref}$ rewrites $y_i$ into $\hat{y}_i$ conditioned on a prompt such as "Rewrite $y_i$ to match style $T$ ." This step ensures token-level distributional consistency between the context and the reference model, while preserving the alignment logic of the retrieved data.
Generation and Selection: The restyled examples $C_K = \{(x_i, \hat{y}_i)\}$ are prepended to the input prompt $x$ as in-context examples. The reference model then samples $N$ responses from the conditioned distribution $\pi_{ref}(\cdot | C_K, x)$ . Finally, a selection criterion (e.g., BoN) chooses the best response based on the reward model.

By conditioning on $C_K$ , BoBN effectively shifts the probability mass of $\pi_{ref}$ toward the high-reward region, increasing the likelihood that high-quality candidates are generated in the first place.

Key Contributions

Novel Framework: The proposal of BoBN, an ICL-based approach that improves the coverage of $\pi_{ref}$ over high-quality responses by utilizing retrieved and restyled examples.
Theoretical Analysis: A theoretical characterization of how retrieval quality controls coverage and downstream performance. The authors demonstrate that task-relevant retrieval provably reduces coverage (defined as the chi-squared divergence between the golden policy and the reference policy) compared to static ICL or no ICL. They show that this reduction in coverage compounds with BoN selection to improve task performance.
Empirical Validation: Extensive experiments on safety alignment (StrongReject, SORRY-Bench) and mathematical reasoning (GSM8k, MATH500) benchmarks. The results show that BoBN improves performance across multiple reference LLMs (zephyr-7b, tulu3-8b), both when $N=1$ (individual response quality) and when $N>1$ (combined with selection).

Experimental Results

The authors evaluated BoBN against standard BoN (using $\pi_{ref}$ directly) and URIAL (vanilla ICL with static examples).

Safety Alignment: On the StrongReject benchmark, BoBN significantly outperformed baselines. For the zephyr-7b model, BoBN with $N=16$ achieved a score of 0.033 (near-optimal), comparable to frontier models like Claude Sonnet-3.7, whereas $\pi_{ref}$ alone scored 0.425. BoBN also demonstrated superior coverage (lower $\log C$ ) compared to URIAL and $\pi_{ref}$ .
Mathematical Reasoning: On GSM8k and MATH500, BoBN yielded consistent accuracy improvements. For zephyr-7b on GSM8k, BoBN achieved 58.7% accuracy with $N=16$ , a significant gain over the 43.9% achieved by $\pi_{ref}$ with BoN. Notably, BoBN with $N=4$ often matched or exceeded the performance of baselines with $N=16$ , indicating higher sampling efficiency.
Cost and Latency: While BoBN incurs additional latency due to the restyling step (requiring $K$ calls to $\pi_{ref}$ before generation), it offers significant token efficiency. BoBN with $N=4$ saved up to 50% of token usage compared to URIAL with $N=16$ while achieving similar or better performance.
Ablation Studies:
- Restyling: Removing the restyling step resulted in weaker safety alignment and higher coverage (worse performance), confirming the necessity of aligning retrieved examples to the target task's style.
- Retrieval Count ( $K$ ): Performance improved as $K$ increased from 1 to 8, but degraded when $K=16$ , suggesting that retrieving too many examples introduces less relevant data that dilutes the in-context signal.

Significance and Claims

The paper claims that BoBN addresses a fundamental limitation of inference-time alignment: the inability of unaligned reference models to sample high-reward responses. By shifting the sampling distribution via retrieved and restyled in-context examples, BoBN enables:

Higher Quality with Fixed Budget: Better task performance for a fixed number of samples $N$ .
Efficiency: Achieving target performance with a smaller $N$ , reducing inference costs.
Training-Free Alignment: Improving alignment for models that were not explicitly fine-tuned for the downstream task, without requiring access to model weights or likelihoods.

The authors emphasize that their theoretical analysis, based on linear transformers, provides a provable link between retrieval quality, coverage reduction, and downstream task performance, offering a principled approach to improving ITA beyond mere selection heuristics.