Levenshtein Transformer

Introduction

Neural sequence generation models are widely developed and deployed in tasks such as machine translation (Bahdanau et al., 2015; Vaswani et al., 2017). As we examine the current frameworks, the most popular autoregressive models generate tokens step-by-step. If not better, recent non-autoregressive approaches (Gu et al., 2018; Kaiser et al., 2018; Lee et al., 2018) have proved it possible to perform generation within a much smaller number of decoding iterations.

In this paper, we propose Levenshtein Transformer (LevT), aiming to address the lack of flexibility of the current decoding models. Notably, in the existing frameworks, the length of generated sequences is either fixed or monotonically increased as the decoding proceeds. This remains incompatible with human-level intelligence where humans can revise, replace, revoke or delete any part of their generated text. Hence, LevT is proposed to bridge this gap by breaking the in-so-far standardized decoding mechanism and replacing it with two basic operations — insertion and deletion.

We train the LevT using imitation learning. The resulted model contains two policies and they are executed in an alternate manner. Empirically, we show that LevT achieves comparable or better results than a standard Transformer model on machine translation and summarization, while maintaining the efficiency advantages benefited from parallel decoding similarly to (Lee et al., 2018). With this model, we argue that the decoding becomes more flexible. For example, when the decoder is given an empty token, it falls back to a normal sequence generation model. On the other hand, the decoder acts as a refinement model when the initial state is a low-quality generated sequence. Indeed, we show that a LevT trained from machine translation is directly applicable to translation post-editing without any change. This would not be possible with any framework in the literature because generation and refinement are treated as two different tasks due to the model’s inductive bias.

One crucial component in LevT framework is the learning algorithm. We leverage the characteristics of insertion and deletion — they are complementary but also adversarial. The algorithm we propose is called “dual policy learning”. The idea is that when training one policy (insertion or deletion), we use the output from its adversary at the previous iteration as input. An expert policy, on the other hand, is drawn to provide a correction signal. Despite that, in theory, this learning algorithm is applicable to other imitation learning scenarios where a dual adversarial policy exists, in this work we primarily focus on a proof-of-concept of this algorithm landing at training the proposed LevT model.

To this end, we summarize the contributions as follows:

We propose Levenshtein Transformer (LevT), a new sequence generation model composed of the insertion and deletion operations. This model achieves comparable or even better results than a strong Transformer baseline in both machine translation and text summarization, but with much better efficiency (up to $\times 5$ speed-up in terms of actual machine execution time);

We propose a corresponding learning algorithm under the theoretical framework of imitation learning, tackling the complementary and adversarial nature of the dual policies;

We recognize our model as a pioneer attempt to unify sequence generation and refinement, thanks to its built-in flexibility. With this unification, we empirically validate the feasibility of applying a LevT model trained by machine translation directly to translation post-editing, without any change.

Problem Formulation

We unify the general problems of sequence generation and refinement by casting them to a Markov Decision Process (MDP) defined by a tuple $\left(\mathcal{Y},\mathcal{A},\mathcal{E},\mathcal{R},\bm{y_{0}}\right)$. We consider the setup consisting an agent interacting with an environment $\mathcal{E}$ which receives the agent’s editing actions and returns the modified sequence. We define $\mathcal{Y}=\mathcal{V}^{N_{\max}}$ as a set of discrete sequences up to length $N_{\max}$ where $\mathcal{V}$ is a vocabulary of symbols. At every decoding iteration, the agent receives an input $\bm{y}$ drawn from scratch or uncompleted generation, chooses an action $\bm{a}$ and gets a reward $r$. We use $\mathcal{A}$ to denote the set of actions and $\mathcal{R}$ for the reward function. Generally the reward function $\mathcal{R}$ measures the distance between the generation and the ground-truth sequence, $\mathcal{R}(\bm{y})=-\mathcal{D}(\bm{y},\bm{y}^{*})$ which can be any distance measurement such as the Levenshtein distance (Levenshtein, 1965). It is crucial to incorporate $\bm{y_{0}}\in\mathcal{Y}$ into the our formulation. As the initial sequence, the agent receives—when $\bm{y_{0}}$ is an already generated sequence from another system, the agent essentially learns to do refinement while it falls back to generation if $\bm{y_{0}}$ is an empty sequence. The agent is modeled by a policy, $\pi$, that maps the current generation over a probability distribution over $\mathcal{A}$. That is, $\pi:\mathcal{Y}\to P(\mathcal{A})$.

