Medusa Heads

Medusa 头是轻量级解码头，添加在大型语言模型的最后一个隐藏层之上；每个头都经过训练，用来预测比模型自身输出层更靠后的一个词元。模型只需一次前向传播，就能生成常规的下一个词元预测，以及对后续几个位置的猜测；接下来的一次前向传播中，树注意力步骤会检查这些猜测，并保留模型认可的最长连续序列。Medusa 是一种推测解码方法，不需要单独的草稿模型，因此无需训练、托管或保持第二个网络同步。Tianle Cai、Tri Dao 及其同事于 2024 年 1 月在论文《Medusa：具有多个解码头的简单 LLM 推理加速框架》中提出了该方法。

解码头如何预测未来词元#

标准的文本生成步骤会将位置 t 的最后一个隐藏状态输入语言模型输出头，由它为位置 t + 1 的词元评分。Medusa 的第 k 个头读取同一个隐藏状态，并为位置 t + k + 1 的词元评分。每个头都是一个单层前馈层，使用 SiLU 激活函数和残差连接，随后再投影到词表。投影层一开始会复制语言模型输出头，而前馈层的初始值为零，因此训练前每个头都会复现基础模型的下一个词元预测。论文发现五个头就足够了。下面这个 PyTorch 示例构建了三个头，并比较它们与基础模型输出头在同一个隐藏状态上的输出：

import torch
import torch.nn.functional as F

hidden_size, vocab_size, num_heads = 256, 1000, 3
lm_head = torch.nn.Linear(hidden_size, vocab_size, bias=False)  # 冻结模型的输出层

# 每个 Medusa 头：一个零初始化的残差 SiLU 层和语言模型输出头的副本
w1 = [torch.nn.Linear(hidden_size, hidden_size) for _ in range(num_heads)]
w2 = [torch.nn.Linear(hidden_size, vocab_size, bias=False) for _ in range(num_heads)]
for layer, out in zip(w1, w2):
    torch.nn.init.zeros_(layer.weight)
    torch.nn.init.zeros_(layer.bias)
    out.weight.data.copy_(lm_head.weight.data)

# 所有头都读取同一个最后隐藏状态；第 k 个头预测位置 t + k + 1
h = torch.randn(1, hidden_size)
head_logits = [out(h + F.silu(layer(h))) for layer, out in zip(w1, w2)]
print([torch.equal(logits, lm_head(h)) for logits in head_logits])

输出 [True, True, True] 表明，零初始化的残差层会让每个新头在初始化时与基础模型的预测完全相同。随后，训练会让第 k 个头转而预测它对应的未来位置，而 lm_head 和主干网络保持不变。

树注意力与接受机制#

各个头给出的最高概率猜测会组合成候选续写。取第 1 个头概率最高的两个词元和第 2 个头概率最高的三个词元，它们的笛卡尔积会生成六个共享前缀、构成一棵树的候选：

在下一次前向传播中，所有候选词元会一次性附加到输入中，而注意力掩码会让每个词元只能看到树中属于自己的祖先节点，因此模型可以并行评估所有分支。模型接受的最长分支会被保留，其最后一个隐藏状态会输入各个头，用于下一轮预测。由于所有头都读取同一个隐藏状态，因此无论第 1 个头选择哪个分支，第 2 个头提出的词元都相同。作者发现，使用最准确组合构建的 64 节点稀疏树，比 256 节点的密集树能带来更高的加速率。

使用贪心解码时，候选词元与模型的最高概率选择一致即可被接受。对于采样，Medusa 提供典型接受机制：只要候选词元在基础模型下的概率超过一个与分布熵相关的阈值，就会予以保留。这种方法能接受更多词元，但代价是无法精确匹配原始分布。

Medusa-1 和 Medusa-2 的训练#

Medusa-1 只训练添加在冻结主干网络之上的新头，因此模型输出质量保持不变；论文报告其生成速度提升超过 2.2 倍。使用 60,000 段 ShareGPT 对话，在一块 A100 GPU 上训练 Vicuna-7B 的解码头耗时约五小时。这与 freeze=23 设置的方案相同：在微调 Ultralytics YOLO26时，第 0 至 22 层保持冻结，只有第 23 层的检测头使用新数据进行学习。

Medusa-2 会将解码头和主干网络一起微调，使解码头的准确度更高，并将加速比提升到 2.3–3.6 倍；它使用的训练方案能保护主干网络原有的能力。如果无法获得原始训练数据，自蒸馏会使用模型自行生成的文本训练解码头。

Medusa 属于几种直接从目标模型生成草稿词元的方法之一。EAGLE 使用一个小型自回归层，根据模型的内部特征生成草稿；DeepSeek-V3 使用的多词元预测，则会在预训练期间训练额外的预测模块，供服务引擎之后用作草稿生成器。这三种方法在交互式聊天机器人的小型批次大小下最能降低推理延迟，因为每次前向传播时，GPU大部分时间都在等待内存。