SlowGuess/NAACL_2021_OCR huggingface dataset:
Ethan A. Chi
Stanford University\* ethanchi@cs.stanford.edu
Julian Salazar
Amazon AWS AI
julsal@amazon.com
Katrin Kirchhoff
katrinki@amazon.com
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, AlignRefine, 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  $14 \times$  decoding speedup; on LibriSpeech, we reach an LM-free testother 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.
Transformer encoder-decoder models (Vaswani et al., 2017) have achieved high performance in sequence-to-sequence tasks like neural machine translation (NMT; Edunov et al., 2018) and end-to-end automatic speech recognition (ASR; Karita et al., 2019). 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, non-autoregressive (NAR) models decode all output tokens independently and in parallel. When combined with self-attention, one gets fast, constant-time inference in NMT (Gu et al., 2018) and end-to-end ASR (Salazar et al., 2019). However, these models underperform their autoregressive counter
Figure 1: In Align-Refine, a Transformer encoder labels each input frame to give a latent alignment. The refiner, a non-causal Transformer decoder, improves the alignment conditioned on the encoder. After a bounded # of iterations, the result is collapsed into the output.
parts, as the conditional independence between output tokens results in globally inconsistent outputs. $^{1}$
To mitigate these issues, infilling methods like Mask-Predict (Ghazvininejad et al., 2019) 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 (Chen et al., 2020), Imputer (Chan et al., 2020), and Mask-CTC (Higuchi et al., 2020). 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 (Lee et al., 2018); here, full proposals are predicted and trained on at each iteration, with no masking required. This
<sup>1</sup>Gu et al. call this the multimodality problem, as it is induced by the highly multimodal distribution of target translations (or in the case of ASR, frame-level alignments).
Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 1920-1927 June 6-11, 2021. ©2021 Association for Computational Linguistics
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 (Graves et al., 2006), 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 $^2$ , 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 at  $< 1/4$ th the effective layers (and thus estimated RTF) of Imputer. Our work suggests that iterative realignment is a promising direction for other sequence-to-sequence tasks, such as NMT.
2.1 Connectionist Temporal Classification
CTC (Graves et al., 2006) is a strategy for defining latent monotonic alignments from an input sequence  $\mathbf{x}$  to a shorter output sequence  $\mathbf{y}$ . Let  $\underline{\cdot}_{-}^{\cdot}$ , termed a blank, be an additional possible output token; then a CTC alignment is reduced to an output sequence by collapsing repeated labels then removing blanks, e.g., AB\_BB\_A  $\mapsto$  ABBA. Since this is a many-to-one process, to calculate  $p(\mathbf{y}|\mathbf{x})$ , we marginalize over all alignments  $\psi (\mathbf{y})$  mapping to an output  $\mathbf{y}$ . Assuming alignment labels are
conditionally independent:
$$
p (\mathbf {y} | \mathbf {x}) = \sum_ {a \in \psi (\mathbf {y})} p _ {\theta} (a | \mathbf {x}) = \sum_ {a \in \psi (\mathbf {y})} \prod_ {t = 1} ^ {| \mathbf {x} |} p _ {\theta} (a _ {t} | \mathbf {x}),
$$
making  $J_{\mathrm{CTC}} = -\log p(\mathbf{y}|\mathbf{x})$  differentiable and efficiently computable via dynamic programming.
2.2 Existing Approaches
Iterative refinement methods non-autoregressively refine an initial proposal  $\mathbf{y}^0$  of length  $L$ , with full re-predicted proposals  $\mathbf{y}^k$  conditioned on previous proposals  $\mathbf{y}^{k - 1}$  and the input sequence  $\mathbf{x}$ :
$$
p (\mathbf {y} ^ {k} | \mathbf {x}) = \prod_ {i} p (y _ {i} ^ {k} | \mathbf {x}, \mathbf {y} ^ {k - 1}).
$$
By contrast, infilling methods such as Mask-Predict (Ghazvininejad et al., 2019) mask part of the previous proposal every iteration. Only masked positions are re-predicted, conditioned on all unmasked tokens and the inputs  $\mathbf{x}$ :
$$
p (\mathbf {y} _ {\mathrm {m a s k}} ^ {k} | \mathbf {x}) = \prod_ {i} p (y _ {i} ^ {k} | \mathbf {x}, \mathbf {y} ^ {k - 1} \setminus \mathbf {y} _ {\mathrm {m a s k}} ^ {k - 1}).
$$
Typically, at each iteration a decreasing proportion of high-confidence tokens are unmasked such that the full budget of  $K$  iterations is always required. Furthermore, state-of-the-art methods for English ASR—Mask-CTC (Higuchi et al., 2020, Figure 2) and Imputer (Chan et al., 2020)—enforce an added constraint that no decisions may be reversed at all:  $p(y_{i}^{k}|\mathbf{x}) = \mathbf{1}[y_{i}^{k} = y_{i}^{k - 1}]$  when  $y_{i}^{k - 1}\neq <\mathrm{MASK}>$ .
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 Lee et al. (2018), our proposals are latent CTC alignments  $\mathbf{a}_{1:|\mathbf{x}|}^{k}$ , not outputs  $\mathbf{y}_{1:L}^{k}$ . 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_{\mathrm{enc}}(\mathbf{y}^{0}|L,\mathbf{x})$  is generated, either explicitly (sampling from a modelled length distribution  $p(L|\mathbf{x})$ ; Lee et al., 2018) or implicitly (collapsing CTC outputs before infilling; Higuchi et al., 2020). Either way, the decoder cannot fix insertions/deletions.
3In Mask-CTC, this is done by the decoder. In Imputer, full encoder passes are required, as there is no decoder.
2The decoding time divided by the length of audio.
Enc	···F···	···S···RE···LY···EL···	TO····L···	H· SAID	__SSA_FFFO_ID	H'S DIRECTLY	RRELATED	TOO	HEALTH	HHE	SAAIDD_
k=1	···F···	H·S·D·RE···HLY···EL···D	TO·H···L···	H· SAID	__SSA__ FFO_ID	HIS	DIRECTLY	RELATED	TO	HEALTH	HE SAAID_
k=2	···F···	HAS·DIRE···HLY·RELAX···D	TO·H·AL···	H· SAID	__SSA__ FFO_ID	HIS	DIRECTLY	RELATED	TO	HEALTH	HE SAAID_
k=3	S···F···	HAS DIRE·THLY·RELAT-D	TO·H·ALT	H· SAID	__SSAI_FFO_ID	HIS	DIRECTLY	RELATED	TO	HEALTH	HE SAAID_
k=4	S···FO···	HAS DIRE·THLY RELATED	TO H·ALT	HE SAID
k=5	SE FOET HAS DIRECTLY	RELATED TO HEALT	HE SAID
End	SE FOET HAS DIRECTLY	RELATED TO HEALT	HE SAID
Figure 2: WSJ dev93 utterance as decoded by our models. Mask-CTC's masks are denoted with  $\bullet$ . Mask-CTC gives HEALT as it has no space for H from the very beginning, and outputs DIRECTLY as it cannot undo HLY's emission at iteration  $k = 1$ . Meanwhile, Align-Refine makes the mistake HEALT immediately, but corrects it over two steps: at  $k = 1$  it deletes HHE  $\mapsto$  HE, and at  $k = 2$  it sees the new space and inserts HEALT  $\mapsto$  HEALTH.
By contrast, since CTC alignments are by nature fixed-length, Align-Refine easily handles insertions/deletions by placing/replacing blanks and spaces (Figure 2). At a given iteration  $k$ , we have:
$$
p (\mathbf {y} | \mathbf {x}) = \mathbb {E} _ {\mathbf {a} ^ {k - 1}} \left[ \sum_ {\mathbf {a} ^ {k} \in \psi (\mathbf {y})} p _ {\text {r e f}} \left(\mathbf {a} ^ {k} \mid \mathbf {a} ^ {k - 1}, \mathbf {x}\right) \right].
$$
This expectation over previous alignments is intractable, so like Lee et al. (2018) we take a deterministic lower bound: sampling the mode  $\hat{\mathbf{a}}^{k - 1}$ . After we marginalize over all  $\mathbf{a}^k$ , the loss is  $J_{\mathrm{CTC}}(\hat{\mathbf{a}}^{k - 1},\mathbf{y};\theta_{\mathrm{ref}},\theta_{\mathrm{enc}})$ . Since we don't know a priori which  $k$  is final, we apply the loss at  $k = 1,\ldots ,K$  with weights  $w_{1},\ldots ,w_{K}$  for some hyperparameter  $K$ . For  $k = 0$  we get  $J_{\mathrm{CTC}}(\mathbf{x},\mathbf{y};\theta_{\mathrm{enc}})$ .
In summary, we take the greedy alignment at each iteration and apply the CTC loss, as shown in Figure 1 for  $K = 2$ . In practice, we upweight the encoder and first iteration terms with weights  $\lambda$  and  $w_{1}$ , then sum to give the total loss. For this and other training details, consult Appendix B, C.
Data. We evaluate on two English ASR benchmarks: WSJ (81 hours; Paul and Baker, 1992) and LibriSpeech (960 hours; Panayotov et al., 2015). For WSJ, we run at the character level, matching Mask-CTC; for LibriSpeech, we build a 400-token BPE vocabulary, matching Imputer. We use 80-dim. filter banks and SpecAugment (Park et al., 2019).
Model. For WSJ we use a 12-layer encoder and 6-layer decoder, as in Higuchi et al. (2020); each layer has 4 heads over 256 units. For LibriSpeech, we use 8 heads over 512 units. With CNN frontends these are 27M and 71M parameters. Unless stated otherwise, we do  $K = 4$  training iterations.
Decoding. We evaluate with decoding iterations  $k$  from  $\{0,1,3,5,10\}$ , exiting early on convergence (consecutive iterations are identical). The
final CTC alignment is collapsed to give the result. To match previous non-autoregressive ASR work, we do not use a language model (LM).
WSJ results (Table 1). Since Mask-CTC and Align-Refine share an identical architecture, the difference lies solely in training and evaluation. For both models, joint training with the refinement objective improves the encoder's performance as a standalone CTC model ( $k = 0$ ) to a similar degree. However, from  $k = 1$  onwards, Align-Refine outperforms Mask-CTC, improving the initial encoder proposal by  $1.9\%$  absolute in just one iteration:
Reference: SEAFOOD IS DIRECTLY RELATED TO HEALTH HE SAID
Model	# passes	WER	RTF
Enc	Dec	dev93	eval92
Autoregressive baseline (Higuchi et al., 2020)
CTC+ATTENTION	1	L	14.4	11.3	0.97*
+ beam search	1	> L	13.5	10.9	4.62*
Previous work (Higuchi et al., 2020; Chan et al., 2020)
CTC	1	-	22.2	17.9	0.03*
MASK-CTC	1	0	16.3	12.9	0.03*
1	1	15.7	12.5	0.04*
1	5	15.5	12.2	0.05*
1	#mask	15.4	12.1	0.13*
IMPUTER (8-LYR) †	8	-	-	12.7	-
Our work
CTC	1	-	18.6	15.0	0.036
MASK-CTC	1	5	15.3	12.8	0.072
ALIGN-REFINE	1	0	16.2	13.5	0.037
1	1	14.1	11.6	0.048
1	3	13.9	11.5	0.066
1	5	13.7	11.4	0.068
1	10	13.7	11.4	0.068
Table 1: Non-autoregressive ASR on WSJ. No LMs are used. For CERs, see Table 4 (Appendix A).*: RTFs from Higuchi et al., which are lower than ours for corresponding models. †: No SpecAugment.
By  $k = 5$ , Align-Refine closes the performance gap with a comparable autoregressive model at 1/14th the RTF. As for the 8-layer Imputer, it seems un
Mask-CTC, 5 iterations
Align-Refine (char.),  $\leq 5$  iterations
Enc	[WHEN__[DI [DI_CKI__[CAME__[DOWN__[HIS__[A__UNUNTT__[S__IGHT_LY__[S[SLA__PP_T__[HIM
k=1	[WHEN__[DI [DI_CKIEE[AME__[DOWN__[HIS__[A__UNTT__[SL_IGHT_LY__[SLA__PP_ED__[HIM
k=2	[WHEN__[DI [DI_CKI_E[AME__[DOWN__[HIS__[A__UN_T__[SL_IGHT_LY__[SLA__PP_ED__[HIM
k=3	(identical, so collapse early)
...	↓
End	[WHEN [DICKIE [CAME [DOWN [HIS [AUNT [SLIGHTLY [SLAPPED [HIM
likely that re-introducing SpecAugment would outperform this augmented 12+6-layer autoregressive baseline; even if performances did match, RTFs would be higher than Align-Refine's (Section 5).
LibriSpeech results (Table 2). Align-Refine gives  $9.0\%$  WER on test-other with no LM, outperforming published non-autoregressive models by  $2.1\%$  WER absolute. This is 5.1 points better than training the encoder with CTC only, even after SpecAugment is used (compare with 1.9 points from 16-layer CTC to Imputer). In Figure 3 we see the refiner make output-conditional edits that would be difficult for greedy CTC inference; we attribute our outsized gain on test-other to this LM-like behavior. Future work could start from stronger CTC encoders like QuartzNet (Kriman et al., 2020) to achieve even better results.
Model	# passes	WER (test)
Enc	Dec	clean	other
Autoregressive models (Han et al., 2020)
LAS (360M)	1	L	2.6	6.0
RNN-T ( CONTEXTNET-L)	1	L	2.1	4.6
Previous work (Chan et al., 2020; Kriman et al., 2020)
CTC (16-LYR) †	1	-	4.6	13.0
IMPUTER (16-LYR) †	8	-	4.0	11.1
CTC (JASPER DR 10x5)	1	-	4.3	11.8
CTC (QUARTZNET 15x5)	1	-	3.9	11.2
Our work
CTC	1	-	5.1	14.1
ALIGN-REFINE	1	0	4.6	11.5
1	1	3.8	9.5
1	3	3.6	9.0
1	5	3.6	9.0
Table 2: Non-autoregressive ASR on LibriSpeech. No LMs are used. †: No SpecAugment.
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
Figure 3: LibriSpeech test-other utterance; reference matches the prediction. At  $k = 1$  three separate corrections are made, two of which (DICKI→DICKIE; SLAPPT→SLAPPED) cannot be done from the audio. In  $k = 2$  the multimodal prediction E E is resolved into _E, though the repetition-collapsed transcript would be correct regardless.
Figure 4: Proportion of utterances in each LibriSpeech test set whose alignments become fixed at iteration  $k$ .  $\infty$  denotes utterances which cycle between alignments indefinitely.
edits (Figures 2 and 3). We conclude by discussing some properties, tradeoffs, and limitations:
Number of iterations. In Figure 4 we graph how many utterances became fixed at each iteration  $k$  (i.e., where  $\hat{\mathbf{a}}^{k + 1} = \hat{\mathbf{a}}^k \neq \hat{\mathbf{a}}^{k - 1}$ ). First, we see that the CTC alignment is always revised by the refiner, evidence that CTC's non-autoregressive greedy mode is fundamentally different from the conditionally dependent mode annealed to by the refiner. We see that upweighting  $w_{1}$  and training with  $K = 4$  largely confined the "fixed point" iteration index  $k$  to 4 or less.
One interesting phenomenon is the small fraction of utterances which never reach a fixed point
Align-Refine (subword), up to 5 iterations
Figure 5: Bar plot of iterations at which each alignment became fixed, as a function of alignment length (4× downsampled from regular audio frames), over both test-clean and test-other. For comparability we normalize each bucket to 1.0; the total number of utterances in each bucket is above the chart.
(labeled  $\infty$ ); rather, they cycle through two or more alignments repeatedly. In these cases, the transcript is finalized except on a local set of tokens where the model flips back and forth, e.g., -0US versus -US. One could mitigate this by stopping after an edit distance of 1 between alignments is reached (similar to Lee et al., 2018), or by comparing with proposals e.g., two iterations prior.
Length independence. In Figure 5, we empirically validate the length independence of our method by plotting the "fixed point" iteration  $k$  versus the alignment length. The medians increase from  $k = 1$  to  $k = 2$  by lengths  $>300$ , but no further (for the lengths in LibriSpeech). One observation is that as an utterance gets longer, the chance that "multimodal" corrections like [DI_CKI__]→[DI_CKIEE (Figure 3) are made increases, which a further iteration then resolves ([DI_CKIEE]→[DI_CKI_E] independent of whether it affects the collapsed transcript. Interestingly, a few short alignments (e.g., lengths  $< 100$ ) use from 3 all the way up to 8 iterations, perhaps due to lack of linguistic context the refiner can use for disambiguation.
Speed and decoder depth. By factoring out refinement from feature processing, we can effectively adapt to variable inference budgets by adjusting the number of decoding iterations. While Imputer requires a separate run of the entire model
for every iteration, for Align-Refine, only the decoder must be rerun. Another advantage is that since we have a full proposal at every iteration, in the vast majority of cases we exited early once the proposals stabilized.
Kasai et al. (2021) critiqued this factorization in non-autoregressive models by showing autoregressive models could shift layers from decoder to encoder to reduce the speed gap; however, we show that Align-Refine benefits from a similar reallocation (Table 3):
Model-(# enc)-(# dec)	WER (test-other) for each k
0	1	2	3	5
ALIGN-REFINE-12-6	12.5	10.0	9.4	9.3	9.3
ALIGN-REFINE-15-3	10.5	9.5	9.4	9.3	9.3
ALIGN-REFINE-17-1	10.4	10.0	10.0	10.0	10.0
Table 3: Results with various encoder-decoder splits on LibriSpeech. We take  $K = 2$  to speed up training.
The refiner needs some depth for best performance: 17-1 underperforms on test-other (though within 0.1 points on test-clean). However, 15-3 performs as well as 12-6 at  $k = 1$  onwards, despite passing through half the number of decoder layers. At  $k = 3$ , our LibriSpeech model's RTF was 0.171, while 15-3's RTF is 0.136. In this configuration, we pass through 15 encoder and at most  $3 \times 3 = 9$  decoder layers. By contrast, Imputer inference passes through  $8 \times 16 = 128$  encoder layers with the same alignment-length inputs and layer size.
Limitations. The refiner sometimes makes edits which do not affect post-collapse outputs (Figure 3,  $k = 2$ ); variants that use repetition tokens (ASG; Collobert et al., 2016) or prohibit repetition collapse (Chan et al., 2020) may mitigate this behavior. The initial downsampling also restricts what edits can be done in one step (Figure 2,  $k = 1, 2$ ).
We found that Align-Refine performed worse when there were more alignment labels per word (Appendix A). In general, word-level errors like SAI FOOD HIS (Figure 2) may require more coordinated mechanisms to fix. Future work could integrate non-autoregressive LMs like BERT (Devlin et al., 2019) via alignment synthesis, fusion, or otherwise (e.g., Salazar et al., 2020). Some concurrent works (Inaguma et al., 2021; Tian et al., 2021) also propose reranking NAR candidates with a jointly-trained autoregressive decoder.
We are grateful to the AWS Speech Science team and the Stanford NLP Group for their advice and feedback, to Yosuke Higuchi and William Chan for their helpful correspondence, and to the anonymous reviewers for their thoughtful feedback.
William Chan, Chitwan Sahara, Geoffrey E. Hinton, Mohammad Norouzi, and Navdeep Jaitly. 2020. Imputer: Sequence modelling via imputation and dynamic programming. In ICML.
Nanxin Chen, Shinji Watanabe, Jesús Villalba, and Najim Dehak. 2020. Listen and fill in the missing letters: Non-autoregressive transformer for speech recognition. CoRR, abs/1911.04908v2.
Ronan Collobert, Christian Puhrsch, and Gabriel Synnaeve. 2016. Wav2letter: an end-to-end convnet-based speech recognition system. CoRR, abs/1609.03193v2.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019. BERT: Pre-training of deep bidirectional transformers for language understanding. In *NAACL-HLT* (1), pages 4171–4186. Association for Computational Linguistics.
Linhao Dong, Shuang Xu, and Bo Xu. 2018. Speech-Transformer: A no-recurrence sequence-to-sequence model for speech recognition. In ICASSP, pages 5884-5888. IEEE.
Sergey Edunov, Myle Ott, Michael Auli, and David Grangier. 2018. Understanding back-translation at scale. In EMNLP, pages 489-500. Association for Computational Linguistics.
Marjan Ghazvininejad, Omer Levy, Yinhan Liu, and Luke Zettlemoyer. 2019. Mask-Predict: Parallel decoding of conditional masked language models. In EMNLP/IJCNLP, pages 6111-6120. Association for Computational Linguistics.
Alex Graves, Santiago Fernández, Faustino J. Gomez, and Jürgen Schmidhuber. 2006. Connectionist temporal classification: Labelling unsegmented sequence data with recurrent neural networks. In ICML, volume 148 of ACM International Conference Proceeding Series, pages 369-376. ACM.
Jiatao Gu, James Bradbury, Caiming Xiong, Victor O. K. Li, and Richard Socher. 2018. Nonautoregressive neural machine translation. In *ICLR (Poster)*. OpenReview.net.
Wei Han, Zhengdong Zhang, Yu Zhang, Jiahui Yu, Chung-Cheng Chiu, James Qin, Anmol Gulati, Ruoming Pang, and Yonghui Wu. 2020. ContextNet: Improving convolutional neural networks for automatic speech recognition with global context. In INTERSPEECH. ISCA.