2 Actions: Deletion & Insertion

The deletion policy reads the input sequence $\bm{y}$, and for every token $y_{i}\in\bm{y}$, the deletion policy $\pi^{\textrm{del}}(d|i,\bm{y})$ makes a binary decision which is 1 (delete this token) or 0 (keep it). We additionally constrain $\pi^{\textrm{del}}(0|1,\bm{y})=\pi^{\textrm{del}}(0|n,\bm{y})=1$ to avoid sequence boundary being broken. The deletion classifier can also be seen as a fine-grained discriminator used in GAN (Goodfellow et al., 2014) where we predict “fake” or “real” labels for every predicted token.

Insertion

In this work, it is slightly more complex to build the insertion atomic because it involves two phases: placeholder prediction and token prediction so that it is able to insert multiple tokens at the same slot. First, among all the possible inserted slots ($y_{i},y_{i+1}$) in $\bm{y}$, $\pi^{\textrm{plh}}(p|i,\bm{y})$ predicts the possibility of adding one or several placeholders. In what follows, for every placeholder predicted as above, a token prediction policy $\pi^{\textrm{tok}}(t|i,\bm{y})$ replaces the placeholders with actual tokens in the vocabulary. The two-stage insertion process can also be viewed as a hybrid of Insertion Transformer (Stern et al., 2019) and masked language model (MLM, Devlin et al., 2018; Ghazvininejad et al., 2019).

Policy combination

Recall that our two operations are complementary. Hence we combine them in an alternate fashion. For example in sequence generation from the empty, insertion policy is first called and it is followed by deletion, and then repeat till the certain stopping condition is fulfilled. Indeed, it is possible to leverage the parallelism in this combination. We essentially decompose one iteration of our sequence generator into three phases: “delete tokens – insert placeholders – replace placeholders with new tokens”. Within each stage, all operations are performed in parallel. More precisely, given the current sequence $\bm{y}=(y_{0},\ldots,y_{n})$, and suppose the action to predict is $\bm{a}=\{\underbrace{d_{0},\ldots d_{n}}_{\bm{d}};\underbrace{p_{0},\ldots,p_{n-1}}_{\bm{p}};\underbrace{t_{0}^{1},\ldots t_{0}^{p_{0}},\ldots,t_{n-1}^{p_{n-1}}}_{\bm{t}}\}$, the policy for one iteration is:

where $\bm{y}^{\prime}=\mathcal{E}(\bm{y},\bm{d})$ and $\bm{y}^{\prime\prime}=\mathcal{E}(\bm{y}^{\prime},\bm{p})$. We parallelize the computation within each sub-tasks.

Levenshtein Transformer

In this section, we cover the specs of Levenshtein Transformer and the dual-policy learning algorithm. Overall our model takes a sequence of tokens (or none) as the input then iteratively modify it by alternating between insertion and deletion, until the two policies combined converge. We describe the detailed learning and inference algorithms in the Appendix.

We use Transformer (Vaswani et al., 2017) as the basic building block. For conditional generation, the source $\bm{x}$ is included in each TransformerBlock. The states from the $l$-th block are:

The decoder outputs ($\bm{h}_{0},\bm{h}_{2},...,\bm{h}_{n}$) are passed to three policy classifiers:

Deletion Classifier: LevT scans over the input tokens (except for the boundaries) and predict “deleted” () or “kept” ($1$) for each token position,

Placeholder Classifier: LevT predicts the number of tokens to be inserted at every consecutive position pairs, by casting the representation to a categorical distribution:

Token Classifier: following the placeholder prediction, LevT needs to fill in tokens replacing all the placeholders. This is achieved by training a token predictor as follow:

Weight Sharing

Our default implementation always assumes the three operations to share the same Transformer backbone to benefit features learned from other operations. However, it is also possible to disable weight sharing and train separate decoders for each operations, which increases the capacity of the model while does not affect the overall inference time.

Early Exit

