跳转至

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

会议: ACL 2025 (Outstanding Paper)
arXiv: 2504.15754
代码: 无
领域: 模型压缩 / 高效推理 / 投机解码
关键词: Token Recycling, 投机解码, 邻接矩阵, BFS树构建, 无训练加速

一句话总结

提出Token Recycling——一种无需额外训练的投机解码方法，将解码过程中产生的候选token存入邻接矩阵，通过BFS算法构建draft tree并用tree attention验证，仅需<2MB额外存储即在所有规模LLM上实现约2倍加速，超越现有无训练方法30%和有训练方法25%。

背景与动机

LLM参数规模的快速增长使推理延迟成为核心瓶颈。投机解码（speculative decoding）通过"猜测-验证"范式提供无损加速，但现有方法要么需要额外的draft模型（需训练），要么依赖检索库（存储大、检索慢、适应性差）。关键洞察：解码过程中被拒绝的候选token虽然在当前步骤不正确，但高概率会在未来序列中出现——这些"废弃"token就是可以回收利用的"宝藏"。

核心问题

如何设计一种无训练、低存储、高适应性的投机解码方法，利用解码过程中"自然产生"的候选token来加速后续推理？

方法详解

整体框架

Token Recycling维护一个轻量级邻接矩阵，记录token之间的后继关系。在每步解码中： 1. 收集：将当前步骤中被采样但未被选中的候选token存入邻接矩阵 2. 构建：从当前token出发，在邻接矩阵上执行BFS搜索，构建包含多条可能路径的draft tree 3. 验证：通过tree attention一次性验证draft tree中的所有路径 4. 更新：验证后新产生的候选token再次更新邻接矩阵

关键设计

  1. 邻接矩阵存储
  2. 矩阵大小为词表大小×k（k为每个token保留的最大后继数）
  3. 存储每个token最可能的后继token及其概率
  4. 仅需<2MB存储（对比检索库方法可能需要GB级别）

  5. 随解码过程动态更新，自然适应当前任务/领域

  6. BFS树构建

  7. 从当前已接受的token出发，在邻接矩阵中做广度优先搜索
  8. 每层展开概率最高的若干后继节点
  9. 生成的draft tree覆盖多条并行猜测路径

  10. 树的深度和宽度可根据计算预算灵活调整

  11. Tree Attention验证

  12. 将整棵draft tree通过一次前向传播送入目标LLM
  13. 使用特殊的tree attention mask确保因果性
  14. 接受与目标LLM分布一致的token路径（无损保证）

训练策略

无需任何训练——完全在推理时自组织。邻接矩阵从零开始，随解码过程逐步丰富。

实验关键数据

对比方法 Token Recycling优势
vs 无训练方法（检索/n-gram） +30% 加速比
vs 有训练方法（draft model） +25% 加速比
存储开销 <2MB（极低）
所有LLM规模 均实现~2x加速

消融实验要点

  • 邻接矩阵在解码100+ token后即收敛到有效状态
  • BFS树的宽度和深度存在最优平衡点——太宽浪费验证计算，太深匹配概率低
  • 方法对不同任务类型（对话、代码、翻译）均有效，因为邻接矩阵自适应更新

亮点

  • "回收废弃token"的洞察极为简洁优雅：解码过程中自然产生的候选token包含丰富的next-token概率信息，但一直被白白丢弃。这篇论文将这些"垃圾"变为加速推理的"宝藏"
  • 零训练+极低存储：<2MB邻接矩阵即可实现2x加速，对比需要额外训练draft model或维护大型检索库的方法，实用性大幅提升
  • 自适应性强：邻接矩阵随解码过程在线更新，自然适应当前上下文、任务和领域，无需预构建领域特定的检索库
  • 与tree attention的完美结合：BFS构建的多路径draft tree通过tree attention高效并行验证

局限性 / 可改进方向

  • 对话开始时邻接矩阵为空，需要"冷启动"若干步
  • 邻接矩阵基于token级别的局部后继关系，无法捕捉更长距离的依赖
  • 在极度开放性生成（高温度采样）时，token复现率可能降低，影响加速效果
  • 未探索与基于模型的投机解码方法的结合

与相关工作的对比

  • vs Medusa/EAGLE等draft-model方法：后者需要额外训练小模型作为drafter；Token Recycling完全无训练，更通用
  • vs 检索增强投机解码（REST等）：后者需构建大型检索库；Token Recycling仅需<2MB，且自适应更新
  • vs N-gram匹配方法：后者依赖静态n-gram表；Token Recycling的邻接矩阵动态更新，适应性更强

启发与关联

  • 与NSA（Native Sparse Attention）的共同主题：充分利用"已有计算"的副产品——NSA利用压缩注意力分数做块选择，Token Recycling利用被拒绝的候选token做未来猜测
  • 邻接矩阵的在线学习思路可迁移到其他需要预测模式的场景

评分

  • 新颖性: ⭐⭐⭐⭐⭐ "回收废弃token"的洞察新颖且直觉，方法简洁有力
  • 实验充分度: ⭐⭐⭐⭐ 多规模LLM验证，但部分细节需更多消融
  • 写作质量: ⭐⭐⭐⭐ 方法描述清晰，结构合理
  • 价值: ⭐⭐⭐⭐⭐ 即插即用的推理加速方法，极高实用价值