TETRIS: Optimal Draft Token Selection for Batch Speculative Decoding

📄 arXiv: 2502.15197v2 📥 PDF

作者: Zhaoxuan Wu, Zijian Zhou, Arun Verma, Alok Prakash, Daniela Rus, Bryan Kian Hsiang Low

分类: cs.CL, cs.AI

发布日期: 2025-02-21 (更新: 2025-05-30)

备注: 17 pages, 11 figures, 5 tables

💡 一句话要点

TETRIS：面向批量推测解码的最优草稿令牌选择方法

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

关键词: 批量推测解码 大型语言模型 推理优化 令牌选择 资源利用率

📋 核心要点

  1. 现有批量推测解码方法在多请求场景下资源利用率低，导致推理速度受限。
  2. TETRIS通过主动选择最有可能被接受的草稿令牌，减少资源浪费，提升吞吐量。
  3. 实验表明，TETRIS在接受率和资源利用率上优于现有方法，实现更高效的批量推理。

📝 摘要（中文）

我们提出了TETRIS，一种新颖的方法，旨在优化多请求场景下批量推测解码的总吞吐量。与现有方法不同，TETRIS并非针对单个请求或一组请求进行整体优化，而是主动选择在并行验证时最有可能被接受的草稿令牌（针对批次中的每个请求），从而减少被拒绝的令牌数量，进而减少计算资源的浪费。这种有效利用资源以实现大型语言模型（LLM）快速推理的方法，对于推理能力有限的服务提供商尤为重要。与基线推测解码相比，TETRIS始终产生更高的接受率，并更有效地利用有限的推理能力。我们从理论和实验上证明，TETRIS优于基线推测解码和现有的动态选择草稿令牌的方法，从而在LLM中实现更高效的批量推理。

🔬 方法详解

问题定义：论文旨在解决多请求场景下，批量推测解码过程中计算资源利用率不高的问题。现有方法要么针对单个请求优化，要么将一组请求视为整体，无法充分利用并行验证的优势，导致大量草稿令牌被拒绝，浪费计算资源。

核心思路：TETRIS的核心思路是，在每个批次中，针对每个请求，主动选择最有可能被验证器接受的草稿令牌。通过这种精细化的选择，最大化被接受的令牌数量，从而减少被拒绝的令牌数量，提高计算资源的利用率，最终提升整体吞吐量。

技术框架：TETRIS的技术框架主要包含以下几个阶段：1. 草稿生成：使用小型语言模型为每个请求生成多个草稿令牌。2. 令牌选择：TETRIS算法根据一定的策略，从所有草稿令牌中选择一个子集，用于并行验证。3. 并行验证：使用大型语言模型并行验证所选的草稿令牌。4. 令牌接受/拒绝：根据验证结果，接受或拒绝相应的草稿令牌。5. 迭代：重复上述过程，直到达到预定的生成长度或满足其他停止条件。

关键创新：TETRIS的关键创新在于其令牌选择策略。与现有方法不同，TETRIS并非简单地选择置信度最高的令牌，而是综合考虑了每个令牌被接受的概率以及对整体吞吐量的影响。这种策略能够更有效地利用有限的计算资源，实现更高的吞吐量。

关键设计：TETRIS的具体令牌选择策略未知，论文中可能涉及概率模型、优化算法等技术细节，用于评估每个草稿令牌被接受的概率，并选择最优的令牌组合。损失函数的设计目标是最大化整体吞吐量，同时考虑计算资源的限制。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

论文通过实验证明，TETRIS在接受率和资源利用率上优于基线推测解码和现有的动态选择草稿令牌的方法。具体性能提升数据未知，但论文强调TETRIS能够更有效地利用有限的推理能力，实现更高效的批量推理。

🎯 应用场景

TETRIS可应用于各种需要高效批量推理的大型语言模型服务场景，例如在线翻译、文本摘要、代码生成等。通过提高推理吞吐量，TETRIS能够降低服务延迟，提升用户体验，并降低服务提供商的运营成本。该研究对于推动大型语言模型的实际应用具有重要意义。

📄 摘要（原文）

We propose TETRIS, a novel method that optimizes the total throughput of batch speculative decoding in multi-request settings. Unlike existing methods that optimize for a single request or a group of requests as a whole, TETRIS actively selects the most promising draft tokens (for every request in a batch) to be accepted when verified in parallel, resulting in fewer rejected tokens and hence less wasted computing resources. Such an effective resource utilization to achieve fast inference in large language models (LLMs) is especially important to service providers with limited inference capacity. Compared to baseline speculative decoding, TETRIS yields a consistently higher acceptance rate and more effective utilization of the limited inference capacity. We show theoretically and empirically that TETRIS outperforms baseline speculative decoding and existing methods that dynamically select draft tokens, leading to a more efficient batch inference in LLMs.