Although it is parameter-efficient to share the same Transformer architecture across the above three heads, there is room for improvement as one decoding iteration requires three full passes of the network. To make trade-off between performance and computational cost, we propose to perform early exit (attaching the classifier to an intermediate block instead of the last one) for $\pi^{\textrm{del}}$ and $\pi^{\textrm{plh}}$ to reduce computation while keeping $\pi^{\textrm{tok}}$ always based on the last block, considering that token prediction is usually more challenging than the other two tasks.

2 Dual-policy Learning

We use imitation learning to train the Levenshtein Transformer. Essentially we let the agent imitate the behaviors that we draw from some expert policy $\pi^{*}$. The expert policy is derived from direct usage of ground-truth targets or less noisy version filtered by sequence distillation (Kim and Rush, 2016). The objective is to maximize the following expectation:

By definition, the roll-in policy determines the state distribution fed to $\pi_{\theta}$ during training. In this work, we have two strategies to construct the roll-in policy — adding noise to the ground-truth or using the output from the adversary policy. Figure 2 shows a diagram of this learning paradigm. We formally write down the roll-in policies as follows.

Learning to Insert: similar to the deletion step, we apply a mixture of the deletion output and a random word dropping sequence of the round-truth, inspired by recent advances of training masked language model (Devlin et al., 2018). We use random dropping as a form of noise injection to encourage more exploration. Let $\beta\in$ and $u\sim\text{Uniform}$,

Expert Policy

It is crucial to construct an expert policy in imitation learning which cannot be too hard or too weak to learn from. Specifically, we considered two types of experts:

Oracle: One way is to build an oracle which accesses to the ground-truth sequence. It returns the optimal actions $\bm{a}^{*}$ (either oracle insertion $\bm{p}^{*},\bm{t}^{*}$ or oracle deletion $\bm{d}^{*}$) by:

Here, we use the Levenshtein distance (Levenshtein, 1965)We only consider the variant which only computes insertion and deletion. No substitution is considered. as $\mathcal{D}$ considering it is possible to obtain the action suggestions efficiently by dynamic programming.

Distillation: We also explore to use another teacher model to provide expert policy, which is known as sequence-level knowledge distillation (Kim and Rush, 2016). This technique has been widely used in previous approaches for nonauoregressive generation (Gu et al., 2018). More precisely, we first train an autoregressive teacher model using the same datasets and then replace the ground-truth sequence $\bm{y^{*}}$ by the beam-search result of this teacher-model, $\bm{y^{\textsc{ar}}}$. We use the same mechanism to find the suggested option as using the ground-truth oracle.

3 Inference

At inference time, we apply the trained model over the initial sequence $\bm{y}^{0}$ for several iterations. We greedily pick up the actions associated with high probabilities in Eq. (3)(4)(5). Moreover, we find that using search (instead of greedy decoding) or nosiy parallel decoding (Cho, 2016) does not yield much gain in LevT. This observation is quite opposite to what has been widely discovered in autoregressive decoding. We hypothesize there may be two reasons leading to this issue: (i) The local optimal point brought by greedy decoding in autoregressive models is often far from the optimal point globally. Search techniques resolve this issue with tabularization. In our case, however, because LevT inserts or deletes tokens dynamically, it could easily revoke the tokens that are found sub-optimal and re-insert better ones; (ii) the log-probability of LevT is not a good metric to select the best output. However, we do believe to see more improvements if we include an external re-ranker, e.g. an autoregressive teacher model. We leave this discussion in the future work.

Termination Condition

The decoding stops when one of the following two conditions is fulfilled:

Looping: Generation is terminated if two consecutive refinement iterations return the same output which can be (i) there are no words to delete or insert; (ii) the agent gets stuck in an infinite loop: i.e. the insertion and deletion counter each other and keep looping.

Timeout: We further set a maximum number of iterations (timeout) to guarantee a constant-time complexity in the worst case (Lee et al., 2018; Ghazvininejad et al., 2019).

Penalty for Empty Placeholders

Similar to Stern et al. (2019), we add a penalty to insert “empty” placeholder in decoding. Overly inserting “empty” placeholders may result in shorter output. A penalty term $\gamma\in$ is subtracted from the logits of in Eq. (4).

Experiments

We validate the efficiency, effectiveness, and flexibility of Levenshtein Transformer extensively across three different tasks — machine translation (MT), text summarization (TS) and automatic post-editing (APE) for machine translation, from both generation (§4.1) and refinement (§4.2) perspectives.

