Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Inference Time Alignment with Reward-Guided Tree Search

Chia-Yu Hung, Navonil Majumder, Ambuj Mehrish, Soujanya Poria

TL;DR

DARWIN is proposed, an inference-time alignment method that leverages the guidance of a reward model to achieve alignment through a reward-guided tree search and achieves performance comparable to preference-tuned models on both benchmarks.

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 algorithms have proven to be effective in boosting the performance of LLMs. These approaches strategically trade increased computational resources for improved model responses. In this work, we proposed DARWIN, an inference-time alignment method that leverages the guidance of a reward model to achieve alignment through a reward-guided tree search. Empirical evidences indicates that our method outperforms 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. We have released our codes at https://github.com/declare-lab/darwin.

Inference Time Alignment with Reward-Guided Tree Search

TL;DR

DARWIN is proposed, an inference-time alignment method that leverages the guidance of a reward model to achieve alignment through a reward-guided tree search and achieves performance comparable to preference-tuned models on both benchmarks.

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 algorithms have proven to be effective in boosting the performance of LLMs. These approaches strategically trade increased computational resources for improved model responses. In this work, we proposed DARWIN, an inference-time alignment method that leverages the guidance of a reward model to achieve alignment through a reward-guided tree search. Empirical evidences indicates that our method outperforms 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. We have released our codes at https://github.com/declare-lab/darwin.
Paper Structure (44 sections, 7 equations, 7 figures, 7 tables, 1 algorithm)

This paper contains 44 sections, 7 equations, 7 figures, 7 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Methodology
  3. Reward model $R_\theta(\cdot)$ as a state evaluator.
  4. Mechanisms of Exploration
  5. Sample N Generations (Sample N).
  6. Instruction Mutation and Response Generation.
  7. Reward Exploitation
  8. Best Reward Selection from N Samples (Best-of-N).
  9. Reward-guided Beam Replacement.
  10. Darwin
  11. Multiple Beams per Mutations.
  12. Experiments
  13. Models and Settings.
  14. Baselines
  15. Evaluation Benchmarks
  16. ...and 29 more sections

Figures (7)

  • Figure 1: The stages of Darwin.
  • Figure 2: An illustration of our proposed framework, Darwin, for inference-time alignment. It executes the following steps in a sequence: (1) an archive of instructions is initialized with the input instruction, (2) a sample from the archive is sampled and modified to create mutated instructions, (3) decode for a replacement period, (4) evaluate the quality of the new tokens w.r.t. the original instruction with a reward model, (5) replace the worse quality generations with the better ones, and (3) decode for a replacement period, and so on. For brevity, we omitted the factor of replacement period $m$ in the state superscript. Thus, $s^{(i)}_j$ in the figure is equivalent to $s^{(im)}_j$ in \ref{['sec:exploitation']}.
  • Figure 3: The impact of Mutation and Replacement on the WR and LC metrics.
  • Figure 4: An illustration of the average (a) Jaccard similarity and (b) rank-biased overlap (RBO) between the sets of the top-$k$ rewarded beams in two consecutive replacement cycles, where $k=3$.
  • Figure 5: A plot of the probability of the winning instruction of a mutation cycle appearing with the top-k beams of a given replacement step in mutation cycle (a) 1, (b) 2, and (c) 3.
  • ...and 2 more figures