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

Authors: zhanhui  zhou
zhixuan  liu
jie  liu
zhichen  dong
chao  yang
yu  qiao
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-13 15:12:53

Abstract: Large language models are usually fine-tuned to align with human preferences. However, fine-tuning a large language model can be challenging. In this work, we introduce $\textit{weak-to-strong search}$, framing the alignment of a large language model as a test-time greedy search to maximize the log-probability difference between small tuned and untuned models while sampling from the frozen large model. This method serves both as (1) a compute-efficient model up-scaling strategy that avoids directly tuning the large model and as (2) an instance of weak-to-strong generalization that enhances a strong model with weak test-time guidance.
Empirically, we demonstrate the flexibility of weak-to-strong search across different tasks. In controlled-sentiment generation and summarization, we use tuned and untuned $\texttt{gpt2}$s to improve the alignment of large models without additional training. Crucially, in a more difficult instruction-following benchmark, AlpacaEval 2.0, we show that reusing off-the-shelf small models (e.g., $\texttt{zephyr-7b-beta}$ and its untuned version) can improve the length-controlled win rates of both white-box and black-box large models against $\texttt{gpt-4-turbo}$ (e.g., $34.4\% \rightarrow 37.9\%$ for $\texttt{Llama-3-70B-Instruct}$ and $16.0\% \rightarrow 20.1\%$ for $\texttt{gpt-3.5-turbo-instruct}$), despite the small models' low win rates $\approx 10.0\%$.

