Hugging Face 用推测解码让 Whisper 推理提速 2 倍

Open AI 的 Whisper 是一个通用的语音转写模型，在多种不同的基准测试和音频条件下均取得了最先进的结果。最新的 large-v3 模型位居 OpenASR 排行榜榜首，被评为英语最佳开源语音转写模型。该模型还展现出强大的多语言性能，在 Common Voice 15 数据集中测试的 58 种语言中，有 42 种实现了低于 30% 的词错误率（WER）。

虽然转写准确率非常出色，但推理时间非常慢。一段 1 小时的音频片段在 16GB T4 GPU 上需要 6 分钟以上才能完成转写，即使已经采用了 flash attention、半精度和 分块等推理优化手段。

在这篇博客文章中，我们演示了如何使用推测解码将 Whisper 的推理时间减少 2 倍，同时在数学上确保模型产生完全相同的输出。因此，该方法可以完美地直接替换现有的 Whisper 流水线，因为它在保持相同准确率的同时提供了免费的 2 倍加速。如需更精简的博客版本（解释更少但包含所有代码），请参阅附带的 Google Colab。

推测解码

推测解码由 Google 的 Yaniv Leviathan 等人在 Fast Inference from Transformers via Speculative Decoding 中提出。其工作原理基于这样一个前提：一个更快的助手模型经常能生成与更大的主模型相同的 token。

首先，助手模型自回归地生成一个包含 N N 个候选 token的序列，y^1:N \hat{\boldsymbol{y}}_{1:N} 。在下图中，助手模型生成了 5 个候选 token 的序列：The quick brown sock jumps。

虽然这些候选 token 生成得很快，但它们可能与主模型预测的 token 不同。因此，在第二步中，候选 token 被传递给主模型进行"验证"。主模型将候选 token 作为输入，执行单次前向传播。主模型的输出是 token 序列中每一步的"正确"token y1:N \boldsymbol{y}_{1:N} 。

在上图中，我们看到主模型预测的前三个 token 与助手模型的预测一致：The quick brown。然而，助手模型的第四个候选 token sock 与主模型的正确 token fox 不匹配。

我们知道，在第一个不匹配之前的所有候选 token 都是正确的（The quick brown），因为这些与主模型的预测一致。然而，在第一个不匹配之后，候选 token 与主模型预测的实际 token 发生了偏离。因此，我们可以将第一个不正确的候选 token（sock）替换为主模型的正确 token（fox），并丢弃此后所有预测的 token，因为它们已经偏离了。修正后的序列 The quick brown fox 现在构成了助手模型的新输入：

推理过程随后重复，助手模型生成一组新的 N N 候选 token，这些 token 由主模型在一次前向传递中验证。

由于我们使用快速助手模型自回归生成，并且仅使用慢速主模型执行验证前向传递，解码过程得到大幅加速。此外，主模型执行的验证前向传递确保获得完全相同的输出，就像单独使用主模型一样。这使得推测解码可以完美地直接替换现有的 Whisper 流程，因为可以确信将达到相同的质量。

为了最大程度降低延迟，助手模型应显著快于主模型，同时尽可能频繁地预测相同的 token 分布。在实践中，这两个属性形成权衡：模型越快，其准确性越低。然而，由于所有预测 token 中的 70-80% 往往是“更容易”的 token，这种权衡严重偏向于选择更快的模型，而不是更准确的模型。因此，助手模型应至少比主模型快 3 倍（越快越好），同时正确预测示例中的所有“容易” token。剩余的 20-30% 更“困难”的 token 则可由更大的主模型验证。

选择助手模型的唯一约束是它必须与主模型共享相同的词表。也就是说，助手模型必须使用与主模型一对一的相同 tokenizer。因此，如果我们想对 Whisper 的多语言变体（例如 large-v2（多语言））使用推测解码，我们需要选择 Whisper 的多语言变体作为助手模型，例如 tiny。反之，如果我们想对仅英语版本的 Whisper（例如 medium.en）使用推测解码，我们需要仅英语版本的模型作为助手模型，例如 tiny.en。目前，Whisper large-v3 是一个例外，因为它是唯一具有扩展词表大小的 Whisper 检查点，因此与之前的 Whisper 检查点不兼容。

