Home / 大语言模型 / Medusa：让一个模型提前猜多步，再一次验证候选树

085｜Medusa：让一个模型提前猜多步，再一次验证候选树

论文：Tianle Cai 等，Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads，2024，ICML。采用论文 v3；代码见官方仓库。文中“质量接近”与“分布完全不变”分别讨论。

图表来源：Tianle Cai 等，2024，原文；CC BY 4.0。本篇展示做了PNG格式转换、按原图面板顺序拼合，图中数据与标签未改动。

先接住前作的问题

投机解码用独立草稿与目标验证减少串行调用，但还要维护第二个模型。Medusa改用同主干的多个未来预测头，继续追问候选怎样组织、验证和接受；它可组合FlashAttention，却不自动继承任意随机采样的精确证明。

五个前置概念

五个前置概念：自回归条件分布依赖真实已接受前缀；草稿预测廉价猜测未来；树形注意力隔离候选祖先关系；接受长度每轮正式提交的token数量；算术强度单位数据搬运对应的计算量。

一、问题设定：共享主干多猜未来位置，怎样减少串行生成轮数

原始隐藏状态 $h_t$ 预测 $t+1$ ，第 $k$ 个附加头预测 $t+k+1$ 时还没看见中间未来词，因此只能提出草稿。需要用主干按候选自身前缀重新验证，才可能一次提交多个token而不把猜测当作已知条件。

越远的位置通常越不确定。若直接输出所有头的结果，早期一个错误就可能让后面内容不成立。验证步骤仍按候选路径的真实前缀运行主干，重新判断下一 token；所以它改变的是执行组织，而不是宣布文本没有顺序依赖。

本文属于推理系统与辅助监督训练。没有环境奖励、PPO 或策略价值函数；所谓接受长度也不是强化学习回报。它可以与量化、GQA 和缓存管理结合，但各自解决的瓶颈不同。

二、区分未来预测、候选树与严格或近似接受

1. 多个轻量头提出草稿

一个头由残差前馈层及词表投影构成，抽象为

$p_t^{(k)}=\operatorname{softmax}\{W_2^{(k)}[h_t+\operatorname{SiLU}(W_1^{(k)}h_t)]\}.$

$p_t^{(k)}$ 是未来位置的词表概率；矩阵维度按使用列向量或行向量相应转置。初始化时中间投影置零、输出投影复制原语言模型头，让新头从已有表示出发。

冻结主干时，只训练这些头：

$L_{\mathrm{M1}}=-\sum_{t,k}\lambda_k\log p_t^{(k)}(y_{t+k+1}),\qquad \lambda_k=0.8^k.$

较远预测的不确定性较大，衰减权重避免它支配训练。Medusa-2 增加主干下一 token 损失，并用不同学习率、头部预热或辅助损失权重调度保护原能力。这属于经验训练配方，不是主干语义能力不变的理论证明。

2. 树形注意力批量验证候选

若第一头取两个候选，第二头取三个，笛卡尔组合有六条路径。树表示共享的前缀，只保存必要节点。节点可以读取原始上下文和自身祖先，不能读兄弟分支；位置编码应依树深度分配，不能随意把展平数组下标当作文本位置。

否则，验证一个候选时会偷看另一个互斥未来，得到的概率不再是该路径的真实条件概率。树形掩码是正确性条件，不能省略成“把所有候选拼在一起跑”。

3. 严格接受与典型接受有不同承诺

贪心验证要求候选与主干在对应前缀下的贪心选择一致。若采用正确的投机采样校正，也可以保持目标采样分布。Medusa 另提出 typical acceptance，以主干概率是否足够合理作为接受条件：

$p(x_{n+k}\mid x_{<n+k})> \min\{\epsilon,\delta e^{-H(p)}\}.$

$H(p)$ 是当前主干分布熵， $\epsilon$ 是固定阈值， $\delta$ 控制熵相关阈值。较高熵允许更多合理候选。论文该分支先贪心接受第一个 token，以保证前进，再接受符合条件的最长候选前缀。

这是一种近似生成规则，目标是质量与速度折中，并不保持原始采样分布。冻结主干的 Medusa-1 若配这种规则，权重虽然没变，输出分布仍然变了；联合训练的 Medusa-2 连主干分布也可能变化。两处“变化”应分别检查。

用两个词的分布，核对“严格无损”到底靠什么

