来自 Flashlight Text（Flashlight ML 框架的一部分）的业界领先速度的束搜索解码现已在 TorchAudio 中获得官方支持，为基于 PyTorch 构建的语音和文本应用带来了高性能的束搜索和文本工具。当前集成支持 CTC 风格的解码，但可用于任何随时间步输出 token 级别概率分布的建模设置。
束搜索简要回顾
在语音和语言处理中，束搜索是一种高效的贪婪算法，它可以使用有效的序列的可选约束（例如，词典）、可选的外部评分（例如，为有效序列评分的语言模型）以及针对特定序列的其他得分调整，将连续值序列（即概率或得分）转换为图或序列（即 token、word-piece、单词）。
在接下来的示例中，我们将考虑一个 token 集合 {ϵ, a, b}，其中 ϵ 是一个特殊 token，我们可以想象它表示单词之间的空格或语音中的停顿。此处及下方的图例均取自 Awni Hannun 关于 CTC 和束搜索的优秀 distill.pub 文章。
采用类似贪婪的方法，束搜索根据现有 token 序列考虑下一个可行的 token——在上面的示例中，a, b, b 是一个有效序列，而 a, b, a 则不是。在束搜索的每个步骤中，我们根据一个评分函数对每个可能的下一个 token 进行排名。评分函数 (s) 通常看起来像
其中 ŷ 是一个潜在的路径/token 序列，x 是输入 (P(ŷ|x) 表示模型随时间步的预测)，𝛼 是语言模型概率的权重 (P(y) 是序列在语言模型下的概率)。一些评分函数会增加 𝜷，它根据预测序列的长度 |ŷ| 调整得分。这种特定的评分函数在 FAIR 先前关于端到端 ASR 的工作中使用，并且根据不同的应用领域，评分函数有许多变体。
给定一个特定序列，为了评估该序列中的下一个可行 token（可能受到一组允许的单词或序列的约束，例如单词词典），束搜索算法会为添加了每个候选 token 的序列评分，并根据这些得分对候选 token 进行排序。为了提高效率，并且由于路径数量随 token 集合大小呈指数级增长，因此会保留得分最高的 top-k 个候选——k 代表束大小。
关于束搜索如何进行，还有许多其他细节：例如，相似的假设序列可以被“合并”。
评分函数可以进一步增强，以提高或降低 token 插入或长单词或短单词的权重。使用更强大的外部语言模型进行评分，虽然会产生计算成本，但也能显著提高性能；这通常被称为LM 融合。解码还有许多其他可调参数——这些参数记录在 TorchAudio 的文档中，并在 TorchAudio 的 ASR 推理教程中进一步探讨。由于解码非常高效，参数可以很容易地进行扫描和调整。
多年来，束搜索在 ASR 中得到了广泛应用，涉及的作品众多，在此不一一列举。它在近期强劲的成果和系统中也得到了应用，包括 wav2vec 2.0 和 NVIDIA 的 NeMo。
为什么选择束搜索？
束搜索仍然是重量级解码方法（例如 RNN-Transducer）的快速竞争者，Google 已投入资金将其部署到设备上，并在常见基准测试中显示出强劲的结果。大规模自回归文本模型也可以从束搜索中受益。除此之外，束搜索提供了
- 灵活的性能/延迟权衡 — 通过调整束大小和外部语言模型，用户可以牺牲延迟换取准确性，或以较小的延迟成本获得更准确的结果。不使用外部语言模型进行解码可以以非常小的性能成本提高结果。
- 无需重新训练即可实现可移植性 — 现有的神经网络模型可以从多种解码设置中受益，并且无需训练或微调即可与外部语言模型即插即用。
- 引人注目的复杂性/准确性权衡 — 将束搜索添加到现有的建模流程中，几乎不会增加额外的复杂性，同时可以提高性能。
性能基准测试
目前最常用的支持外部语言模型集成的束搜索解码库包括 Kensho 的 pyctcdecode 和 NVIDIA 的 NeMo 工具包。我们使用在 100 小时音频上训练的 wav2vec 2.0 基础模型，并在 LibriSpeech dev-other 数据集上，结合官方 KenLM 3-gram 语言模型，对 TorchAudio + Flashlight 解码器进行了基准测试。基准测试在 Intel E5-2698 CPU 的单线程上运行。所有计算都在内存中进行 — KenLM 内存映射被禁用，因为它未得到广泛支持。
进行基准测试时，我们测量达到指定词错误率 (WER) 所需的时间——由于解码算法实现的细微差异以及参数与解码速度之间复杂的关联，不同运行中的一些超参数有所不同。为了公平评估性能，我们首先扫描能达到基线 WER 的参数，并在可能的情况下最小化束大小。
预训练 wav2vec 2.0 模型在 Librispeech dev-other 数据集上的解码性能。在低 WER 下，TorchAudio + Flashlight 解码性能领先一个数量级。
各解码器在优先选择较小束大小时达到指定 WER 所需的时间结果。TorchAudio + Flashlight 解码器在较大束大小和较低 WER 下的扩展性要好得多。
TorchAudio API 和用法
TorchAudio 提供了一个用于 CTC 束搜索解码的 Python API，支持以下功能：
- 基于词典和无词典解码
- KenLM N-gram 语言模型集成
- 字符和 word-piece 解码
- 预训练的 LibriSpeech KenLM 模型示例以及相应的词典和 token 文件
- 各种可自定义的束搜索参数（束大小、剪枝阈值、语言模型权重等…）
要设置解码器，请使用工厂函数 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) 的发射值，解码器将计算并返回一个批次列表的 List，每个列表包含与发射值对应的 nbest 假设。每个假设可以进一步细分为 token、单词（如果提供了词典）、得分和时间步组件。
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 的 nightly build 中提供，并计划在 TorchAudio 0.13 版本中发布。
自回归/seq2seq 解码 — Flashlight Text 还支持自回归模型的序列到序列 (seq2seq) 解码，我们希望为其添加绑定，并将其添加到 TorchAudio 和 TorchText 中，同时提供高效的 GPU 实现。
更好的构建支持 — 为了受益于 Flashlight Text 的改进，TorchAudio 将直接将 Flashlight Text 作为子模块，以便更容易地进行上游修改和改进。这已经在 TorchAudio 的 nightly build 中生效，并计划在 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}
}
