Gumiho: 一种用于推测解码中优先处理早期令牌的混合架构
计算与语言
2025-07-01 v2 人工智能
机器学习
摘要
推测解码（SPD）旨在加速目标大语言模型（LLM）的自回归令牌生成过程。部分方法采用具有多个头的草稿模型来预测未来令牌序列，其中每个头处理序列中的一个令牌。目标 LLM 验证预测序列并接受对齐的令牌，从而实现高效的多令牌生成。然而，现有方法假设序列中的所有令牌同等重要，采用相同的头结构，并依赖单一生成范式（串行或并行）。为此，我们从理论上证明了草稿序列中的初始令牌比后续令牌更重要。基于这一洞察，我们提出了 Gumiho，一种结合串行头和并行头的混合模型。具体而言，考虑到早期令牌的关键重要性，我们在串行配置中采用复杂的 Transformer 架构来处理早期草稿头，以提高准确性。对于后续令牌，我们利用多个轻量级 MLP 头并行运行以提升效率。通过为早期头分配更先进的模型结构和更长的运行时间，Gumiho 实现了整体性能的提升。实验结果表明，我们的方法优于现有方法，充分验证了其有效性。
引用
@article{arxiv.2503.10135,
title = {Gumiho: A Hybrid Architecture to Prioritize Early Tokens in Speculative Decoding},
author = {Jinze Li and Yixing Xu and Haiduo Huang and Xuanwu Yin and Dong Li and Edith C. H. Ngai and Emad Barsoum},
journal= {arXiv preprint arXiv:2503.10135},
year = {2025}
}
备注
Accepted to the 42nd International Conference on Machine Learning (ICML 2025). Code: https://github.com/AMD-AIG-AIMA/Gumiho