Reward-Guided Tree Search for Inference Time Alignment of Large Language Models (2025)

Abstract: Inference-time computation methods enhance the performance of Large Language Models (LLMs) by leveraging additional computational resources to achieve superior results. Common techniques, such as Best-of-N sampling, Majority Voting, and variants of tree-search algorithm have proven to be effective in boosting the performance of LLMs. These approaches strategically trade increased computational resource for improved model responses. In this work, we proposed DARWIN, an inference-time alignment method that leverage the guidance of a reward model to achieve alignment through reward-guided tree search. Empirical evidences indicates that our method outperform other inference-time alignment methods such as Best-of-N and ARGS on two widely accepted alignment benchmarks AlpacaEval 2 and MT-Bench. Furthermore, we show that our inference-time approach achieves performance comparable to preference-tuned models on both benchmarks, highlighting the effectiveness of trading inference-time compute for enhanced performance during inference.

## introduction having llms generate aligned responses-such as adhering to specific output formats, citing sources, avoiding harmful language, and refusing inappropriate questions-has been extensively explored within the training-time framework. in this context, reinforcement learning from human feedback (rlhf) @xcite , direct preference optimization (dpo) @xcite , and their variants have been shown to be effective and, as such, widely adopted for achieving these alignment objectives. the former uses proximal policy optimization (ppo) @xcite with a preference reward model and the latter minimizes dpo-loss to tune llms to enforce the preferences. on the other hand, o1 @xcite demonstrates the effectiveness of increasing inference-time compute to achieve superior performance, especially on reasoning tasks. this raises a natural question: can llms generate more aligned responses through more inference-time compute? perhaps the simplest yet strong inference-time approach would be best-of-n sampling @xcite and in this work, we demonstrate this method generates more aligned responses. other works on inference-time alignment such as, args @xcite and @xcite , modify the conditional probability using a reward model on the vocabulary at a given decoding step to steer the decoding. however, since reward scores are unbounded, having a fixed reward regulation factor may either overwhelm the llm score or be dominated by the llm score in the overall score, leading to generic or unaligned responses, respectively. another issue could be evaluating the reward score at every decoding step may over-optimize the response toward the reward model @xcite , curtailing the expressiveness of the llm. moreover, these works are not evaluated on general alignment, but rather on some special cases, such as helpfulness vs harmlessness and toxicity. to address these limitations, we proposed an inference-time alignment algorithm through framing this problem as reward guided tree search. analogous to several tree search algorithms, balancing exploration and exploitation are the key ingredients to these algorithms. best-of-n can be viewed as a tree search, where n samples explore different branches, and the exploitation step is selecting the highest reward at the end. inspired by this, we proposed instruction mutation that encourages exploration and reward-guided beam replacement that encourages exploitation to improve the tree search process. our evaluation mirrors the stateof-the-art evaluation of the alignment of methods like orpo @xcite , simpo @xcite , etc, making the general comparison with them straightforward. we empirically show that our method darwin outperforms the inference-time alignment method args @xcite on alpacaeval 2 @xcite and mt-bench @xcite alignment benchmarks. overall, our method makes the following contributions:(i) we demonstrated the effectiveness of scaling inference-time compute to achieve more aligned responses (ii) we proposed a novel inference-time alignment algorithm based on reward-guided tree search that outperforms other inference-time methods, as well as surpassing few preference optimization methods. ## methodology we frame the inference-time alignment problem as a reward-guided tree search problem. our reward model serves as a proxy for human preference. this reward model r θ (s, i) acts as an evaluator of state s, given instruction i. the reward models are ranked by rewardbench @xcite which made it simple for us to choose one based on the empirical performance. formally, each node in the tree is a state representing the decoded tokens: @xmath0, o 2 , • • • , o t ], where i is an instruction and o i are the tokens generated by an llm prompted with instruction i. we say a state has reached an end if an end-ofsentence token is generated. given a seed instruction i, we aim to search for s * = [o 1 * , • • • , o t * ] that maximizes r θ (s * , i). we say that if states s 1 and s 2 are reached under the guidance of instructions i 1 and i 2 , respectively, and r θ (s 1 , i) > r θ (s 2 , i), then i 1 ≻ i 2 . we investigated two primary strategies: reward exploitation and exploration in tree search. reward exploration mutates a given instruction into several instructions, guiding independent search processes. this encourages diverse exploration of the state, potentially discovering higher reward states. however, overexploration can be sub-optimal, so we also employ strategies for exploitation, focusing on leveraging high-reward states discovered during the search process and continuing the search process from there. reward model r θ (•) as a state evaluator. the reward model r θ (s, i) serves as a state evaluator that assigns a scalar reward r to each state on how aligned the currently generated sequence s is based on the instruction i. this serves as a heuristic to guide our tree search process to find the state that has the highest reward. suppose, we are given (i, o w , o l ) a tuple of an instruction, an aligned response, and an unaligned response. given states s w := o w and s l := o l , we expect r θ (s w , i) > r θ (s l , i), where r θ (s, i) assigns higher value to states that represent more aligned output. ## related works inference-time compute methods and llm alignment are quite active research area aimed at aligning llms to human intentions, thus making them more useful. ## conclusion in this work, we demonstrate the effectiveness using inference-time compute for general alignment. we proposed darwin, a inference-time alignment technique that employs evolutionary strategies to implement exploration and exploitation aspects for a more-balanced reward optimization of the generated llm responses. the empirical results strongly indicate our approach's supremacy over existing inference-time alignment methods and competitiveness with preference optimization methods.