来自 Flashlight Text（Flashlight 机器学习框架的一部分）的行业领先速度的集束搜索解码现在已在 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）的快速竞争者，谷歌已投入精力将其 部署到设备上，并在 常见基准 上显示出强大的结果。大规模自回归文本模型也可以从集束搜索中受益。除其他外，集束搜索提供了：
- 灵活的性能/延迟权衡——通过调整集束大小和外部 LM，用户可以牺牲延迟来提高准确性，或者以很小的延迟成本获得更准确的结果。在没有外部 LM 的情况下进行解码可以以极小的性能成本提高结果。
- 无需重新训练的便携性 — 现有神经网络模型可以受益于多种解码设置，并与外部 LM 即插即用，无需训练或微调。
- 引人注目的复杂性/准确性权衡——将集束搜索添加到现有建模管道中几乎不会增加额外复杂性，并且可以提高性能。
性能基准
目前最常用的支持外部语言模型集成的集束搜索解码库包括 Kensho 的 pyctcdecode 和 NVIDIA 的 NeMo 工具包。我们使用在 100 小时音频上训练的 wav2vec 2.0 基础模型，并在 LibriSpeech dev-other 上使用官方 KenLM 3-gram LM 对 TorchAudio + Flashlight 解码器进行基准测试。基准测试在 Intel E5-2698 CPU 上以单线程运行。所有计算均在内存中进行——KenLM 内存映射被禁用，因为它未得到广泛支持。
在基准测试时，我们测量的是达到 WER（词错误率）所需的时间——由于解码算法实现的细微差异以及参数与解码速度之间复杂的关系，一些超参数在不同运行中有所不同。为了公平评估性能，我们首先扫描实现基线 WER 的参数，并尽可能减小束大小。
预训练 wav2vec 2.0 模型在 Librispeech dev-other 上的解码性能。TorchAudio + Flashlight 解码器在低 WER 时表现出数量级的优势。
在解码器中，以较小束大小为准的 Time-to-WER 结果。TorchAudio + Flashlight 解码器在较大束大小和较低 WER 下的扩展性要好得多。
TorchAudio API 和用法
TorchAudio 提供用于 CTC 集束搜索解码的 Python API，支持以下功能：
- 有词典和无词典解码
- KenLM N-gram 语言模型集成
- 字符和词元解码
- 示例预训练 LibriSpeech KenLM 模型以及相应的词典和标记文件
- 各种可定制的集束搜索参数（集束大小、剪枝阈值、LM 权重...）
要设置解码器，请使用工厂函数 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 ...
)
给定形状为 (batch, time, num_tokens) 的发射，解码器将计算并返回一个批次列表，每个列表包含与发射对应的 nbest 假设。每个假设可以进一步分解为标记、单词（如果提供了词典）、分数和时间步长组件。
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 中添加支持，允许用户使用自定义的解码器兼容 LM。自定义词级语言模型已在 TorchAudio 每晚构建中可用，并计划在 TorchAudio 0.13 中发布。
自回归/序列到序列解码——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}
}