Source-linked AI summary
Imputer: Sequence Modelling via Imputation and Dynamic Programming
William Chan, Chitwan Saharia, Geoffrey Hinton, Mohammad Norouzi, Navdeep Jaitly
TL;DR
For sequence problems with monotonic alignments, including speech recognition, long outputs make sub-linear inference important, while existing autoregressive and non-autoregressive approaches trade generation speed against output dependencies. The Imputer iteratively refines alignments in a constant number of steps and uses tractable dynamic programming to approximately marginalize over alignments and generation orders. On LibriSpeech, Imputer achieves 11.1 WER, outperforming the CTC baseline at 13.0 WER and seq2seq baseline at 12.5 WER.
Problem
For sequence problems with monotonic alignments, including speech recognition, long outputs make sub-linear inference important, while existing autoregressive and non-autoregressive approaches trade generation speed against output dependencies.
Source: §1 · pp. 1–2Method
Results
On LibriSpeech, Imputer achieves 11.1 WER, outperforming the CTC baseline at 13.0 WER and seq2seq baseline at 12.5 WER.
Source: §4 · p. 6Takeaways & Limitations
For end-to-end speech recognition, the Imputer outperforms prior non-autoregressive approaches and achieves competitive results to autoregressive models.
Source: §6 · p. 9The direct objective has high variance from double sampling, and the choices of masking and alignment distributions are unclear, making naive optimization difficult.
Source: §2.4 · p. 3
Abstract
from arXiv · showhide
This paper presents the Imputer, a neural sequence model that generates output sequences iteratively via imputations. The Imputer is an iterative generative model, requiring only a constant number of generation steps independent of the number of input or output tokens. The Imputer can be trained to approximately marginalize over all possible alignments between the input and output sequences, and all possible generation orders. We present a tractable dynamic programming training algorithm, which yields a lower bound on the log marginal likelihood. When applied to end-to-end speech recognition, the Imputer outperforms prior non-autoregressive models and achieves competitive results to autoregressive models. On LibriSpeech test-other, the Imputer achieves 11.1 WER, outperforming CTC at 13.0 WER and seq2seq at 12.5 WER.
1. Introduction
The paper seeks a balance between autoregressive models, which estimate exact likelihoods but require one generation step per token, and faster non-autoregressive models that often underperform. It targets sequence problems with monotonic alignments, where long outputs make sub-linear inference especially valuable.
1. Introduction
- Autoregressive models enable exact likelihood estimation but need n steps for n tokens; non-autoregressive models can generate in one step but often underperform because of conditional independence assumptions.Intermediate approaches can require logarithmic or constant generation steps while sacrificing almost no performance.Source: §1 · p. 1
- Monotonic alignment is a useful inductive bias for long-output sequence problems such as speech recognition, where sequences are typically an order of magnitude longer than in machine translation.The paper identifies sub-linear inference time as important for these longer sequences.Source: §1 · p. 1
- CTC marginalizes alignments efficiently but assumes conditional independence; RNN-T models output dependencies but requires n steps and typically uses a linear bottleneck for tractability.Seq2seq often performs best empirically but requires n generation steps for n tokens.Source: §1 · pp. 1–2
- The Imputer uses a constant number of generation steps, approximately marginalizes over alignments and generation orders, and supports bidirectional contextualization and local and global dependencies.It superimposes input and output sequences, avoiding the cross-attention mechanism typically used in encoder-decoder models.Source: §1 · p. 2
2. The Imputation Framework
The framework represents the output through a latent monotonic alignment with the input, using blanks so that removing them recovers the output sequence. It contrasts this alignment approach with CTC’s conditionally independent predictions.
2. The Imputation Framework
- A latent alignment has length |x|, is monotonic, and collapses to output y when blank tokens are removed.The framework defines β(y) as the set of possible alignments and β−1(a) as the collapse function; repeated-token handling is treated as a hyperparameter choice.Source: §2 · p. 2
- Figure 1’s example uses four blocks of three tokens, and decoding takes exactly B iterations.Each decoding iteration imputes one token per block.Source: Figure 1 · p. 2, §2 · p. 2
2.1. Connectionist Temporal Classification (CTC)
- CTC marginalizes over all alignments compatible with y using dynamic programming, but its conditional independence assumption limits its ability to model language dependencies and multimodal outputs.That assumption also permits one-step non-autoregressive generation.Source: §2.1 · p. 2
2.2. The Imputer
The Imputer iteratively refines a partially generated alignment, conditioning each new alignment on the previous one. Its block-based inference procedure balances parallel generation with conditional dependencies while requiring a constant number of steps.
2.2. The Imputer
- At each step, the Imputer emits a complete alignment conditioned on a previous partial alignment, then retains selected predictions to refine that partial alignment.Committed tokens remain consistent; conditioning across steps models dependencies, while token predictions within a step use conditional independence for parallel generation.Source: §2.2 · pp. 2–3
2.3. Inference
- The framework also permits fully non-autoregressive generation in one step or fully autoregressive generation in |x| steps, with the latter not constrained to left-to-right order.The block strategy provides an intermediate balance between these extremes.Source: §2.3 · p. 3
- Block inference selects the highest-probability masked slot in each block and commits one token per block per iteration, finishing in exactly B iterations.B=1 gives fully non-autoregressive generation, while B=|a| gives fully autoregressive generation.Source: §2.3 · p. 3
2.4. Marginalization
The training objective samples alignments and maskings to approximate marginalization over alignments and generation orders. The paper develops a dynamic-programming target policy for summing compatible completions, while noting challenges with the direct objective.
2.4. Marginalization
- The objective uses q and r to sample alignments and maskings, then applies Jensen’s inequality to obtain a lower bound while omitting entropy terms that do not affect the gradient.The bound assumes q and r have support everywhere and do not share parameters with pθ.Source: §2.4 · p. 3
- Directly optimizing the objective is difficult because double sampling has high variance and choosing the masking and alignment distributions is unclear.The paper frames training as choosing a roll-in policy and a target policy, then considers imitation learning or dynamic programming for the latter.Source: §2.4 · pp. 3–4
- Dynamic programming marginalizes over all alignments compatible with the conditioning alignment’s masked regions.The roll-in policy samples a masked-out alignment; the algorithm sums compatible completions over those regions.Source: §2.4 · p. 4, Figure 2 · p. 4
2.5. Roll-in Policy
The roll-in policy uses an inexpensive alignment source and a masking scheme designed to match inference conditions. A stale model provides an alternative that more closely matches those conditions.
2.5. Roll-in Policy
- A pseudo-expert policy perturbs an empirical alignment, allowing alignments to be stored offline and cheap noise to be sampled during training.Sampling directly from the model would require costly on-the-fly decoding, while storing expert logits would require large amounts of memory.Source: §2.5 · p. 4
- A stationary distribution from a stale model copy avoids live decoding and backpropagation while more closely matching training and inference conditions.The paper experiments with both this approach and the pseudo-expert policy.Source: §2.5 · p. 4
- Block sampling masks the same number of randomly selected tokens in each block, matching inference, which selects one token per block in parallel.The number masked per block is sampled uniformly from 0 to B−1, where B is the block size.Source: §2.5 · p. 4
2.6. Imitation Learning
Imitation Learning trains the Imputer to imitate a pseudo-expert using a roll-in policy. Dynamic Programming instead marginalizes over compatible alignments, and that marginal can be computed efficiently.
2.6. Imitation Learning
- Imitation Learning uses qφ′ × r as its roll-in policy and teaches the model to imitate the pseudo-expert φ′.The architecture conditions on a previous alignment and uses audio features processed by a convolutional network and Transformer self-attention to generate a new alignment.Source: §2.6 · pp. 4–5, Figure 3 · p. 5
2.7. Dynamic Programming
- Dynamic Programming retains qφ′ × r as the roll-in policy but marginalizes over compatible alignments for its target policy.The compatible-alignment sum is computed efficiently by dynamic programming, using the same algorithm as CTC except that it forces emit symbols.Source: §2.7 · p. 5
2.8. Lower bound
The section relates Imitation Learning and Dynamic Programming objectives to the log probability through lower bounds. Dynamic Programming accounts for compatible alignments and is itself bounded by the imitation objective.
2.8. Lower bound
- Imitation Learning is a lower bound on Dynamic Programming, which in turn lower-bounds the log probability up to a constant.The constant can be dropped during gradient computation.Source: §2.8 · p. 5
- Dynamic Programming sums over alignments compatible with a partially masked alignment, with repeated alignments corrected by a repetition constant.That constant depends on the partial and complete alignments.Source: §2.8 · p. 5
- The imitation objective drops compatible-alignment terms other than the original alignment; because the omitted terms are positive, it is a lower bound.Both objectives use the same roll-in policy, so the importance-sampling correction terms can be ignored.Source: §2.8 · p. 5
3. Model
The Imputer processes audio features and a previous alignment to predict a new alignment through convolutional and Transformer layers.
3. Model
- A stack of convolutional layers processes the audio features, such as filter banks.The resulting features are added to embeddings from the previous alignment.Source: §3 · p. 6
- The combined audio features and previous-alignment embeddings pass through Transformer self-attention layers.Source: §3 · p. 6
- A softmax layer models the new alignment from the Transformer output.Source: §3 · p. 6
4. Experiments
Experiments show that the Imputer’s dynamic-programming (DP) training outperforms its imitation-learning alternative, and its LibriSpeech result beats both CTC and seq2seq baselines. Further analyses examine block size, decoding strategies, masking policies, and roll-in alignments.
4. Experiments
- 11.1 WER on LibriSpeech beats CTC at 13.0 WER and seq2seq at 12.5 WER; on WSJ, Imputer scores 12.7 WER versus CTC’s 15.2 WER.The experiments use 82-hour WSJ and 960-hour LibriSpeech, without data augmentation or external language models.Source: §4 · p. 6
- 14.6 WER for IL versus 11.1 WER for DP on LibriSpeech; the authors hypothesize higher IL gradient variance may contribute to its weaker performance.The authors attribute DP’s advantage over IL to its tighter lower bound on likelihood, while presenting the gradient-variance explanation as a hypothesis.Source: §4 · p. 6
4.1. Analysis
- B = 8 gives the lowest WER for both decoding strategies when training and inference use the same block size; smaller and larger training blocks perform worse.The authors hypothesize that larger-block models are harder to optimize.Source: Figure 5 · p. 7, §4.1 · p. 7
- With training fixed at B = 8, increasing inference block size monotonically lowers WER for Right Most Last, while Alternate Sub-block has a local minimum near B = 16.The authors hypothesize Alternate Sub-block’s behavior reflects its constraint that tokens be predicted in specific sub-blocks at each iteration.Source: Figure 6 · p. 7, §4.1 · p. 8
4.2. Masking Policy
- Imputer is robust to the choice among Bernoulli, Uniform, and Block masking policies.The reported comparison covers LibriSpeech dev-other WER.Source: Table 3 · p. 8, §4.2 · p. 8
4.3. Training with Model Samples
- Updating roll-in alignments with Imputer samples every 50k steps did not improve performance over starting with noisy pre-trained CTC alignments.The authors conclude that CTC alignments suffice for roll-in training in these experiments.Source: §4.3 · p. 8
4.4. Top-k Decoding
- 11.6 WER with greedy top-k decoding versus 11.2 WER with Block inference on LibriSpeech dev-other; both use constant B generation steps.The authors describe the top-k result as slightly worse and consider block inference reasonable for relatively local speech.Source: §4.4 · p. 8
4.5. Simulated Annealing Decoding
- Simulated-annealing-like decoding improved unconverged models but did not improve converged models.The tested strategies deleted token commitments randomly or based on token-level likelihood.Source: §4.5 · p. 8
5. Related Work
The related work contrasts Imputer with non-autoregressive and insertion-based models, and places its convolutional Transformer architecture alongside prior speech-recognition work. The cited comparisons emphasize sequence-length handling, monotonicity, and architectural choices.
5. Related Work
- NAT and Mask-Predict predict target length first and use encoder-decoder cross-attention, while Imputer uses a fixed-size canvas and local prediction structure for monotonic speech tasks.The paper describes the suitability of cross-attention for long monotonic sequence problems as uncertain.Source: §5 · p. 8
- Unlike Imputer’s fixed canvas and constant generation steps, the Insertion Transformer dynamically grows its canvas and generates in logarithmic iterations; its non-monotonic order was difficult to apply to speech.The passage cites the Insertion Transformer’s logarithmic generation result as Stern et al. (2019).Source: §5 · p. 8
- Imputer’s convolutional Transformer draws on prior speech-recognition architectures; concurrent work by Synnaeve et al. (2019) used substantially larger models, data augmentation, and a semi-supervised-learning focus.Mohamed et al. (2019) also combined convolutions for local acoustic features with self-attention for global signals and bidirectional language modeling.Source: §5 · p. 8
6. Conclusion
The Imputer iteratively generates sequences by updating partial alignments, using a tractable dynamic program to approximately marginalize over alignments and generation orders. On speech recognition, it outperforms prior non-autoregressive approaches and achieves comparable performance to autoregressive models.
6. Conclusion
- 11.1 WER on LibriSpeech test-other beats CTC (13.0 WER) and seq2seq (12.5 WER); overall, Imputer outperforms prior non-autoregressive and matches autoregressive approaches.Source: §6 · p. 9
- Imputer conditions on a previous partial alignment to generate a new alignment, requiring a fixed number of iterations independent of input or output length.Source: §6 · p. 9
- Imputer can be trained to approximately marginalize over all possible alignments and generation orders.Source: §6 · p. 9
- A tractable dynamic program marginalizes over all completions given a previous alignment and provides a lower bound on likelihood.Source: §6 · p. 9