1 Introduction↩︎

Inference-time alignment (ITA) is a paradigm which has achieved remarkable performance gains for various downstream tasks performed by large language models (LLMs) without requiring further training [1]. These methods provide an adaptable and easy-to-use alternative to the standard LLM post-training pipeline, where a pretrained LLM may undergo supervised fine-tuning and reinforcement learning [2], [3] to align the model with human preferences and boost performance on downstream tasks. Such tasks may include improved helpfulness and engagement during chat-bot interactions, more accurate reasoning chains for arithmetic tasks, or refusal capabilities when prompted with toxic material.

While the LLM post-training pipeline has proven highly successful for real-world deployment of LLMs [3], [4], it is both impractical and costly due to unavailability of the LLM weights and the high cost of training. ITA presents a training-free substitute by performing a proxy version entirely at inference time, enabling improved performance on downstream tasks without needing LLM weights or likelihood access.

Current ITA methods typically operate in the sample-and-evaluate framework [5]. Specifically, let \(\piref\) be a reference LLM which we wish to align at inference time. Given a prompt \(\bx\), an ITA method will sample responses from \(\piref(\cdot|\bx)\), and then select an output response according to a criterion. This criterion typically depends on a reward model which encodes response quality; larger rewards imply a more desired response for the task of interest. For example, the best-of-\(N\) (BoN) criterion [6], [7] chooses the highest reward response. Within the sample-and-evaluate framework, there are two fundamental assumptions which determine the success of ITA in practice:

1. the responses sampled from \(\piref(\cdot|\bx)\) contain at least one high-quality response, and

2. among the sampled responses, the selection criterion will choose a good response.

Both assumptions need to be satisfied for good performance. If (A1) is not satisfied, then no selection criterion will be able to output an aligned response; if (A2) is not satisfied, then the ITA method will fail to select a good candidate even if good candidates were generated. Therefore, a failure to satisfy (A1) presents a fundamental roadblock in downstream ITA performance. Successfully satisfying (A1), such that high-quality responses are sampled with higher probability, additionally creates an opportunity for compounding benefits: better task performance may be achieved under a fixed sample budget, and less samples may be needed from \(\piref\) to achieve sufficient alignment.

While many efforts in ITA have tried to improve the chances of satisfying (A2) by mitigating reward hacking [5], [8], there has been limited work in improving (A1). The primary challenge of satisfying (A1) is that high-quality responses may almost never be sampled from the reference LLM when it is not aligned to the downstream task, causing all responses from \(\piref\) to fall in the low-reward region (Fig. 1). For example, if we wish to use ITA for safety alignment, a reference LLM that was never explicitly safety trained will almost always sample low-reward, unsafe responses. In other words, the output space of \(\piref\) may fail to cover or overlap with the set of high-reward responses that we desire, thereby making ITA ineffective no matter how good the selection criterion is.

In this paper, we seek to improve the coverage of \(\piref\) over high-quality responses without access to \(\piref\)’s model weights or likelihood access. Equivalently, we seek to shift the response distribution of \(\piref\) under prompt \(\bx\) toward the high-reward region under \(r\). In-context learning (ICL) is a weight- and likelihood-free method that is known to be capable of shifting the output distribution of the reference LLM toward the provided context examples [9], [10] by leveraging the pattern-matching abilities of reference LLMs gained during pretraining. We introduce Best-of-Better-\(N\) (BoBN), a framework which utilizes ICL with retrieved examples along with a restyling step to address the challenges of poor coverage of the reference LLM. Our contributions are the following.

1. We propose BoBN, a novel ICL-based approach to improve coverage of high-quality responses. BoBN shifts the reference LLM’s response distribution toward the in-context examples, increasing the chance a high-quality candidate is generated and later chosen.

2. We theoretically analyze how retrieval quality controls coverage and downstream task performance, showing that task-relevant retrieval provably reduces coverage compared to static or no ICL, and that this improvement compounds with BoN selection.

3. We experimentally study BoBN on real-world LLM benchmarks for safety alignment and mathematical reasoning. We show that across multiple reference LLMs, BoBN improves benchmark performance. Performance improves both when \(N=1\) and when \(N>1\), demonstrating that the pre-aligned BoBN responses help improve performance individually as well as when combined with the BoN criterion. Additionally, we show that coverage empirically improves with benchmark performance.

2 Background and Related Work↩︎

2.1 Inference-Time Alignment↩︎

Let \(\mathcal{V}\) be a vocabulary representing the set of all possible tokens. We represent a large language model (LLM) policy \(\pi\) as a conditional probability distribution \(\pi(\cdot|\bx) \in \mathcal{P}(\mathcal{V}^*)\), where \(\bx \in \mathcal{V}^*\) is an input prompt. We denote \(\by \in \mathcal{V}^*\) as the response generated by \(\pi\) when prompted with \(\bx\), i.e., \(\by \sim \pi(\cdot|\bx)\). We let \(r:\mathcal{V}^* \times \mathcal{V}^* \rightarrow \mathbb{R}\) be a reward model that measures response quality for a given task. There are two primary phases in the sample-and-evaluate framework for ITA [5], which we now define.

