Weak-to-Strong Search: Align Large Language Models via Searching over Small Language Models

Introduction

Learning-based algorithms have become the standard approach for aligning large language models (LLMs) with human preferences . However, fine-tuning large language models is resource-intensive and difficult to implement . These challenges have motivated recent studies on search-based algorithms that keep the large language models frozen and steer their decoding with test-time guidance . Typical examples of search-based algorithms include rejection sampling and Monte Carlo Tree Search . These search-based algorithms are promising as they can reuse the same learned guidance to steer the decoding of any large language model without additional training. However, existing search-based methods either simplify the search over tokens as a bandit problem , which limits their steerability, or require a value function learned from scratch to address preference reward sparsity and reduce search depth , which can be as difficult as fine-tuning a large language model.

To make search-based algorithms better suited for aligning large language models, we introduce weak-to-strong search, a simple algorithm that frames the alignment of a large model as a test-time search over the log-likelihoods of small language models. This algorithm makes two contributions: (1) First, it builds on the theoretical foundation of the token-level MDP for alignment , using the log-likelihood difference between small tuned and untuned language models as both reward and critic to guide the decoding of a large model (Section 4.1). Theoretically, this formulation is suitable for search as it converts the otherwise sparse preference reward to a per-token dense reward, which can be summed up as a value function (critic) . Practically, this formulation allows the reuse of off-the-shelf small tuned and untuned language model pairs as steering forces, avoiding the need to train a reward or critic model from scratch. (2) Second, it introduces a beam search variant, Chunk-level Beam Search (CBS), tailored for optimizing the proposed search objective. CBS guides the large language model towards high-reward regions by alternating between sampling from the frozen large model and expanding promising states as evaluated by the small tuned and untuned models (Section 4.2). Especially, when the small models are weaker than the large model, our method can be viewed as an instance of weak-to-strong generalization that makes the strong model stronger with weak test-time guidance (Figure 1).

Empirically, we verify weak-to-strong search’s flexibility in various tasks (Section 5). First, in controlled-sentiment generation and summarization , our method uses small language models of 124M parameters (i.e., gpt2) to effectively steer much larger language models from the GPT-2 (e.g., gpt2-xl) , Llama-2 and Llama-3 families, at least as effective as existing methods. Then, in a more difficult instruction-following benchmark, AlpacaEval 2.0 , we show reusing off-the-shelf small models (e.g., zephyr-7b-beta and its untuned version) as test-time guidance can significantly improve the length-controlled win rates of both white-box and black-box large models against gpt-4-turbo (e.g, $34.4\rightarrow 37.9$ for Llama-3-70B-Instruct, and $16.0\rightarrow 20.1$ for gpt-3.5-turbo-instruct), despite the small models’ low win rates $\approx 10.0$ (Figure 1).

Related Work

Large unsupervised language models trained on internet-scale corpus acquire broad knowledge and abilities . However, these large pre-trained language models may not always align with human values. To instill the desired behaviors into language models, most existing methods fine-tune these pre-trained language models on human comparisons of model-generated responses . Despite these successes, fine-tuning a large language model requires substantial computational resources and engineering effort. These problems are compounded by the reality that different humans have different values , as it is nearly impossible to train a new large language model from scratch for individual preference. In light of these issues, our work takes a search-based approach, folding as much of the complexity of alignment as possible into the decoding phase. This allows us to keep the large pre-trained language models frozen, steering their outputs at test time with only small models that are easier to obtain.

Framing alignment as a test-time search to maximize a reward function is not a novel formulation. However, most existing works either simplify autoregressive decoding as a bandit problem , which limits their steerability, or require a value function learned from scratch to handle sparse preference rewards and reduce search depth , which can be as difficult as training a large language model from scratch. Our work avoids these issues by parametrizing the reward function with the log-likelihood difference between small tuned and untuned language models . This parametrization not only simplifies the search objective, allowing a simple greedy search algorithm to generate good results, but also reuses off-the-shelf models as steering forces, eliminating the need to train a reward or critic model from scratch.

Concurrently with our work, Rafailov et al. proposes a token-level MDP interpretation for language model alignment, demonstrating that a greedy likelihood search over a trained language model can achieve improvements over regular decoding. Our work builds on their theoretical foundations and proposes a novel greedy search algorithm designed for weak-to-strong guidance.

