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

• URL: http://arxiv.org/abs/2412.15287v1
• Date: Wed, 18 Dec 2024 20:43:47 GMT
• Title: Inference-Aware Fine-Tuning for Best-of-N Sampling in Large Language Models
• Authors: Yinlam Chow, Guy Tennenholtz, Izzeddin Gur, Vincent Zhuang, Bo Dai, Sridhar Thiagarajan, Craig Boutilier, Rishabh Agarwal, Aviral Kumar, Aleksandra Faust,
• Abstract summary: We propose a novel inference-aware fine-tuning paradigm, in which the model is fine-tuned in a manner that directly optimize the performance of the inference-time strategy. 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. Our experiments demonstrate the effectiveness of BoN-aware fine-tuning in terms of improved performance and inference-time compute.
• Score: 80.65242356955231
• License:
• Abstract: 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%.

Related papers

• Evaluation of Best-of-N Sampling Strategies for Language Model Alignment [6.4706370001155955]
  Best-of-N (BoN) sampling with a reward model has been shown to be an effective strategy for aligning Large Language Models (LLMs) with human preferences at the time of decoding. Previous work proposes Regularized BoN sampling (RBoN), a BoN sampling with regularization to the objective, and shows that it outperforms BoN sampling. This paper proposes an extension of the RBoN framework, called RBoN sampling (SRBoN), which is a theoretically guaranteed approach to worst-case RBoN proxy reward.
  arXiv  Detail & Related papers  (2025-02-18T09:18:02Z)
• FORLAPS: An Innovative Data-Driven Reinforcement Learning Approach for Prescriptive Process Monitoring [3.4437362489150254]
  This study introduces an innovative evaluation model, benchmarking its performance against earlier works using nine publicly available datasets. The proposed model, FORLAPS, demonstrated exceptional performance, outperforming existing state-of-the-art approaches in suggesting the most optimal policies or predicting the best next activities within a process trace.
  arXiv  Detail & Related papers  (2025-01-17T20:31:35Z)
• Enhancing Knowledge Distillation for LLMs with Response-Priming Prompting [1.9461727843485295]
  We propose a set of novel response-priming prompting strategies to enhance the performance of student models. Our approach fine-tunes a smaller Llama 3.1 8B Instruct model by distilling knowledge from a quantized Llama 3.1 405B Instruct teacher model. We find that Ground Truth prompting results in a 55% performance increase on GSM8K for a distilled Llama 3.1 8B Instruct.
  arXiv  Detail & Related papers  (2024-12-18T20:41:44Z)
• The Surprising Effectiveness of Test-Time Training for Abstract Reasoning [64.36534512742736]
  We investigate the effectiveness of test-time training (TTT) as a mechanism for improving models' reasoning capabilities. TTT significantly improves performance on ARC tasks, achieving up to 6x improvement in accuracy compared to base fine-tuned models. Our findings suggest that explicit symbolic search is not the only path to improved abstract reasoning in neural language models.
  arXiv  Detail & Related papers  (2024-11-11T18:59:45Z)
• Building Math Agents with Multi-Turn Iterative Preference Learning [56.71330214021884]
  This paper studies the complementary direct preference learning approach to further improve model performance. Existing direct preference learning algorithms are originally designed for the single-turn chat task. We introduce a multi-turn direct preference learning framework, tailored for this context.
  arXiv  Detail & Related papers  (2024-09-04T02:41:04Z)
• Variational Best-of-N Alignment [58.7977683502207]
  Best-of-N (BoN) is a popular and effective algorithm for aligning language models to human preferences. We propose to fine-tune the language model to mimic what BoN does during inference. Our approach is analogous to mean-field variational inference and, thus, we term it variational BoN (vBoN)
  arXiv  Detail & Related papers  (2024-07-08T15:59:44Z)
• Towards Compute-Optimal Transfer Learning [82.88829463290041]
  We argue that zero-shot structured pruning of pretrained models allows them to increase compute efficiency with minimal reduction in performance. Our results show that pruning convolutional filters of pretrained models can lead to more than 20% performance improvement in low computational regimes.
  arXiv  Detail & Related papers  (2023-04-25T21:49:09Z)
• DeBERTa: Decoding-enhanced BERT with Disentangled Attention [119.77305080520718]
  We propose a new model architecture DeBERTa that improves the BERT and RoBERTa models using two novel techniques. We show that these techniques significantly improve the efficiency of model pre-training and the performance of both natural language understanding (NLU) and natural langauge generation (NLG) downstream tasks.
  arXiv  Detail & Related papers  (2020-06-05T19:54:34Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.