We note here that these two phases correspond to assumption (A1) and (A2) respectively. The most well-known ITA method is best-of-\(N\) (BoN), which generates candidates by sampling \(\by_1,\dots,\by_N\) independently, and selects a candidate under the highest reward criterion [6], [7]: \(\mathcal{A}_{\textsf{BoN}}(\by_1,\dots,\by_N) = \argmax_{\by \in \{\by_1,\dots,\by_N\}} r(\bx, \by)\). Other ITA methods that fall under the sample-and-evaluate framework include Soft-BoN [11], [12], InferenceTimePessimism [5], as well as variations that leverage rejection sampling [13], [14], and Markov-chain Monte-Carlo sampling, such as [8]. As mentioned in Sec. 1, the primary motivation of these works is (A2): choosing a better response sampled from \(\piref\). While there is a plethora of work in improving ITA, these methods primarily aim to improve the response selection process and avoid known challenges such as reward hacking [15]. In contrast, our work addresses (A1)—generating better responses prior to selection—and provides theoretical analysis on how ICL supports this aim.

While the above ITA methods operate in the sample-and-evaluate framework, there are several recent works that do not. URIAL [16] utilizes in-context learning capabilities of LLMs to place a curated but static set of examples in-context. A key component of URIAL is manually restyling these examples, which improves the helpfulness, engagement, and safety of the responses. Our method improves upon URIAL’s static in-context examples with retrieved ones that are restyled by \(\piref\).

Another relevant work is TPO [17], which iteratively updates a set of candidate responses using \(\piref\) and reward model values as a learning signal, requiring multiple \(\piref\) and reward model calls that scale with the number of iterations. In contrast, our method only makes a constant number of \(\piref\) and reward model calls. There also exist many ITA methods that directly adjust the next-token distribution of \(\piref\) during sampling. ARGS [18] adjusts next-token likelihoods of \(\piref\) by tilting it with the reward model \(r\). GenARM [19] extends this approach with autoregressive reward models and PAD [20] extends this to personalized preference alignment. In contrast to our method, these ITA methods require LLM likelihood access which makes them unsuitable for black-box settings in which only text responses to \(\piref\) are available. In our work, we focus on the BoN candidate selection criterion since it possesses high empirical performance, it is the most popular ITA method, and it does not require weight or likelihood access of \(\piref\).

Training-time alignment. Traditionally, LLM alignment is done via reinforcement learning with human feedback (RLHF) using a learned reward model over a preference dataset [2]. Variations of this framework include DPO [21] which uses a preference dataset but does not require a reward model. However, training-time alignment requires LLM weight access and additionally suffers from alignment tax [3], [22], where applying RLHF can hurt the LLM’s performance on general task performance. In contrast, ITA methods are more flexible at inference time, where different reward models or preference datasets may be adaptively applied depending on the task. There exist several works that can be classified as ITA-aware finetuning [23], [24], where \(\piref\) is finetuned to either simulate ITA zero-shot or improve ITA at inference time. While these methods do address (A1) and improve \(\piref\) response quality, they require \(\piref\) finetuning; in contrast, our work improves \(\piref\) response quality without any finetuning.

In-context learning. ICL is the practice of providing demonstrations to the LLM within its context in order to enable the LLM to respond to an input query in a similar manner [10], [25], [26]. These capabilities are often leveraged by augmenting the LLM with external knowledge via retrieval from data sources, a practice known as retrieval-augmented generation. In contrast, alignment pertains to eliciting desired behaviors from an LLM rather than augmenting incomplete knowledge.

2.2 Coverage of the Reference LLM↩︎

We now formally define coverage, and describe how coverage and the criterion \(\mathcal{A}\) are intertwined in characterizing downstream task performance. Let \(r^*\) denote the golden reward function encoding response quality for a given downstream task. In practice, \(r^*\) can be interpreted as the query-wise benchmark or evaluation performance of some task. For example, in mathematical reasoning, \(r^*(\bx, \by)\) could be an indicator function for whether or not the response \(\by\) contains the correct answer of a given mathematical problem \(\bx\). On the other hand, \(r\) is a reward model trained on preference data that serves as a proxy for \(r^*\), as \(r^*\) is not available during inference. Let \(\pi^*\) be a golden LLM policy with respect to \(r^*\), such that it achieves optimal task performance \(\EE_{\by \sim \pi^*(\cdot|\bx)} [r^*(\bx, \by)]\). The coverage, defined below, determines the overlap between two policies \(\mu\) and \(\nu\).

We note that \(\mathsf{C}(\mu||\nu) \geq 1\), with equality if and only if \(\mu = \nu\): in this case, \(\nu\) perfectly covers \(\mu\). Given prompt \(\bx\), the coverage of \(\piref(\cdot|\bx)\) with respect to \(\pi^*(\cdot|\bx)\) is \(\mathsf{C}(\pi^*(\cdot|\bx)||\piref(\cdot|\bx))=\EE_{\by \sim \pi^*(\cdot |\bx)} \bracket*{\frac{\pi^*(\by|\bx)}{\piref(\by|\bx)}}\). The coverage provides a measure for how well \(\piref\) is able to sample high-quality responses generated by the golden LLM policy \(\by \sim \pi^*(\cdot|\bx)\). It decreases when \(\piref\) is able to assign higher probability to responses in the high-reward region under \(r^*\).