The idea of using small language models to align large language models has arisen in many recent works. The most related is proxy or emulated fine-tuning , which uses the distributional difference of a small tuned and untuned model pair to modify the output distribution of a large model, approximating the output of the directly tuned large model. However, these methods require that both small and large models share the same vocabulary, limiting their practical applications. In contrast, our approach does not modify the sampling distribution of the large model at the token level. Instead, we perform a tree search that periodically selects the most promising states for further expansion (as evaluated by the small models) while sampling from the frozen large model’s distribution. Thus our approach does not require shared vocabulary and is applicable to black-box language models.

Preliminaries

In this section, we introduce the mathematical formulation of alignment (Section 3.1) and describe the duality between language models and reward functions (Section 3.2).

The alignment of language models is typically cast as a KL-constrained optimization problem :

2 Duality between Language Models and Reward Functions

The analytical solution to Eq. 1 can be obtained through the following Lagrangian :

which has a well-known closed-form solution that expresses a duality between the reward function $r({\mathbf{x}},{\mathbf{y}})$ and the optimal language model $\pi^{*}({\mathbf{y}}\mid{\mathbf{x}})$ :

where $Z({\mathbf{x}})=\sum_{\mathbf{y}}{\pi_{\text{ref}}({\mathbf{y}}\mid{\mathbf{x}})}\exp\left(\frac{1}{\beta}r({\mathbf{x}},{\mathbf{y}})\right)$ denotes the partition function. One takeaway from this duality is that we can always express a reward function using tuned and untuned language models: (1) If a reward function is given , we can first obtain the optimally tuned language model under this reward function with any learning-based algorithms, and then use the tuned and untuned models $(\pi^{*},\pi_{\text{ref}})$ to reparametrize the reward function (Eq. 3); (2) If a dataset is given from which the reward function can be derived, we can then directly parametrize the reward function with the tuned and untuned language models $(\pi^{*},\pi_{\text{ref}})$ during reward modeling (Eq. 3).

Weak-to-Strong Search

In this section, we introduce weak-to-strong search, a search-based algorithm that aligns a large language model by searching over the log-likelihood difference between small tuned and untuned language models. First, we discuss how using language models to (re)parametrize the reward function from Eq. 1 makes the challenging objective solvable by a simple greedy search algorithm (e.g., beam search) (Section 4.1). Then, we introduce a practical beam search variant, called Chunk-level Beam Search (CBS) (Section 4.2), for optimizing the proposed objective, which is applicable to steering both white-box and black-box large language models.

One practical challenge for search-based alignment algorithms is the sparsity of the preference reward signal. The preference reward function $r({\mathbf{x}},{\mathbf{y}})$, based on the Bradley-Terry model , only emits a terminal reward when the model response is complete. Search-based algorithms often struggle without any intermediate rewards or a critic model predicting future returns . However, if we parameterize this sparse reward function with language models (Section 3.2), we can obtain both a dense reward function and a critic function simultaneously.

To obtain a dense reward function, we leverage the duality between the sparse preference reward and the dense language model probability (Eq. 3). By explicitly factorizing the log-likelihood of a complete response ${\mathbf{y}}$ under the language models, we can obtain a sum-of-rewards style formulation for Eq. 3:

where ${\mathbf{y}}_{<t}$ denotes the response tokens from $1$ to $t-1$, and the last response token $y_{|{\mathbf{y}}|}$ is always the EOS token. Combining Eq. 1 and 4, we can rewrite the original objective with a per-token reward:

Setting aside the KL-constraint (Eq. 5b) for now (detailed in Section 4.2), we can apply existing search algorithms like beam search to optimize Eq. 5a. However, it is tempting to argue that beam search is prone to local optima , as it greedily retains promising states midway through generation $({\mathbf{x}},{\mathbf{y}}^{\prime})$ (${\mathbf{y}}^{\prime}$ is incomplete)${\mathbf{y}}$ denotes a complete response, while ${\mathbf{y}}^{\prime}$ denotes a response that can be either incomplete or complete. based on partial return $\log\pi^{*}({\mathbf{y}}^{\prime}\mid{\mathbf{x}})-\log\pi_{\text{ref}}({\mathbf{y}}^{\prime}\mid{\mathbf{x}})$, but high partial return may tell little about the overall return. Although this argument holds for most MDPs, appealing to the token-level MDP framework , we argue that the opposite is true here:

