Adaptive Inference-Time Compute: LLMs Can Predict if They Can Do Better, Even Mid-Generation

📄 arXiv: 2410.02725v1 📥 PDF

作者: Rohin Manvi, Anikait Singh, Stefano Ermon

分类: cs.CL, cs.AI, cs.LG

发布日期: 2024-10-03

💡 一句话要点

提出自评估生成方案，LLM可预测自身生成质量并自适应调整计算量，提升推理效率。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大语言模型 推理优化 自适应计算 生成式评估 计算效率 模型剪枝

📋 核心要点

  1. 现有Best-of-N采样方法计算成本高昂，需要外部奖励模型和生成多个样本。
  2. 提出一种生成式自评估方案，LLM可预测重启生成是否能提升性能，无需外部奖励模型。
  3. 实验表明，该方法能显著提升模型性能，并大幅减少平均采样数，降低计算成本。

📝 摘要（中文）

本文提出了一种新的生成式自评估方案，旨在自适应地减少大型语言模型（LLM）在推理时生成的样本数量，同时保持甚至提高性能。该方案使用生成式奖励模型，允许LLM在生成过程中预测重启生成是否会产生更好的响应。这种预测无需外部奖励模型，可用于决定是否生成更多样本、提前剪枝不理想的样本或选择最佳样本。该方法成本很低，只需生成一个预定义的token。使用真实LMSYS用户提示构建的数据集进行训练后，Llama 3.1 8B在AlpacaEval上对GPT-4的胜率从21%提高到34%（使用16个样本），在GSM8K上的数学性能从84%提高到91%。通过仅在LLM认为有利时进行采样并自适应调整温度退火，仅需平均1.2个样本即可实现使用16个样本带来的74%的性能提升。此外，还证明了可以在生成早期剪枝50-75%的样本，而性能下降最小。总而言之，本文方法能够更高效、更可扩展地利用LLM推理时的计算资源。

🔬 方法详解

问题定义：论文旨在解决大型语言模型推理时计算效率低下的问题。现有的Best-of-N采样方法虽然能提升性能，但需要生成大量的候选样本，并依赖外部奖励模型进行评估，计算成本很高，限制了其在资源受限场景下的应用。

核心思路：论文的核心思路是让LLM具备自我评估能力，即在生成过程中预测自身生成质量，并根据预测结果自适应地调整计算量。具体来说，LLM预测重启生成是否能产生更好的结果，从而决定是否继续采样或提前剪枝不理想的样本。

技术框架：整体框架包含以下几个阶段：1) LLM生成token序列；2) 在生成过程中，LLM预测一个特殊的token，该token代表重启生成后获得更好结果的概率；3) 基于该概率，决定是否继续生成更多样本，或者提前停止生成并选择当前最佳样本；4) 使用真实用户提示数据训练LLM，使其能够准确预测生成质量。

关键创新：最重要的创新点在于提出了生成式自评估机制，LLM无需外部奖励模型即可预测自身生成质量。与传统的依赖外部评估的方法不同，该方法将评估过程融入到生成过程中，降低了计算成本，并提高了效率。

关键设计：关键设计包括：1) 使用一个预定义的token来表示重启生成后获得更好结果的概率；2) 设计合适的损失函数，训练LLM准确预测该token；3) 自适应调整温度退火策略，平衡探索和利用；4) 设计剪枝策略，提前停止生成不理想的样本。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，使用该方法后，Llama 3.1 8B在AlpacaEval上对GPT-4的胜率从21%提高到34%（使用16个样本），在GSM8K上的数学性能从84%提高到91%。更重要的是，仅需平均1.2个样本即可实现使用16个样本带来的74%的性能提升，并且可以在生成早期剪枝50-75%的样本，而性能下降最小。

🎯 应用场景

该研究成果可应用于各种需要高效LLM推理的场景，例如移动设备上的本地部署、低延迟的在线服务以及资源受限的边缘计算环境。通过自适应地调整计算量，该方法能够在保证性能的同时显著降低计算成本，提高LLM的实用性。

📄 摘要（原文）

Inference-time computation is a powerful paradigm to enhance the performance of large language models (LLMs), with Best-of-N sampling being a widely used technique. However, this method is computationally expensive, requiring both (1) an external reward model and (2) the generation of multiple samples. In this work, we introduce a new generative self-evaluation scheme designed to adaptively reduce the number of generated samples while maintaining or even improving performance. We use a generative reward model formulation, allowing the LLM to predict mid-generation the probability that restarting the generation will yield a better response. These predictions are obtained without an external reward model and can be used to decide whether or not to generate more samples, prune unpromising samples early on, or to pick the best sample. This capability is very inexpensive as it involves generating a single predefined token. Trained using a dataset constructed with real unfiltered LMSYS user prompts, Llama 3.1 8B's win rate against GPT-4 on AlpacaEval increases from 21% to 34% with 16 samples and math performance on GSM8K improves from 84% to 91%. By sampling only when the LLM determines that it is beneficial to do so and adaptively adjusting temperature annealing, we demonstrate that 74% of the improvement from using 16 samples can be achieved with only 1.2 samples on average. We further demonstrate that 50-75% of samples can be pruned early in generation with minimal degradation in performance. Overall, our methods enable more efficient and scalable compute utilization during inference for LLMs.