WeNet更新：U2++ 双向 re-score

写文章标题时候正好是 6月7日 的上午，隔壁的校园里正在进行2021年高考的语文考试，而我也距离高考整整九年（现在去参加高考，怕啥也考不上了）。愿今年参加高考的这批孩子，未来接触语音识别时，第一个用到的工具就是 WeNet 😊。

这次 WeNet 带来的更新是 U2++ 双向建模，其核心思想是同时利用标注序列的前向和后向信息训练模型，在解码时同时利用双向的 decoder 进行 re-score。实验证明，该方法在各个数据集上都能取得一致性的5%～8%的相对错误率的下降。该工作目前已经整理成论文，详请参考

https://arxiv.org/pdf/2106.05642.pdf。

本文是关于该方法的概览的中文介绍和说明，首先我们先回顾一下 WeNet 中 U2 模型和的解码方式。

回顾

WeNet 中采用的 U2 模型，如下图所示，该模型使用 Joint CTC/AED 的结构，训练时使用 CTC 和Attention Loss 联合优化，并且通过 dynamic chunk 的训练技巧，使 Shared Encoder 能够处理任意大小的 chunk（即任意长度的语音片段）。在解码的时候，先使用 CTC Decoder 产生得分最高的多个候选结果，再使用 Attention Decoder 对候选结果进行重打分(Re-scoring)，并选择重打分后得分最高的结果作为最终识别结果。解码时，当设定 chunk 为无限大的时候，模型需要拿到完整的一句话才能开始做解码，该模型适用于非流式场景，可以充分利用上下文信息获得最佳识别效果；当设定 chunk 为有限大小（如每0.5秒语音作为1个chunk）时，Shared Encoder 可以进行增量式的前向运算，同时 CTC Decoder 的结果作为中间结果展示，此时模型可以适用于流式场景，而在流式识别结束时，可利用低延时的 Re-scoring 算法修复结果，进一步提高最终的识别率。可以看到，依靠该结构，我们同时解决了流式问题和统一模型的问题。

U2 存在的不足

U2 建模存在的问题

WeNet 里的 U2 做到了流式和非流式模型的统一，Re-scoring 解码方式与自回归解码方式相比，在维持相近甚至更好的模型性能的情况下，展示了更好的实时率以及对长语音解码的鲁棒性。但是这种建模方法是完美的吗，是否完美适用于 Re-scoring 解码方法呢，它存在哪些不足呢？Re-scoring 过程是否可以进一步控制延时呢？此次 WeNet 更新的 U2++ 将给出这两个问题的答案。首先我们思考第一个问题，由于 U2 decoder 部分在训练时是使用了标准的 Transformer 的 teacher-forcing 的训练方式即输入序列开始加上字符<sos>；输出序列的末尾加上字符<eos>，这样训练的原因是为了做自回归生成式的解码，并且使用了 left-to-right mask 来限制当前时刻只能看到之前的信息，便于训练时可以并行计算。受限于这样的建模与训练的影响，U2 decoder Re-scoring 解码时也需要 left -to-right mask 来保证这样的限制（和训练一致）。但是这样的限制对于 Re-scoring 解码是有必要的吗？Re-scoring 解码的本质是 decoder 对 CTC 解出的候选集（hypothesis） 进行重新打分，此时候选集中的每条结果都已经是完整的 token 序列，也就是说，每个时刻都可以知道自己之后的时刻是什么。那么这时如果 decoder 每个时刻仍然只能看到其之前的信息岂不是浪费？那我们如何能让当前时刻可以看到其之后时刻的信息呢？

潜在的解决方案

回顾一下 NLP 里百家争鸣的预训练模型，有 GPT 这种 left-to-right LM 的方式（和U2 decoder的部分建训练与建模方式相同），也有 BERT 这种 masked LM 的方式，即根据两边（当前时刻之前 和当前时刻之后）预测中间（当前时刻，实际上通常是当前连续的几个时刻）还有打乱顺序的某某模型等等。直觉上 BERT 这种正是我们需要的方式，那我们可以尝试将原有的 decoder 换成 BERT 的方式进行训练。但是实际上我们并没有选择这种方式。因为我们对候选的序列进行重新打分时需要得到每个时刻的概率，为了在解码时得到这样的概率，我们在训练时需要创造这样的任务即对任意时刻：mask 掉该时刻，然后用该时刻前后的信息预测当前时刻。而 BERT 在训练时，只能 mask 一个序列中的一部分，无法对任意时刻都进行 mask 的操作（因为如果任意时刻都 mask，模型将无法进行有效训练），那么难道 BERT 带来的双向建模看似这么契合 Re-scoring 解码的方式真的没办法解决这个问题了吗？答案也是可以的，但是这个方案会极大增加训练和解码时的计算量，读者可以自己先行思考。即使可以忽略训练和解码时的计算量的问题，我们也暂时不会考虑 BERT 的方式了，因为我们有更简单的方式，并且可以在 Re-scoring 解码时更灵活。