In Appendix A, we show that (inspired heavily by )

where $V^{*}({\mathbf{x}},{\mathbf{y}}^{\prime})$ denotes the value function, predicting the expected terminal reward under the optimal $\pi^{*}$ in the original KL-constrained sparse reward setting. Although $V^{*}({\mathbf{x}},{\mathbf{y}}^{\prime})$ is not necessarily achievable by the searched policy, it approximates how good the state $({\mathbf{x}},{\mathbf{y}}^{\prime})$ is in the long run. In other words, continuing from the state $({\mathbf{x}},{\mathbf{y}}^{\prime})$ of high partial return $\log\pi^{*}({\mathbf{y}}^{\prime}\mid{\mathbf{x}})-\log\pi_{\text{ref}}({\mathbf{y}}^{\prime}\mid{\mathbf{x}})$ is likely to generate a complete response ${\mathbf{y}}$ with high overall return $\log\pi^{*}({\mathbf{y}}\mid{\mathbf{x}})-\log\pi_{\text{ref}}({\mathbf{y}}\mid{\mathbf{x}})$.

2 Chunk-level Beam Search (CBS)

After analyzing the feasibility of optimizing Eq. 5a with greedy search algorithms (e.g., beam search), we introduce a practical beam search variant that optimizes the dense reward objective (Eq. 5a) while ensuring the KL-constraint from $\pi_{\text{base}}$ (Eq. 5b).

The core algorithm providing the foundation of our method, Chunk-level Beam Search (CBS), is detailed in Algorithm 1 and illustrated in Figure 2. The key insight is that our beam search operates at the level of chunk. The search starts at the prompt and always maintains a hypothesis set $\mathcal{H}=\{({\mathbf{x}},{\mathbf{y}}^{\prime})_{i}\}^{W}_{i=1}$ of $W$ states. For each state $({\mathbf{x}},{\mathbf{y}}^{\prime})$ in $\mathcal{H}$, CBS samples $K$ continuation chunks ${\mathbf{y}}_{L}$ of length $L$ from $\pi_{\text{base}}$. This results in $W\cdot K$ successor states. Among these successors, only the top-$W$ successors with the highest partial return $\log\pi^{*}({\mathbf{y}}^{\prime}\circ{\mathbf{y}}_{L}\mid{\mathbf{x}})-\log\pi_{\text{ref}}({\mathbf{y}}^{\prime}\circ{\mathbf{y}}_{L}\mid{\mathbf{x}})$ are stored in $\mathcal{H}$ and expanded further. Finally, the terminal state $({\mathbf{x}},{\mathbf{y}})$ with the highest overall return $\log\pi^{*}({\mathbf{y}}\mid{\mathbf{x}})-\log\pi_{\text{ref}}({\mathbf{y}}\mid{\mathbf{x}})$ is selected, from which the complete response ${\mathbf{y}}$ is extracted.

Notably, CBS is a unified framework that encompasses several search-based algorithms: (1) CBS with $W=1$, $K=N$, $L=\infty$, is equivalent to BoN sampling with $\log\pi^{*}({\mathbf{y}}\mid{\mathbf{x}})-\log\pi_{\text{ref}}({\mathbf{y}}\mid{\mathbf{x}})$ as the scoring function. (2) CBS with $K=\infty$, $L=1$ (directly inspecting the log-likelihoods of all possible next tokens from the vocabulary) is equivalent to vanilla token-level beam search.

However, we always ensure finite chunk length and limited successor exploration via sampling (even when $\bm{L=1}$) to achieve the best of both worlds: (1) Using a finite chunk length allows CBS to discard bad states earlier and focus computational resources on expanding promising states, enhancing steerability more efficiently compared to BoN. (2) Sampling from $\pi_{\text{base}}$ with limited successor exploration implicitly enforces the KL-constraint from $\pi_{\text{base}}$ (Eq. 5b). Without this limit, integrating the KL-constraint into the objective (Eq. 5a) as in Eq. 2 would be necessary, but this can be challenging, especially when vocabularies of $\pi_{\text{base}}$ and $(\pi^{*},\pi_{\text{ref}})$ differ or with black-box language models $\pi_{\text{base}}$ whose log-likelihoods are inaccessible. In addition, it can be beneficial to keep more than one promising state $\bm{W>1}$ because the partial return only approximately correlates with the overall return (Section 4.1). While the partial return provides useful long-term guidance, we can still not strictly guarantee achieving the global optima by following the best local estimates.

