借助来自 Flashlight Text（Flashlight ML 框架的组成部分）的业界领先速度，集束搜索解码现已在 TorchAudio 中获得官方支持，为基于 PyTorch 构建的语音和文本应用带来了高性能集束搜索和文本实用工具。当前的集成支持 CTC 风格的解码，但它可以用于任何在时间步长上输出令牌级别概率分布的建模设置。
集束搜索简要回顾
在语音和语言设置中，集束搜索是一种高效的贪婪算法，可以使用可选约束（例如词典）来约束有效序列、可选外部评分（例如对有效序列进行评分的语言模型）以及其他针对特定序列的分数调整，将连续值（即概率或分数）序列转换为图或序列（即令牌、词片、单词）。
在接下来的示例中，我们将考虑一个令牌集 {ϵ, a, b}，其中 ϵ 是一个特殊令牌，我们可以想象它表示单词之间的空格或语音中的停顿。此处及下方的图表均取自 Awni Hannun 关于 CTC 和集束搜索的精彩 distill.pub 文章。
通过类似贪婪的方法，集束搜索在给定现有令牌序列的情况下考虑下一个可行的令牌——在上述示例中，a、b、b 是一个有效序列，但 a、b、a 不是。我们在集束搜索的每一步，根据评分函数对每个可能的下一个令牌进行排名。评分函数 (s) 通常如下所示：
其中 ŷ 是潜在的路径/令牌序列，x 是输入（P(ŷ|x) 表示模型随时间的预测），𝛼 是语言模型概率（P(y) 是序列在语言模型下的概率）的权重。一些评分函数添加 𝜷，它根据预测序列的长度 |ŷ| 调整分数。这个特定的评分函数在 FAIR 之前关于端到端 ASR 的工作中使用，评分函数有许多变体，在不同的应用领域可能会有所不同。
给定一个特定序列，为了评估该序列中的下一个可行令牌（可能受一组允许的单词或序列（例如词汇表）的约束），集束搜索算法会用每个候选令牌添加后对序列进行评分，并根据这些分数对令牌候选进行排序。为了效率，并且由于路径的数量与令牌集大小呈指数关系，保留得分最高的 top-k 候选者——k 代表 集束大小。
集束搜索的进展还有许多其他细微之处：例如，相似的假设序列可以“合并”。
评分函数可以进一步增强，以增加/减少令牌插入或长短单词的权重。使用更强的外部语言模型进行评分，虽然会产生计算成本，但也可以显著提高性能；这通常被称为语言模型融合。解码还有许多其他可调参数——这些在 TorchAudio 的文档中有所记载，并在 TorchAudio 的 ASR 推理教程中进行了进一步探讨。由于解码效率很高，参数可以轻松扫描和调整。
多年来，集束搜索在 ASR 中得到了广泛应用，文献多得无法一一引用，并且在强大的最新成果和系统中，包括 wav2vec 2.0 和 NVIDIA 的 NeMo。
为什么选择集束搜索？
集束搜索仍然是比诸如 RNN-Transducer 等更重的解码方法更快的竞争对手，谷歌已经投入将其 部署到设备上，并在 常见基准上显示出强大的结果。大规模的自回归文本模型也可以从集束搜索中受益。除其他外，集束搜索提供了：
- 灵活的性能/延迟权衡——通过调整集束大小和外部语言模型，用户可以牺牲延迟以提高准确性，或者以少量延迟成本换取更准确的结果。在没有外部语言模型的情况下进行解码可以以很小的性能成本改进结果。
- 无需重新训练的便携性——现有神经模型可以受益于多种解码设置，并与外部语言模型即插即用，而无需训练或微调。
- 引人注目的复杂性/准确性权衡——将集束搜索添加到现有建模管道中，几乎不会增加额外复杂性，并且可以提高性能。
性能基准
目前最常用的支持外部语言模型集成的集束搜索解码库包括 Kensho 的 pyctcdecode 和 NVIDIA 的 NeMo 工具包。我们使用在 100 小时音频上训练的 wav2vec 2.0 基础模型，并在 LibriSpeech dev-other 上使用官方 KenLM 3-gram 语言模型评估，与它们进行基准测试。基准测试在 Intel E5-2698 CPU 的单线程上运行。所有计算都在内存中进行——KenLM 内存映射被禁用，因为它不支持广泛。
在基准测试时，我们测量 time-to-WER（词错误率）——由于解码算法实现的细微差异以及参数与解码速度之间的复杂关系，一些超参数在不同运行中有所不同。为了公平地评估性能，我们首先扫描实现基线 WER 的参数，如果可能，最小化集束大小。
预训练 wav2vec 2.0 模型在 Librispeech dev-other 上的解码性能。TorchAudio + Flashlight 解码在低 WER 下表现优于一个数量级。
不同解码器在较小集束大小下的 WER 时间结果。TorchAudio + Flashlight 解码器在较大集束大小和较低 WER 下的扩展性要好得多。
TorchAudio API 和用法
TorchAudio 提供了一个用于 CTC 集束搜索解码的 Python API，支持以下功能：
- 词典和无词典解码
- KenLM n-gram 语言模型集成
- 字符和词片解码
- 预训练的 LibriSpeech KenLM 模型以及相应的词典和令牌文件示例
- 各种可定制的集束搜索参数（集束大小、修剪阈值、语言模型权重等）
要设置解码器，请使用工厂函数 torchaudio.models.decoder.ctc_decoder
from torchaudio.models.decoder import ctc_decoder, download_pretrained_files
files = download_pretrained_files("librispeech-4-gram")
decoder = ctc_decoder(
lexicon=files.lexicon,
tokens=files.tokens,
lm=files.lm,
nbest=1,
... additional optional customizable args ...
)
给定形状为 (批次, 时间, 令牌数) 的输出，解码器将计算并返回一个批次列表的列表，每个列表包含与输出对应的 n 个最佳假设。每个假设可以进一步分解为令牌、单词（如果提供了词典）、分数和时间步长组件。
emissions = acoustic_model(waveforms) # (B, T, N)
batch_hypotheses = decoder(emissions) # List[List[CTCHypothesis]]
# transcript for a lexicon decoder
transcripts = [" ".join(hypo[0].words) for hypo in batch_hypotheses]
# transcript for a lexicon free decoder, splitting by sil token
batch_tokens = [decoder.idxs_to_tokens(hypo[0].tokens) for hypo in batch_hypotheses]
transcripts = ["".join(tokens) for tokens in batch_tokens]
请参阅文档了解更多 API 详细信息，并参阅教程（ASR 推理解码）或示例推理脚本以获取更多用法示例。
即将推出的改进
完整的 NNLM 支持——使用大型神经语言模型（例如 Transformer）进行解码在大规模应用中仍未得到充分探索。Flashlight 已经支持，我们计划在 TorchAudio 中添加支持，允许用户使用自定义的解码器兼容语言模型。自定义词级别语言模型已在 TorchAudio 每夜构建版本中提供，并计划在 TorchAudio 0.13 中发布。
自回归/seq2seq 解码——Flashlight Text 还支持自回归模型的 序列到序列 (seq2seq) 解码，我们希望为其添加绑定，并将其添加到 TorchAudio 和 TorchText 中，并实现高效的 GPU 实现。
更好的构建支持——为了从 Flashlight Text 的改进中受益，TorchAudio 将直接子模块化 Flashlight Text，以便更容易地进行上游修改和改进。这已在 TorchAudio 每夜构建版本中生效，并计划在 TorchAudio 0.13 中发布。
引用
要引用此解码器，请使用以下内容：
@inproceedings{kahn2022flashlight,
title={Flashlight: Enabling innovation in tools for machine learning},
author={Kahn, Jacob D and Pratap, Vineel and Likhomanenko, Tatiana and Xu, Qiantong and Hannun, Awni and Cai, Jeff and Tomasello, Paden and Lee, Ann and Grave, Edouard and Avidov, Gilad and others},
booktitle={International Conference on Machine Learning},
pages={10557--10574},
year={2022},
organization={PMLR}
}
@inproceedings{yang2022torchaudio,
title={Torchaudio: Building blocks for audio and speech processing},
author={Yang, Yao-Yuan and Hira, Moto and Ni, Zhaoheng and Astafurov, Artyom and Chen, Caroline and Puhrsch, Christian and Pollack, David and Genzel, Dmitriy and Greenberg, Donny and Yang, Edward Z and others},
booktitle={ICASSP 2022-2022 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)},
pages={6982--6986},
year={2022},
organization={IEEE}
}