Inference Strategies for Machine Translation with Conditional Masking (2020)

Abstract: Conditional masked language model (CMLM) training has proven successful for non-autoregressive and semi-autoregressive sequence generation tasks, such as machine translation. Given a trained CMLM, however, it is not clear what the best inference strategy is. We formulate masked inference as a factorization of conditional probabilities of partial sequences, show that this does not harm performance, and investigate a number of simple heuristics motivated by this perspective. We identify a thresholding strategy that has advantages over the standard “mask-predict” algorithm, and provide analyses of its behavior on machine translation tasks.

## introduction the widely successful masked language modeling paradigm popularized by bert @xcite has recently been adapted to conditional masked language model (cmlm) training for semiautoregressive sequence generation @xcite , where model predictions are conditioned on the complete input sequence and the observed (non-masked) portion of the output sequence. the cmlm's simplicity and its clear links to the very active field of linguistic representation learning are advantages over its semiautoregressive competitors, such as iterative refinement of token sequences @xcite , refinement of non-linguistic intermediate representations @xcite and learning to predict parallel edit operations @xcite . it is not obvious how to best perform inference with the cmlm. starting from a partially-observed output sequence, the optimal choice to complete it within a single step would be to generate the most likely token at each unobserved (masked) position independently. however, it is less clear how to progress from an initial, completely masked sequence to a final hypothesis semi-autoregressively over a number of steps, with each successive step unmasking new context for the next. this requires not only ordering the tokens for generation, but also making decisions about how many tokens to simultaneously predict in each step. @xcite propose the maskpredict algorithm, which iteratively generates fresh model predictions for all masked positions, and then unmasks a predefined number of the most likely predictions. given a fixed number of iterations, a decaying schedule determines how many predictions to unmask in each iteration. each successive iteration provides mode-breaking @xcite context for the next. by fixing the number of iterations, this approach allows for constant-time semi-autoregressive decoding. the fixed-iteration strategy is very practical and has yielded empirical success in a range of machine translation experiments, but there is no guarantee that it is optimal. the tokens to be unmasked on a given iteration are all predicted independently, and therefore might contain repeated words, or words with low model confidence. these issues can be mitigated by later re-masking a token to repair it @xcite or by adapting the model to incorrect contexts @xcite . we instead adopt a fully probabilistic view of the masked prediction sequence, which we enable by simply disallowing the re-masking of previously unmasked tokens. this view guides us to a heuristic inference schedule that selects sets of unmasked tokens according to a threshold on the product of their conditionally independent model probabilities. this heuristic naturally slows down in the situations mentioned above, and speeds up in the presence of high confidence, which allows us to achieve favorable quality-to-speed trade-offs. we focus on strengthening the cmlm inference (section 3) while leaving its training algorithm unchanged (section 2), and maintaining much of the structure of the original inference strategy. for our experiments on machine translation (section 4), we compare inference heuristics in terms of their quality-speed trade-offs. we analyze the development of quality over iterations, and the influence of sentence length. with examples of unmasking schedules we furthermore illustrate the role of mode breaking through choosing the right contexts. ## cmlm model and training the cmlm is a model for p(y mask |y obs , x), the probability of masked tokens y mask given a partially observed output sequence y obs and an input sequence x. y mask and y obs are sets of tokens at specified positions that together form a complete output sequence y : y @xmath0. the model is implicitly conditioned on output sequence length @xmath1, and the tokens in y mask are conditionally independent: during training, masks are placed randomly: first, the mask size @xmath2, . . . , n } is sampled from a uniform distribution, then s positions are randomly chosen to define the subsets y obs and y mask . cross-entropy loss is incurred via p(y i |y obs , x) for each y @xmath3. an additional classifier on top of encoder representations is trained to predict the output length n . ## cmlm inference inference starts with a context of only mask tokens. until a stop condition is met, decoder predictions iteratively replace a subset of these in selected positions ("unmasking"). with a single iteration, inference is non-autoregressive; when the number of iterations t is less than the sentence length n it is semi-autoregressive; and when @xmath4. due to the use of a uniform distribution over reference contexts, training is agnostic to these different regimes. in general, we seek to minimize t without trading off too much quality. the challenge in doing so is to identify the subset of predictions that are most likely to provide suitable conditioning context for future iterations @xcite . structural or linguistic dependencies in the output may also play an important role for resolving linguistic ambiguities @xcite . ## experiments our cmlm is implemented with a base transformer @xcite built on a tensor-flow implementation of @xcite . the input to the decoder is y obs , with mask tokens at masked positions, and the output is y mask , predictions for all masked positions without future attention masking. we use data from wmt14 en↔de @xcite and wmt17 zh↔en @xcite with a sentence piece vocabulary of 32k, focusing mainly on @xmath5, and providing results for all pairs in appendix a. the cmlm is trained on distilled training data from an autoregressive transformer and initialized with its parameters. figure 2 shows the performance of the update strategies described in section 3.1 versus length beam b. all strategies use the mask-predict heuristic with a fixed 10-iteration limit. as beam size increases past 2, the update-masked strategies increasingly dominate, indicating that their scores are more reliable for choosing among length hypotheses. there is no significant difference between the two variants of update-masked. this suggests that our probabilistic factorization constraint (updatemasked-sub) does not hurt in practice. ## conclusion we investigated inference strategies for machine translation based on cmlm with a focus on the trade-off between generation speed and quality. we introduce a perspective which views generation sequences as probabilistic factorizations of the final output sequence, and use it to analyze and extend previous heuristics. our new heuristics achieve better speed/quality balance by flexibly adjusting the number of total iterations, and by taking the probabilities of sets of tokens into account. for future work we would like to explore if their success transfers to other generation tasks with mlms where inference efficiency is a concern.