3 BoBN: Improving Sampling with Coverage↩︎

BoBN seeks to improve ITA task performance by improving the coverage of \(\piref\), i.e., reducing \(\coverage\), using ICL with \(K\) retrieved examples \(C_K\). We first provide an illustrative example showing how BoBN’s retrieved examples \(C_K\) shifts the response distribution under \(\piref\) toward the high-reward region.

Reward distribution shift under BoBN. We compare the distributions of the reward \(r^*(\bx, \by)\) of responses \(\by\) generated by \(\piref(\cdot|\bx)\) and \(\piref(\cdot|C_K, \bx)\) (BoBN). The prompts \(\bx\) are drawn from the StrongReject benchmark [27], and the reference LLM used is zephyr-7b [28].

For StrongReject, \(r^*(\bx, \by) = 1-\mathsf{score} \in [0,1]\), where \(\mathsf{score}\) is the three-metric combination defined in [27]; smaller \(\mathsf{score}\) implies better safety alignment. Shown in Fig. 2, we see that samples from \(\piref(\cdot|\bx)\) are highly concentrated in the low-reward region. This is consistent with the fact that \(\texttt{zephyr-7b}\) is not explicitly finetuned for safety. On the other hand, \(\piref(\cdot|C_K,\bx)\) shifts the \(r^*\) distribution toward the high-reward region, significantly increasing the chance that the BoN criterion will choose a safety-aligned response and thereby coverage. For example, under prompt 2, even if the number of responses \(N\) is very large, the probability that any reward will be greater than 0.1 under \(\piref\) is very small. In contrast, the probability that a single (i.e., \(N=1\)) response reward will be greater than 0.5 is roughly 50% under BoBN. A similar effect on math reasoning is also shown on the GSM8k dataset with \(\texttt{tulu3-8b}\) [29], where \(r^*\) is a binary indicator on whether the LLM answered the math question correctly. Next, we describe how BoBN forms the context \(C_K\) to drive this reward distribution shift.

3.1 Retrieval from Prompt-Response Datasets↩︎

We assume that the ITA algorithm has access to a prompt-response dataset \(\prefdata = \{(\bx_i, \by_i)\}_{i=1}^{M}\) containing responses \(\by_i\) to prompts \(\bx_i\). Training-time finetuning and alignment methods typically require such datasets \(\prefdata\); in ITA, we still wish to perform alignment using \(\prefdata\), except we do not modify or have access to \(\piref\) weights.

BoBN begins by retrieving the \(K\) most similar examples according to prompt similarity under an embedding model \(\phi : \mathcal{V}^* \rightarrow \mathbb{R}^d\). That is, for a given input prompt \(\bx\), the retrieval step returns \(\prefdataretr := \{(\bx_i, \by_i)\}_{i \in I}\) where the indices \(I\) are chosen via \(I = \argmax_{S \subseteq \{1,\dots,M\}, |S|=K} \sum_{i \in S} s(\phi(\bx), \phi(\bx_i))\) and \(s(\ba, \bb) = \frac{\ba \cdot \bb}{\|\ba\| \|\bb\|}\) is cosine similarity. While an unaligned \(\piref\) may not know how to refuse a dangerous prompt or how to structure a reasoning chain, the prompt-response dataset contains demonstrations for these high-reward behaviors, effectively guiding \(\piref\) into the high-reward region of the response space \(\mathcal{V}^*\). Compared to vanilla ICL, which uses static in-context examples, retrieving \(K\) examples with similar prompts enables more fine-grained alignment within specific domains. For example, these could include adapting to the style of refusing queries on how to perform internet fraud, or mathematical reasoning traces that involve simplifying fractions. Rather than provide broad examples of how we want the LLM to behave, the examples in \(\prefdataretr\) provide query-specific demonstrations for preferred behaviors. These benefits over vanilla ICL are discussed analytically with linear transformers in Sec. 5 and experimentally on practical LLM benchmarks in Sec. 4.

3.2 Restyling of Retrieved Responses↩︎

While BoBN could directly place \(\prefdataretr\) in the context of \(\piref\), this may not result in the best performance. This is because while retrieved examples encode a desired alignment logic, they may not exactly match the desired form, format, style, or persona that the query and downstream task require. For example, the query or task may require certain constraints on the final response, such as formatting requirements like JSON or putting an arithmetic result in a boxed format. Additionally, chat interactions may also reward a more polite or engaging tone than what \(\prefdata\) possesses. To mitigate this potential mismatch, BoBN incorporates a restyling step in which \(\piref\) itself may be used to restyle the responses in \(\prefdataretr\). Specifically, the retrieved responses \(\by_i\) for each \(i \in I\) are restyled into \(\hat{\by}_i \sim \piref(\cdot|\text{``Rewrite \by_i to match style T''})\) where \(T\) represents the downstream task. This enables \(\hat{\by}_i\) to retain the alignment logic encoded within the data, yet maintain consistency with the stylistic and structural constraints of the downstream task. This restyling step widens the applicability of any dataset when used in BoBN, as the restyling step curates the \(\prefdata\) to each task, reducing the need to collect task-specific datasets. Moreover, since the raw retrieved responses \(\by_i\) were likely written by a different LLM than \(\piref\), the restyling step ensures token-level distributional consistency between \(\hat{\by}_i\) and \(\piref\), which is known to improve downstream LLM task performance [30], [31].

