T ree B o N : Enhancing Inference-Time Alignment with Speculative Tree-Search and Best-of-N Sampling (2025)

Abstract: Inference-time alignment enhances the performance of large language models without requiring additional training or fine-tuning but presents challenges due to balancing computational efficiency with high-quality output. Best-of-N (BoN) sampling, as a simple yet powerful approach, generates multiple responses and selects the best one, achieving improved performance but with a high computational cost. We propose TreeBoN, a novel framework that integrates a speculative tree-search strategy into Best-of-N (BoN) Sampling. TreeBoN maintains a set of parent nodes, iteratively branching and pruning low-quality responses, thereby reducing computational overhead while maintaining high output quality. Our approach also leverages token-level rewards from Direct Preference Optimization (DPO) to guide tree expansion and prune low-quality paths. We evaluate TreeBoN using AlpacaFarm, UltraFeedback, GSM8K, HH-RLHF, and TutorEval datasets, demonstrating consistent improvements. Specifically, TreeBoN achieves a 65% win rate at maximum lengths of 192 and 384 tokens, outperforming standard BoN with the same computational cost. Furthermore, TreeBoN achieves around a 60% win rate across longer responses, showcasing its scalability and alignment efficacy.

## introduction aligning large language models (llms) with human values is essential for ensuring their outputs reflect human intentions and ethical standards. when data on human preferences is available, a pretrained llm can be fine-tuned to align with these preferences. one popular approach for fine-tuning is reinforcement learning from human feedback (rlhf), where a reward model is trained on a human-labeled preference dataset, followed by reinforcement learning to fine-tune the llm as a policy model @xcite . alternative methods such as direct preference optimization @xcite and its variants @xcite @xcite enable direct alignment via fine-tuning using a contrastive loss, eliminating the need for a separate reward model. this paper focuses on optimizing inference-time alignment of large language models (llms). by leveraging inference-time search, the capability of llms is enhanced during the generation process, improving real-time decision-making. various techniques, such as monte carlo tree search (mcts), have been effectively applied to reasoning, planning, and accelerated decoding tasks @xcite @xcite , demonstrating the potential for better decoding outcomes @xcite . in this work, we aim to explore tree search strategies to further capitalize on decoding-time alignment. our goal is to enhance the quality of alignment while simultaneously reducing the computational cost of inference, providing a more efficient and aligned llm experience. a most simple, yet powerful inference-time alignment method is the best-of-n (bon) method. we start our discussion with bon to motivate our development of more efficient solutions. bon generates multiple sample responses and chooses the best one based on a reward function r(y|x) which characterizes how well-aligned a generated response y is with respect to the given prompt x. more formally, bon aims to approximate the solution to the following optimization problem: max y r(y|x) where the only access to y is through auto-regressively sampling the next token y t from the base policy π base (•|x, y 1:t-1 ), conditioned on the previous tokens. bon generates n samples and selects the response from y 1 , y 2 , . . . , y n that achieves the highest reward model score. due to its simplicity and effectiveness, best-of-n sampling and its variants are widely studied to align llm outputs with human preferences @xcite @xcite @xcite @xcite . also, best-of-n sampling is commonly used in expert iteration and iterative finetuning @xcite , which plays an important role in the alignment of llama2 @xcite and @xcite . in detail, llama2 @xcite combines rejection sampling with proximal policy optimization (ppo) in an iterative fine-tuning process to align llama 2 with human preferences. more recently, llama3 @xcite uses rejection sampling to generate high-quality data for alignment in an iterative process. while best-of-n sampling has proven effective, it has a significant drawback: efficiency. naively implementing bon requires generating n separate responses and the total inference flops scales linearly with n . this not only demands n times more computation but also potentially leads to n times longer latency. the computational overhead can be prohibitively expensive for llms with billions of parameters, particularly when real-time or low-latency responses are needed. some potential solutions involve more intelligent sampling strategies such as pruning to improve efficiency. speculative best-of-n (sbon) @xcite alleviates the problem by continuing the generation of high-quality responses and rejecting the low-quality responses at an early stage of the generation. cascade reward sampling(cards) @xcite use rejection sampling to iteratively generate small semantic segments to form such prefixes, based on rewards computed using incomplete sentences. these accelerated methods are based on the hypothesis that utterances receiving high/low rewards early on in the generation process are likely to yield high/low rewards in the final complete response. however, this hypothesis is too good to be true. in fact, off-the-shelf reward models are typically trained on complete responses, and therefore the score of partial completions by the reward model is usually chaotic and doesn't accurately predict the final output's quality, especially for long responses. our analysis confirmed that rewards of partial completions are not necessarily positively correlated with the final reward (see our experiment results in section 4.2.4 and appendix g). to enable faster, efficient inference-time alignment, we propose to incorporate a tree search strategy into bon sampling, in order to improve the alignment quality as well as reduce the overall in-ference cost. our treebon method maintains an active set of nodes, and actively grows a tree via branching and pruning. in other words, treebon would sample more frequently from good parent nodes but prunes nodes with low predicted rewards. this tree search strategy makes it possible to efficiently explore the search space. another design feature of treebon is the use of implicit reward from dpo-aligned models for guidance of the tree research. dpo @xcite states that the dpo policy model can provide an implicit reward. @xcite further points out that dpo training implicitly learns a token-level reward function. thus, we design treebon to be able to leverage any off-the-shelf dpo model for inference-time decoding of the target model. our extensive experiments show that a weighted combination of implicit dpo rewards would lead to superior, robust performance. our observation is consistent with the fact that one can detect safety levels of the full response using the first few tokens @xcite . our experiments show that under the same computing budget, treebon achieves better performance than bon extensively and stably, with the highest win-rates of 65% on tutoreval @xcite , 63% on alpacafarm @xcite with length 192 and 384, above 60% across hh-rlhf @xcite and @xcite , and increased pass@1 solve rate on gsm8k @xcite as well. by choosing a smaller n, treebon could achieve better performance and improve efficiency at the same time. with only 6.3% of the compute, treebon still maintains a 55% win-rate against bon. on the other hand, sbon can be viewed as a special example of our method with a two-layer tree whose children number is equal to one and bon can be viewed as a two-layer tree with the children number equal to n. treebon has the potential to further improve efficiency than expected by taking advantage of the key-value cache which is especially beneficial to the tree structure since the keys and values of parent tokens can be cached and shared by children. the main contributions of this paper are as follows: ## method in this section, we introduce treebon, a novel inference-time algorithm that enhances alignment quality and efficiency by incorporating a speculative tree-search structure into the best-of-n (bon) sampling framework. treebon iteratively expands high-reward partial responses, pruning low-quality candidates at early stages. the algorithm leverages a weighted implicit reward from a direct preference optimization (dpo) policy model to improve the quality of partial response evaluation. below, we describe the key steps involved in treebon. prompt: were unicorns easily caught in medieval times? ## additional results compared to beam search we also compare treebon with simple beam search. under the same compute, treebon, using the sft model to decode and a dpo-aligned model to provide partial reward, outperforms naive beam search that uses the own probabilities to guide decoding of the same dpo-aligned model in treebon. we conduct an additional experiment that compare treebon using the sfr model with @xmath0. the win-rate is 55.33 ± 2.88% on max length 192, and 58.00 ± 2.85% on 384. ## conclusion treebon is a novel framework that combines the speculative tree-search strategy with best-of-n (bon) sampling and token-level reward guidance modified from dpo implicit reward. through extensive experiments, we show that treebon not only has robust alignment improvements but also maintains efficiency, which provides a potential solution for efficient inference and alignment of llms.