令目标下一词分布为 $p$ 、草稿分布为 $q$ 。标准投机校正先从 $q$ 采样词 $x$ ，再以 $\alpha(x)=\min(1,p(x)/q(x))$ 接受。被直接接受的概率质量为 $q(x)\alpha(x)=\min(p(x),q(x))$ ；拒绝后从归一化的正残差 $[p(x)-q(x)]_+$ 采样，其中 $[u]_+=\max(u,0)$ 。接受质量加残差恰好补成 $p$ 。这套耦合采样与“两个模型各自独立抽一个词，再比较是否相等”不是同一算法。

原创教学例。 词表只有甲、乙， $p=(0.6,0.4)$ 、 $q=(0.8,0.2)$ 。甲被草稿选中的概率0.8，再以0.75接受，贡献0.6；乙总是接受，贡献0.2。剩余0.2的拒绝质量全部补给乙，最终恰为 $(0.6,0.4)$ 。若 $p=q$ ，两词的接受率都为1；独立抽样比较却可能不相等。因此原文 §2.3.1 的同分布拒绝说法需要保留疑点，不能成为严格投机算法的性质。

这个算例证明的是单步标准校正，不是替Medusa的任意候选树提供一份正确性证明。实际树分支选择、草稿概率和残差处理也必须匹配其算法。Typical acceptance 省去这种精确分布承诺，其速度收益应连同质量评价一起阅读。

拒绝处之后的候选不能跳过错误继续接受，因为条件前缀已不成立。缓存也必须跟随选中的路径整理，否则后续生成会受未选分支污染。

候选树中相同深度不表示相同上下文

若节点A、B是同父兄弟，C是A的子节点，验证C时可见共同提示与A，却不可见B。即使A、B具有相同位置编号，也不是两个已经共同发生的词。树形mask应按祖先关系构造，再与每条路径独立运行的logit对照；仅套普通三角mask会把排在前面的兄弟误当历史。

教法联系：Lilian Weng的Large Transformer Model Inference Optimization。算例为本文自构。

三、训练预测头并用祖先掩码验证连续前缀

自构例子。 已有前缀“首都是”，头部提出“北/南”“京/海”。树包含“北京”“北海”“南京”“南海”等路径，主干根据每条自身前缀验证。如果接受“北”但拒绝其后的“海”，只能提交“北”，不能把后续某个更远候选直接接上。假设平均接受3个token，而每轮成本是普通解码的1.5倍，则理论近似速度为2倍，尚未计提示处理和调度等额外开销。

输入：已生成前缀、主干、多个未来预测头、预设候选树
先处理提示并获得最后隐藏状态
直到结束：
    原始头提出下一 token，附加头并行提出更远候选
    根据树结构选择候选节点，避免无节制笛卡尔展开
    用只允许祖先可见的注意力掩码，一次运行主干验证树
    按所选规则接受连续前缀：
        严格分支使用正确贪心匹配或概率校正
        近似分支使用典型接受阈值
    输出被接受 token，仅保留所选路径的有效 KV
    复用新隐藏状态进入下一轮
输出：生成文本，以及实际耗时与接受长度统计

四、把接受长度除以单步开销后再比较速度

候选树增加接受长度也增加验证成本；严格校正与typical acceptance要连同训练方式、温度和质量指标分别读。

项目	论文设置与需要保留的条件
主模型	Vicuna v1.5 7B/13B；另外考察 Vicuna 33B、Zephyr 7B
Medusa 头	训练 5 个单层头；推理可只用其中3—4个，依树配置决定
Vicuna 数据	ShareGPT，主实验训练2 epoch；§2.2.1 给出7B头部训练60K样本、单A100约5小时的示例
联合训练	附录 B：全局批量64、主干峰值学习率 $5\times10^{-4}$ 、头部 $2\times10^{-3}$ ，40步warmup；4bit主干和QLoRA
自蒸馏	33B/Zephyr分别约100K生成样本，长度2,048、批量128；通过原模型概率的KL蒸馏保护主干
LoRA	附录 B.2：所有线性层含语言模型头，秩32、 $\alpha=16$ 、dropout 0.05
推理场景	主要批量1；基线为默认Hugging Face实现，任务为MT-Bench，附录补AlpacaEval
质量判定	GPT-4为MT-Bench评审；并非所有任务都有人工盲评，也未给全部设置训练重复区间
硬件范围	附录G测A100-80GB-PCIe、A40、A6000算子；正文并未把每张主结果表完整逐行映射到所有硬件配置

