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

arXiv: https://arxiv.org/abs/2412.15287 · HTML: https://mlsanity.com/paper/2412.15287/

## Key facts

- Cited by 89 papers in the map (data through 14 Oct 2026).
- Placed in Test-time compute and efficient thinking, one of 17 directions in LLM reasoning.
- Builds on 31 of its 49 references that are in the map.
- First posted on arXiv 18 Dec 2024.

## Placed in

- LLM reasoning / [Test-time compute and efficient thinking](https://mlsanity.com/lr/test-time-compute-and-efficient-thinking/)

## Builds on (most cited first, 30 of 31)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [Program Synthesis with Large Language Models](https://mlsanity.com/paper/2108.07732/) (2021)
- [Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context](https://mlsanity.com/paper/2403.05530/) (2024)
- [Learning to summarize from human feedback](https://mlsanity.com/paper/2009.01325/) (2020)
- [High-Dimensional Continuous Control Using Generalized Advantage Estimation](https://mlsanity.com/paper/1506.02438/) (2015)
- [Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters](https://mlsanity.com/paper/2408.03314/) (2024)
- [Large Language Monkeys: Scaling Inference Compute with Repeated Sampling](https://mlsanity.com/paper/2407.21787/) (2024)
- [STaR: Bootstrapping Reasoning With Reasoning](https://mlsanity.com/paper/2203.14465/) (2022)
- [A General Theoretical Paradigm to Understand Learning from Human Preferences](https://mlsanity.com/paper/2310.12036/) (2023)
- [Generative Verifiers: Reward Modeling as Next-Token Prediction](https://mlsanity.com/paper/2408.15240/) (2024)
- [Training Language Models to Self-Correct via Reinforcement Learning](https://mlsanity.com/paper/2409.12917/) (2024)
- [Beyond Human Data: Scaling Self-Training for Problem-Solving with Language Models](https://mlsanity.com/paper/2312.06585/) (2023)
- [V-STaR: Training Verifiers for Self-Taught Reasoners](https://mlsanity.com/paper/2402.06457/) (2024)
- [Preference Fine-Tuning of LLMs Should Leverage Suboptimal, On-Policy Data](https://mlsanity.com/paper/2404.14367/) (2024)
- [Stream of Search (SoS): Learning to Search in Language](https://mlsanity.com/paper/2404.03683/) (2024)
- [From Decoding to Meta-Generation: Inference-time Algorithms for Large Language Models](https://mlsanity.com/paper/2406.16838/) (2024)
- [Theoretical guarantees on the best-of-n alignment policy](https://mlsanity.com/paper/2401.01879/) (2024)
- [BoNBoN Alignment for Large Language Models and the Sweetness of Best-of-n Sampling](https://mlsanity.com/paper/2406.00832/) (2024)
- [Differentiable Ranks and Sorting using Optimal Transport](https://mlsanity.com/paper/1905.11885/) (2019)
- [Beyond A*: Better Planning with Transformers via Search Dynamics Bootstrapping](https://mlsanity.com/paper/2402.14083/) (2024)
- [Functional Benchmarks for Robust Evaluation of Reasoning Performance, and the Reasoning Gap](https://mlsanity.com/paper/2402.19450/) (2024)
- [BOND: Aligning LLMs with Best-of-N Distillation](https://mlsanity.com/paper/2407.14622/) (2024)
- [Value-Incentivized Preference Optimization: A Unified Approach to Online and Offline RLHF](https://mlsanity.com/paper/2405.19320/) (2024)
- [MathOdyssey: Benchmarking Mathematical Problem-Solving Skills in Large Language Models Using Odyssey Math Data](https://mlsanity.com/paper/2406.18321/) (2024)
- [Provable Bayesian Inference via Particle Mirror Descent](https://mlsanity.com/paper/1506.03101/) (2015)

## Built on by (most cited first, 30 of 89)

- [ReasoningBank: Scaling Agent Self-Evolving with Reasoning Memory](https://mlsanity.com/paper/2509.25140/) (2025)
- [RewardBench 2: Advancing Reward Model Evaluation](https://mlsanity.com/paper/2506.01937/) (2025)
- [Optimizing Test-Time Compute via Meta Reinforcement Fine-Tuning](https://mlsanity.com/paper/2503.07572/) (2025)
- [A Survey of Frontiers in LLM Reasoning: Inference Scaling, Learning to Reason, and Agentic Systems](https://mlsanity.com/paper/2504.09037/) (2025)
- [Is Best-of-N the Best of Them? Coverage, Scaling, and Optimality in Inference-Time Alignment](https://mlsanity.com/paper/2503.21878/) (2025)
- [Pass@K Policy Optimization: Solving Harder Reinforcement Learning Problems](https://mlsanity.com/paper/2505.15201/) (2025)
- [Scaling Test-Time Compute Without Verification or RL is Suboptimal](https://mlsanity.com/paper/2502.12118/) (2025)
- [Rewarding the Unlikely: Lifting GRPO Beyond Distribution Sharpening](https://mlsanity.com/paper/2506.02355/) (2025)
- [AgentGym-RL: Training LLM Agents for Long-Horizon Decision Making through Multi-Turn Reinforcement Learning](https://mlsanity.com/paper/2509.08755/) (2025)
- [Tapered Off-Policy REINFORCE: Stable and efficient reinforcement learning for LLMs](https://mlsanity.com/paper/2503.14286/) (2025)
- [Autellix: An Efficient Serving Engine for LLM Agents as General Programs](https://mlsanity.com/paper/2502.13965/) (2025)
- [Optimizing Language Models for Inference Time Objectives using Reinforcement Learning](https://mlsanity.com/paper/2503.19595/) (2025)
- [Exploring the Limit of Outcome Reward for Learning Mathematical Reasoning](https://mlsanity.com/paper/2502.06781/) (2025)
- [Sample, Scrutinize and Scale: Effective Inference-Time Search by Scaling Verification](https://mlsanity.com/paper/2502.01839/) (2025)
- [POPE: Learning to Reason on Hard Problems via Privileged On-Policy Exploration](https://mlsanity.com/paper/2601.18779/) (2026)
- [Weight Ensembling Improves Reasoning in Language Models](https://mlsanity.com/paper/2504.10478/) (2025)
- [Self-Improving Language Models for Evolutionary Program Synthesis: A Case Study on ARC-AGI](https://mlsanity.com/paper/2507.14172/) (2025)
- [InfAlign: Inference-aware language model alignment](https://mlsanity.com/paper/2412.19792/) (2024)
- [Thinking vs. Doing: Agents that Reason by Scaling Test-Time Interaction](https://mlsanity.com/paper/2506.07976/) (2025)
- [Inference-Time Reward Hacking in Large Language Models](https://mlsanity.com/paper/2506.19248/) (2025)
- [Reasoning on a Budget: A Survey of Adaptive and Controllable Test-Time Compute in LLMs](https://mlsanity.com/paper/2507.02076/) (2025)
- [Scalable Power Sampling: Unlocking Efficient, Training-Free Reasoning for LLMs via Distribution Sharpening](https://mlsanity.com/paper/2601.21590/) (2026)
- [Roll the dice & look before you leap: Going beyond the creative limits of next-token prediction](https://mlsanity.com/paper/2504.15266/) (2025)
- [Reinforce LLM Reasoning through Multi-Agent Reflection](https://mlsanity.com/paper/2506.08379/) (2025)
- [AI-Slop to AI-Polish? Aligning Language Models through Edit-Based Writing Rewards and Test-time Computation](https://mlsanity.com/paper/2504.07532/) (2025)
- [Tina: Tiny Reasoning Models via LoRA](https://mlsanity.com/paper/2504.15777/) (2025)
- [Differential Smoothing Mitigates Sharpening and Improves LLM Reasoning](https://mlsanity.com/paper/2511.19942/) (2025)
- [SofT-GRPO: Surpassing Discrete-Token LLM Reinforcement Learning via Gumbel-Reparameterized Soft-Thinking Policy Optimization](https://mlsanity.com/paper/2511.06411/) (2025)
- [Representation-Based Exploration for Language Models: From Test-Time to Post-Training](https://mlsanity.com/paper/2510.11686/) (2025)
- [The Best of N Worlds: Aligning Reinforcement Learning with Best-of-N Sampling via max@k Optimisation](https://mlsanity.com/paper/2510.23393/) (2025)

## Abstract (arXiv)

Recent studies have indicated that effectively utilizing inference-time compute is crucial for attaining better performance from large language models (LLMs). In this work, we propose a novel inference-aware fine-tuning paradigm, in which the model is fine-tuned in a manner that directly optimizes the performance of the inference-time strategy. We study this paradigm using the simple yet effective Best-of-N (BoN) inference strategy, in which a verifier selects the best out of a set of LLM-generated responses. We devise the first imitation learning and reinforcement learning~(RL) methods for BoN-aware fine-tuning, overcoming the challenging, non-differentiable argmax operator within BoN. We empirically demonstrate that our BoN-aware models implicitly learn a meta-strategy that 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 Bo32 performance of Gemma 2B on Hendrycks MATH from 26.8% to 30.8%, and pass@32 from 60.0% to 67.0%, as well as the pass@16 on HumanEval from 61.6% to 67.1%.