3.3 Generating and Selecting Pre-Aligned Responses↩︎

Having constructed a set of task-aligned, high-reward demonstrations, the next phase leverages them to steer the reference policy \(\piref\) toward the high-reward region of the output space. BoBN modifies the candidate sampling distribution by prepending the restyled context \(C_K := \{(\bx_i, \hat{\by}_i)\}_{i \in I}\) consisting of the \(K\) retrieved prompt-response pairs that have been rewritten to match the target task’s constraints. By conditioning on \(C_K\), we guide \(\piref\) to mimic the safety boundaries and reasoning patterns inherent in the retrieved examples. The full end-to-end flow is shown in Fig. 3. The candidate generation process proceeds by sampling \(\by_1, \dots, \by_N \sim \piref(\cdot|C_K, \bx)\). This is followed by applying a candidate selection criterion \(\mathcal{A}\). Since \(C_K\) contains examples where \(\hat{\by}_i^+\) is both high-reward (due to retrieval) and in-distribution (due to the restyling step), the probability mass of \(\piref(\cdot|C_K, \bx)\) is significantly more concentrated on valid responses compared to \(\piref\). Rather than relying on the reward model to filter through potentially poor, un-aligned candidates, we use ICL to ensure the candidate pool is rich with high-potential, pre-aligned options instead.

4 Experimental Study↩︎

Experimental setup. We conduct our experiments on safety alignment and mathematical reasoning benchmarks. We use zephyr-7b [28] and tulu3-8b [29] as \(\piref\). We use \(K=8\) retrieved examples, and vary the number of responses \(N\). We compare our method BoBN with BoN using responses from \(\piref\) and URIAL [16], which performs vanilla ICL with static in-context examples. For reward models, we use ArmoRM [32] for math, and a finetuned Llama-3.1-8b-Instruct [33] using Bradley-Terry reward modelling on the HH-RLHF preference dataset [2] for safety alignment. Results are reported as mean \(\pm\) 1 standard deviation across three random seeds.

Empirical coverage estimates. We estimate the coverage, which allows comparison between task performance and coverage of \(\piref\), corresponding to Thm. 4. We choose \(\pi^*\) to be an LLM that achieves near-optimal performance on each benchmark. If \(\pi_0\) is a sampler of response candidates, we estimate the coverage of \(\pi_0\) w.r.t. \(\pi^*\), normalized by response length, as \(\underline{\mathsf{C}}(\pi^*(\cdot|\bx) || \pi_0(\cdot|\bx)) = \mathop{\EE}_{\by \sim \pi^*(\cdot|\bx)} \bracket*{\paran*{\frac{\pi^*(\by|\bx)}{\pi_0(\by|\bx)}}^{\frac{1}{T^*}}}\) using an empirical average over samples from \(\pi^*(\cdot|\bx)\), where \(T^*\) is the number of tokens of response \(\by\) under \(\pi^*\). This allows normalization of coverage across prompts within a dataset, as coverage scales with response length. Below, we report the mean and standard deviation of \(\log \underline{\mathsf{C}}(\pi^*(\cdot|\bx) || \pi_0(\cdot|\bx))\) over all queries for each benchmark. These are reported for each response sampler \(\pi_0 \in \{\piref, \text{URIAL}, \text{BoBN}\}\); changing \(N\) does not affect \(\logC\). Further experimental details can be found in Sec. 7.

4.1 Safety Alignment↩︎

4pt

For safety alignment, we use the StrongReject benchmark [27] and SORRY-Bench [34]; the score ranges from \([0,1]\), with \(0\) indicating fully safe and \(1\) as toxic; see Sec. 7.3. For \(\prefdata\), we use the HH-RLHF preference dataset [2] for BoBN. Shown in Tab. 1, we see that for zephyr-7b, single-response sampling from \(\piref\) with unsafe prompts \(\bx\) result in poor safety alignment; this is expected as zephyr-7b was not explicitly post-trained for safety. On the other hand, applying URIAL improves safety alignment, which outperforms BoN selection on top of \(\piref\) candidates. However, BoBN-1 performs best for the single-response samplers as well as when BoN selection is applied: BoBN-16 achieves a near-optimal score of 0.033 on StrongReject (same score as Claude Sonnet-3.7), and also achieves 0.361 on SORRY-Bench (same score as GPT-4o-mini [34]). This shows that BoBN is capable of aligning small language models that were not explicitly safety trained to the refusal performance of frontier LLMs, despite no finetuning. The coverage estimates \(\logC\), which are estimated using Llama-3.3-70b-Instruct as \(\pi^*\), decrease as performance improves, validating Thm. 4 in that task performance improves with decreasing coverage. Varying \(N\) for BoBN and comparing with \(\piref\), URIAL further validates Thm. 4, as decreasing coverage enables similar or better performance with less \(N\). We observe similar trends for tulu3-8b in Tab. 1, except \(\piref\) itself is much more safety-aligned compared to zephyr-7b.

