投机解码实现 Whisper 推理速度提升 2×

TL;DR

Hugging Face 表明，将投机解码应用于 OpenAI 的 Whisper 模型可将推理时间大幅降低约 2×（例如，在一个小基准上从 73 秒降至 33 秒），且转录质量毫无损失，因为该方法保证输出的 token 与原模型完全相同。

什么是投机解码？

投机解码由 Leviathan et al.（2022）提出，它将一个 快速助手模型 与一个 更大的主模型 配对。助手生成一段简短的候选 token 序列（例如，五个 token）。随后主模型在一次前向传播中验证该序列，接受直到首次不匹配的所有 token，其余的则被丢弃。校正后的前缀成为助手的新上下文，循环重复。由于助手的运行速度远快于主模型，整体解码速度提升，而验证步骤确保 最终输出与仅使用主模型时完全相同。

关键要求：

  • 助手必须与主模型共享完全相同的分词器/词表。
  • 它的速度应至少比主模型快 3×，并且能够正确预测大多数（≈70‑80 %）的“简单” token。

Whisper 基准推理速度

作者在一个包含 73 条样本的 LibriSpeech 验证集（总计约 9 MB）上对 Whisper large‑v2 进行基准测试，使用了：

  • float16 精度、FlashAttention（attn_implementation="sdpa"）以及低 CPU 内存加载。
  • 按样本测量生成时间。

结果： 总计 72.99 秒（≈每样本 1 秒），词错误率（WER）为 3.5 %。

将投机解码应用于 Whisper

助手模型选择

  • Distil‑Whisper distil‑large‑v2 被选为助手模型。它保留 Whisper 的编码器，但仅使用 32 层解码器中的 2 层，提供 6× 的速度优势，同时在分布外数据上保持与完整模型相差不超过 1 % 的 WER。
  • 编码器可以在主模型和助手之间共享，因此显存开销相对较小（约额外 8 %）。

实现细节

assistant_model = AutoModelForCausalLM.from_pretrained(
    "distil-whisper/distil-large-v2",
    torch_dtype=torch_dtype,
    low_cpu_mem_usage=True,
    use_safetensors=True,
    attn_implementation="sdpa",
).to(device)

# Generation with assistant
outputs = model.generate(**inputs, assistant_model=assistant_model, **kwargs)

assistant_model 参数在 🤗 Transformers 中激活 辅助生成 策略，实现投机解码。

加速结果（英文）

  • 投机解码时间： 总计 32.70 秒（≈每样本 0.45 秒）。
  • 加速比： 2.2×，比基准快。
  • WER： 与基准相同，为 3.5 %，确认输出完全保持一致。

相同的方法也适用于高层 pipeline API，只需传入 generate_kwargs={"assistant_model": assistant_model}。

多语言转录

Distil‑Whisper 检查点仅支持英语，因此在多语言场景下，作者使用最小的多语言 Whisper 检查点（tiny）作为助手。

在 VoxPopuli 的 73 条荷兰语样本上进行基准测试：

  • 基准（large‑v2）： 116.5 秒，WER 12.8 %。
  • 投机解码（assistant = tiny）： 62.1 秒，WER 12.8 %。
  • 加速比： 1.9×。

通过向 generate 提供相应的 language 和 task 参数，该方法同样适用于转录和翻译任务。

提升效率的策略

选择助手模型

  • 目标是相较于主模型速度提升 ≥ 3×，且 token 一致率 ≥ 70‑80 %。
  • 针对特定语言，可将大型 Whisper 模型（例如 large‑v3）微调为主模型，并对相同架构进行蒸馏以生成快速助手。这会使 token 分布保持一致，并提升两者的 WER。

批量大小考虑

  • 投机解码在 batch size = 1 时收益最大。
  • 更大的批量需要批内所有候选与主模型匹配；不匹配会导致提前丢弃，削弱加速效果。
  • 经验表明，加速效果可维持至 batch size 4；超过此值，开销大于收益（参见 Distil‑Whisper 论文，Sec. D.3）。

实际要点

  • 投机解码是现有 Whisper 流程的 即插即用 替代方案：只需添加 assistant_model 参数。
  • 它提供 免费 2× 推理加速，同时保证转录质量不变。
  • 该方法与模型无关；任何基于 Transformer 的序列模型，只要存在共享相同分词器的更快助手，都可受益。

“投机解码为现有 Whisper 流程提供了完美的即插即用替代方案，因为它在保持相同准确度的同时提供免费 2× 加速。” — Sanchit Gandhi, Hugging Face 博客 (2023‑12‑20)

致谢：本文感谢 Patrick von Platen、Pedro Cuenca 和 Joao Gante 的反馈，以及 🤗 Transformers 中的辅助生成实现。

Sources