Proposal for the integration of Tree-constrained Pointer Generator and Minimum Biasing Word Error (MBWE) training for contextual ASR

🚀 The feature

I’d like to propose the integration of tree-constrained pointer generator (TCPGen) [1] and Minimum Biasing Word Error (MBWE) training [2] for contextual biasing into torchaudio package.

The implementation of TCPGen consists of 2 parts. The first part organises the biasing list and constructs a wordpiece-level prefix tree. The second part performs the prefix-tree search and calculates relevant pointer generator distributions in RNNT. During training, at the beginning of each minibatch, a biasing list will be organised by finding biasing words in the current utterance and adding a certain number of distractors. This biasing list is then converted into a prefix tree. Then, at each RNNT predictor step, a prefix-tree search is performed according to the wordpiece history, and a set of valid wordpieces are selected. For each combination of the encoder and predictor steps in RNNT, in addition to the model output distribution, we calculate another distribution over these valid wordpieces (TCPGen distribution). An interpolation between the model output distribution and this TCPGen distribution is then performed to achieve contextual biasing. The interpolation weight is also predicted by the TCPGen indicating how much contextual biasing is needed at each step. Inference follows the procedure as training.

For implementation, I would put the first part (organising prefix-trees) under some supporting function directories, such as utils/. I would put the second part (tree search and calculating relevant distributions) in the RNNT model implementation, before the joint network. It takes the combination of predictor and encoder output vectors and predicts a distribution for each of the combinations, which is of the same size as the logits output. The output vector of TCPGen, which is a (weighted) sum of wordpiece embeddings, can be also fed into the joint network to achieve deep biasing.

For MBWE training, we need to perform a batched beam search algorithm to get n-best hypotheses for each utterance in the minibatch, and then calculate the expected word error (WE) and biasing word error (BWE) for each utterance. For the batched beam search, we adopted a one-step constrained beam search which constrained the output at each encoder step can only be 0 or 1 wordpiece token hence all forward computation could be parallelised. Once the n-best hypotheses are obtained, the WE and the BWE will be calculated. The BWE is calculated as the edit-distance between the sequence of biasing words that appeared in the hypothesis and reference. The WE and BWE for each hypothesis are summed up and then multiplied with the normalised probabilities to get the final MBWE loss to be optimised. The batched beam search could potentially go into the rnnt_decoder.py file and the calculation of MBWE loss could either be directly implemented as a class method of the RNNT model or be separately put into some supporting function directories.

There are mainly 2 issues with the implementation. First, the interpolation is between 2 distributions, but the RNNT loss function takes logits as arguments, it is necessary to modify the RNNT loss to include such interpolation. The second issue is that in our experiments, wordpiece models are suffix-based models that mark word boundary at the end of a word. The benefit is that it is clear when to get back to the root of the tree, whereas, for the prefix-based one, any starting wordpieces could be valid which influences the effect of both our method and deep biasing.

Please refer to the following papers for more details.

[1] Guangzhi Sun, Chao Zhang, Philip C Woodland. Tree-constrained Pointer Generator for End-to-end Contextual Speech Recognition. In Proc. ASRU. 2021. Link: https://arxiv.org/abs/2109.00627.

[2] Guangzhi Sun, Chao Zhang, Philip C Woodland. Minimising Biasing Word Errors for Contextual ASR with the Tree-Constrained Pointer Generator. Link: https://arxiv.org/pdf/2205.09058.pdf.

Motivation, pitch

Contextual biasing, which integrates contextual knowledge into an ASR system, has become increasingly important to many applications. The tree-constrained pointer generator (TCPGen) is an effective contextual biasing component that combines the benefit of deep biasing and shallow fusion-based contextual biasing. The biasing list which contains a list of words or entities that are likely to appear in a given context could be customised depending on the use case. In addition, the minimum biasing word error (MBWE) training provides an effective way of training contextual ASR systems by particularly emphasising the biasing word errors.

Alternatives

No response

Additional context

No response