使用推测解码 (Speculative Decoding) 使 Whisper 实现 2 倍的推理加速

Open AI 推出的 Whisper 是一个通用语音转录模型，在各种基准和音频条件下都取得了非常棒的结果。最新的 large-v3 模型登顶了 OpenASR 排行榜，被评为最佳的开源英语语音转录模型。该模型在 Common Voice 15 数据集的 58 种语言中也展现出了强大的多语言性能，在 42 种语言上的单词错误率 (WER) 低于 30％。

尽管转录准确度非常优秀，但推理速度非常缓慢。即使利用 flash attention 、半精度和 分块 等优化推理技术，1 小时长度的音频在 16GB T4 GPU 上也需要超过 6 分钟的转录时间。

在本文中，我们将演示如何运用推测解码将 Whisper 的推理时间缩减 2 倍，同时在数学上确保完全取得与原模型 相同的输出。因此，这种方法可以完美地替换现有的 Whisper 流水线，因为它可以在不降低准确性的情况下免费获得 2 倍的加速。想要看附带有更简洁解释的全部代码，请参阅配套的 Google Colab。

推测解码

推测解码由 Yaniv Leviathan 等人在 Fast Inference from Transformers via Speculative Decoding 中提出。其思想是，一个更快的 辅助模型 通常会生成和更大的 主模型 相同的 token。

首先，辅助模型会通过自回归生成 个 候选 token 序列: 。在下图中，辅助模型生成了一个包含 5 个候选 token 的序列: The quick brown sock jumps 。

尽管这些候选 token 可以快速生成，但它们可能与主模型预测的 token 不同。因此，在第二步中，候选 token 被传入主模型以进行“验证”。主模型将候选 token 作为输入，并执行 单次前馈传播。主模型的输出是每个步骤中“正确”token 的序列 。

在上图中，我们看到主模型预测的前三个 token 与辅助模型的 token 一致: <span style="color:green"> The quick brown 但是，辅助模型的第四个候选 token: “ <span style="color:red"> sock”与主模型的正确 token: “ <span style="color:green"> fox”不一致。

我们知道，所有候选 token 一直到第一个不匹配之前都是正确的 ( <span style="color:green"> The quick brown)，因为这些与主模型的预测一致。但是，在第一个不匹配之后，候选 token 开始偏离主模型实际预测的 token。因此，我们可以用主模型的正确 token ( <span style="color:green"> fox) 替换第一个不正确的候选 token ( <span style="color:red"> sock)，并放弃之后所有预测的 token，因为这些已经逐渐偏离主模型的预测。经过校正的序列 The quick brown fox 现在成为辅助模型的新输入:

然后，辅助模型再次通过自回归推理，生成一组新的 个候选 token，这些 token 再次通过主模型的单次前馈传播进行验证。

由于我们在生成的时候使用的快速的辅助模型进行自回归，并且缓慢的主模型仅用于验证前馈传播，解码过程将大大加快。此外，经过主模型前馈传播验证后可以确保与仅使用主模型时获得完全相同的输出。这使得推测解码可以完美地替换现有的 Whisper 流水线，因为我们可以确定会取得相同质量的输出。

为了最大限度地减少延迟，辅助模型应该比主模型快得多，同时尽可能频繁地预测相同的 token 分布。实际上，这两个属性之间需要权衡: 模型越快，其准确度越低。然而，由于所有预测 token 中的 70-80％ 往往是“较易”的 token，此权衡倾向于选择一个更快的模型，而不是一个更准确的模型。因此，辅助模型应该至少比主模型快 3 倍 (越快越好)，同时在示例中正确预测所有较“易”token。剩余的 20-30％ 更“难”的 token 可以由更大的主模型进行验证。

选择辅助模型的唯一约束是它必须与主模型使用相同的词汇表。也就是说，辅助模型必须使用与主模型完全一对一相同的分词器。因此，如果我们想对诸如 large-v2 (多语言) 的 Whisper 多语言版本使用推测解码，我们需要选择诸如 tiny 的 Whisper 多语言版本作为辅助模型。而如果我们想对诸如 medium.en 的 Whisper 英文版本使用推测解码，我们需要选择诸如 tiny.en 的 Whisper 英文版本作为辅助模型。目前，large-v3 是唯一一个扩展了词汇量的 Whisper 检查点，因此与以前的 Whisper 检查点不兼容。

现在我们已经了解了推测解码背后的原理，我们准备实际实现它。在