Turning Trash into Treasure: Accelerating Inference of Large Language Models with Token Recycling

📄 arXiv: 2408.08696v3 📥 PDF

作者: Xianzhen Luo, Yixuan Wang, Qingfu Zhu, Zhiming Zhang, Xuanyu Zhang, Qing Yang, Dongliang Xu

分类: cs.CL, cs.LG

发布日期: 2024-08-16 (更新: 2025-05-25)

备注: Accepted by ACL2025. Code is here. Token Recycling has already merged into SpecBench

💡 一句话要点

提出Token Recycling，加速大语言模型推理，无需额外训练。

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

关键词: 大语言模型 推理加速 推测解码 Token重用 免训练

📋 核心要点

  1. 现有大语言模型推理速度慢，推测解码虽能加速，但依赖额外训练或检索，存在存储和效率问题。
  2. Token Recycling通过存储和重用解码过程中的候选token，构建草稿树，加速推理过程。
  3. 实验表明，Token Recycling仅需少量额外存储，即可实现约2倍加速，显著优于现有免训练和训练方法。

📝 摘要（中文）

大语言模型（LLM）的巨大参数量使得推理延迟成为一个根本瓶颈。推测解码是一种通过猜测-验证范式来加速推理的无损方法。一些方法依赖于额外的架构来猜测草稿token，这需要额外的训练。或者，基于检索的免训练技术从预先存在的语料库或通过n-gram生成构建库。然而，它们面临着诸如大存储需求、耗时的检索和有限的适应性等挑战。观察到解码过程中生成的候选token很可能在未来的序列中再次出现，我们提出了Token Recycling。它将候选token存储在邻接矩阵中，并采用类似广度优先搜索（BFS）的算法来构建草稿树，然后通过树注意力进行验证。来自解码过程的新候选token然后用于更新矩阵。Token Recycling需要小于2MB的额外存储，并在所有尺寸的LLM上实现大约2倍的加速。它显著优于现有的免训练方法30％，甚至优于广泛认可的训练方法25％。

🔬 方法详解

问题定义：大语言模型推理延迟是实际应用中的主要瓶颈。现有的推测解码方法，如需要额外训练的架构或基于检索的方法，存在额外的训练成本、存储需求和检索延迟等问题，限制了其应用范围。

核心思路：论文的核心思路是观察到大语言模型在解码过程中生成的候选token，在后续的解码过程中有很高的概率被重复使用。因此，可以通过有效地存储和重用这些候选token来加速推理过程，避免重复计算。

技术框架：Token Recycling方法主要包含以下几个阶段：1）候选token存储：使用邻接矩阵存储解码过程中生成的候选token及其关系。2）草稿树构建：采用类似广度优先搜索（BFS）的算法，利用邻接矩阵构建草稿树，该树表示可能的token序列。3）树注意力验证：使用树注意力机制验证草稿树的正确性，并选择最优的token序列。4）矩阵更新：使用新的候选token更新邻接矩阵，以适应新的解码上下文。

关键创新：Token Recycling的关键创新在于其免训练的特性和高效的token重用机制。与需要额外训练的方法不同，Token Recycling可以直接应用于现有的预训练大语言模型，无需任何额外的训练步骤。此外，通过邻接矩阵和BFS算法，可以高效地存储和检索候选token，从而加速推理过程。

关键设计：邻接矩阵的大小是一个关键参数，需要根据实际应用场景进行调整，以平衡存储需求和token重用率。树注意力机制的设计也至关重要，需要能够有效地验证草稿树的正确性，并选择最优的token序列。论文中具体参数设置未知。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，Token Recycling在各种尺寸的大语言模型上实现了约2倍的推理加速，且仅需小于2MB的额外存储。与现有的免训练方法相比，Token Recycling的性能提升了30％，甚至优于一种广泛认可的训练方法25％，证明了其有效性和优越性。

🎯 应用场景

Token Recycling可广泛应用于各种需要加速大语言模型推理的场景，例如在线对话系统、文本生成、机器翻译等。该方法降低了推理延迟，提升了用户体验，并降低了计算成本，具有重要的实际应用价值。未来，该方法可以进一步扩展到其他类型的序列生成任务中。

📄 摘要（原文）

Massive parameters of LLMs have made inference latency a fundamental bottleneck. Speculative decoding represents a lossless approach to accelerate inference through a guess-and-verify paradigm. Some methods rely on additional architectures to guess draft tokens, which need extra training before use. Alternatively, retrieval-based training-free techniques build libraries from pre-existing corpora or by n-gram generation. However, they face challenges like large storage requirements, time-consuming retrieval, and limited adaptability. Observing that candidate tokens generated during the decoding process are likely to reoccur in future sequences, we propose Token Recycling. It stores candidate tokens in an adjacency matrix and employs a breadth-first-search (BFS)-like algorithm to construct a draft tree, which is then validated through tree attention. New candidate tokens from the decoding process are then used to update the matrix. Token Recycling requires \textless2MB of additional storage and achieves approximately 2x speedup across all sizes of LLMs. It significantly outperforms existing train-free methods by 30\% and even a widely recognized training method by 25\%.