Inference-Aware Fine-Tuning for Best-of-N Sampling in Large Language Models
Fuente:
arXiv
Saved in:
| Main Authors: | , , , , , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2024
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866908676163371008 |
|---|---|
| author | Chow, Yinlam Tennenholtz, Guy Gur, Izzeddin Zhuang, Vincent Dai, Bo Thiagarajan, Sridhar Boutilier, Craig Agarwal, Rishabh Kumar, Aviral Faust, Aleksandra |
| author_facet | Chow, Yinlam Tennenholtz, Guy Gur, Izzeddin Zhuang, Vincent Dai, Bo Thiagarajan, Sridhar Boutilier, Craig Agarwal, Rishabh Kumar, Aviral Faust, Aleksandra |
| contents | 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%. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2412_15287 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | Inference-Aware Fine-Tuning for Best-of-N Sampling in Large Language Models Chow, Yinlam Tennenholtz, Guy Gur, Izzeddin Zhuang, Vincent Dai, Bo Thiagarajan, Sridhar Boutilier, Craig Agarwal, Rishabh Kumar, Aviral Faust, Aleksandra Computation and Language Artificial Intelligence Machine Learning 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%. |
| title | Inference-Aware Fine-Tuning for Best-of-N Sampling in Large Language Models |
| topic | Computation and Language Artificial Intelligence Machine Learning |
| url | https://arxiv.org/abs/2412.15287 |