Discussions on computation costs and steerability-KL tradeoff. In practice, CBS samples $W\cdot K$ continuation chunks in parallel and evaluates new states by calling $(\pi^{*},\pi_{\text{base}})$ every $L$ tokens. Larger $W\cdot K$ and smaller $L$ enhance steerability at the cost of increased computations, whereas smaller $W\cdot K$ and larger $L$ sacrifice steerability for more efficient computations. Note that high steerability, while beneficial, is not always ideal as it may lead to large KL deviation and over-optimization .

3 Application: Model Up-Scaling and Weak-to-Strong Generalization

The most practical use of CBS occurs when the tuned and untuned models, $(\pi^{*},\pi_{\text{ref}})$, are smaller than the model to steer, $\pi_{\text{base}}$. (1) First, this instance serves as a model up-scaling strategy, directly tuning a small model $\pi_{\text{ref}}\rightarrow\pi^{*}$, by which the large model decoding can then be guided, to achieve similar outcomes as directly tuning the large model. (2) Second, since the small models $(\pi^{*},\pi_{\text{ref}})$ are usually weaker than the large model to steer $\pi_{\text{base}}$, this instance also exemplifies weak-to-strong generalization , enhancing the strong model with only weak test-time guidance. We refer to this instance of CBS as weak-to-strong search, which is the main focus of our study.

Experiments

In this section, we empirically evaluate weak-to-strong search’s ability to align large language models using only test-time guidance from small language models. First, in controlled-sentiment generation and summarization , we tune gpt2 to model the desired behaviors in each task and then use tuned and untuned gpt2 to steer larger models of various scales (Section 5.1). Next, in a more difficult instruction-following benchmark, AlpacaEval 2.0 , instead of tunning small models, we reuse off-the-shelf open-source 7B models and their untuned versions to steer a series of large models, including open-source 70B models and a black-box model (Section 5.2).

In addition to weak-to-strong search, we evaluate several existing test-time approaches that steer large language models $\pi_{\text{base}}$ using small tuned and untuned language models $(\pi^{*},\pi_{\text{ref}})$: (1) Base: we explore regular decoding from the frozen large language model with n-shot prompting (see Appendix C.1.6 for prompt details). (2) Best-of-N Sampling (BoN) : BoN uses $r=\log\pi^{*}({\mathbf{y}}\mid{\mathbf{x}})-\log\pi_{\text{ref}}({\mathbf{y}}\mid{\mathbf{x}})$ to select the highest-scoring responses among the $N$ independent responses from the frozen large language model. Since weak-to-strong search (CBS) samples $W\cdot K$ response chunks in parallel, for fair computational comparisons, we always ensure $N=W\cdot K$. (3) Emulated Fine-Tuning (EFT) : EFT approximates the results of directly fine-tuning the large language model by sampling from $\log\pi_{\text{EFT}}(y_{t}\mid{\mathbf{x}},{\mathbf{y}}_{<t})\propto\log\pi_{\text{base}}(y_{t}\mid{\mathbf{x}},{\mathbf{y}}_{<t})+\beta^{-1}(\log\pi^{*}(y_{t}\mid{\mathbf{x}},{\mathbf{y}}_{<t})-\log\pi_{\text{ref}}(y_{t}\mid{\mathbf{x}},{\mathbf{y}}_{<t}))$, where $\beta$ is the hyperparameter from Eq. 2. Note that EFT is only applicable when all models share the same vocabulary (which is necessary for composing output distributions from different models). Whenever possible, we also compare test-time methods against directly fine-tuning the large models in the same way small models are tuned.

1 Controlled-Sentiment Generation & Summarization

