Align-Refine: Non-Autoregressive Speech Recognition via Iterative Realignment (2021)

Abstract: Non-autoregressive encoder-decoder models greatly improve decoding speed over autoregressive models, at the expense of generation quality. To mitigate this, iterative decoding models repeatedly infill or refine the proposal of a non-autoregressive model. However, editing at the level of output sequences limits model flexibility. We instead propose *iterative realignment*, which by refining latent alignments allows more flexible edits in fewer steps. Our model, Align-Refine, is an end-to-end Transformer which iteratively realigns connectionist temporal classification (CTC) alignments. On the WSJ dataset, Align-Refine matches an autoregressive baseline with a 14x decoding speedup; on LibriSpeech, we reach an LM-free test-other WER of 9.0% (19% relative improvement on comparable work) in three iterations. We release our code at https://github.com/amazon-research/align-refine .

## introduction transformer encoder-decoder models @xcite have achieved high performance in sequence-to-sequence tasks like neural machine translation (nmt; @xcite and endto-end automatic speech recognition (asr; @xcite . however, like their recurrent predecessors, these models are autoregressive at inference: tokens are generated sequentially, with each token conditioned on all previous tokens. this makes decoding time linear in output sequence length, which is slow for long sequences. by contrast, nonautoregressive (nar) models decode all output tokens independently and in parallel. when combined with self-attention, one gets fast, constanttime inference in nmt @xcite and endto-end asr @xcite . however, these models underperform their autoregressive counterparts, as the conditional independence between output tokens results in globally inconsistent outputs. 1 to mitigate these issues, infilling methods like mask-predict @xcite refine an initial non-autoregressive proposal, repeatedly predicting a masked subset of low-confidence proposal tokens in a fixed number of decoding passes. in asr, most iterative non-autoregressive methods use infilling, like a-fmlm @xcite , imputer @xcite , and mask-ctc @xcite . however, during training, infilling requires partial proposals to be simulated by synthetically masking ground truths or samples from an expert. the resulting train-test mismatch leads to poor-quality generation. an alternative proposed in nmt is iterative refinement @xcite ; here, full proposals are predicted and trained on at each iteration, with no masking required. this reduces mismatch but still lacks flexibility: among other problems, working at the output sequence level constrains every iteration to the initial length l predicted by the model, making it difficult to correct insertions or deletions. in this work, we propose iterative realignment, a variation on iterative refinement where latent alignments are edited instead. by working at the alignment level, we avoid length prediction and enable more powerful edits, while preserving the flexibility and reduced train-test mismatch of iterative refinement. our model, align-refine, demonstrates this on asr: after a transformer encoder first produces a noisy non-autoregressive proposal, the refiner (a non-causal transformer decoder) repeatedly conditions on both the previous proposal and the initial encoder representation to produce a better proposal. both the encoder and refiner are supervised with ctc @xcite , a loss defined between sequences and their latent monotonic alignments. unlike past methods, align-refine requires no token masking or expert policies. we validate our approach on two english asr benchmarks, improving on state-of-the-art infilling methods-mask-ctc and imputer-in word error rate (wer) and/or inference time (as measured by real-time factor foot_1 , or rtf). on wsj, we close the wer gap with an autoregressive baseline at 1/14th the rtf, outperforming mask-ctc after a single iteration. on librispeech, we improve on published (lm-free) nar results by 2.1% wer absolute on test-other @xmath0. our work suggests that iterative realignment is a promising direction for other sequence-to-sequence tasks, such as nmt. ## methods we propose align-refine, a variation of iterative refinement which, like its predecessor, always keeps the proposal fully formed; this permits flexibility in decoding (as iteration can be stopped at any time) and potential speedups (errors seen and fixed in parallel; easy utterances refined in fewer steps). however, unlike @xcite , our proposals are latent ctc alignments a k 1:|x| , not outputs y k 1:l . in previous work, working at the output sequence level is another source of irreversibility: the length l must be predicted even before the initial proposal p enc (y 0 |l, x) is generated, either explicitly (sampling from a modelled length distribution p(l|x); @xcite or implicitly (collapsing ctc outputs before infilling; @xcite . either way, the decoder cannot fix insertions/deletions. ## results wsj results ( librispeech results ( ## analysis in all, we have seen that align-refine improves performance over infilling methods like imputer and mask-ctc across tokenizations and dataset sizes (tables 1 and 2 ). we saw improvements qualitatively as parallel and multi-stage insertion/deletion 0 1 2 3 4 5 6 7 ∞ edits (figures 2 and 3 ). we conclude by discussing some properties, tradeoffs, and limitations: