本文提出 EAGLE-3 方法,通过移除特征预测约束和多层特征融合技术,显著提高了大语言模型的推理加速比,并在实验中实现了高达 6.5 倍的无损速度提升。
Large Language Model, Inference Acceleration, Speculative Sampling, Feature Fusion, Token Prediction, Scaling Laws
Yuhui Li, Fangyun Wei, Chao Zhang, Hongyang Zhang
Peking University, Microsoft Research, University of Waterloo, Vector Institute
Generated by grok-3-mini-latest
Background Problem
大语言模型(LLMs)的顺序生成特性使其推理过程缓慢且成本高昂,影响用户体验和应用部署。投机采样方法已被证明能有效加速推理,但现有方法如EAGLE在增加训练数据时改进有限,这是由于EAGLE的特征预测约束限制了模型的表达能力。本文的工作起点是解决这一问题,通过改进投机采样框架来实现更好的推理加速,同时确保在数据规模扩展时性能持续提升。
Method
- 核心思想: EAGLE-3 放弃了特征预测,转而直接预测 token,并通过一种名为 ‘训练时测试’ 的技术融合目标模型的多层特征(低层、中层和高层),以捕获更丰富的语义信息,从而提高草稿模型的接受率和加速比。
- 实现方式: 在推理管道中,EAGLE-3 交替进行草稿生成和验证阶段。草稿阶段使用融合后的特征(如 ,通过全连接层将低、中、高层特征concatenate后降维)作为输入,结合之前 token 的嵌入和草稿模型输出,进行多步 token 预测。训练时,通过 ‘训练时测试’ 模拟多步生成过程,调整注意力掩码(例如,使用对角线掩码或向量点积计算注意力分数),移除特征预测损失 ,并确保草稿模型适应不同输入分布。EAGLE-3 还兼容 EAGLE-2 的动态草案树技术,实现上下文感知的树状草稿生成。
- 主要步骤: 1. 训练草稿模型:使用训练数据模拟多步生成,预测 token 而非特征。2. 推理阶段:重用目标模型的特征融合结果,生成草稿 token;验证阶段使用目标模型并行验证草稿。整个过程不修改目标模型权重,仅在草稿模型中引入少量计算开销。
Experiment
- 实验设置: 本文在多个开源模型上(如 Vicuna 13B、LLaMA-Instruct 3.1 8B、LLaMA-Instruct 3.3 70B 和 DeepSeek-R1-Distill-LLaMA 8B)评估 EAGLE-3,任务包括多轮对话(MT-bench)、代码生成(HumanEval)、数学推理(GSM8K)、指令遵循(Alpaca)和摘要生成(CNN/Daily Mail)。使用 AdamW 优化器,学习率 5e-5,训练数据包括 ShareGPT 和 UltraChat-200K 等。实验指标为加速比、平均接受长度 τ 和接受率 n-α,确保无损加速(即不降低生成质量)。
- 结果分析: EAGLE-3 实现了最高 6.5 倍加速比,平均比 EAGLE-2 提升 20%-40%。例如,在温度=0 时,Vicuna 13B 上 MT-bench 任务加速比达 5.58x,τ 为 6.65;消融实验(表 2)显示移除特征预测约束和使用特征融合分别贡献显著提升,验证了设计合理性。接受率实验(图 7)表明 EAGLE-3 的 n-α 几乎不受自预测输入影响,而 EAGLE 急剧下降,符合预期。在 SGLang 和 vLLM 框架下,大批量时 EAGLE-3 仍保持吞吐量提升(批量大小 64 时 38% 改进),证明了其在实际部署中的鲁棒性。整体结果匹配预期,展示了 EAGLE-3 在数据规模扩展下的新 scaling law。
Further Thoughts
EAGLE-3 的 ‘训练时测试’ 技术可能启发其他 AI 领域,如计算机视觉中的多层特征提取或强化学习中的模拟训练,潜在地减少训练-推理不一致问题;此外,与 DeepSeek-v3 的多 token 预测结合,可能进一步提升模型泛化能力,但需注意特征融合在不同任务中的鲁棒性,以及如何在资源受限设备上优化计算开销。