For these two tasks, we follow the synthetic setups from , assuming access to a gold reward model $r_{\text{gold}}$. For controlled-sentiment generation, $r_{\text{gold}}$ encourages positive continuations of movie reviews, while for summarization, it encourages high-quality summaries of Reddit posts (details in Appendix C.1.4). We generate synthetic preference datasets $\mathcal{D}=\{({\mathbf{x}},{\mathbf{y}}_{w},{\mathbf{y}}_{l})_{i}\}_{i=1}^{N}$ from $r_{\text{gold}}$ with $p({\mathbf{y}}_{1}\succ{\mathbf{y}}_{2}\mid{\mathbf{x}})=\sigma(r_{\text{gold}}({\mathbf{x}},{\mathbf{y}}_{1})-r_{\text{gold}}({\mathbf{x}},{\mathbf{y}}_{2}))$ to mimic human feedback .

To obtain the small language models, we optimize gpt2 (124M parameters) using the standard DPO pipeline : (1) we first obtain the reference model $\pi_{\text{ref}}$ through supervised fine-tuning on both chosen and rejected responses from the synthetic preference dataset, then (2) we apply DPO on the synthetic preference dataset with $\pi_{\text{ref}}$ as the reference policy to obtain the optimal language model $\pi^{*}$. Note that the first stage primarily informs the language model of the desired response format, with most of the tuning occurring in the second DPO stage.

Given the tuned and untuned (un-DPO-tuned) gpt2 pair $(\pi^{*},\pi_{\text{ref}})$, we use them to steer the large pre-trained language models without additional training. The large pre-trained language models we study fall into two categories based on whether they share the same vocabulary as the small models: (1) same vocabulary: gpt2-large (774M), gpt2-xl (1.5B) and (2) cross vocabulary: Llama-2-7b, Llama-3-8B. Eventually, since we have access to the gold reward model, language model responses can be fairly evaluated on the test split of prompts using this gold reward model.

Figure 3 demonstrates weak-to-strong search’s great flexibility and steerability in both tasks. For summarization, weak-to-strong search consistently outperforms other test-time methods by large margins. For controlled-sentiment generation, weak-to-strong search is second only to EFT with a carefully selected hyperparameter ($\beta^{*}=1/4$) when EFT is applicable. We hypothesize that token-level adjustments from EFT are sufficient for controlled-sentiment generation, which primarily requires minor stylistic changes at the token level (e.g., “hate” $\rightarrow$ “love”). However, in the more complex task of summarization, where broader sequence-level manipulations are essential, weak-to-strong search excels. Please refer to Appendix E for quantitative comparisons of samples from different methods. We need to mention that we do not meaningfully tune weak-to-strong search (CBS)’s hyperparameters to obtain the results in Figure 3 (we use a fixed set of hyperparameters of $(4,4,5)$ for $W,K,L$ across all models), which may underestimate the performance of our method. In addition, our method enables consistent weak-to-strong generalization in the harder task of summarization: most large pre-trained models (except for gpt2-large) are stronger than the tuned gpt2 in summarizing long text, but the weak models are still able to improve the strong models through test-time guidance, nearly matching the results of direct fine-tuning. The phenomenon of weak-to-strong generalization will be further studied in Section 5.2.

We perform ablations to understand how CBS hyperparameters (beam width $W$, successors per state $K$, and chunk length $L$) influence performance. Figure 4 displays the ablation results for $W,K$. With the same computation budget (i.e., $W\cdot K$), the optimal trade-off between $W$ and $K$ varies by tasks: for controlled-sentiment generation, the best results come from retaining the most promising state and concentrating computational efforts on expanding from it $(W,K=1,16)$; in contrast, for summarization, maintaining multiple hypotheses $(W,K=8,2)$ yields the best results by avoiding local optima. Figure 5 displays the ablation results for $L$ where smaller $L$ benefits controlled-sentiment generation, while an intermediate $L$ is optimal for summarization. These results are consistent with our findings in Figure 3, suggesting that the simple nature of controlled-sentiment generation makes token-level manipulation sufficient and partial return a more reliable indicator of overall return. See Appendix D.1 for extended ablations.

2 Instruction Following

