Gumiho: A Hybrid Architecture to Prioritize Early Tokens in Speculative Decoding

📄 arXiv: 2503.10135v2 📥 PDF

作者: Jinze Li, Yixing Xu, Haiduo Huang, Xuanwu Yin, Dong Li, Edith C. H. Ngai, Emad Barsoum

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

发布日期: 2025-03-13 (更新: 2025-06-30)

备注: Accepted to the 42nd International Conference on Machine Learning (ICML 2025). Code: https://github.com/AMD-AIG-AIMA/Gumiho

💡 一句话要点

Gumiho：一种混合架构，通过优先处理推测解码中的早期token来加速LLM推理。

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

关键词: 推测解码 大型语言模型 模型加速 混合架构 Transformer MLP 自回归生成

📋 核心要点

  1. 现有推测解码方法假设所有token同等重要，忽略了早期token对后续生成的影响，限制了效率。
  2. Gumiho提出混合架构，对早期token采用更复杂的串行Transformer头，对后续token采用轻量级并行MLP头。
  3. 实验表明，Gumiho优于现有方法，验证了其有效性，证明了优先处理早期token的策略的优越性。

📝 摘要（中文）

推测解码(SPD)旨在加速目标大型语言模型(LLM)的自回归token生成过程。一些方法采用具有多个头的draft模型来预测未来token序列，其中每个头处理序列中的一个token。目标LLM验证预测的序列并接受对齐的token，从而实现高效的多token生成。然而，现有方法假设序列中的所有token同等重要，采用相同的头结构，并依赖于单一生成范式，无论是串行还是并行。为此，我们从理论上证明了draft序列中的初始token比后面的token更重要。基于这一洞察，我们提出了Gumiho，一种结合串行和并行头的混合模型。具体来说，考虑到早期token的关键重要性，我们采用复杂的Transformer架构，以串行配置用于早期draft头，以提高准确性。对于后面的token，我们利用多个轻量级MLP头并行运行，以提高效率。通过为早期头分配更先进的模型结构和更长的运行时间，Gumiho实现了改进的整体性能。实验结果表明，我们的方法优于现有方法，充分验证了其有效性。

🔬 方法详解

问题定义：现有推测解码方法在加速LLM推理时，通常假设序列中的所有token同等重要，采用相同的模型结构处理每个token。然而，早期token的预测准确性对后续token的生成至关重要，直接影响整个序列的验证效率。现有方法未能区分token的重要性，导致资源分配不合理，影响整体性能。

核心思路：Gumiho的核心思路是根据token在序列中的位置，动态分配计算资源。早期token对后续生成影响更大，因此采用更复杂的模型结构（Transformer）和串行处理方式，以提高预测准确性。后续token的重要性相对较低，采用轻量级的MLP模型和并行处理方式，以提高效率。这种混合架构旨在在准确性和效率之间取得平衡。

技术框架：Gumiho的整体架构包含一个目标LLM和一个draft模型。Draft模型由多个头组成，用于预测token序列。与现有方法不同，Gumiho的draft模型采用混合架构：早期头采用串行的Transformer结构，负责预测序列中的前几个token；后续头采用并行的MLP结构，负责预测序列中的剩余token。目标LLM验证draft模型生成的序列，并接受对齐的token。

关键创新：Gumiho的关键创新在于其混合架构，它根据token的重要性动态分配计算资源。通过对早期token采用更复杂的串行Transformer头，Gumiho提高了早期token的预测准确性，从而提高了整体性能。与现有方法相比，Gumiho不再假设所有token同等重要，而是根据token的位置和重要性，采用不同的模型结构和处理方式。

关键设计：Gumiho的关键设计包括：1) 确定早期Transformer头和后续MLP头的数量，这需要根据具体任务和模型大小进行调整。2) Transformer头的具体结构，例如层数、注意力头数等。3) MLP头的具体结构，例如层数、隐藏层大小等。4) 损失函数的设计，需要平衡早期token的准确性和整体序列的生成效率。具体参数设置未知，需要根据实验结果进行调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，Gumiho在多个基准测试中优于现有推测解码方法。具体性能提升数据未知，但摘要中提到“我们的方法优于现有方法，充分验证了其有效性”。Gumiho通过优先处理早期token，实现了更高的生成效率和准确性。

🎯 应用场景

Gumiho可应用于各种需要加速LLM推理的场景，例如机器翻译、文本摘要、对话生成等。通过提高LLM的推理速度，Gumiho可以降低计算成本，提高用户体验，并促进LLM在资源受限设备上的部署。未来，Gumiho的混合架构思想可以推广到其他序列生成任务中。

📄 摘要（原文）

Speculative decoding (SPD) aims to accelerate the auto-regressive token generation process of a target Large Language Model (LLM). Some approaches employ a draft model with multiple heads to predict a sequence of future tokens, where each head handles a token in the sequence. The target LLM verifies the predicted sequence and accepts aligned tokens, enabling efficient multi-token generation. However, existing methods assume that all tokens within a sequence are equally important, employing identical head structures and relying on a single-generation paradigm, either serial or parallel. To this end, we theoretically demonstrate that initial tokens in the draft sequence are more important than later ones. Building on this insight, we propose Gumiho, a hybrid model combining serial and parallel heads. Specifically, given the critical importance of early tokens, we employ a sophisticated Transformer architecture for the early draft heads in a serial configuration to improve accuracy. For later tokens, we utilize multiple lightweight MLP heads operating in parallel to enhance efficiency. By allocating more advanced model structures and longer running times to the early heads, Gumiho achieves improved overall performance. The experimental results demonstrate that our method outperforms existing approaches, fully validating its effectiveness.