U2++ 双向 re-score

模型&训练

实际上解决上述问题的方法也很简单，我们只需要再次增加一个 right-to-left（R2L） 的 decoder 即可。此次更新的 U2++ 模型如下图，即在原有U2的基础上，增加了新的 Right to Left Attention decoder。

这个增加的 decoder 可以根据右侧信息（当前时刻之后的信息）来预测当前时刻。因此增加的 left to right decoder 可以对右侧信息进行建模，U2 原有的 decoder 可以对左侧信息建模，那么二者结合起来就以达到对完整的序列进行双向建模的能力。如何实现这个 R2L decoder 呢？

L2R decoder 实现

一种方式是引入 R2L mask （上图 b）来完成 R2L decoder 的训练，此时输入与输出的处理与 U2 原有 L2R decoder 相同（不考虑 batch training 问题，实际由于为了 batch training 的 padding ，需要很多额外的处理，这个方式的实现可以参考 diwu-bidirectional-rescore 分支）。另外一种方式是我们使用与 U2 原有 L2R decoder 相同的 mask （上图 a），但是将输入和输出序列分别反向即可，这种方式的实现更为简单。这二种实现的方式的建模思想是一致的，区别仅在于输入 token 位置信息不同。经过对比这两种实现方式的结果很接近，为了简洁，我们选择了后者的实现方式。在训练时一共存在三个任务即 CTC loss、L2R AED loss 和新增的 R2L AED loss。将这三个 loss 加权及得到最终的 loss。α 和 λ 在训练时通常都设为 0.3。

解码

Re-scoring 解码的过程与 U2 一致，只是现在将会有两个不同的 L2R 和 R2L decoder，这两个 decoder 会根据 encoder out 对 CTC 产生的 hypothesis 重新打分，并且使用 α 和 λ 对这两个 decoder 的得分 SR2L，SL2R 与 SCTC 得分进行加权，选取加权之后得分最高的一条序列作为最终识别结果。此时的权重系数 α 和 λ 与训练时固定的数值不同，在解码时可以进行超参数搜索，来获得更优的结果。

实验结果

在 AISHELL-1 和 AISHELL-2 数据集上进行了测试，并且为了更好的展现 U2++ 双向建模的优势，我们在 Transformer 和 Conformer 模型上都进行了测试。为了公平比较，控制了 U2++ 两个 decoder（各3层）总的层数与 U2 相同（6层）。结果如下所示，其中 AISHELL-1 的数据做了 speed perturb 数据增强。U2++ 的实验也都增加了 SpecSub 数据增强 。在这两个数据集上，相同规模的参数量且使用字建模的情况下，U2++ 甚至比目前已知使用了 LM 的模型结果还要好。

我们对比了 Re-scoring （RS） 解码方式和自回归的解码方法（AR）。在 AISHELL-1 数据集上，U2++ RS 解码方法仍然好过自回归的解码方法 AR（下图）。其中 U2++ 的自回归解码由于有两个 decoder，所以在解码时我们采用的策略是：1）两个 decoder 分别进行自回归解码，2）然后选取二者产生所有路径中的最大路径作为最终结果。

此外，不管是 RS 解码方式还是 AR 解码方法，在使用相同的 decoder 层数的情况下，U2++ 都要好过 U2，甚至 U2++ 只使用其中的一个 decoder （三层）做 RS 和 AR 解码，其结果也比 U2 6层 decoder 的 baseline 要好。

更灵活的解码

上文说到不再选择 BERT 而选择当前 U2++ 这样的方式的另外一个原因是 U2++ 在解码时可以更加灵活，可玩性更高。并且也可以利用 U2++ 来解决上文的第二个问题---优化 re-score 的延迟。U2++ 现在有两个 decoder，那么在 Re-scoring 解码时，我们可以根据实际情况，只选择其中的一个 decoder 进行 re-score，这时相比较于原来 U2 decoder 6层参数的前向计算，U2++ decoder 部分前向运算只剩下三层，在模型结果更好的情况下，减少了近一半的运算，进而减少了 re-score 过程由于计算带来的延迟。

为此在实际使用过程中，可以根据实际情况来选择是否需要同时使用两个 decoder 做 re-score。我们期待更好的结果的时候可以选择使用两个 decoder，期待更少的 re-score 延迟的时候可以选择只使用一个 decoder。

使用

此次更新的双向 re-score 模块，向前做了兼容，并且 python 和 c++ runtime 都做了更新。对于之前的训练好的模型，小伙伴们可以 pull 新的代码，将模型重新导出即可。训练部分的配置可以参考 aishell/s0/conf/train_u2++_conformer.yaml 来进行模型的训练。在跑 python 或 runtime 解码时可以搜索 λ（reverse weight）的值，通常 0.4 或 0.5 时最优。

总结

此次 WeNet 更新了 U2++ 双向 re-score 建模，更完善的建模方法，更好的模型结果，更为灵活的解码。