现在我们了解了推测解码背后的背景，准备深入探讨实际实现。在 🤗 Transformers 库中，推测解码被实现为“辅助生成”推理策略。有关实现的更多细节，建议读者阅读 Joao Gante 关于辅助生成的优秀博客文章。

英语语音转录

基线实现

我们首先对 Whisper large-v2 进行基准测试，以获得推理速度的基线数值。我们可以通过便捷的 AutoModelForSpeechSeq2Seq 和 AutoProcessor 类加载主模型及其对应的处理器。我们将以 float16 精度加载模型，并通过传入 low_cpu_mem_usage=True 确保加载时间尽可能短。此外，我们希望通过传入 use_safetensors=True 确保模型以 safetensors 格式加载。最后，我们将传入参数 attn_implementation="sdpa"，以通过 PyTorch 的 SDPA attention kernel 获得 Flash Attention 加速：

import torch
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor

device = "cuda:0" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

model_id = "openai/whisper-large-v2"

model = AutoModelForSpeechSeq2Seq.from_pretrained(
    model_id,
    torch_dtype=torch_dtype,
    low_cpu_mem_usage=True,
    use_safetensors=True,
    attn_implementation="sdpa",
)
model.to(device)

processor = AutoProcessor.from_pretrained(model_id)

让我们加载将用于基准测试的英语语音转录数据集。我们将从 LibriSpeech ASR 验证清理数据集中加载一个包含 73 个样本的小型数据集。这大约为 9MB 的数据，非常轻量，可在设备上快速下载：

from datasets import load_dataset

dataset = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")

对于基准测试，我们只想测量生成时间，所以让我们编写一个简短的辅助函数来衡量这一步骤。以下函数将返回解码后的词元和运行模型所花费的时间：

import time

def generate_with_time(model, inputs, **kwargs):
    start_time = time.time()
    outputs = model.generate(**inputs, **kwargs)
    generation_time = time.time() - start_time
    return outputs, generation_time

现在我们可以遍历数据集中的音频样本，并累加总体生成时间：

from tqdm import tqdm

all_time = 0
predictions = []
references = []

for sample in tqdm(dataset):
    audio = sample["audio"]
    inputs = processor(audio["array"], sampling_rate=audio["sampling_rate"], return_tensors="pt")
    inputs = inputs.to(device=device, dtype=torch.float16)
    
    output, gen_time = generate_with_time(model, inputs)
    all_time += gen_time
    predictions.append(processor.batch_decode(output, skip_special_tokens=True, normalize=True)[0])
    references.append(processor.tokenizer._normalize(sample["text"]))

print(all_time)

输出：

100%|██████████| 73/73 [01:37<00:00,  1.33s/it]
72.99542546272278

好的！我们看到转写这73个样本花了73秒。让我们检查预测的词错误率（WER）：

from evaluate import load

wer = load("wer")
print(wer.compute(predictions=predictions, references=references))

输出：

0.03507271171941831

我们的最终基线数字是73秒，词错误率为3.5%。

推测解码

现在让我们加载用于推测解码的辅助模型。在这个例子中，我们将使用Whisper的蒸馏变体，distil-large-v2。该蒸馏模型复制了Whisper的整个编码器，但仅复制了32个解码器层中的2个。因此，它运行速度比Whisper快6倍，同时在分布外测试集上的词错误率（WER）表现仅相差1%以内。这使得它成为辅助模型的完美选择，因为它既具有高转录准确性，又能快速生成1{}^1。

由于Distil-Whisper使用与Whisper模型完全相同的编码器，我们可以在主模型和辅助模型之间共享编码器。然后我们只需从Distil-Whisper加载2层解码器作为“仅解码器”模型。我们可以通过方便的AutoModelForCausalLM自动类来实现这一点。实际上，与仅使用主模型相比，这只会使VRAM增加8%。

from transformers import AutoModelForCausalLM

assistant_model_id = "distil-whisper/distil-large-v2"

assistant_model = AutoModelForCausalLM.from_pretrained(
    assistant_model_id,
    torch_dtype=torch_dtype,
    low_cpu_mem_usage=True,
    use_safetensors=True,
    attn_implementation="sdpa",
)

assistant_model.to(device)

1{}^1 我们打算发布一个改进版的Distil-Whisper，其在标记分布上具有更强的对齐，这将进一步提高推测解码的性能。请关注Distil-Whisper仓库以获取更新。

