Research map / Large language models / Instruction and fine-tuning
Inference-Aware Fine-Tuning for Best-of-N Sampling in Large Language Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes inference-aware fine-tuning that optimizes LLMs for Best-of-N (BoN) inference, deriving SFT and RL objectives (BoN-SFT, BoN-RL-V/S, BoN-RLB, BoN-RLB(P)) with a variational form of the BoN policy, evaluated on Gemma 2B/9B for MATH, MBPP/HumanEval.
Problem
Standard fine-tuning (SFT, RL) ignores the inference-time strategy actually used, so models are not optimized for Best-of-N sampling; BoN's argmax is non-differentiable, making direct training hard.
Approach
Derive a variational (energy-based) approximation of the BoN policy, pi_BoN ∝ pi·exp(λ_N Q_pi), that yields differentiable gradients for training the base policy so that its BoN outputs are optimized, covering imitation and RL settings, and for binary rewards a closed-form weighted REINFORCE (BoN-RLB) with positive-only variant.
Key results
- BoN-RL-V raises Bo32 accuracy of Gemma 2B on MATH from 26.8% to 30.8%.
- BoN-RL-S improves pass@32 on MATH from 60.0% to 67.0% for Gemma 2B.
- BoN-RLB(P) raises pass@16 on HumanEval from 61.6% to 67.1%, while standard RL (N'=1) drops it to 59.8%.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Train with BoN-aware RL using verifier feedback and N'=32 | Gemma 2B · Hendrycks MATH | Bo-32 accuracy: 26.8% -> 30.8% | Helped |
| Train with BoN-aware RL using environment reward as verifier | Gemma 2B · Hendrycks MATH | pass@32 accuracy: 60.0% to 67.0% | Helped |
| Train with BoN-RLB(P), using binary reward and positive examples only | Gemma 2B · HumanEval | pass@16 accuracy: 61.6% to 67.1% | Helped |
| Use standard RL fine-tuning with N'=1 instead of pass@N-aware training | Gemma 2B · HumanEval | pass@16 accuracy: reducing pass@$16 to 59.8% | Hurt |
| Use standard RL with N'=1 instead of BoN-aware training | Gemma 2B · Hendrycks MATH | pass@32 accuracy | Hurt |
| Train with BoN-aware SFT rather than base SFT | Gemma 2B · Hendrycks MATH | BoN accuracy | Helped |
| Use base SFT with N'=1 rather than the pretrained base model | Gemma 2B · Hendrycks MATH | BoN accuracy | Hurt |
| Train BoN-RL-V to optimize verifier robustness rather than pass@N | Gemma 2B · Hendrycks MATH | pass@N accuracy | No clear effect |
| Train BoN-RL-V with N'=32 | Gemma 2B · Hendrycks MATH | direct N=1 accuracy: from 22% to 26% | Helped |
Surprising findings
- Standard RL with verifier feedback (RL-V, N'=1) performs poorly, attributed to reward hacking.
- Standard RL fine-tuning (N'=1) slightly degrades pass@32 on MATH, indicating over-optimization hurts exploration.
Limitations
- BoN-RL is sample inefficient for large N because each update requires N samples from the current policy.
- Experiments limited to Gemma 2B and 9B on math and code; authors defer larger-scale problems to future work.
- Gradient weights for BoN-RLB can become very large for hard problems, requiring clipping for stability.
Experimental setup
Datasets: MATH MBPP
Benchmarks: MATH HumanEval Functional MATH MathOdyssey
Metrics: BoN accuracy pass@N
Baselines: STaR RL-V RL-S SFT Base model
Builds on: STaR REINFORCE IPO ReST-EM
Closest papers by meaning
- On the Limits of Test-Time Compute: Sequential Reward Filtering for Better Inference (ICML 2026 · similarity 0.90)
- T1: Advancing Language Model Reasoning through Reinforcement Learning and Inference Scaling (ICML 2025 · similarity 0.90)
- Is Best-of-N the Best of Them? Coverage, Scaling, and Optimality in Inference-Time Alignment (ICML 2025 · similarity 0.89)
- Amortizing intractable inference in large language models (ICLR 2024 · similarity 0.89)
- BoNBoN Alignment for Large Language Models and the Sweetness of Best-of-n Sampling (NeurIPS 2024 · similarity 0.89)
- Variational Best-of-N Alignment (ICLR 2025 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.