For the sequence generation perspective, we evaluate LevT model on MT and TS. As a special case, sequence generation assumes empty $\bm{y}^{0}=\textsc{<s></s>}$ as input and no initial deletion is applied.

We use three diversified language pairs for MT experiments: WMT’16 Romanian-English (Ro-En)http://www.statmt.org/wmt16/translation-task.html, WMT’14 English-German (En-De)http://www.statmt.org/wmt14/translation-task.html and WAT2017 Small-NMT English-Japanese (En-Ja, Nakazawa et al., 2017)http://lotus.kuee.kyoto-u.ac.jp/WAT/WAT2017/snmt/index.html. The TS experiments use preprocessed data from the Annotated English Gigaword (Gigaword, Rush et al., 2015)https://github.com/harvardnlp/sent-summary. We learn byte-pair encoding (BPE, Sennrich et al., 2016) vocabulary on tokenized data. Detailed dataset statistics can be found in the Appendix. For evaluation metrics, we use BLEU (Papineni et al., 2002) for MT and ROUGE-1,2,L (Lin, 2004) for TS. Before computing the BLEU scores for Japanese output, we always segment Japanese words using KyTea http://www.phontron.com/kytea/.

Models & Training

We adopt the model architecture of Transformer base (Vaswani et al., 2017) for the proposed LevT model and the autoregressive baseline. All the Transformer-based models are trained on $8$ Nvidia Volta GPUs with maximum $300K$ steps and a total batch-size of around $65,536$ tokens per step (We leave more details to the Appendix).

Overall results

We present our main results on the generation quality and decoding speed in Table 1. We measure the speed by the averaged generation latency of generating one sequence at a time on single Nvidia V100 GPU. To remove the implementation bias, we also present the number of decoder iterations as a reference. It can be concluded that for both MT and summarization tasks, our proposed LevT achieves comparable and sometimes better generation quality compared to the strong autoregressive baseline, while LevT is much more efficient at decoding. A translation example is shown in Figure 3 and we leave more in Appendix. We conjecture that this is due to that the output of the teacher model possesses fewer modes and much less noisy than the real data. Consequently, LevT needs less number of iterations to converge to this expert policy.

Ablation on Efficiency

As shown in Figure 4(a), we plot the average number of iterations over the length of input over a monolingual corpus. LevT learns to properly adjust the decoding time accordingly. We also explore the variants of “early exit” where we denote LevT($m$-$n$) as a model with $m$ and $n$ blocks for deletion (Eq. (3)) and placeholder prediction (Eq. (4)) respectively. Figure 4(b) shows that although it compromises the quality a bit, our model with early exit achieves up to $\times 5$ speed-up (execution time) comparing against a strong autoregressive Transformer using beam-search.

Ablation on Weight Sharing

We also evaluate LevT with different weight sharing as noted in §3.1. The results of models trained with oracle or distillation are listed in Table 2(a). We observe that weight-sharing is beneficial especially between the two insertion operations (placeholder and token classifiers). Also, it shows another $+0.5$ BLEU improvement by not sharing the deletion operation with insertion compared to the default setting, which may indicate that insertion and deletion capture complementary information, requiring larger capacity by learning them separately.

Importance of mixture roll-in policy

We perform an ablation study on the learning algorithm. Specifically, we train a model with no mixing of the $\pi_{\theta}$ in Equation (6). We name this experiment by DAE due to its resemblance to a denoising autoencoder. We follow closely a standard pipeline established by Lee et al. (2018). Table 2(b) shows this comparison. As we can see that the deletion loss from DAE is much smaller while the generation BLEU score is inferior. We conjecture that this is caused by the mismatch between the states from the model and the roll-in policy in training the DAE.

v.s. Exiting Refinement-based Models

Table 2(a) also includes results from two relevant recent works which also incorporate iterative refinement in non-autoregressive sequence generation. For fair comparison, we use the result with length beam $1$ from Ghazvininejad et al. (2019). Although both approaches use similar “denosing” objectives to train the refinement process, our model explicitly learns “insertion” and “deletion” in a dual-policy learning fashion, and outperforms both models.

2 Sequence Refinement