4.2 Mathematical Reasoning↩︎

For math, we use GSM8k [35] and MATH500 [36] benchmarks. GSM8k contains grade-school math problems, and MATH500 contains more difficult high-school competition questions. These are evaluated using accuracy averaged over the dataset. The retrieval dataset \(\prefdata\) used in BoBN is Math-Step-DPO-10K [37], which contains correct reasoning traces for chosen responses and erroneous ones for rejected responses. Shown in Tab. 2 for zephyr-7b, tulu3-8b, we see that for URIAL, especially when combined with BoN, does not improve the accuracy over \(\piref\) (with BoN). Unlike safety alignment, where URIAL yielded consistent gains over \(\piref\), in math reasoning, the fixed set of examples in URIAL is not sufficiently relevant to the reasoning task. On the other hand, BoBN yields consistent accuracy improvement over \(\piref\) (with and without BoN), with BoBN gaining up to 37% and 11% accuracy gains compared to directly prompting zephyr-7b for GSM8k and MATH500 respectively, without finetuning. Similar trends follow for tulu3-8b, however with less significant gains, as tulu3-8b is a stronger reference LLM for reasoning than zephyr-7b and there is less room for improvement. Similar to safety alignment, improved coverage enables better performance for a fixed \(N\), and a smaller \(N\) to achieve better performance: BoBN with \(N=4\) performs similar or better than BoN with URIAL and \(\piref\) with \(N=16\).

4.3 Cost and Latency Analysis↩︎

We perform a cost analysis of BoBN and compare to URIAL and \(\piref\). Shown in Tab. 3, we report the latency per sample, as well as token statistics. Input, output, and total token usage with respect to \(\piref\) is reported per sample, with all metrics aggregated over StrongReject and GSM8k with zephyr-7b as \(\piref\). For URIAL and \(\piref\), there will be \(N\) calls to \(\piref\) per sample; for BoBN, there are up to \(K+N\) due to the restyling step. We observe that latency is roughly similar across \(N=16\) for \(\piref\) and URIAL and BoBN-4; however, BoBN-4 saves up to 50% token usage compared to URIAL, while achieving similar or better performance. This showcases the cost benefit of the improved sampling efficiency afforded by BoBN due to its superior coverage. While BoBN-16 is slower and uses considerably more tokens, its performance is significantly better than URIAL and BoN, despite no direct finetuning.

4.4 Ablation Study↩︎

r0.45

The effect of the restyling step is shown in Tab. 4, where we compare the effect of restyling on zephyr-7b on safety alignment. We see that BoBN-1 exhibits weaker safety alignment and worse coverage when directly using the retrieved examples compared to when the retrieved examples undergo restyling first, further supporting the effectiveness of restyling in improving coverage of \(\piref\) to the downstream task.

r0.45

We also provide an ablation study on the number of retrieved examples \(K\) in BoBN. Tab. 5 shows how task performance and coverage \(\logC\) behave as \(K\) increases, for zephyr-7b on StrongReject and GSM8k. This is conducted on a held-out validation set for each benchmark. We see that performance and \(\logC\) improves as \(K\) increases initially, but decreases when \(K\) goes to 16; \(K=8\) appears to be a sweet spot. This may be attributed to retrieved preference examples that are less task-relevant as \(K\) grows.

5 Coverage Gains in BoBN: A Linear-Transformer Case Study↩︎

In order to provide insight on how ICL, which forms a core component of BoBN, interacts with coverage and task performance, we now provide a case study of these components in linear transformers. We theoretically analyze the coverage of pretrained linear transformers with ICL, as this is a well-established framework that enables precise analysis of the mechanisms of ICL [26], [38]. We precisely show how factors such as pretraining context length, number of retrieved examples, and task-relevance of retrieval affect coverage and downstream performance, particularly when BoN is applied.

One-layer linear transformers. We first define the one-layer linear transformer [26], [39]. Let \(\be_1,\dots,\be_t\) be a sequence of input token embeddings, and let \(W_Q, W_K, W_V\) be the query, key, and value matrices. Then the output of a one-layer linear transformer \(f_{\mathsf{LT}}\) at the \((t+1)\)-th token is \([f_{\mathsf{LT}}(X)]_{t+1} = \be_t + W_V XX^\top W_K^\top W_Q\be_{t+1}\), where \(X = [\be_1,\dots,\be_t]\). For pretraining/inference, we use the standard linear-regression ICL setting for one-layer linear transformers. The following result, originally from [39]–[41], shows that after being optimally-pretrained on random linear regression tasks, the transformer then equivalent to performing one step of gradient descent on in-context examples provided during inference.

In other words, \(f^*_{\mathsf{LT}}\) predicts \(\hat{y} = \wgd^\top \bx\), where \(\wgd := \frac{\eta}{K} \sum_{i=1}^K y_i \bx_i\) is the vector that regresses \(\{(\bx_i, y_i)\}_{i=1}^K\), and \(\eta\) is a scaling factor due to optimal pretraining on length-\(n\) sequences. Next, we analyze BoBN coverage and the factors affecting task performance.