我们可以为推测解码基准测试定义一个修改后的函数。与之前函数的唯一区别是在调用.generate时传入了辅助模型：

def assisted_generate_with_time(model, inputs, **kwargs):
    start_time = time.time()
    outputs = model.generate(**inputs, assistant_model=assistant_model, **kwargs)
    generation_time = time.time() - start_time
    return outputs, generation_time

让我们使用Distil-Whisper作为Whisper的辅助模型，运行带有推测解码的基准测试：

all_time = 0
predictions = []
references = []

for sample in tqdm(dataset):
    audio = sample["audio"]
    inputs = processor(audio["array"], sampling_rate=audio["sampling_rate"], return_tensors="pt")
    inputs = inputs.to(device=device, dtype=torch.float16)
    
    output, gen_time = assisted_generate_with_time(model, inputs)
    all_time += gen_time
    predictions.append(processor.batch_decode(output, skip_special_tokens=True, normalize=True)[0])
    references.append(processor.tokenizer._normalize(sample["text"]))

print(all_time)

输出：

100%|██████████| 73/73 [00:38<00:00,  1.88it/s]
32.69683289527893

使用推测解码，推理时间仅为33秒，比之前快了2.2倍！让我们验证词错误率（WER）是否相同：

print(wer.compute(predictions=predictions, references=references))

输出：

0.03507271171941831

完美！词错误率再次为3.5%，因为我们得到的输出与单独使用主模型时完全相同。

推测解码也可以与简单的🤗 Transformers pipeline API一起用于推理。下面，我们使用模型和处理器实例化pipeline，然后用它来转录玩具数据集中的第一个样本。这可以扩展到转写任意长度的音频样本，包括使用批处理：

from transformers import pipeline

pipe = pipeline(
    "automatic-speech-recognition",
    model=model,
    tokenizer=processor.tokenizer,
    feature_extractor=processor.feature_extractor,
    max_new_tokens=128,
    chunk_length_s=15,
    batch_size=4,
    generate_kwargs={"assistant_model": assistant_model},
    torch_dtype=torch_dtype,
    device=device,
)

sample = dataset[0]["audio"]
result = pipe(sample)
print(result["text"])

输出：

 Mr. Quilter is the apostle of the middle classes and we are glad to welcome his gospel.

关于使用Whisper和Distil-Whisper进行推测解码的端到端代码片段，可以在Distil-Whisper模型卡上找到。它将本笔记本中涵盖的推理阶段合并为一个单一的代码示例。

多语言语音转录

Distil-Whisper 是英语语音转录的完美助手模型，因为它在短音频和长音频样本上的性能与原始 Whisper 模型的 WER 差距在 1% 以内，同时速度提升了 6 倍。然而，官方的 Distil-Whisper 检查点仅支持英语，这意味着它们无法用于多语言语音转录。

要将推测解码用于多语言语音转录，可以使用官方多语言 Whisper 检查点之一，或使用 Whisper 的微调变体。撰写本文时，HF Mirror Hub 上有超过 5,000 个微调 Whisper 检查点，涵盖 100 多种语言。这些为选择在单一语言上表现极佳的助手 Whisper 检查点提供了绝佳的起点。在本示例中，我们将使用最小的官方多语言检查点 Whisper tiny。欢迎尝试用你的语言微调过的不同检查点！

让我们为新的助手模型 Whisper tiny 加载权重。由于 Whisper tiny 中的编码器与 large-v2 中的不同，这次我们将使用 AutoModelForSpeechSeq2Seq 类加载编码器和解码器：

assistant_model_id = "openai/whisper-tiny"

assistant_model = AutoModelForSpeechSeq2Seq.from_pretrained(
    assistant_model_id,
    torch_dtype=torch_dtype,
    low_cpu_mem_usage=True,
    use_safetensors=True,
    attn_implementation="sdpa",
)

assistant_model.to(device);

对于我们的基准测试数据集，我们将从 VoxPopuli 数据集的荷兰语（"nl"）划分中加载 73 个样本：

dataset = load_dataset("sanchit-gandhi/voxpopuli_dummy", "nl", split="validation")

