Scaling LLM Inference with Optimized Sample Compute Allocation

📄 arXiv: 2410.22480v1 📥 PDF

作者: Kexun Zhang, Shang Zhou, Danqing Wang, William Yang Wang, Lei Li

分类: cs.CL, cs.AI

发布日期: 2024-10-29

🔗 代码/项目: GITHUB

💡 一句话要点

提出OSCA算法，通过优化采样计算分配显著提升大语言模型推理效率。

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

关键词: 大语言模型 推理优化 采样算法 计算资源分配 Agentic工作流

📋 核心要点

  1. 大语言模型推理中采样计算资源分配不合理，导致计算效率低下，难以在有限资源下扩展推理规模。
  2. OSCA算法将采样配置选择和样本数量确定建模为学习问题，通过优化不同推理配置的组合来提升效率。
  3. 实验表明，OSCA在代码生成和推理任务上显著降低了计算成本，并在Agentic工作流中提升了准确率。

📝 摘要（中文）

本文提出了一种名为OSCA的算法，旨在优化大语言模型（LLM）推理时采样计算资源的分配。采样是LLM推理算法中的基本操作。在计算资源有限的情况下，为了高效地扩展推理规模，找到采样计算预算的最佳分配至关重要，即确定使用哪些采样配置（模型、温度、语言等），以及在每种配置中生成多少个样本。本文将这些选择建模为一个学习问题，并通过寻找不同推理配置的最佳组合来优化采样计算分配。实验结果表明，通过学习到的混合分配，在代码生成任务上，OSCA的准确率优于最佳单一配置，且计算量减少了128倍；在四个推理任务上，计算量减少了25倍。此外，OSCA在超越单轮任务的Agentic工作流中也表现出有效性，在SWE-Bench上的准确率优于默认配置，且计算量减少了3倍。代码和生成结果已开源。

🔬 方法详解

问题定义：现有的大语言模型推理方法在采样时，通常采用固定的计算资源分配策略，即对所有采样配置（例如不同的模型、温度、语言等）分配相同的计算预算。这种策略忽略了不同配置对最终结果的影响差异，导致计算资源的浪费和推理效率的降低。尤其是在计算资源有限的情况下，如何合理分配采样计算预算，以最大化推理性能，是一个亟待解决的问题。

核心思路：OSCA的核心思路是将采样计算资源的分配问题建模为一个学习问题。通过学习不同采样配置对最终结果的贡献，动态地调整每种配置的计算预算。具体来说，OSCA旨在找到一个最优的采样配置组合，以及每种配置下生成样本的数量，从而在有限的计算资源下，最大化推理性能。

技术框架：OSCA算法的整体框架包含以下几个主要步骤：1) 定义采样配置空间，包括不同的模型、温度、语言等；2) 收集不同配置下的采样数据，并评估其对最终结果的贡献；3) 利用学习算法（具体算法未知）优化采样计算资源的分配，即确定每种配置的计算预算；4) 根据优化后的分配策略进行采样，并生成最终结果。

关键创新：OSCA的关键创新在于将采样计算资源的分配问题建模为一个学习问题，并提出了一种自动优化分配策略的算法。与传统的固定分配策略相比，OSCA能够根据不同配置的贡献动态地调整计算预算，从而更有效地利用计算资源，提升推理效率。

关键设计：论文中没有详细描述具体的学习算法、损失函数、网络结构等技术细节。但是，可以推测，OSCA可能采用了某种强化学习或贝叶斯优化算法，以学习不同配置的贡献，并优化计算资源的分配。此外，如何有效地评估不同配置的贡献，也是一个关键的设计问题。具体的评估方法未知。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

OSCA算法在代码生成任务上，以128倍更少的计算量，达到了优于最佳单一配置的准确率。在四个推理任务上，计算量减少了25倍。在SWE-Bench的Agentic工作流中，OSCA以3倍更少的计算量，实现了更高的准确率。这些实验结果表明，OSCA能够显著提升大语言模型的推理效率，并降低计算成本。

🎯 应用场景

OSCA算法可广泛应用于各种需要高效利用计算资源的大语言模型推理场景，例如代码生成、知识推理、对话系统等。通过优化采样计算分配，OSCA能够显著降低推理成本，提高推理效率，从而加速大语言模型在实际应用中的部署和推广。此外，OSCA还可以应用于Agentic工作流，提升智能体的决策能力和效率。

📄 摘要（原文）

Sampling is a basic operation in many inference-time algorithms of large language models (LLMs). To scale up inference efficiently with a limited compute, it is crucial to find an optimal allocation for sample compute budgets: Which sampling configurations (model, temperature, language, etc.) do we use? How many samples do we generate in each configuration? We formulate these choices as a learning problem and propose OSCA, an algorithm that Optimizes Sample Compute Allocation by finding an optimal mix of different inference configurations. Our experiments show that with our learned mixed allocation, we can achieve accuracy better than the best single configuration with 128x less compute on code generation and 25x less compute on 4 reasoning tasks. OSCA is also shown to be effective in agentic workflows beyond single-turn tasks, achieving a better accuracy on SWE-Bench with 3x less compute than the default configuration. Our code and generations are released at https://github.com/LeiLiLab/OSCA.