5.1 Effect of ICL Retrieval Quality on Coverage↩︎

Suppose that \(f^*_{\mathsf{LT}}\) is optimally-pretrained. This transformer represents the reference LLM that we would like to apply ITA to. Let \(\bw_* \in \mathbb{R}^d\) be a downstream task of interest. The context \(C_K := \{(\bx_i, y_i)\}_{i=1}^K\) is drawn from a potentially misspecified task \(\bw'\), i.e., \(y_i = \bw'^\top \bx_i\), and \(\bx_i \sim \mathcal{N}(\bzero, I_d)\). In the best-case scenario, \(\bw' = \bw_*\), and BoBN retrieves high-quality examples as if they were generated by the golden LLM. In practice, retrieval is imperfect, and we factor in the retrieval misspecification \(\Delta := \bw_* - \bw'\) within the subsequent analysis.

We denote the golden stochastic policy \(\pi^*(y|\bx) = \mathcal{N}(\bw_*^\top \bx, \sigma^2)\), and the policy induced by \(f^*_{\mathsf{LT}}\) with context \(C_K\) be \(\piref(y|C_K,\bx) = \mathcal{N}(\wgd^\top \bx, \sigma^2)\); analogously, the policy without context is \(\piref(y|\bx) = \mathcal{N}(0, \sigma^2)\). Without context, the coverage of \(\piref\) w.r.t. \(\pi^*\) satisfies \(\mathsf{C}(\pi^*(\cdot|\bx)||\piref(\cdot|\bx)) = \exp\paran*{\frac{(\bw_*^\top \bx)^2}{\sigma^2}}\), where we use the fact that \(\chi^2\paran*{\mathcal{N}(\mu_1, \sigma^2), \mathcal{N}(\mu_2, \sigma^2)} = \exp\paran*{\frac{(\mu_1-\mu_2)^2}{\sigma^2}}\). We first show a general bound on the in-context coverage \(\mathsf{C}(\pi^*(\cdot|\bx)||\piref(\cdot|C_K, \bx))\) in Thm. 2. We defer full proofs to Appendix 10.

In Fig. 5, we provide a heatmap of \(\coverage\) from a pretrained \(f^*_{\mathsf{LT}}\) on a fixed task \(\bw_*\) and input \(\bx\), and show that the relationship between coverage, \(\|\Delta\|\), \(n\), and \(K\) indeed follow the bound in 1 . Next, we use the coverage bounds to predict downstream task performance as a function of the number of responses \(N\) and \(\|\Delta\|\) for large \(n\) and \(K\).

5.2 Characterization of BoN Task Performance with ICL↩︎

The following result characterizes the regret, or gap to optimality, of the downstream task performance \(\EE_{\by \sim \piref(\cdot | C_K,\bx)}[r^*(\bx, \by)]\) under the BoN criterion.

Discussion. Thms. 2, 4 show that coverage and task performance are controlled by \(\|\Delta\|\). We now relate this to three ICL strategies studied in this work. Note that coverage depends on \(\wgd = \frac{\eta}{K}\sum_i y_i \bx_i\), which is well-defined for any set of in-context pairs regardless of their origin, and that \(\|\bw_* - \wgd\| \rightarrow \|\Delta\|\) as \(n, K \rightarrow \infty\). Without in-context examples, \(\wgd = \bzero\), so \(\|\bw_* - \wgd\| = \|\bw_*\|\). Under vanilla ICL, \(\wgd\) is determined by a fixed set of examples independent of the query, giving some \(\|\bw_* - \wgdstatic\| \leq \|\bw_*\|\) provided \(\wgdstatic^\top \bw_* \geq \frac{\|\wgdstatic\|^2}{2}\), i.e., the examples are non-adversarial. Under retrieval-based ICL, the examples adapt to the query \(\bx\), and since retrieval can always replicate any fixed set of examples by ignoring the query, \(\|\bw_* - \wgdretrieval\| \leq \|\bw_* - \wgdstatic\|\). This gives the ordering \(\|\bw_* - \wgdretrieval\| \leq \|\bw_* - \wgdstatic\| \leq \|\bw_*\|\), which implies better coverage and performance from Thms. 2, 4 when stronger retrieval is used.

Compounding benefits of improving coverage with \(N\). These results also exemplify the compounding benefit of improving coverage and applying BoN. Thm. 4 shows that the regret’s leading term reduces exponentially with \(\|\Delta\|\). This means that better retrieval quality results in both better performance at fixed \(N\) as well as smaller \(N\) required to reach a target task performance. The left side of Fig. 4, which shows the bound in 2 for \(R_{\mathsf{max}}=1, \epsilon_{\mathsf{RM}}=0.1\), illustrates this compounding effect with coverage \(\mathsf{C}_1:= \exp \paran*{\frac{\|\bx\|^2 \|\Delta\|^2}{\sigma^2}}\). For a fixed \(N\), task performance improves with decreasing coverage, and the \(N\) required to achieve a fixed regret decreases linearly with \(\mathsf{C}_1\); when \(N\) is large, reward hacking occurs, and performance decreases due to reward model noise. While other ITA work, such as [5], attempt to improve this regret bound using improved selection criteria to mitigate reward hacking, reward hacking is not problematic for most regimes of \(N\); a more direct way to improve task performance across most regimes of \(N\) in practice is to simply reduce \(\piref\)’s coverage, which is what BoBN accomplishes.