Main paper: 
 introduction: 
 Learning-based algorithms [1, 2, 3, 4, 5] have become the standard approach for aligning large language models (LLMs) with human preferences [3, 6, 7, 8, 9, 10] . However, fine-tuning large language models is resource-intensive and difficult to implement [4] . 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 [11, 12, 13, 14, 15, 16, 17] . Typical examples of search-based algorithms include rejection sampling [16, 17] and Monte Carlo Tree Search [18, 19] . These searchbased algorithms are promising as they can reuse the same guiding signal 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 [16, 17] , which limits their steerability, or require a value function learned from scratch to address preference reward sparsity and prune search space [13, 18, 14] , 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-probabilities of small language models. This algorithm makes two contributions:(1) First, it builds on the theoretical foundation of the token-level MDP for alignment [20] , using the 38th Conference on Neural Information Processing Systems (NeurIPS 2024). zephyr-7b-beta (tuned) mistral-7b-sft-beta (untuned) Base Weak-to-strong search Figure 1 : Weak-to-strong search enhances the alignment of large models through test-time guidance from small models (dashed lines). This method is applicable to white-box models that use the same or different vocabularies as the small models, as well as to black-box models. We present the results for the instruction-tuned models from each family (e.g., Llama2-7B denotes Llama-2-7b-chat).log-probability difference between small tuned and untuned language models as both reward and value [4, 20] to guide the decoding of a large model (Section 4.1). Theoretically, this formulation is suitable for search as it converts the otherwise sequence-level sparse preference reward function to a per-token dense reward function, which can be summed up as a value function [20] . 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 value 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 [21] 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 [22] and summarization [2] , 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) [23] , Llama-2 [7] and Llama-3 [24] families, at least as effective as existing methods. Then, in a more difficult instruction-following benchmark, AlpacaEval 2.0 [25] , 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% → 37.9% for Llama-3-70B-Instruct, and 16.0% → 20.1% for gpt-3.5-turbo-instruct), despite the small models' low win rates ≈ 10.0% (Figure 1 ).The following table lists the models and their corresponding links.When sampling from Llama-3-8B-Instruct and Llama-3-70B-Instruct, we use the T = 0.6, top-k = 1.0 and top-p = 0.9 as per the official Llama-3 generation configuration. For other models, we default to temperature T = 0.7, top-k = 50 and top-p = 1.0. Specific hyperparameters for each method are detailed in Tables 3 and 4 .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.We show the extended CBS hyperparameters (W, K, L) ablations in Figures 8 and 7 . beam width (W ) beam width (W ) With the same computation budget (i.e., same W K), the optimal hyperparameters differ by tasks.In addition to gpt-4-turbo evaluations, we assess language model responses using two top-ranking reward models from RewardBench [42] : UltraRM-13b [55] and Starling-RM-34B [56] . 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 [20] , which does not consistently improve over direct sampling for instruction following (Table 2 ).Llama- In Section 5, we have demonstrated that by converting weak language models to a per-token dense reward function (Eq. 4), we can further improve strong models at test time without additional training. This section presents preliminary experiments showing that this reparametrized dense reward function can also benefit RL fine-tuning. Using the same setup from Section 5, we start with a dense reward function parametrized by tuned and unturned gpt2 pair (π * , π ref ) and then we tune two larger models under this dense reward on sentiment-controlled generation using PPO [57] . Figure 9 shows the fine-tuning results with different reward sparsity, where we find that dense reward function does benefit RL fine-tuning, and weak-to-strong generalization [21] also occurs consistently during training when the reward signal comes from weaker language models (note that π * only achieve a 4.6 gold reward on sentiment-controlled generation from Figure 3 ). a huge fan of the Godzilla series, I do enjoy watching all of the films in the franchise, and I enjoyed watching this one with a good friend. The film has a very good story, with some great action scenes. The movie was very well [truncated for brevity] 5.16Direct Tuning to be "old" at this point in my life, I still love this movie. It's a very entertaining film. 5.52Base to be a fan of the James Bond series, I do like the movie series. The most recent is a film with a highly complex plot, but it is not the best of the series, and it's not a good idea to watch it [truncated for brevity] 2.04EFT (β * ) NA BoN (16)a fan of the franchise, I have seen most of the movies in this series, so I could appreciate this one to a certain extent. However, I'm not sure if it is a good movie. The characters are not as well [truncated for brevity] -1.65Weak-to-strong search (4, 4, 5) a fan of the original trilogy, I do think this is a fantastic remake. It's essentially the same story, yet it's different enough to be enjoyable. There's also a lot of nods to the [truncated for brevity] 4.89Direct Tuning a huge fan of romantic comedies,Ï do consider this film to be a delightful movie. While I can see how people might not like this film, I can also see why people would like it. I think it's [truncated for brevity] 4.73Base an intellectual, I do consider myself a fan of cerebral movies, and this is easily the most intelligent movie I've ever seen. I'm not sure if I'd call this film uplifting, but it is without a doubt inspiring and thought-provoking [truncated for brevity] 5.16EQUATIONa huge fan of either comic book or of super heroes, I have been very impressed by the overall quality of the Marvel's cinematic universe. This movie is a great addition to the series. The story is complex and multi-layered, with the characters [truncated for brevity] 5.33Weak-to-strong search (4, 4, 5) a horror fan, I do love a good story which is why I thoroughly enjoyed this movie. The story is told well, very well. The acting is great, and the special effects are great too. I thought the story was very well done. [truncated for brevity] 5.41Direct Tuning a ḧorror fan, I enjoy a very wide variety of genres. However, it is my belief that few genre films are ever truly ḧorrorör cinematic films; they are all variations on a theme. I think To give a brief background, my boyfriend and I have been dating for about 9 months. He's a physician in an intense fellowship program and generally very stressedtiredbusy, I'm currently in school, and stay pretty busy with an internship and working. We spend almost everyday together, support each other, have a lot of fun and both refer to the other as our best friend (that we enjoy touching inappropriately.)For the first couple months we were very happy together. The issues didn't start until I brought up the dreaded question, So what are we? ¨. Then began a bit of a struggle. I always wanted more from the relationship and was greeted by hesitancy, by 3 months I asked for exclusivity and he eventually agreed but only because I was going to leave if he didn't (I wanted something serious.) by 6-7 months I said I love you, he didn't say it back. This is totally fine, painful, but fine. I said I'd be patient with him until he is ready. He still hasn't said it back, but now at 9 months he said that he feels like he is falling in love, but isn't in love yet. This progress has given me hope, but not being on the same page has made me feel a bit insecure about our relationship. He has mentioned that he doesn't want to fall in love unless he knows that we will have a serious future together, so he doesn't get hurt like he did in his 2 prior relationships. :(We have been fighting a lot recently because: I still haven't met his family (they live 5 minutes from him), he doesn't want to share our relationship on Facebook (not the biggest deal!), he isn't in love with me, we don't really talk about the future, and the quality of time spent together isn't the best due to the nature of his work and how stressed and tired he is all of the time. After expressing my concerns about these things, he said he doesn't want to lose me and will try to open up to me. The fighting is really getting to us both. , with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 
related work: 
 Large unsupervised language models trained on internet-scale corpus acquire broad knowledge and abilities [26, 27, 28] . 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 [1, 2, 3, 4, 7, 24, 6, 29] . 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 [3, 30, 13, 31, 32, 33] , 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 [16, 17] , which limits their steerability, or require a value function learned from scratch to handle sparse preference rewards and prune search space [13, 18] , which can be as difficult as training a large language model from scratch. Our work avoids these issues by parametrizing the sparse preference reward function with the log-probability difference between small tuned and untuned language models [4] . 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. [20] proposes a token-level MDP interpretation for language model alignment, demonstrating that a greedy probability search over a trained language model can achieve improvements over regular decoding. Our work builds on their theoretical foundations and proposes a practical 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 [34, 12, 11, 35] , 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 prioritizes 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 [1] :EQUATIONwhere p(x) is a distribution of prompts, y is the complete language model response, r is a preference reward function that encourages human-preferred responses, and D KL limits how far the optimized language model π can deviate from the reference (untuned) model π ref .There are two main categories of alignment algorithms: (1) search-based algorithms that optimize Eq. 1 with graph-based search during inference [16, 13, 18, 19, 11, 12] , and (2) learning-based algorithms that optimize Eq. 1 through gradient descent, aiming for a parametrized optimal language model [1, 36, 4, 37] . Our work falls in the first category, proposing a search-based algorithm capable of using small language models to guide the decoding of a large language model to align with human preferences.The analytical solution to Eq. 1 can be obtained through the following Lagrangian [38, 39] :EQUATIONwhich has a well-known closed-form solution that expresses a duality between the reward function r(x, y) and the optimal language model π * (y | x) [40, 41] :EQUATIONwhere Z(x) = y π ref (y | x) exp 1 β r(x, y) 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 [1, 2, 3] , 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 (π * , π ref ) to reparametrize the reward function [42] ; (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 (π * , π ref ) during reward modeling [4] . 
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-probability difference between small tuned and untuned language models. First, we discuss how using language models to parametrize the preference reward function (Eq. 1) makes the reward-maximization problem solvable by a simple greedy search algorithm (e.g., beam search) (Section 4.1). Then, we introduce a practical beam search method, Chunk-level Beam Search (CBS) (Section 4.2), that balances reward maximization and KL minimization, 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(x, y), based on the Bradley-Terry model [43] , only emits a terminal reward when the model response is complete. Search-based algorithms often struggle without any intermediate rewards or a value function providing intermediate guidance [44, 45] . However, if we parameterize this sparse reward function with language models (Section 3.2), we can obtain both a dense reward function and a value function simultaneously.Language models as a dense reward function. To obtain a dense reward function, we leverage the duality between the sparse preference reward function and the dense language model probability (Eq. 3). By explicitly factorizing the log-probability of a complete response y under the language models, we obtain a sum-of-rewards style formulation for Eq. 3:EQUATIONwhere y <t denotes the response tokens from 1 to t -1, and the last response token y |y| is the EOS token. Combining Eq. 1 and 4, we rewrite the original objective with a per-token reward function:EQUATIONwhere β and Z(x) are omitted as they do not influence the optimal solution. It is important to note that the reference model that parametrizes the reward function (π ref ) (Eq. 5a) and the reference model that constrains the test-time search space (π base ) (Eq. 5b) can be different. Practically, decoupling the reference models is useful as it allows using a tuned and untuned language model pairnamely (π * , π ref ) -to steer the decoding of any base language model π base without retraining.Setting aside the KL constraint (Eq. 5b) for now, we can apply existing search algorithms like beam search [45, 20] to optimize Eq. 5a. Beam search is often criticized for leading to myopic solutions [45] , as it tends to greedily prioritize states (x, y ′ ) (y ′ is incomplete) 1 with high cumulative reward log π * (y ′ | x)log π ref (y ′ | x) midway through generation, which is generally viewed as poorly correlated with the overall return we care about. While this criticism is valid for most MDPs, we argue that in the token-level MDP [20] of our case, the cumulative reward mid-generation is actually a reliable indicator of the long-term value, making beam search less myopic.Cumulative reward under language models as a value function [20] . Appendix A shows that:EQUATIONwhere V * (x, y ′ ) denotes the value function, predicting the expected terminal reward under the optimal π * in the original KL-constrained sparse reward problem. Although V 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 π 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 model to steer π base , tuned model π * , and untuned model π ref .H = {(x, y ′ ) i } W i=1 of W states.3: Output: optimal terminal state (x, y)4: Initialize H = {(x, y ′ = ∅) i } W i=1 5: while ∃(x, y ′ ) ∈ H such that y ′ is incomplete do 6: Initialize C = {} 7: for each (x, y ′ ) ∈ H do 8: Y ← {(y L ) i } K i=1 i.i.d. ∼ π base (• | x, y ′ ) // y L = ∅ if y ′ is complete 9: C ← C ∪ {(x, y ′ • y L ) | y L ∈ Y} 10:end for 11:H ← Top-W (x,y ′ •y L )∈C (log π * (y ′ • y L | x) -log π ref (y ′ • y L | x)) 12: end while 13: return arg max (x,y)∈H (log π * (y | x) -log π ref (y | x))CBS is a unified framework that encompasses several search-based algorithms: (1) CBS with W = 1, K = N , L = ∞ (i.e., infinite chunk length) is equivalent to BoN sampling with log π * (y | x)log π ref (y | x) as the scoring function, and (2) CBS with K = ∞, L = 1 (i.e., exploring 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 to achieve the best of both worlds: (1) Using a finite chunk length allows CBS to prune bad states during generation, enhancing steerability more efficiently compared to BoN. (2) Sampling from π base with limited successor exploration implicitly enforces the KL-constraint from π base (Eq. 5b); otherwise, integrating the KL-constraint into the objective (Eq. 5a) would be necessary for token-level search, but this can be challenging, especially when vocabularies of models differ or with black-box language base models π base whose log-probabilities are inaccessible.Computation costs. In practice, CBS samples W K continuation chunks in parallel from the frozen base model π base and prune states by calling tuned and untuned model pair (π * , π ref ) every L tokens. Larger W K and smaller L enhance steerability at the cost of increased computations. Note that high steerability, while beneficial, is not always ideal as it may lead to large KL deviation and over-optimization [16] .The most practical use of CBS occurs when the tuned and untuned models, (π * , π ref ), are smaller than the model to steer, π base . (1) First, this instance serves as a model up-scaling strategy, directly tuning a small model π ref → π * , 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 (π * , π ref ) are usually weaker than the large model to steer π base , this instance also exemplifies weak-to-strong generalization [21] , 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 [22] and summarization [2] , 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 [25] , 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).Baselines. In addition to weak-to-strong search, we evaluate several existing test-time approaches that steer a large language model π base using small tuned and untuned language models (π * , π ref ):(1) Base: we explore regular decoding from the frozen large language model with n-shot prompting (see Appendix B.1.6 for prompt details). ( 2) Best-of-N Sampling (BoN) [16, 17] : BoN uses r = log π * (y | x)log π ref (y | 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 K response chunks in parallel, for fair computational comparisons, we always ensure N = W K.(3) Emulated Fine-Tuning (EFT) [34, 12, 11, 35] : EFT approximates the results of directly finetuning the large language model by sampling from log π EFT (yt | x, y <t ) ∝ log π base (y t | x, y <t ) + β -1 (log π * (y t | x, y <t )-log π ref (y t | x, y <t )), where β 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.Setup. For these two tasks, we follow the synthetic setups from [16, 46, 4] , assuming access to a gold reward model r gold . For controlled-sentiment generation, r gold encourages positive continuations of movie reviews, while for summarization, it encourages high-quality summaries of Reddit posts (details in Appendix B.1.4). We generate synthetic preference datasets D = {(x, y w , y l ) i } N i=1 from r gold with p(y 1 ≻ y 2 | x) = σ(r gold (x, y 1 )r gold (x, y 2 )) to mimic human feedback [47] .To obtain the small language models, we optimize gpt2 (124M parameters) using the standard DPO pipeline [4] : (1) we first obtain the reference model π 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 π ref as the reference policy to obtain the optimal language model π * . 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.GPT2-large GPT2-xl Llama2-7B Llama3-8B 0 1 2 3 4 5 Gold RM (r gold ) Controlled-Sentiment Generation GPT2-large GPT2-xl Llama2-7B Llama3-8B -2 -1 0 1 2 Summarization GPT2-DPO (π * ) GPT2-SFT (π ref ) Base (π base )Weak-to-strong search (4, 4, 5) BoN ( 16) Given the tuned and untuned (un-DPO-tuned) gpt2 pair (π * , π 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.EFT (β * ) Directly tuned. 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 (β * = 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" → "love"). However, in the more complex task of summarization, where broader subsequence-level manipulations are essential, weak-to-strong search excels. Please refer to Appendix D 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.Chunk-level Beam Search ablations. We perform additional 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 and K. With the same computation budget (i.e., W 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 probably because it helps avoid 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 from Figure 3 , suggesting that the simple nature of controlled-sentiment generation makes token-level manipulation sufficient and cumulative reward mid-generation a more reliable indicator of overall return. See Appendix C.1 for extended ablations on more models. beam width (W ) ). We show the mean rewards (± standard deviations) across three random seeds.Setup. Next, we evaluate weak-to-strong search on a standard single-turn instruction-following benchmark, AlpacaEval 2.0 [25] , 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 [48] , 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; and (3) black box: gpt-3.5-turbo-instruct. As it is nearly impossible to reproduce the exact training pipeline for these small models (π ref → π * ), 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.Results. Experimental results with Zephyr and Tulu guidance are shown in Figure 6 (detailed hyperparameters in Appendix B.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 π 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 over the strong models without external guidance [20] but we found no obvious improvements compared with regular decoding (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-tostrong 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 relatively long chunk length of 100, as the black-box language model APIs are stateless and do not retain activation caches, making repeated context embedding costly. Despite the long chunk length, our method still effectively improves the alignment of black-box models, significantly outperforming BoN, a special case of weak-to-strong search (CBS) with infinite chunk length.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes] Justification: See Appendix B. We also submit our code in a zip file along with the instructions to reproduce the experimental results.Guidelines:• The answer NA means that paper does not include experiments requiring code. We run experiments across multiple random seeds for the two synthetic experiments and report the means and standard deviations (Figure 3 ). In instruction-following experiments, where GPT-4 is used as the evaluator, we can not afford multiple tests (each test costs about 5 dollars). However, we verified the GPT-4 evaluations using other complementary metrics (Tables 3 and 4 ). Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [Yes] Justification: The positive societal impact is that our method is able to better align language models with human values, which include helpfulness and harmlessness. We do not expect negative societal impacts as a direct result of the contributions in our paper. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
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-probability difference between small tuned and untuned language models can serve both as a dense reward function and a value function, and then 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 [21] that makes strong models stronger with only weak test-time guidance. Empirically, this approach is effective in controlled-sentiment generation, summarization, and instruction following.Limitations & Future Work. While our work focuses on aligning with human preferences, weakto-strong search could also apply to tasks like reasoning [49, 50] and coding [51] , 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 π ref due to the two-stage nature of preference learning; however, in single-stage fine-tuning tasks, does weak-to-strong search still work with a pre-trained model serving as the untuned model π ref ? (2) Although our method shows consistent weak-to-strong generalization across diverse alignment tasks, it is also critical to probe its potential failure modes [52] . 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? (3) Additionally, while our work mainly focuses on how dense language model reward function (Eq. 4) benefits language model decoding at test time, it's also worth exploring the potential benefits of this reward parametrization for RL tuning. Although Appendix C.3 presents some promising preliminary results, we leave further analysis for future work.A Mathematical Derivations for Eq. 6As introduced in Section 3, the alignment objective under the traditional contextual bandit framing is given by (Eq. 2):EQUATIONwhere p(x) is a distribution of prompts, y is the complete language model response, r is a reward function that encourages human-preferred responses, and D KL limits how far the optimized language model π can deviate from the reference model π ref .Given the autoregressive nature of language models, we can frame language model alignment as solving a token-level MDP [20] . This token-level MDP is define by the tuple (S, A, f, r(s t , a t )).Here, the state s t := (x, y <t ) ∈ S consists of the prompt and all response tokens generated so far; the action a := y t determines the next token to generate from the vocabulary A; the dynamics f is a deterministic function that updates the state by concatenating the current state and action s t+1 := (s t , a t ); r(s t , a t ) is a sparse reward that equals r(x, y) if a t is EOS and 0 otherwise. We use ρ π (s t ) to denote the state marginals of the trajectory distribution induced by the policy π. Under the token-level MDP, the objective from Eq. 7 can be written asEQUATIONwhere T specifies the response length. The solution of Eq. 8 is given by [40] as:EQUATIONwhere the optimal Q-function and V-function satisfiesEQUATIONEQUATIONHere, V * predicts the expected future return (the terminal reward in this sparse reward setting) penalized with the future KL constraint starting from the state s t , under the optimal language model π * . Combining Eq. 9 and Eq. 10, we haveEQUATIONNote that (1) r(s t , a t ) is a sparse reward that is non-zero if a t is EOS, and (2) V * (s t+1 ) = 0 if a t is EOS. Then, summing Eq. 12 from timestep 1 to H yieldEQUATIONwhere V * (s H+1 ) = V * ((s H , a H )) due to the deterministic transition. Now, returning back to the sequence-level MDP (Section 4), where we define y as a complete response, y ′ as a response that can be either complete or incomplete, we have that s 1 = (x, y <1 ) = (x, ∅) = x and (s H , a H ) = ((x, y <H ), y H ) = (x, y <H+1 ). Thus, we can rewrite Eq. 13, with a slight abuse of notations, asEQUATION)Here is a movie review from imdb: {prompt}For summarization, we use a two-shot prompt (the exemplars are selected arbitrarily):{examplar [1] .prompt}TL;DR: {examplar [1] .response} {examplar [2] .prompt}TL;DR: {examplar [2] .response} {prompt}TL;DR:B.2 Instruction Following B.2. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)? Answer: [NA] Justification: This study does not involve crowdsourcing nor research with human subjects. Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.y denotes a complete response, while y ′ denotes a response that can be either incomplete or complete. 
 
 Appendix:  


Decision: This is a solid paper that proposes to align larger language models by using the difference in logits from smaller (tuned vs. untuned) models. The idea is interesting, and backed up by solid empirical experiments. There were some initial concerns as to the use of stronger baselines (in particular, those that incorporating search during decoding), as well as its position relative to related work. These have been largely addressed through the rebuttal. 

After reading the paper, I recommend that the authors further discussion the limitations of the proposed method (e.g., this method clearly would not work if the difference in capabilities of the two models is too large (e.g., n-gram model vs. a neural LM).

Reviews: review: paper_summary: paper_summary: This paper addresses the alignment of large language models without the need for fine-tuning. It conceptualizes the alignment process as a search problem, leveraging the log-likelihood difference between small tuned and untuned language models as both a reward and a critic. By transforming a sparse preference reward into a per-token dense reward, the method achieves weak-to-strong generalization. The effectiveness of the proposed approach is empirically validated through controlled sentiment generation, summarization, and an instruction-following benchmark.
strength_weakness: strength_weakness: 1. The paper does not include comparisons with existing decoding-based alignment baselines [1,2].

2. The token-based MDP formulation is largely derived from existing work [3]. The main innovation of this paper lies in using the log-likelihood difference between two small models for guidance based on the existing formulation.


[1] ARGS: Alignment as Reward-Guided Search, ICLR 2024

[2] Controlled Decoding from Language Models, ICML 2024

[3] From r to Q*: Your Language Model is Secretly a Q-Function.
questions: questions: What is the difference between the formulation in section 4.1 and the results in [3]?
limitations: limitations: The paper has discussed limitations.
paper_summary: paper_summary: The paper proposes a search/decoding method "weak-to-strong search" for improving LLM's performance at test time by leveraging the log likelihood from small language models. Specifically, it utilizes the log likelihood differences between tuned and untuned small models to guide the decoding of larger models. This approach avoids the computationally intensive process of fine-tuning large models by using smaller, more manageable models for alignment at inference time, guiding the large model to generate better responses. The paper demonstrates the effectiveness of this method across various tasks, including controlled-sentiment generation, summarization, and instruction-following tasks.
strength_weakness: strength_weakness: * Although the method avoids the computational costs associated with fine-tuning, the paper does not explain clearly the overhead introduced during inference. How does weak-to-strong search impact memory usage and inference speed?
* How does the proposed method compare to other inference-time techniques, such as in-context learning (using few-shot examples) and prompting methods like CoT?
* The method appears to be broadly applicable and not strictly limited to a weak-to-strong approach. For smaller models, such as Llama3-8B as used in the experiments, why would this method be preferred over techniques like LoRA fine-tuning, which can be performed at a low cost and yield better performance? Additionally, if the target model has already been fine-tuned, can the proposed method still enhance its performance? 
* (minor) It would be interesting to explore how weak-to-strong search could enhance a model's reasoning abilities, for example on math datasets like GSM8K.
questions: questions: See weaknesses above.
limitations: limitations: Yes, the authors discussed limitations in the paper.
paper_summary: paper_summary: The paper introduces the "weak-to-strong search" method for aligning a stronger large language models (LLMs) by leveraging two weaker LLMs during test-time without requiring fine-tuning of the large models. This approach aims to improve alignment by maximizing the log-likelihood difference between tuned and untuned small models while using the frozen large model for decoding. The method demonstrates effectiveness in various tasks, including sentiment-controlled generation, summarization, and instruction following.

I generally think this is an inspiring paper towards acceptance. The reason I rate it a weak accept rather than accept is I personally think an ablation study of using the log-likelihood difference to conduct token/chunk-wise PPO is necessary. It can make the paper more complete but lacking it is kinda fine though.
strength_weakness: strength_weakness: 1. Obtaining a tuned weaker model is still not a very trivial thing, and the performance of $\pi_\text{base}$ on downstream tasks could heavily depend on the tuned weaker model as well. There is not too much discussion regarding it.
2. As we can formulate a partial reward function using a tuned/untuned model pairs, apart from using it to guide decoding, a natural next step/ablation study is using it to conduct PPO. It is worth trying this experiment and compare the downstream performance as well as the overall cost.
questions: questions: N/A no questions here; the paper is presented very clearly.
limitations: limitations: The paper includes a limitation section at the end.
paper_summary: paper_summary: In this paper, they frame alignment as a decoding time problem, allowing the large language model to be frozen. They do this by parametrizing a reward function with the difference in the log-likelihood between small untuned and tuned language models. Interestingly, their approach does not require shared vocabulary and is applicable to black-box LLMs. They perform experiments on controlled-sentiment generation, summarization, and instruction following.
strength_weakness: strength_weakness: Although there’s a related work section, I believe the comparison between the proposed method and other methods, e.g., the ones that rely on small language models to aid the alignment of LLMs, should be described in more detail. In particular, the novelty of the proposed method should be made more evident; for non-experts, I think it’s a bit difficult to understand if some parts are novel or if they already exist in previous work. I also think that some citations might be missing (see my questions below).
questions: questions: Questions: 
- Is it \pi_base in L181?
- What about the vocabulary of the small tuned and tuned models? Does it need to be shared?
- How does your work compare with that of Li et al. (ACL 2023)? What about Zhao et al. (2024)?

Minor comments: 
- BoN is defined in L204 but it’s used before (e.g., L166).

References:

- Contrastive Decoding: Open-ended Text Generation as Optimization (Li et al., ACL 2023)
- Weak-to-Strong Jailbreaking on Large Language Models (Zhao et al., 2024)
limitations: limitations: Yes.

Comments: title: Thank you for your rebuttal
comment: Thanks for your response and feedback!
Thanks for your response. I appreciate the rebuttal and have read other reviews. I maintain my score as 6 and will engage in further discussion with other reviewers and AC if needed.
Thanks for your response and feedback!
Thank you for answering my questions. I have read the rebuttal and the other reviews. I believe that incorporating some of this discussion in the updated version of the paper is a good idea.
I totally get that it's too early to finalize the rating. Thanks again for your thoughtful comments and suggestions!
It's too early to talk about the rating. I'll reach out if I need further clarifications. I might recalibrate my rating during the AC-reviewer discussion period but it depends on the discussion and other reviewers as well.
Thank you for your prompt response! And I wanted to kindly remind you that you rated our paper as "weak accept" due to some of the weaknesses you mentioned. If we have satisfactorily addressed these weaknesses in our rebuttal, would you be open to increasing your score? If not, are there any other clarifications I can provide to assist in your evaluation?
Thanks for your response and additional response! It currently looks good to me.

I'll follow up if there are further questions. So far so good.