渐进树草稿法让大模型推理提速2倍

背景

自回归语言模型（如 GPT、Llama 系列）的推理过程受到严重的内存带宽瓶颈限制：每次只生成一个 token，且模型的利用率极低。推测解码（Speculative Decoding）是一种缓解该瓶颈的有效技术，它通过引入一个轻量级的草稿模型来并行生成多个候选 token，再由目标模型一次性验证，从而提升吞吐量。然而，传统的推测解码通常依赖外部辅助草稿模型，这不仅带来了额外的训练和部署开销，还引入了模型间的通信延迟。近年来，一些方法尝试在目标模型自身内部生成草稿（例如使用模型自身的浅层或共享注意力），但这些方法往往缺乏结构化的协同机制，无法充分挖掘模型内在的并行潜力。

核心内容

本文提出渐进树草稿（Progressive Tree Drafting, PTD），一种无需外部草稿模型、完全在目标模型内部运行的结构化并行草稿策略。PTD 的核心思路是：在单次前向传递（forward pass）中，引导语言模型同时探索多条语义路径，从而生成一组多样且连贯的候选草稿 token 序列。具体而言，PTD 包含两个关键组件：

1. 渐进树结构（Progressive Tree Structure）：草稿生成过程被建模为一个动态扩展的树。树的根节点是当前上下文，每一层对应一个待预测的 token 位置，分支代表不同的语义候选。与传统固定长度的线性草稿不同，PTD 允许树在深度和宽度上灵活增长，确保草稿覆盖更多可能的连续路径。
2. 逐步剪枝机制（Stepwise Pruning Mechanism）：随着树逐层扩展，基于模型自身的置信度（logits）对分支进行实时剪枝，保留高概率的路径，丢弃低质量的候选。这种机制在保证草稿多样性的同时，控制计算开销，避免树过度膨胀。

通过将上述树结构和剪枝机制耦合，PTD 在单次前向传递中实现了“多条路径并行生成 — 逐层筛选 — 最终验证”的流程。目标模型仅需一次完整的 Transformer 前向传播，即可验证整个树的所有候选路径，从而大幅降低逐 token 推理的序列化延迟。

实验在多个标准基准（如 MT-Bench、AlpacaEval 等）上进行，结果表明 PTD 相比标准自回归解码实现了高达 2 倍 的加速，并且无需额外训练、不依赖模型架构（model-agnostic）。代码已开源。

关键要点

• 无需训练：PTD 不依赖任何辅助草稿模型或额外微调，完全利用目标模型自身的并行能力。
• 模型无关：可应用于任意自回归语言模型（如 GPT、Llama、Mistral 等），无需修改模型结构。
• 单次前向传递：整个草稿生成和验证过程仅需一次完整的模型前向传播，大幅减少内存访问次数。
• 渐进树 + 剪枝：动态树结构保证草稿的多样性与覆盖范围，逐步剪枝控制计算开销，平衡质量和效率。
• 2倍加速：在多种基准上达到约 2 倍的解码速度提升，且不损失生成质量（因验证机制保证分布正确）。
• 开源：代码已公开，便于社区复现和使用。

意义与影响

PTD 提出了一种轻量级、训练免费、模型无关的推测解码方案，直接挑战了传统依赖外部草稿模型的主流范式。其核心贡献在于通过结构化的树状草稿和剪枝机制，有效释放了自回归模型中固有的“隐藏并行性”——即模型在单次前向传播中本可同时计算多条路径的概率，只是因自回归约束而被串行化。PTD 使得这一并行潜力得以实际利用，为 LLM 低延迟推理（如对话系统、实时生成）提供了一种实用且易于部署的加速手段。此外，由于不需要额外的模型或训练，该方法尤其适合资源受限的环境或需要快速迭代的部署场景。未来工作可进一步探索更复杂的树拓扑结构以及与其他推理优化技术（如 KV cache、量化）的协同。