We additionally provide experimental results verifying exact regret bounds as a function of exact coverage and \(N\). We fix a downstream task \(\bw_*\), construct no-context, static-context, and retrieved-context policies with controlled misspecification \(\Delta=\bw_*-\wgdretrieval\), and apply BoN (see Sec. 7.4). The resulting regret curves in Fig. 4 (middle) show the same behavior predicted by Thm. 4’s bounds: better retrieval reduces coverage, improves performance at fixed \(N\), and reduces the number of samples required to reach a target regret. For large \(N\), reward hacking occurs. Finally, Fig. 4 (right) shows that as misspecification increases, coverage worsens, and the minimum \(N\) required to achieve a fixed, target regret increases roughly linearly with coverage.

The relationship between \(N\), \(\|\Delta\|\), and regret enables practitioners to tune retrieval and \(N\) for their inference setup. Despite the fact that our analysis has been focused on linear transformers, it provides a precise characterization of how \(n, K\), and \(\|\Delta\|\) control downstream task performance through their effect on coverage. These are corroborated by our results in Sec. 4 within practical LLM benchmarks, where BoBN’s retrieval achieves lower coverage than vanilla ICL and \(\piref\) and corresponding improvements in task performance.

6 Conclusion↩︎

We introduced a preference-retrieval framework which overcomes the coverage limitations of inference-time alignment. BoBN is shown to effectively shift the response distribution toward high-reward, aligned behaviors, with no training required. Future work may investigate: different criterions \(\mathcal{A}\) (e.g., SoftBoN), analyzing token- or segment-level coverage for use with process reward models, and ICL analysis with nonlinear transformers.

6.0.0.1 Disclaimer.

This paper was prepared for informational purposes by the Global Technology Applied Research center of JPMorgan Chase & Co. This paper is not a product of the Research Department of JPMorgan Chase & Co. or its affiliates. Neither JPMorgan Chase & Co. nor any of its affiliates makes any explicit or implied representation or warranty and none of them accept any liability in connection with this paper, including, without limitation, with respect to the completeness, accuracy, or reliability of the information contained herein and the potential legal, compliance, tax, or accounting effects thereof. This document is not intended as investment research or investment advice, or as a recommendation, offer, or solicitation for the purchase or sale of any security, financial instrument, financial product or service, or to be used in any way for evaluating the merits of participating in any transaction.

7 Further Experimental Details↩︎

7.1 BoBN Experimental Details↩︎

7.1.0.1 Embedding model.

The embedding model \(\phi\) used for the retrieval step is nvidia/llama-embed-nemotron-8b.

7.1.0.2 Restyling details.

The system prompt used to restyle \(\prefdataretr\) with \(\piref\) for safety alignment is shown below:

For mathematical reasoning, the restyling step does not use \(\piref\) but instead reformats the retrieved responses. Specifically, within each preference example response, the “Step \(k\):” labels are removed, and the “The answer is:” is reformatted into “Answer:” in order to match the desired formatting of the corresponding test harness; see Sec. 7.3.

7.2 Coverage Estimation↩︎

We estimate coverage for a single prompt \(\bx\) by generating \(N_{\mathsf{est}}=100\) samples from \(\pi^*(\cdot|\bx)\) and estimating \(\underline{\mathsf{C}}\) using a sample mean: \[\underline{\mathsf{C}}(\pi^*(\cdot|\bx) || \pi_0(\cdot|\bx)) \approx \frac{1}{N_{\mathsf{est}}}\sum_{i=1}^{N_{\mathsf{est}}} \paran*{\frac{\pi^*(\by_i|\bx)}{\pi_0(\by_i|\bx)}^{\frac{1}{T^*_i}}}.\] In the above, \(\by_i \sim \pi^*(\cdot|\bx)\) are independently drawn, and \(T^*_i\) is the number of tokens of \(\by_i\) when tokenized using \(\pi^*\)’s tokenizer. In practice, since we have access to token-wise log probabilities of \(\pi^*\) and \(\pi_0\), we estimate \[\underline{\mathsf{C}}(\pi^*(\cdot|\bx) || \pi_0(\cdot|\bx)) \approx \frac{1}{N_{\mathsf{est}}}\sum_{i=1}^{N_{\mathsf{est}}} \exp\cbracket*{{\frac{1}{T^*_i} \paran*{\sum_{t=1}^{T^*_i} \log \pi^*(\by_{i, t}|\by_{i, <t}, \bx) - \sum_{t=1}^{T_{i,0}} \log \pi_0(\by_{i, t}|\by_{i, <t}, \bx)}}},\] where \(T_{i,0}\) is the number of tokens of \(\by_i\) when tokenized by \(\pi_0\).

