Inference-Aware Fine-Tuning for Best-of-N Sampling in Large Language Models

Authors: Yinlam Chow, Guy Tennenholtz, Izzeddin Gur, Vincent Zhuang, Bo Dai, Aviral Kumar, Rishabh Agarwal, Sridhar Thiagarajan, Craig Boutilier, Aleksandra Faust
Primary area: foundation or frontier models, including LLMs
Keywords: Best-of-N sampling, Reinforcement Learning, Language models

Abstract:
Recent studies indicate that effectively utilizing inference-time compute is crucial for attaining good performance from large language models (LLMs). Specifically, the Best-of-N (BoN) inference strategy, where an LLM generates multiple responses and a verifier selects the best, has shown strong empirical performance. Motivated by this, we develop a novel inference-aware fine-tuning paradigm, which encompasses the BoN-aware inference framework as a special case. We devise the first imitation learning and reinforcement learning (RL) methods for fine-tuning LLMs using BoN, overcoming the challenging, non-differentiable argmax operator in BoN. We empirically demonstrate that our BoN-aware models implicitly learn a per-example "meta-strategy", which interleaves best responses with more diverse responses that might be better suited to a test-time input—a process reminiscent of the exploration-exploitation trade-off in RL. Our experiments demonstrate the effectiveness of BoN-aware fine-tuning in terms of improved performance and inference-time compute. In particular, we show that our methods improve the BoN performance of Gemma 2B on Hendrycks MATH from 26.8% to 30.8%, and Pass@K from 60% to 67%.

Decision: accept (poster)
Reviewer scores: 5, 6, 6

Reviews and discussion:
<thread>
creation date:
1730035621219

thread_type:
official_review

<thread_object>
creator: reviewer_gx5h

<content>
summary:
This paper proposes an inference aware fine-tuning strategy for Best-of-N (BoN) to overcome the non-differentiable argmax operator for BoN. The proposed strategy incentivizes the model to balance the best and diverse responses for fine-tuning, achieves better performance compared to the base Gemma 2B model on the MATH dataset.

strengths:
1. The incorporating of inference strategy into training is a compelling and well-founded idea.
2. The proposed BoN-aware fine-tuning to balance the trade-off between the best and diverse responses is well-motivated.