太好了！我们现在可以像之前一样为基线 Whisper large-v2 模型重新运行基准测试。我们唯一改变的是将 language 和 task 参数传递给 generate 函数，以确保我们执行的是语音转录（而非语音翻译）。推测解码完全兼容语音转录和翻译任务。只需根据需要设置下面的 task 参数：

all_time = 0
predictions = []
references = []

for sample in tqdm(dataset):
    audio = sample["audio"]
    inputs = processor(audio["array"], sampling_rate=audio["sampling_rate"], return_tensors="pt")
    inputs = inputs.to(device=device, dtype=torch.float16)
    
    output, gen_time = generate_with_time(model, inputs, language="nl", task="transcribe")
    all_time += gen_time
    predictions.append(processor.batch_decode(output, skip_special_tokens=True, normalize=True)[0])
    references.append(processor.tokenizer._normalize(sample["normalized_text"]))

wer_result = wer.compute(predictions=predictions, references=references)

print("Time:", all_time)
print("WER:", wer_result)

输出：

100%|██████████| 73/73 [02:05<00:00,  1.72s/it]
Time: 116.50992178916931
WER: 0.127190136275146

没错！我们得到的基线时间为 117 秒，WER 为 12.8%。让我们使用推测解码重新运行生成过程：

all_time = 0
predictions = []
references = []

for sample in tqdm(dataset):
    audio = sample["audio"]
    inputs = processor(audio["array"], sampling_rate=audio["sampling_rate"], return_tensors="pt")
    inputs = inputs.to(device=device, dtype=torch.float16)

    output, gen_time = assisted_generate_with_time(model, inputs, language="nl", task="transcribe")
    all_time += gen_time
    predictions.append(processor.batch_decode(output, skip_special_tokens=True, normalize=True)[0])
    references.append(processor.tokenizer._normalize(sample["normalized_text"]))

wer_result = wer.compute(predictions=predictions, references=references)

print("Time:", all_time)
print("WER:", wer_result)

输出：

100%|██████████| 73/73 [01:08<00:00,  1.06it/s]
Time: 62.10229682922363
WER: 0.127190136275146

同样，我们达到了 12.8% 的 WER，但这次推理时间仅为 62 秒，实现了 1.9 倍的加速。鉴于加载助手模型的开销很低，以及能够获得完全相同输出的数学特性，推测解码为现有的 Whisper 流水线提供了完美的即插即用替代方案。

高效推测解码的策略

在最后一节中，我们将介绍两种策略，以确保使用推测解码时获得最快的推理时间。

助手模型

我们的目标是选择一个助手模型，它比主模型至少快 3 倍并且至少能正确转录 70-80% 的预测 token，通常是示例中“较容易”的 token。如果你有想要转录的特定语言，一个有效的策略是训练两个不同大小的 Whisper 模型，并将其中一个用作另一个的助手：

• 首先，微调 Whisper large-v3 作为你的主模型
• 其次，在同一数据集上蒸馏 Whisper large-v3 作为快速助手模型

微调和蒸馏可以提升主模型和助手模型在你所选语言上的 WER 性能，同时最大化 token 分布的对齐。关于 Whisper 微调的完整指南可以在这里找到，蒸馏指南在这里。

批大小

值得注意的是，投机性解码在批大小为 1 时速度提升最大。对于批量投机性解码，整个批次中的所有候选 token 都必须与验证 token 匹配，token 才会被接受。如果批次中某个位置的 token 不一致，则该位置之后的所有候选 token 都会被丢弃。因此，投机性解码更有利于较小的批大小。在实践中，我们发现投机性解码在批大小不超过 4 时能够提供加速。当批大小超过 4 时，投机性解码的推理速度比单独使用主模型更慢。完整结果请参阅 Distil-Whisper 论文的 D.3 节。

结论

在这篇博客文章中，我们介绍了投机性解码的推理策略，并将其应用于语音转录的 Whisper 模型。我们展示了如何实现 2 倍加速，同时在数学上确保输出与单独使用原始模型相同。鉴于使用额外辅助模型的开销很低，并且能够保证相同的转录结果，我们鼓励你尝试将投机性解码作为现有 Whisper 流水线的直接替代方案。

致谢

博客文章由 Sanchit Gandhi 撰写。非常感谢 Patrick von Platen 和 Pedro Cuenca 提出的建设性意见，以及 Joao Gante 在 🤗 Transformers 中的辅助生成实现。