自蒸馏不是把生成文字照单全收后随便继续微调。论文观察仅用生成样本的硬标签调整主干会损害质量，因此以原始输出概率作教师；用 LoRA 时关闭适配器即可得到教师，节省一份完整权重副本，但额外教师前向和训练计算仍存在。

证据一｜图3、表1：墙钟收益需要扣掉验证开销。 Vicuna-7B 的 Medusa-1 为2.18倍、Medusa-2为2.83倍；后者平均每步解码3.47个token，每步开销比为1.22。按论文定义，速度约为 $3.47/1.22\approx2.84$ ，与表中2.83的舍入结果接近。不能将3.47误写成直接速度倍数。树搜索图同时画加速率与 tokens/s：候选树更大时接受更多 token，单步计算也更贵，故接受长度增加仍可能使实际速度下降。

Medusa在不同模型规模和任务类别上的生成加速 — 原文图形摘引：Tianle Cai 等，2024，图3，PDF第7页；核对原文。

证据二｜表1、表2：质量相近是有限基准的经验结果。 Vicuna-7B基线MT-Bench为6.17，直接联合微调5.925，Medusa-2为6.18。13B的Medusa-2为6.43，相比原模型下降0.14；Zephyr下降0.07。支持训练保护策略比直接微调更稳妥，但“不影响任何能力”的说法超过这些证据，细小差异也没有完整重复区间帮助判断。

证据三｜图4：更多候选可能让系统更慢。 优化后的64节点稀疏树，接受长度优于256节点稠密树；节点继续增多时计算开销抬升，速度下降。树选择利用校准数据估计各头候选正确率，并假设部分独立性，真实错误相关时这种估计可能偏差，所以它是实用搜索启发，不是所有数据上的无条件最优树。

Medusa候选树节点数对接受长度和实际解码速度的影响 — 原实验图摘引：Tianle Cai 等，2024，图4；核对原文。原版权保留。

证据四｜图5与附录表4：近似规则和任务分布会改变收益。 图5只在写作、角色扮演类别、温度0.7下调典型接受阈值：阈值变严格，质量趋好、接受长度下降。AlpacaEval中Vicuna-13B从29.01到91.54 token/s，即3.16倍，和MT-Bench的2.83不同。这说明速度需按工作负载报告，而不是取最大数字覆盖所有场景。

Medusa后验阈值改变时接受率与回答评价分数的权衡 — 原文图形摘引：Tianle Cai 等，2024，图5，PDF第8页；核对原文。

表3给出逐项系统组合约1.5、1.9、2.2、2.8倍的变化，帮助理解头、树、优化树和联合训练的贡献；它不是严格控制所有变量且带置信区间的完整因果分解。

五、冻结主干不自动保证任意接受规则分布无损

相比独立草稿模型，部署组件更少；相比一次只猜一条续写，候选树提供了更好的准确率与成本折中。

需要保留一个数学边界：§2.3.1 用“草稿与目标分布相同也可能拒绝”解释采样损失效率。若指标准投机采样的接受式 $\min(1,p(x)/q(x))$ ，当 $p=q$ 时接受率应为1，不能以独立采样不一致替代该算法。因此这段动机不能当作标准拒绝校正的普遍性质；典型接受的实际价值应依据它自身的效率/质量实验。标准算法可对照Leviathan等原论文。

六、把多token监督与可选投机推理解耦比较

DeepSeek-V3将多token训练监督与可选投机使用分开。它与本篇共享预测未来的想法，但是否改主干、使用哪个未来输入和如何接受候选，仍须按各自公式核对。

自测：三个“无损”条件能否分开

1. 冻结主干是否足以保证输出分布不变？

答案：不足。还需要接受和采样规则正确保持目标分布。典型接受按阈值挑选合理token，会改变分布，即使主干权重完全冻结。Medusa-2另外更新主干，因此还需评估模型本身的能力变化。

2. 为什么树节点只能看祖先，且拒绝后不能继续接更远token？

答案：主干验证的条件必须等于该候选路径的真实前缀。看到兄弟分支会泄露互斥未来；跳过被拒绝位置会改变后续条件，使此前的验证概率失效。两者都会破坏自回归语义，缓存管理也必须遵守同一选中路径。

3. 平均每步接受4个token，但单步开销为2.5倍，速度是多少？

答案：按论文的简化口径约为 $4/2.5=1.6$ 倍，而非4倍。端到端还要计提示处理、缓存和调度。若增大树只把接受长度从4提高到4.2，却把开销推到3，速度反而降为1.4倍，因此候选数量要依据硬件与工作负载选择。