weaknesses:
1. The experiments are conducted on a base model (Gemma 2B) and an evaluation dataset (MATH). More evaluation results on different base models and evaluation datasets are needed.
2. What is the underlying rationale for the BoN-SFT loss function presented in Equation (7)? The function appears to penalize y', even in scenarios where the sampled y' demonstrates high quality, as evidenced by r(x, y') approaching r(x, y). This raises a question: If the sampled y' exhibits desirable characteristics, wouldn't it be more beneficial for the policy model to learn from these high-quality responses?
3. It would be valuable to conduct a comparative analysis of BoN-SFT against several straightforward baseline approaches. These could include:

(1) Fine-tuning the base model using the highest-quality response selected from N generated samples.

(2) Fine-tuning the base model utilizing the ground truth responses.

(3) Fine-tuning the base model with all the sampled N responses.

(4) Fine-tuning the base model with the combination of all the sampled N responses and the ground truth responses.

questions:
See "Weaknesses".

I would be happy to discuss further with the authors and reassess my score based on the rebuttal stage.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730604134519

thread_type:
official_review

<thread_object>
creator: reviewer_vbde

<content>
summary:
This paper presents a finetuning approach for large language models (LLMs) that incorporates the inference strategy directly into the model’s training process, with a particular focus on Best-of-N (BoN) inference. The authors argue that there is a disparity between the best-of-N policy typically used at inference time and the policy learned through standard supervised fine-tuning (SFT). To address this, they formalize an inference-aware fine-tuning problem that integrates the inference strategy into the training objective. Specifically, they develop a supervised BoN-aware fine-tuning method and a BoN-aware reinforcement learning (RL) with binary rewards, using policy gradients derived only from positive examples to mitigate data inefficiency. The proposed methods are evaluated on the Gemma 2B model, primarily using the Hendrycks MATH benchmark, with an emphasis on understanding the relationship between the number of samples and temperature.

strengths:
- The paper addresses an important and timely problem, as best-of-N inference with verifiers has become increasingly popular for reasoning tasks.
- The paper provides a solid theoretical framework supporting the proposed methods with a clear formalization of inference-aware training objectives.

weaknesses:
- The method is only evaluated on a single model and a single task, which limits insights into its broader applicability.
- Although BoN-aware fine-tuning is a key contribution, the experiments predominantly examine the relationship between the number of samples and temperature. This leans more toward an analysis of the proposed methods than a comparative assessment against strong baselines, leaving questions about its competitive advantages unanswered.

questions:
1. Will inference-aware SFT and inference-aware RL reduce the model’s generalizability? How does the model perform when using beam search alone, without a verifier?
2. For inference-aware SFT experiments, what type of verifier is used? How does the method perform if there’s a mismatch between the training and testing verifiers?
3. In Figure 1, how is the empirical frequency calculated? What defines the best BoN performance? Also, what are the definitions of "easy" and "difficult" problems, given that the figure is based on the same set of MATH benchmark problems?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730697693392

thread_type:
official_review

<thread_object>
creator: reviewer_4vnu

<content>
summary:
The paper proposes a novel inference-aware fine-tuning paradigm, aiming to enhance the inference performance when scaling the compute. 

To overcome the non-differential argmax operator within best-of-N (BoN), the paper proposes the first imitation learning and reinforcement learning methods for fine-tuning language models with BoN. 

The experimental results show the BoN improvement of Gemma 2B on Hendrycks MATH from 26.8% to 30.8%, and Pass@K from 60% to 67%.

strengths:
1. The studied problem is interesting and important: how to improve best-of-N sampling in the test time through advanced training methods applied to fine-tune the models. 

2. The proposed method is reasonable and could be sound and useful if the method can be verified carefully and thoroughly.

weaknesses:
- The paper writing can be improved: it is not easy to follow the core idea of the proposed methods. A visual Figure showing the core mechanism of the method is strongly suggested. 

- The experiments are not sound enough to demonstrate the effectiveness of the proposed method, specifically:

**Regarding the co-scaling behavior of sample number N and temperature**: this analysis was conducted using only a single policy model and a single reward model.

How might the curves in Figure 2 appear if different language models or stronger reward models were used?

The conclusions drawn here cannot be readily generalized to other language models or reward models.


**The results presented in Figure 4 make it challenging to conclude that the proposed inference-aware fine-tuning significantly improves performance.**

One potential issue lies in the weakness of the reward model: the solution selected as having the highest reward is often inaccurate, and fine-tuning solely on this data may result in suboptimal model performance.

Selecting a broader range of solutions could benefit fine-tuning, as a correct answer might be found within these additional options.

Therefore, the observed improvement may stem from identifying a correct solution through an increased sample selection rather than from the inference-aware fine-tuning itself.


- I would also suggest conducting experiments on the alignment tasks and testing the models on AlpacaEval or Arena-Hard test sets.

questions:
1. How is the accuracy of the used reward model? 

2. How does best-of-N sampling relate to majority vote in math reasoning?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732344204105

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response to All Reviewers 1

comment:
We thank the reviewers for their insightful feedback and constructive criticism. We appreciate the reviewers acknowledging the importance and novelty of our work in addressing the gap between training and inference in large language models, particularly for Best-of-N (BoN) sampling. We have carefully considered all comments and made significant revisions to strengthen the paper. Specifically, we have:

1. Improved clarity and presentation by substantially revising the paper to enhance clarity and readability, as suggested by Reviewer 4VnU.  We have streamlined the presentation of our core ideas and provided additional explanations to make the paper more accessible. In the attached version of the paper we mark in red edits to the paper. We’ve also moved much mathematical notations to the appendix, to improve overall readability.

2. Expanded experiments by significantly expanding our experimental evaluation to address the concerns regarding limited model and task diversity. As requested, we have:
- Conducted further **co-scaling** experiments on **majority voting**, and with **9B Gemma model**. (Figure 9 and 10 in paper)
- Conducted **additional Distillation SFT experiments** using various baseline training datasets, as suggested by Reviewer gx5H. These include training on (a) the best of N samples, (b) all N samples as individual targets, and (c) N samples weighted by verifier scores. These experiments provide a more comprehensive analysis of our proposed BoN-SFT method.  (Figure 11 in paper)
- Added experiments illustrating how doing BoN-aware fine-tuning with the presence of **verifier mismatch** can impact BoN performance (Figure 12 in paper) 
- Added results for the Hendrycks MATH tasks using a **9B parameter Gemma** model, showcasing the scalability of our approach to larger LLMs. (Figure 13 in paper)
- Added an additional **Fractional MATH and MATH Odyssey** benchmark for both **Gemma 2B and 9B** experiments. (Figure 14, 15, 16, 17 in paper)
- Included experiments on code generation, using the **MBPP and HumanEval coding datasets**, demonstrating the broader applicability of our method beyond mathematical reasoning. (Figure 18 in paper)

We summarize the highlighted results from the aforementioned new experiments below. Plots of all these all the new results are added to the updated paper in Appendix D3.

# Coding Benchmark
**Coding results trained on MBPP and tested on HumanEval**

| Metric   | Base model | RL-S N'=1 | BoN-RL-S N'=8 | BoN-RLBP N'=8 | BoN-RLB N'=8 |
|----------|----------------|----------------|--------------------|---------------------|-------------------|
| Pass@1  | 40.09%         | 39.37%         | 38.99%            | 41.12%              | 39.37%            |
| Pass@2  | 46.82%         | 45.94%         | 46.55%            | 48.34%              | 46.02%            |
| Pass@4  | 52.61%         | 51.15%         | 53.39%            | 54.98%              | 52.00%            |
| Pass@8  | 57.36%         | 55.57%         | 59.65%            | 61.09%              | 57.54%            |
| Pass@16 | 61.59%         | 59.76%         | 66.46%            | 67.07%              | 62.80%            | 


# Gemma 9B on Hendrycks MATH

**BoN Accuracy**

|  | Base | SFT (N=1) | BoN-SFT N=8 | BoN-SFT N=32 |
|---|---|---|---|---|
| **N=1** | 42.5% | 44.5% | 49.5% | 43.1% |
| **N=5** | 51.5% | 53.5% | 55.3% | 55.5% |
| **N=10** | 53% | 54.5% | 55.8% | 56.3% |
| **N=20** | 53.5% | 54.5% | 55.8% | 56.2% |
| **N=30** | 53.5% | 54.3% | 55.7% | 56% |

|  | Base-model | RL N=1 | BoN-RLV N=8 | BoN-RLS N=8 |
|---|---|---|---|---|
| **N=1** | 42.5% | 46.5% | 47.5% | 49.5% |
| **N=5** | 51.5% | 54.5% | 57.5% | 56% |
| **N=10** | 53% | 55% | 57.8% | 57% |
| **N=20** | 53.5% | 55.3% | 58% | 57% |
| **N=30** | 53.5% | 55% | 58% | 56.8% | 

**Pass@N Accuracy**

|  | Base | SFT (N=1) | BoN-SFT N=8 | BoN-SFT N=32 |
|---|---|---|---|---|
| **N=1** | 43% | 45.5% | 50% | 45.5% |
| **N=5** | 59% | 66% | 69.5% | 68.5% |
| **N=10** | 66% | 71% | 73.5% | 73% |
| **N=20** | 72% | 74% | 76% | 76% |
| **N=30** | 74.5% | 75.5% | 77% | 77.5% | 

|  | Base-model | RL N=1 | BoN-RLV N=8 | BoN-RLS N=8 | BoN-RLBP N=8 |
|---|---|---|---|---|---|
| **N=1** | 43% | 46% | 48% | 49% | 47.5% |
| **N=5** | 59% | 63.5% | 66% | 68.5% | 67.5% |
| **N=10** | 66% | 70% | 71.5% | 74.5% | 73.5% |
| **N=20** | 72% | 75% | 75.5% | 78% | 77.5% |
| **N=30** | 74.5% | 77.5% | 78% | 79.5% | 79% |
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper presents an approach for finetuning LLMs such that they are inference-strategy aware, and find that policies resulting from this strategy are more amenable to inference-time scaling. 

On the plus side, this paper presents a rigorous formulation of an important problem (and quite timely, given the focus on inference-time compute these days) along with a sensible relaxation of the objective that enables tractable learning. On the negative side, experiments are only conducted on basic mathematical reasoning benchmarks, which may not generalize to more realistic settings.