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),将其存储并用于后续推测解码的加速。
- 邻接矩阵存储:TR 使用一个邻接矩阵(adjacency matrix)存储每个词的 top-k 候选词,矩阵初始为空或通过‘热启动’(hot start)继承先前矩阵以提高初始预测能力。矩阵规模较小(<2MB),支持高效并行处理。
- 草稿树构建:在每次解码前,基于当前内容最后一个词,通过类似广度优先搜索(BFS)的算法从矩阵中检索候选词,构建一个静态且不平衡的草稿树(draft tree)。树结构预定义,优先分配计算资源给概率较高的路径。
- 验证与更新:利用树注意力(tree attention)机制并行验证草稿树中的多个序列,选择最长的正确序列作为输出;随后用解码过程中新生成的候选词(包括接受和拒绝的词)更新矩阵,使检索空间动态适应当前内容。 批判性思考:虽然 TR 的无训练特性降低了资源需求,但静态树结构可能限制了方法对复杂序列的适应性。此外,矩阵更新中对重复词的并行处理策略(CUDA合并)可能导致更新结果不一致,论文虽称影响较小,但未提供足够证据证明长期使用下的稳定性。邻接矩阵的存储效率在超大词汇表场景下也值得进一步验证。
Experiment
实验在 SpecBench(涵盖多轮对话、翻译、总结、问答、数学推理和检索增强生成)和 MBPP(代码生成数据集)上进行,使用 Vicuna(通用任务)和 Code Llama(代码任务)模型,规模覆盖 7B、13B 和 33B/34B 参数。实验设置聚焦于贪婪解码和批量大小为1的场景,评估指标包括平均接受词数(MAT)、每秒词数(Ts/s)和加速比(Speedup ratio)。
- 结果:TR 在 SpecBench 上实现了约 2 倍加速比,相较于其他无训练方法(Lookahead、PLD、REST)提升超过 30%,甚至比需训练的 Medusa 方法高出约 25%。在代码领域(MBPP),TR 加速比达到 2.3 倍,表现尤为突出。TR 在大多数子任务(如翻译、数学推理)上表现优异,但在总结任务上略逊于 PLD,可能是由于总结任务重复内容较多,而 TR 对新词生成优化更强。
- 实验设计分析:实验覆盖了多种任务和模型规模,设置较为全面,但贪婪解码和单批量设置可能限制了结果的普适性,未反映复杂采样策略或多批量场景下的表现。此外,论文未充分探讨树结构参数(如深度和广度)对不同任务的影响,仅在 MT-Bench 上做了初步分析,缺乏跨任务的系统性对比。
- 批判性思考:虽然结果显示 TR 效果显著,但部分任务(如总结)表现不佳可能暗示方法对内容重复性高的场景优化不足。此外,实验未涉及推理延迟对响应质量的潜在影响,尤其是在树注意力与自回归解码选择不同词的罕见情况下,论文仅轻描淡写地提及质量影响很小,缺乏数据支持。
Further Thoughts
Token Recycling 的方法提供了一个有趣的视角,即如何利用解码过程中的‘废弃’信息来提升效率,这种思路或许可以扩展到其他生成式 AI 领域,如图像生成或多模态模型中,探索是否能通过类似‘候选特征回收’的方式加速生成过程。此外,TR 的动态更新机制让我联想到联邦学习(Federated Learning)中的模型更新策略,是否可以借鉴联邦学习中的隐私保护技术来优化矩阵更新过程,避免潜在的信息泄露风险,尤其是在矩阵继承‘热启动’时可能包含敏感模式的情况下。另一方面,论文中静态树结构的局限性提示我们,是否可以引入自适应树结构或结合强化学习(Reinforcement Learning)来动态调整树的参数,以更好地适应不同任务需求?这种改进可能进一步提升 TR 的泛化能力,但也可能增加计算开销,如何权衡是一个值得深入研究的方向。