Next, we evaluate weak-to-strong search on a standard single-turn instruction-following benchmark, AlpacaEval 2.0 , which consists of 805 prompts from various open-source datasets. Unlike the previous section where we steer large pre-trained language models (e.g., Llama-2-7b), we now steer large instruction-tuned language models (e.g., Llama-2-7b-chat). This is because (1) instruction-tuned models often require further alignment to match human preferences , and (2) to study weak-to-strong generalization in instruction-following, the models must be proficient at following instructions before steering.

For small language models, we reuse two high-ranking 7B model pairs from the AlpacaEval 2.0 leaderboard as guidance: (1) Zephyr guidance: zephyr-7b-beta and its untuned version mistral-7b-sft-beta; (2) Tulu guidance: tulu-2-dpo-7b and its untuned version tulu-2-7b. All four models use the Llama-2 tokenizer. The large instruction-tuned language models we aim to further align fall into three categories: (1) same vocabulary: Llama-2-7b-chat, Llama-2-70b-chat; (2) cross vocabulary: Llama-3-8B-Instruct, Llama-3-70B-Instruct; (3) black box: gpt-3.5-turbo-instruct. As it is nearly impossible to reproduce the exact training pipeline for these small models ($\pi_{\text{ref}}\rightarrow\pi^{*}$), we do not test the baseline results of directly fine-tuning the large models as in Figure 3. Language model responses are evaluated by their length-controlled win rates (LC WR) against gpt-4-turbo, with gpt-4-turbo serving as the judge.

Experimental results with Zephyr and Tulu guidance are shown in Figure 6 (detailed hyperparameters in Appendix C.2.2). Weak-to-strong search consistently outperforms other test-time baselines with great margins. There are two crucial takeaways worth mentioning: (1) Weak-to-strong search makes strong models stronger with only weak test-time guidance. Take Zephyr guidance for an example (Figure 6, left), even if most large instruction-tuned models $\pi_{\text{base}}$ are stronger than zephyr-7b-beta before steering, weak-to-strong search is still able to enhance their performances using weak models as guidance. Conversely, EFT and BoN mainly interpolate between weak and strong models, resulting in limited, if any, improvements over the strong models. We also tested beam search without external guidance but we found no obvious improvements (Table 2), probably because the latent reward functions behind these language models are not well aligned with the human preference that gpt-4-turbo approximates. The same observations apply to Tulu guidance, even though the tuned tulu-2-dpo-7b is weaker than all the large instruction-tuned language models by significant margins (Figure 6, right). (2) Weak-to-strong search applies to black-box language models. Our method, requiring only sampling from large language models, is also effective for black-box models like gpt-3.5-turbo-instruct. For weak-to-strong search with gpt-3.5-turbo-instruct, we use a chunk length of 100, as the black-box APIs are stateless and do not retain activation caches, making repeated context embedding costly. Despite the long chunk length, our method effectively aligns black-box models, significantly outperforming BoN, a special case of weak-to-strong search (CBS) with infinite chunk length.

Discussion

We have presented weak-to-strong search, an alignment method that keeps the large language model frozen while steering its decoding through a test-time greedy search over small language models. This method builds on the insight that the log-likelihood difference between small tuned and untuned language models can serve both as a dense reward function and a critic, and introduces a novel beam search algorithm designed for balancing reward maximization and KL minimization. This method offers a compute-efficient model up-scaling strategy that eliminates the complexity of directly fine-tuning the large models, and exemplifies weak-to-strong generalization that makes strong models stronger with only weak test-time guidance. Empirically, this approach is effective in controlled-sentiment generation, summarization, and instruction following.

While our work focuses on aligning with human preferences, weak-to-strong search could also apply to tasks like reasoning and coding , where ground truth answers exist. This is because any pair of tuned and untuned language models can act as test-time steering forces, without necessarily being trained on preferences. This then raises several questions beyond the scope of our current study: (1) In our study, we consistently use SFTed policy as the untuned model $\pi_{\text{ref}}$ due to the two-stage nature of preference learning. In general single-stage fine-tuning tasks, does weak-to-strong search still work with a pre-trained model serving as the untuned model $\pi_{\text{ref}}$? (2) Can weak-to-strong search enhance language models in tasks where ground truth answers exist, beyond merely tailoring their knowledge and skills to human preferences?

References

Appendix A Mathematical Derivations for Eq. 6