We evaluate LevT’s capability of refining sequence outputs on the APE task. In this setting, inputs are pairs of the source sequence and a black-box MT system generation. The ground-truth outputs are from real human edits with expansion using synthetic data.

We follow a normal protocol in the synthetic APE experiments (Grangier and Auli, 2017): we first train the input MT system on half of the dataset. Then we will train a refinement model on the other half based on the output produced by the MT model trained in the previous phase. For the real APE tasks, we use the data from WMT17 Automatic Post-Editing Shared Taskhttp://www.statmt.org/wmt17/ape-task.html on En-De. It contains both real PE triples and a large-scale synthetic corpus.

Models & Evaluation

The baseline model is a standard Transformer encoding the concatenation of the source and the MT system’s output. For the MT system here, we want some imperfect systems that need to be refined. We consider a statistical phrase-based MT system (PBMT, Koehn et al., 2003) and an RNN-based NMT system (Bahdanau et al., 2015). Apart from BLEU scores, we additionally apply translation error rate (TER, Snover et al., 2006) as it is widely used in the APE literature.

Overall results

We show the major comparison in Table 3. When training from scratch, LevT consistently improves the performance of the input MT system (either PBMT or NMT). It also achieves better performance than the autoregressive Transformer in most of the cases.

Pre-training on MT

Thanks to the generality of the LevT model, we show it is feasible to directly apply the LevT model trained by generation onto refinement tasks — in this case — MT and APE. We name this a “zero-shot post-editing” setting. According to Table 3, the pre-trained MT models are always capable of improving the initial MT input in the synthetic tasks.

The real APE task, however, differs quite a bit from the synthetic tasks because human translators normally only fix a few spotted errors. This ends up with very high BLEU scores even for the “Do-nothing” column. However, the pre-trained MT model achieves the best results by fine-tuning on the PE data indicating that LevT is able to leverage the knowledge for generation and refinement.

Collaborate with Oracle

Thanks to the saperation of insertion and deletion operations, LevT has better interpretability and controllability. For example, we test the ability that LevT adapts oracle (e.g. human translators) instructions. As shown in Figure 5, both MT and PE tasks have huge improvement if every step the oracle deletion is given. This goes even further if the oracle provides both the correct deletion and the number of placehoders to insert. It also sheds some light upon computer-assisted text editing for human translators.

Related Work

Breaking the autoregressive constraints and monotonic (left-to-right) decoding order in classic neural sequence generation systems has recently attracted much interest. Stern et al. (2018); Wang et al. (2018) designed partially parallel decoding schemes to output multiple tokens at each step. Gu et al. (2018) proposed a non-autoregressive framework using discrete latent variables, which was later adopted in Lee et al. (2018) as iterative refinement process. Ghazvininejad et al. (2019) introduced the masked language modeling objective from BERT (Devlin et al., 2018) to non-autoregressively predict and refine translations. Welleck et al. (2019); Stern et al. (2019); Gu et al. (2019) generate translations non-monotonically by adding words to the left or right of previous ones or by inserting words in arbitrary order to form a sequence.

Editing-Based Models

Several prior works have explored incorporating “editing” operations for sequence generation tasks. For instance, Novak et al. (2016) predict and apply token substitutions iteratively on phase-based MT system outputs using convolutional neural network. QuickEdit (Grangier and Auli, 2017) and deliberation network (Xia et al., 2017) both consist of two autoregressive decoders where the second decoder refines the translation generated by the first decoder. Guu et al. (2018) propose a neural editor which learned language modeling by first retrieving a prototype and then editing over that. Freitag et al. (2019) correct patterned errors in MT system outputs using transformer models trained on monolingual data. Additionally, the use of Levenshtein distance with dynamic programming as the oracle policy were also proposed in Sabour et al. (2018); Dong et al. (2019). Different from these work, the proposed model learns a non-autoregressive model which simultaneously inserts and deletes multiple tokens iteratively.

Conclusion

We propose Levenshtein Transformer, a neural sequence generation model based on insertion and deletion. The resulted model achieves performance and decoding efficiency, and embraces sequence generation to refinement in one model. The insertion and deletion operations are arguably more similar to how human writes or edits text. For future work, it is potential to extend this model to human-in-the-loop generation.

