Skip to content
Go back 2408.08696 arXiv logo

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

Published:  at  11:14 AM
91.73 🤔

Token Recycling 提出了一种无训练的推测解码方法,通过回收候选词并利用邻接矩阵构建草稿树,实现大型语言模型推理约 2 倍加速,相较于其他无训练方法提升超 30%。

Large Language Model, Efficiency, Generative AI, Prediction

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

Harbin Institute of Technology, Du Xiaoman (Beijing) Science Technology Co., Ltd.

Generated by grok-3

Background Problem

大型语言模型(LLM)由于其自回归解码策略(auto-regressive decoding),在推理过程中每次只能生成一个词(token),导致推理延迟高,成为应用中的主要瓶颈。这种延迟主要源于每次解码步骤中将大量参数从高带宽内存传输到加速器缓存,而非计算本身。推测解码(speculative decoding)作为一种无损加速技术,通过‘猜测与验证’(guess-and-verify)范式在单次解码步骤中生成多个词,但现有方法要么依赖额外训练的模型架构,要么依赖存储需求大、检索耗时且适应性差的检索库。论文提出了一种新方法,旨在解决现有推测解码方法中候选词被浪费的问题,通过回收解码过程中生成的候选词来构建动态检索库,从而提高推理效率。

Method

Token Recycling (TR) 的核心思想是回收大型语言模型在解码过程中生成的候选词(candidate tokens),将其存储并用于后续推测解码的加速。

Experiment

实验在 SpecBench(涵盖多轮对话、翻译、总结、问答、数学推理和检索增强生成)和 MBPP(代码生成数据集)上进行,使用 Vicuna(通用任务)和 Code Llama(代码任务)模型,规模覆盖 7B、13B 和 33B/34B 参数。实验设置聚焦于贪婪解码和批量大小为1的场景,评估指标包括平均接受词数(MAT)、每秒词数(Ts/s)和加速比(Speedup ratio)。

Further Thoughts

Token Recycling 的方法提供了一个有趣的视角,即如何利用解码过程中的‘废弃’信息来提升效率,这种思路或许可以扩展到其他生成式 AI 领域,如图像生成或多模态模型中,探索是否能通过类似‘候选特征回收’的方式加速生成过程。此外,TR 的动态更新机制让我联想到联邦学习(Federated Learning)中的模型更新策略,是否可以借鉴联邦学习中的隐私保护技术来优化矩阵更新过程,避免潜在的信息泄露风险,尤其是在矩阵继承‘热启动’时可能包含敏感模式的情况下。另一方面,论文中静态树结构的局限性提示我们,是否可以引入自适应树结构或结合强化学习(Reinforcement Learning)来动态调整树的参数,以更好地适应不同任务需求?这种改进可能进一步提升 TR 的泛化能力,但也可能增加计算开销,如何权衡是一个值得深入研究的方向。



Previous Post
Learning to Think: Information-Theoretic Reinforcement Fine-Tuning for LLMs
Next Post
Concise Reasoning, Big Gains: Pruning Long Reasoning Trace with Difficulty-Aware Prompting