As introduced in Section 3, the alignment objective under the traditional contextual bandit framing is given by (Eq. 2):

Given language models’ autoregressive nature, we can also view language model alignment as solving a token-level MDP . This token-level MDP is define by the tuple $(\mathcal{S},\mathcal{A},f,r({\mathbf{s}}_{t},{\mathbf{a}}_{t}))$. Here, the state ${\mathbf{s}}_{t}:=({\mathbf{x}},{\mathbf{y}}_{<t})\in\mathcal{S}$ consists of the prompt and all response tokens generated so far; the action ${\mathbf{a}}:=y_{t}$ determines the next token to generate from the vocabulary $\mathcal{A}$; the dynamics $f$ is a deterministic function that updates the state by concatenating the current state and action ${\mathbf{s}}_{t+1}:=({\mathbf{s}}_{t},{\mathbf{a}}_{t})$; $r({\mathbf{s}}_{t},{\mathbf{a}}_{t})$ is a sparse reward that equals $r({\mathbf{x}},{\mathbf{y}})$ if ${\mathbf{a}}_{t}$ is EOS and otherwise. We use $\rho_{\pi}({\mathbf{s}}_{t})$ to denote the state marginals of the trajectory distribution induced by the policy $\pi$. Under the token-level MDP, the objective from Eq. 7 can be written as

where $T$ specifies the response length. The solution of Eq. 8 is given by as:

where the optimal Q-function and V-function satisfies

Here, $V^{*}$ predicts the expected future return (the terminal reward in this sparse reward setting) penalized with the future KL constraint starting from the state ${\mathbf{s}}_{t}$, under the optimal language model $\pi^{*}$. Combining Eq. 9 and Eq. 10, we have

Note that (1) $r({\mathbf{s}}_{t},{\mathbf{a}}_{t})$ is a sparse reward that is non-zero if ${\mathbf{a}}_{t}$ is EOS, and (2) $V^{*}({\mathbf{s}}_{t+1})=0$ if ${\mathbf{a}}_{t}$ is EOS. Then, summing Eq. 12 from timestep $1$ to $H$ yield

where $V^{*}({\mathbf{s}}_{H+1})=V^{*}(({\mathbf{s}}_{H},{\mathbf{a}}_{H}))$ due to the deterministic transition. Eventually, in the context of sequence-level MDP (Section 4), where we define ${\mathbf{y}}$ as a complete response, ${\mathbf{y}}^{\prime}$ as a response that can be either complete or incomplete, we have that ${\mathbf{s}}_{1}=({\mathbf{x}},{\mathbf{y}}_{<1})=({\mathbf{x}},\varnothing)={\mathbf{x}}$ and $({\mathbf{s}}_{H},{\mathbf{a}}_{H})=(({\mathbf{x}},{\mathbf{y}}_{<H}),{\mathbf{y}}_{H})=({\mathbf{x}},{\mathbf{y}}_{<H+1})$. Thus, we can rewrite Eq. 13, with a slight abuse of notations, as

Appendix B Changelog

v1: We corrected an evaluation inconsistency: using $\text{top-k}=50$ for regular decoding (i.e., Base) and EFT, but $\text{top-k}=\infty$ for BoN and weak-to-strong search, which had slightly underestimated results for weak-to-strong and BoN. This version standardizes $\text{top-k}=50$ and updates results from the previous version (the version under review).

Appendix C Further Details on the Experimental Setup

The following table lists the models and their corresponding links.

C.1.2 Hyperparameters Specification

We use fixed hyperparameters across all tested models. We use temperature $T=0.7$, $\text{top-k}=50$ and $\text{top-p}=1.0$ when sampling from the language models. For weak-to-strong search (CBS), we use $W,K,L=4,4,5$ ($W$: beam width, $K$: successors per state, $L$: chunk length). For BoN, we use $N=16$ for fair computational comparison with weak-to-strong search (i.e., $W\cdot K=N$). For EFT, we report the best results among $\beta\in\{1/4,1/2,1,2,4\}$.

C.1.3 Compute Resources Specification

Models are evaluated over 1000 test prompts, on one single NVIDIA A100 GPU.

C.1.4 Gold Reward Model Details

We follow the synthetic setup in which we use the gold reward models to play the roles of humans and provide binary preference labels .