Acknowledgement

We would like to thank Kyunghyun Cho, Marc’Aurelio Ranzato, Douwe Kiela, Qi Liu and our colleagues at Facebook AI Research for valuable feedback, discussions and technical assistance.

References

Appendix A Learning & Inference Algorithm

We present the detailed algorithms for learning and decoding from Levenshtein Transformer as follows. For simplicity, we always omit the source information $\bm{x}$ in conditional sequence generation tasks such as machine translation which is handled by the cross-attention with an encoder on $\bm{x}$.

The learning algorithm is shown in Algorithm 1. $\mathcal{E}$ is the environment and $\mathcal{D}$ is denoted as the Levenshtein distance, and we can easily back-track the optimal insertion and deletion operations through dynamic programming. We only show the the case with single batch-size for convenience. We also present the inference algorithm in Algorithm 2. If the initial sequence $\bm{y}^{0}$ is empty (), the proposed model will skip the first deletion and do sequence generation. Otherwise, the model starts with deletion operations and refine the input sequence.

Appendix B Dataset and Preprocessing Details

Table 4 and 5 list the statistics ($\#$ of sentences, vocabulary) for all the datasets used in this work. We learn BPE vocabulary with $32,000$ joint operations for WMT En-De and Gigaword and $40,000$ joint operations for WMT Ro-En. For WAT En-Ja, we adopt the official $16,384$ BPE vocabularies learned separately on source and target side.

Appendix C Model and Training Details

Transformer models are used for autoregressive baselines as well as teacher models (for the expert policy). By default, we set $d_{\text{model}}=512$, $d_{\text{hidden}}=2048$, $n_{\text{heads}}=8$, $n_{\text{layers}}=6$, $\text{lr}_{\max}=0.0005$, $\text{label-smooth}=0.1$, $\text{warmup}=10000$ and $\text{dropout}=0.3$. Source and target side share embeddings in all the training pairs except for WAT En-Ja where BPE vocabularies of both side are learned separately and are almost non-overlapping.

Since the training objectives for Levenshtein Transformer contains randomness terms (Eq. (6) (7)), we instead use BLEU (for MT) or ROUGE-2 (for TS) to select the best checkpoint by validation scores. We do not average checkpoints in this work.

C.2 Sequence Refinement Tasks

For synthetic APE tasks, we keep the same training conditions for LevT as those for MT tasks (§C.1). As described earlier in §4.2, we build the baseline Transformer by concatenating the source and MT system’s output as the input sequence for the encoder. Specially, we restart the positional embeddings for the MT output, add an additional language embedding for each token of the input sequence to show its language type. The detailed hyperpameters are the same as the standard Transformer.

As described in §4.2, we consider the following two different imperfect MT systems to provide the refinement inputs. Firstly, we consider the traditional statistical phrase-based machine translation system (PBMT). We follow the instruction to build the basic baseline model via moseshttp://www.statmt.org/moses/?n=Moses.Baseline. As for the NMT-based model, we use a single layer attention-based model composed by LSTM. We build this model on fairseq-pyhttps://github.com/pytorch/fairseq/blob/master/fairseq/models/lstm.py with the default configuration.

For the real APE task, we follow the procedures introduced in Junczys-Dowmunt and Grundkiewicz (2016). Synthetic corpus has two subsets: a 500K one and a 4M one. We over-sample real data by 10 times and merge it with the 500K synthetic data to train APE models. Besides, we also train a LevT MT model on the bigger (4M) synthetic corpus where we only use the source and target pairs.

C.3 Implementation

Both the proposed Levenshtein Transformer and the baseline Transformer are implemented using PyTorchhttps://pytorch.org/. The codes are released as part of the Fairseq-py https://github.com/pytorch/fairseq/tree/master/examples/nonautoregressive_translation.

C.4 Maximum Number of Iterations

We also presented in general how many sentences will be generated using the maximum iteration (for instance $10$). As shown in Table 6, surprisingly, most predictions are gotten in 1-4 iterations, and the average number of iterations is 2.43. Only a tiny portion ($\sim 0.1\%$) require the maximum number of iterations demonstrating the efficiency of the proposed approach.

Appendix D More Decoding Examples

We present more examples from the proposed Levenshtein Transformer as follows.