7.3 Benchmarks↩︎

7.3.0.1 StrongReject.

StrongReject contains harmful prompts and evaluates LLM responses in terms of their refusal/compliance [27]. However, instead of just measuring refusal rate, StrongReject uses a 3-metric score, which is based on LLM-as-a-judge, and additionally evaluates how helpful the response would be to a malicious user. It computes the score as \[\mathsf{score} = (1-\mathsf{refusal}) \times \frac{\mathsf{specificity} - \mathsf{convincingness}}{2},\] where \(\mathsf{refusal}\) is either 0 or 1, corresponding to refusal or compliance, \(\mathsf{specificity}\) is a [0, 1] value originally from a \(\{1,\dots,5\}\) Likert scale measuring how detailed and actionable the response is; a vague answer gets a low score; a step-by-step tutorial gets a high score. Lastly, \(\mathsf{convincingness}\), also a [0, 1] value originally from a \(\{1,\dots,5\}\) Likert scale, measures how persuasive or realistic the response is to the reader. We report the average \(\mathsf{score}\), where smaller \(\mathsf{score}\) is better as it indicates better safety alignment. The corresponding task-specific golden reward \(r^*\) is \(1 - \mathsf{score}\). We follow the evaluation harness provdied in https://github.com/dsbowen/strong_reject.

7.3.0.2 SORRY-Bench.

SORRY-Bench [34] evaluates LLM responses using LLM-as-a-judge. It uses a smaller, fine-tuned Mistral-7B-Instruct-v0.2 model to classify whether a response refused or complied with a harmful request. It outputs 0 for a refusal to a harmful request (i.e., safe) and 1 when the response complies with the harmful request (i.e., unsafe). The reported score is the average compliance rate over all prompts in SORRY-Bench, where lower is better, as it indicates better safety alignment. Similar to StrongReject, the corresponding task-specific golden reward \(r^*\) is 1 if classified as refused and 0 if complied. We follow the evaluation harness provided in https://github.com/SORRY-Bench/sorry-bench.

7.3.0.3 GSM8k.

GSM8k is a benchmark dataset designed to evaluate the multi-step mathematical reasoning capabilities of large language models [35]. It contains 1319 test examples, consisting of diverse grade school word problems that typically require between two to eight steps to solve. Unlike simple arithmetic tasks, GSM8k specifically tests a model’s ability to chain logic together to derive a final answer. It includes problems ranging from basic arithmetic, fractions, pre-algebra, logic, and basic geometry. String matching is used to determine if the response was correct or not, with the average accuracy reported. The corresponding task-specific golden reward \(r^*\) is simply 1 if correct and 0 otherwise. We follow the evaluation harness provided in https://github.com/QwenLM/Qwen/tree/main/eval.

7.3.0.4 MATH500.

MATH500 [36] is a benchmark dataset comprising 500 challenging mathematical problems selected from the larger MATH dataset. While GSM8k focuses on elementary logic, MATH500 targets competition-level mathematics, evaluating a LLM’s ability to handle complex, multi-step reasoning and advanced problem-solving techniques. It is frequently used to evaluate how well models can generate reasoning solutions that often require producing LaTeX-formatted equations. It covers topics such as algebra, number theory, geometry, combinatorics, and pre-calculus. Similar to GSM8k, string matching is used to determine if the response was correct or not, with the average accuracy reported. Similar to GSM8k, the corresponding task-specific golden reward \(r^*\) is simply 1 if correct and 0 otherwise. We follow the evaluation harness provided in https://github.com/openai/simple-evals.

7.4 Linear Transformer BoN Details↩︎

Regarding Fig. 4 (middle and right), these experiments were conducted by first sampling a downstream task vector \(\bw_*\). For each test query \(\bx\), we use the golden policy \(\pi^*(y | \bx)=\mathcal{N}(\bw_{*}^\top \bx,\sigma^2)\). We then compare three policies for BoN with distributions \(\mathcal{N}(\wgd^\top \bx,\sigma^2)\). For no ICL, we have that \(\wgd = \bzero\). For static ICL, context is drawn from a fixed task \(\bw_{\mathsf{static}}\), resulting in \(\wgd=\wgdstatic\). Finally, for retrieval ICL, context examples are drawn from some task \(\bw'\) closer to \(\bw_*\), and we assume that the misspecification \(\Delta = \bw_* - \bw'\) is fixed.

For BoN, we use the true reward model \(r^*(\bx, y) = -(y-\bw_*^\top \bx)^2 / (2s^2)\) for a fixed \(s > 0\), and a proxy reward \(r(\bx, y) = r^*(\bx, y) + \xi\), where \(\xi \sim \mathcal{N}(0, \tau^2)\). Regret is measured with \(1-r\).

8 Additional Results↩︎

9 Reward Distribution Shift Under BoBN↩︎

The prompts used for the reward distribution shift figures can be found in Tab. 6 and Tab. 7.

10 Transformer Coverage Analysis and Proofs↩︎

10.1 Proof of Theorem 2↩︎

We first restate the theorem below.

10.2 Proof of Theorem 4↩︎

The result follows by combining Thm. 2 and [5].

References↩︎