For controlled-sentiment generation, we reuse the publicly available distilbert-imdb as the gold reward model $r_{\text{gold}}$, which is a fine-tuned distilbert-base-uncased on the imdb dataset for classifying movie review sentiments. We define the gold reward $r_{\text{gold}}$ as $\log p(\text{positive}\,|\,x,y)-\log p(\text{negative}\,|\,x,y)$ to encourage positive review. Synthetic preferences are collected using the truncated movie reviews as prompts ${\mathbf{x}}$, and pairwise completions from gpt2-imdb, ranked with $p({\mathbf{y}}_{1}\succ{\mathbf{y}}_{2}\mid{\mathbf{x}})=\sigma(r_{\text{gold}}({\mathbf{x}},{\mathbf{y}}_{1})-r_{\text{gold}}({\mathbf{x}},{\mathbf{y}}_{2}))$, as preferences.

For summarization, we fit a reward model on the summarize_from_feedback dataset as the gold reward model $r_{\text{gold}}$. Specifically, this reward model is fine-tuned from Llama-2-7b with a linear projection head and binary cross entropy loss, using a batch size of $32$, a learning rate of 1e-5 for the projection head, and 5e-6 for other parameters, over one epoch with a cosine learning rate schedule. Synthetic preferences are generated by relabeling pairwise responses in the original dataset with $p({\mathbf{y}}_{1}\succ{\mathbf{y}}_{2}\mid{\mathbf{x}})=\sigma(r_{\text{gold}}({\mathbf{x}},{\mathbf{y}}_{1})-r_{\text{gold}}({\mathbf{x}},{\mathbf{y}}_{2}))$.

Both gold reward models show high validation accuracies, $0.928$ and $0.736$, demonstrating strong correlation with human judgments.

C.1.5 Direct Tuning Details

Direct tuning on the synthetic preferences $\mathcal{D}=\{({\mathbf{x}},{\mathbf{y}}_{w},{\mathbf{y}}_{l})_{i}\}_{i=1}^{N}$ involves two stages: Supervised Fine-Tuning (SFT) and Direct Preference Optimization (DPO) . During SFT, models are trained on both selected and rejected responses using a batch size of $64$, a learning rate of 2e-5, and a cosine learning rate schedule over one epoch. During DPO, we use a $\beta=0.1$, batch size of $256$, a learning rate of 1e-6, and a cosine learning rate schedule over one epoch.

C.1.6 Prompt Template for Sampling from Base Models

When sampling from large pre-trained models, it is crucial to provide clear task-specific instructions. For sentiment-controlled generation, we use a zero-shot prompt:

Here is a movie review from imdb: {prompt}

For summarization, we use a two-shot prompt (the exemplars are selected arbitrarily):

{examplar.prompt}TL;DR: {examplar.response} {examplar.prompt}TL;DR: {examplar.response} {prompt}TL;DR:

C.2 Instruction Following

The following table lists the models and their corresponding links.

C.2.2 Hyperparameters Specification.

When sampling from Llama-3-8B-Instruct and Llama-3-70B-Instruct, we use the $T=0.6$, $\text{top-k}=1.0$ and $\text{top-p}=0.9$ as per the official Llama-3 generation configuration. For other models, we default to temperature $T=0.7$, $\text{top-k}=50$ and $\text{top-p}=1.0$. Specific hyperparameters for each method are detailed in Tables 3 and 4.

C.2.3 Compute Resources Specification.

Models are evaluated on 805 test prompts. Model inference takes place on one single NVIDIA A100 GPU for 7B&8B and black-box models, and on four NVIDIA A100 GPUs for 70B models.

Appendix D Extended Experimental Results

We show the extended CBS hyperparameters ($W,K,L$) ablations in Figures 8 and 7.

D.2 Evaluation Results for Instruction Following

In addition to gpt-4-turbo evaluations, we assess language model responses using two top-ranking reward models from RewardBench : UltraRM-13b and Starling-RM-34B . Table 3 and 4 show that weak-to-strong search consistently outperform other methods across all metrics. We also test vanilla beam search without any external guidance , which does not consistently improve over direct sampling for instruction following (Table 2).

Appendix E Sample Generations

E.2 Summarization Sample Generations