Parallel Iterative Edit Models for Local Sequence Transduction
Abhijeet Awasthi, Sunita Sarawagi, Rasna Goyal, Sabyasachi Ghosh, Vihari Piratla
Introduction
In local sequence transduction (LST) an input sequence needs to be mapped to an output sequence where the and sequences differ only in a few positions, is close to , and come from the same vocabulary . An important application of local sequence transduction that we focus on in this paper is Grammatical error correction (GEC). We contrast local transduction with more general sequence transduction tasks like translation and paraphrasing which might entail different input-output vocabulary and non-local alignments. The general sequence transduction task is cast as sequence to sequence (seq2seq) learning and modeled popularly using an attentional encoder-decoder (ED) model. The ED model auto-regressively produces each token in the output sequence conditioned on all previous tokens . Owing to the remarkable success of this model in challenging tasks like translation, almost all state-of-the-art neural models for GEC use it Zhao et al. (2019); Lichtarge et al. (2019); Ge et al. (2018b); Chollampatt and Ng (2018b); Junczys-Dowmunt et al. (2018).
We take a fresh look at local sequence transduction tasks and present a new parallel-iterative-edit (PIE) architecture. Unlike the prevalent ED model that is constrained to sequentially generating the tokens in the output, the PIE model generates the output in parallel, thereby substantially reducing the latency of sequential decoding on long inputs. However, matching the accuracy of existing ED models without the luxury of conditional generation is highly challenging. Recently, parallel models have also been explored in tasks like translation Stern et al. (2018); Lee et al. (2018); Gu et al. (2018); Kaiser et al. (2018) and speech synthesis van den Oord et al. (2018), but their accuracy is significantly lower than corresponding ED models. The PIE model incorporates the following four ideas to achieve comparable accuracy on tasks like GEC in spite of parallel decoding.
1. Output edits instead of tokens: First, instead of outputting tokens from a large vocabulary, we output edits such as copy, appends, deletes, replacements, and case-changes which generalize better across tokens and yield a much smaller vocabulary. Suppose in GEC we have an input sentence: fowler fed dog. Existing seq2seq learning approaches would need to output the four tokens Fowler, fed, the, dog from a word vocabulary whereas we would predict the edits {Capitalize token 1, Append(the) to token 2, Copy token 3}.
2. Sequence labeling instead of sequence generation: Second we perform in-place edits on the source tokens and formulate local sequence transduction as labeling the input tokens with edit commands, rather than solving the much harder whole sequence generation task involving a separate decoder and attention module. Since input and output lengths are different in general such formulation is non-trivial, particularly due to edits that insert words. We create special compounds edits that merge token inserts with preceding edits that yield higher accuracy than earlier methods of independently predicting inserts Ribeiro et al. (2018); Lee et al. (2018).
3. Iterative refinement: Third, we increase the inference capacity of the parallel model by iteratively inputting the model’s own output for further refinement. This handles dependencies implicitly, in a way reminiscent of Iterative Conditional Modes (ICM) fitting in graphical model inference Koller and Friedman (2009). Lichtarge et al. (2019) and Ge et al. (2018b) also refine iteratively but with ED models.
4. Factorize pre-trained bidirectional LMs: Finally, we adapt recent pre-trained bidirectional models like BERT Devlin et al. (2018) by factorizing the logit layer over edit commands and their token argument. Existing GEC systems typically rely on conventional forward directional LM to pretrain their decoder, whereas we show how to use a bi-directional LM in the encoder, and that too to predict edits.
Novel contributions of our work are as follows:
Recognizing GEC as a local sequence transduction (LST) problem, rather than machine translation. We then cast LST as a fast non-autoregressive, sequence labeling model as against existing auto-regressive encoder-decoder model.
Our method of reducing LST to non-autoregressive sequence labeling has many novel elements: outputting edit operations instead of tokens, append operations instead of insertions in the edit space, and replacements along with custom transformations.
We show how to effectively harness a pre-trained language model like BERT using our factorized logit architecture with edit-specific attention masks.
The parallel inference in PIE is 5 to 15 times faster than a competitive ED based GEC model like Lichtarge et al. (2019) which performs sequential decoding using beam-search. PIE also attains close to state of the art performance on standard GEC datasets. On two other local transduction tasks, viz., OCR and spell corrections the PIE model is fast and accurate w.r.t. other existing models developed specifically for local sequence transduction.
Our Method
Existing seq2seq ED models factorize to capture the full dependency between a and all previous as . An encoder converts input tokens to contextual states and a decoder summarizes to a state . An attention distribution over contextual states computed from determines the relevant input context and the output token distribution is calculated as . Decoding is done sequentially using beam-search. When a correct sequence corpus is available, the decoder is pre-trained on a next-token prediction loss and/or a trained LM is used to re-rank the beam-search outputs Zhao et al. (2019); Chollampatt and Ng (2018a); Junczys-Dowmunt et al. (2018).
1 Overview of the PIE model
We move from generating tokens in the output sequence using a separate decoder, to labelling the input sequence with edits . For this we need to design a function Seq2Edits that takes as input an pair in and outputs a sequence of edits from an edit space where is of the same length as in spite of and being of different lengths. In Section 2.2 we show how we design such a function.
We invoke Seq2Edits on and learn the parameters of a probabilistic model to assign a distribution over the edit labels on tokens of the input sequence. In Section 2.3 we describe the PIE architecture in more detail. The correct corpus , when available, is used to pre-train the encoder to predict an arbitrarily masked token in a sequence in , much like in BERT. Unlike in existing seq2seq systems where is used to pre-train the decoder that only captures forward dependencies, in our pre-training the predicted token is dependent on both forward and backward contexts. This is particularly useful for GEC-type tasks where future context can be approximated by .
Given an input , the trained model predicts the edit distribution for each input token independent of others, that is , and thus does not entail the latency of sequential token generation of the ED model. We output the most probable edits . Edits are designed so that we can easily get the edited sequence after applying on . is further refined by iteratively applying the model on the generated outputs until we get a sequence identical to one of the previous sequences upto a maximum number of iterations .
2 The Seq2Edits Function
Given a and where may not be equal to , our goal is to obtain a sequence of edit operations such that applying edit on the input token at each position reconstructs the output sequence . Invoking an off-the-shelf edit distance algorithm between and can give us a sequence of copy, delete, replace, and insert operations of arbitrary length. The main difficulty is converting the insert operations into in-place edits at each . Other parallel models Ribeiro et al. (2018); Lee et al. (2018) have used methods like predicting insertion slots in a pre-processing step, or predicting zero or more tokens in-between any two tokens in . We will see in Section 3.1.4 that these options do not perform well. Hence we design an alternative edit space that merges inserts with preceding edit operations creating compound append or replace operations. Further, we create a dictionary of common q-gram insertions or replacements observed in the training data. Our edit space () comprises of copy (c) , delete (d) , append (a) a q-gram after copying , replace (r) with a q-gram . For GEC, we additionally use transformations denoted as which perform word-inflection (e.g. arrive to arrival). The space of all edits is thus:
We present our algorithm for converting a sequence and into in-place edits on using the above edit space in Figure 1. Table 1 gives examples of converting (, ) pairs to edit sequences. We first invoke the Levenshtein distance algorithm (Levenshtein, 1966) to obtain diff between and with delete and insert cost as 1 as usual, but with a modified substitution cost to favor matching of related words. We detail this modified cost in the Appendix and show an example of how this modification leads to more sensible edits. The diff is post-processed to convert substitutions into deletes followed by inserts, and consecutive inserts are merged into a q-gram. We then create a dictionary of the most frequent q-gram inserts in the training set. Thereafter, we scan the diff left to right: a copy at makes = c, a delete at makes , an insert at and a flips the into a if is in , else it is dropped, an insert at and a flips the into a if a match found, else a replace if is in , else it is dropped.
The above algorithm does not guarantee that when is applied on we will recover for all sequences in the training data. This is because we limit to include only the most frequently inserted q-grams. For local sequence transduction tasks, we expect a long chain of consecutive inserts to be rare, hence our experiments were performed with . For example, in NUCLE dataset Ng et al. (2014) which has roughly 57.1K sentences, less than 3% sentences have three or more consecutive inserts.
3 The Parallel Edit Prediction Model
We next describe our model for predicting edits on an input sequence . We use a bidirectional encoder to provide a contextual encoding of each . This can either be multiple layers of bidirectional RNNs, CNNs or deep bidirectional transformers. We adopt the deep bidirectional transformer architecture since it encodes the input in parallel. We pre-train the model using much like in BERT pre-training recently proposed for language modelingDevlin et al. (2018). We first give an overview of BERT and then describe our model.
3.1 A Default use of BERT
Since we have cast our task as a sequence labeling task, a default output layer would be to compute as a softmax over the space of edits from each . If denotes the softmax parameter for edit , we get:
The softmax parameters, in Equation 2, have to be trained from scratch. We propose a method to exploit token embeddings of the pre-trained language model to warm-start the training of edits like appends and replaces which are associated with a token argument. Furthermore, for appends and replaces, we provide a new method of computing the hidden layer output via alternative input positional embeddings and self-attention. We do so without introducing any new parameters in the hidden layers of BERT.
3.2 An Edit-factorized BERT Architecture
The first term in the RHS of above equations captures edit specific score. The second term captures the score for copying the current word to the output. The third term models the influence of a new incoming token in the output obtained by a replace or append edit. For replace edits, score of the replaced word is subtracted from score of the incoming word. For transformation we add the copy score because they typically modify only the word forms, hence we do not expect meaning of the transformed word to change significantly.
The above equation provides insights on why predicting independent edits is easier than predicting independent tokens. Consider the append edit (). Instead of independently predicting at and at , we jointly predict the tokens in these two slots and contrast it with not inserting any new after in a single softmax. We will show empirically (Sec 3.1.4) that such selective joint prediction is key to obtaining high accuracy in spite of parallel decoding.
Finally, loss for a training example () is obtained by summing up the cross-entropy associated with predicting edit at each token .
Experiments
We compare our parallel iterative edit (PIE) model with state-of-the-art GEC models that are all based on attentional encoder-decoder architectures. In Section 3.2 we show that the PIE model is also effective on two other local sequence transduction tasks: spell and OCR corrections. Hyperparameters for all the experiments are provided in Table 13, 13 of the Appendix.
We use Lang-8 Mizumoto et al. (2011), NUCLE Ng et al. (2014) and FCE Yannakoudakis et al. (2011) corpora, which jointly comprise 1.2 million sentence pairs in English. The validation dataset comprises of 1381 sentence pairs from CoNLL-13 Ng et al. (2013) test set. We initialize our GEC model with the publicly available BERT-LARGEhttps://github.com/google-research/bert model that was pre-trained on Wikipedia (2500M words) and Book corpus Zhu et al. (2015) (800M words) to predict 15% randomly masked words using its deep bidirectional context. Next, we perform 2 epochs of training on a synthetically perturbed version of the One-Billion-word corpus Chelba et al. (2013). We refer to this as synthetic training. Details of how we create the synthetic corpus appear in Section A.3 of the appendix. Finally we fine-tune on the real GEC training corpus for 2 epochs. We use a batch size of 64 and learning rate 2e-5. The edit space consists of copy, delete, appends, replaces and transformations and their inverse. Arguments of Append and Replace operations mostly comprise punctuations, articles, pronouns, prepositions, conjunctions and verbs. Transformations perform inflections like add suffix s, d, es, ing, ed or replace suffix s to ing, d to s, etc. These transformations were chosen out of common replaces in the training data such that many replace edits map to only a few transformation edits, in order to help the model better generalize replaces across different words. In Section 3.1.4 we see that transformations increase the model’s recall. The complete list of transformations appears in Table 10 of the Appendix. We evaluate on score over span-level corrections from the MaxMatch (M2) scorer Dahlmeier and Ng (2012) on CONLL-2014-test. Like most existing GEC systems, we invoke a spell-checker on the test sentences before applying our model. We also report Napoles et al. (2016) scores on JFLEG corpus Napoles et al. (2017) to evaluate fluency.
Table 3 compares ensemble model results of PIE and other state of the art models, which all happen to be seq2seq ED models and also use ensemble decoding. For PIE, we simply average the probability distribution over edits from 5 independent ensembles. In Table-2 we compare non-ensemble numbers of PIE with the best available non-ensemble numbers of competing methods. On CoNLL-14 test-set our results are very close to the highest reported by Zhao et al. (2019). These results show that our parallel prediction model is competitive without incurring the overheads of beam-search and slow decoding of sequential models. score, that rewards fluency, is somewhat lower for our model on the JFLEG test set because of parallel predictions. We do not finetune our model on the JFLEG dev set. We expect these to improve with re-ranking using a LM. All subsequent ablations and timing measurements are reported for non-ensemble models.
1.2 Running Time Comparison
Parallel decoding enables PIE models to be considerably faster than ED models. In Figure 3 we compare wall-clock decoding time of PIE with 24 encoder layers (PIE-LARGE, ), PIE with 12 encoder layers (PIE-BASE, ) and competitive ED architecture by Lichtarge et al. (2019) with 6 encoder and 6 decoder layers (T2T, ) on CoNLL-14 test set. All decoding experiments were run and measured on a n1-standard-2https://cloud.google.com/compute/docs/machine-types#standard_machine_types VM instance with a single TPU shard (v-2.8). We observe that even PIE-LARGE is between a factor of 5 to 15 faster than an equivalent transformer-based ED model (T2T) with beam-size 4. The running time of PIE-LARGE increases sub-linearly with sentence length whereas the ED model’s decoding time increases linearly.
1.3 Impact of Iterative Refinement
We next evaluate the impact of iterative refinements on accuracy in Table 4. Out of 1312 sentences in the test set, only 832 sentences changed in the first round which were then fed to the second round where only 156 sentences changed, etc. The average number of refinement rounds per example was 2.7. In contrast, a sequential model on this dataset would require 23.2 steps corresponding to the average number of tokens in a sentence. The score increases from 57.9 to 59.5 at the end of the second iteration.
Table 5 presents some sentences corrected by PIE. We see that PIE makes multiple parallel edits in a round if needed. Also, we see how refinement over successive iterations captures output space dependency. For example, in the second sentence interact gets converted to interacted followed by insertion of have in the next round.
1.4 Ablation study on the PIE Architecture
In this section we perform ablation studies to understand the importance of individual features of the PIE model.
Synthetic Training We evaluate the impact of training on the artifically generated GEC corpus in row 2 of Table 6. We find that without it the score is 3.4 points lower.
Factorized Logits We evaluate the gains due to our edit-factorized BERT model (Section 2.3.2) over the default BERT model (Section 2.3.1). In Table 6 (row 3) we show that compared to the factorized model (row 2) we get a 1.2 point drop in score in absence of factorization.
Inserts as Appends on the preceding word was another important design choice. The alternative of predicting insert independently at each gap with a null token added to performs 2.7 points poorly (Table 6 row 4 vs row 2).
Transformation edits are significant as we observe a 6.3 drop in recall without them (row 5).
Impact of Language Model We evaluate the benefit of starting from BERT’s pre-trained LM by reporting accuracy from an un-initialized network (row 6). We observe a 20 points drop in establishing the importance of LMs in GEC.
Impact of Network Size We train the BERT-Base model with one-third fewer parameters than BERT-LARGE. From Table 6 (row 7 vs 1) we see once again that size matters in deep learning!
2 More Sequence Transduction Tasks
We demonstrate the effectiveness of PIE model on two additional local sequence transduction tasks recently used in Ribeiro et al. (2018).
We use the twitter spell correction dataset Aramaki (2010) which consists of 39172 pairs of original and corrected words obtained from twitter. We use the same train-dev-valid split as Ribeiro et al. (2018) (31172/4000/4000). We tokenize on characters, and our vocabulary and comprises the 26 lower cased letters of English.
We use the Finnish OCR data sethttps://github.com/mpsilfve/ocrpp by Silfverberg et al. (2016) comprising words extracted from Early Modern Finnish corpus of OCR processed newspaper text. We use the same train-dev-test splits as provided by Silfverberg et al. (2016). We tokenize on characters in the word. For a particular split, our vocabulary and comprises of all the characters seen in the training data of the split.
For all the tasks in this section, PIE is a 4 layer self-attention transformer with 200 hidden units, 400 intermediate units and 4 attention heads. No pre-initialization is done. Also, number of iterations of refinements is set to 1.
Table 7 presents whole-word 0/1 accuracy for these tasks on PIE and the following methods: Ribeiro et al. (2018)’s local transduction model (described in Section 4), and LSTM based ED models with hard monotonic attention Aharoni and Goldberg (2017) and soft-attention Bahdanau et al. (2015) as reported in Ribeiro et al. (2018). In addition, for a fair decoding time comparison, we also train a transformer-based ED model referred as Soft-T2T with 2 encoder, 2 decoder layers for spell-correction and 2 encoder, 1 decoder layer for OCR correction. We observe that PIE’s accuracy is comparable with ED models in both the tasks. Table 8 compares decoding speed of PIE with Soft-T2T in words/second. Since more than 90% of words have fewer than 9 tokens and the token vocabulary is small, decoding speed-ups of PIE over ED model on these tasks is modest compared to GEC.
Related Work
is an extensively researched area in NLP. See Ng et al. (2013) and Ng et al. (2014) for past shared tasks on GEC, and thishttps://nlpprogress.com/english/grammatical_error_correction.html website for current progress. Approaches attempted so far include rules Felice et al. (2014), classifiers Rozovskaya and Roth (2016), statistical machine translation (SMT) Junczys-Dowmunt and Grundkiewicz (2016), neural ED models Chollampatt and Ng (2018a); Junczys-Dowmunt et al. (2018); Ge et al. (2018a), and hybrids Grundkiewicz and Junczys-Dowmunt (2018). All recent neural approaches are sequential ED models that predict either word sequences Zhao et al. (2019); Lichtarge et al. (2019) or character sequences Xie et al. (2016) using either multi-layer RNNs Ji et al. (2017); Grundkiewicz and Junczys-Dowmunt (2018) or CNNsChollampatt and Ng (2018a); Ge et al. (2018a) or Transformers Junczys-Dowmunt et al. (2018); Lichtarge et al. (2019). Our sequence labeling formulation is similar to Yannakoudakis et al. (2017) and Kaili et al. (2018) but the former uses it to only detect errors and the latter only corrects five error-types using separate classifiers. Edits have been exploited in earlier GEC systems too but very unlike our method of re-architecting the core model to label input sequence with edits. Schmaltz et al. (2017) interleave edit tags in target tokens but use seq2seq learning to predict the output sequence. Chollampatt and Ng (2018a) use edits as features for rescoring seq2seq predictions. Junczys-Dowmunt et al. (2018) use an edit-weighted MLE objective to emphasise corrective edits during seq2seq learning. Stahlberg et al. (2019) use finite state transducers, whose state transitions denote possible edits, built from an unlabeled corpus to constrain the output of a neural beam decoder to a small GEC-feasible space.
Kaiser et al. (2018) achieve partial parallelism by first generating latent variables sequentially to model dependency. Stern et al. (2018) use a parallel generate-and-test method with modest speed-up. Gu et al. (2018) generate all tokens in parallel but initialize decoder states using latent fertility variables to determine number of replicas of an encoder state. We achieve the effect of fertility using delete and append edits. Lee et al. (2018) generate target sequences iteratively but require the target sequence length to be predicted at start. In contrast our in-place edit model allows target sequence length to change with appends.
is handled in Ribeiro et al. (2018) by first predicting insert slots in using learned insertion patterns and then using a sequence labeling task to output tokens in or a special token delete. Instead, we output edit operations including word transformations. Their pattern-based insert pre-slotting is unlikely to work for more challenging tasks like GEC. Koide et al. (2018) design a special edit-invariant neural network for being robust to small edit changes in input biological sequences. This is a different task than ours of edit prediction. Yin et al. (2019) is about neural representation of edits specifically for structured objects like source code. This is again a different problem than ours.
Conclusion
We presented a parallel iterative edit (PIE) model for local sequence transduction with a focus on the GEC task. Compared to the popular encoder-decoder models that perform sequential decoding, parallel decoding in the PIE model yields a factor of 5 to 15 reduction in decoding time. The PIE model employs a number of ideas to match the accuracy of sequential models in spite of parallel decoding: it predicts in-place edits using a carefully designed edit space, iteratively refines its own predictions, and effectively reuses state-of-the-art pre-trained bidrectional language models. In the future we plan to apply the PIE model to more ambitious transduction tasks like translation.
This research was partly sponsored by a Google India AI/ML Research Award and Google PhD Fellowship in Machine Learning. We gratefully acknowledge Google’s TFRC program for providing us Cloud-TPUs. We thank Varun Patil for helping us improve the speed of pre-processing and synthetic-data generation pipelines.
References
Appendix A Appendix
We keep delete and insert cost as as usual, but for substitutions, we use , where is the absolute difference between the number of characters of replaced and substituted word. We set to . Table 9 shows two minimum edit diffs if the substitution penalty has no such offset. In Diff-1, {.} substitutes {,} and Then substitutes then, followed by insertion of {,}. In Diff-2, {.} is inserted after sat followed by Then substituting {,} , followed by {,} substituting then . In absence of offset in substitution penalty, both the diffs have edit distance of . In presence of offset, Diff-1 has an edit-distance of , while Diff-2 has an edit-distance of , this allows Diff-1 to be preferred over Diff-2. As we observe, offset helps in selection of well aligned minimum edit diffs among multiple minimum edit diffs.
A.2 Suffix transformations
A.3 Artificial Error Generation
Figure 4 shows the algorithm used to introduce artificial errors in clean dataset. Given a sentence, first the number of errors in that sentence is determined by sampling from a multinoulli (over ). Similarly, an error is chosen independently from another multinoulli (over ). The distribution of the number of errors in a sentence and probability of each kind of error was obtained based on the available parallel corpus. For append, replace and delete errors, a position is randomly chosen for the error occurrence. For append error the word in that position is dropped. For delete error a spurious word from a commonly deleted words dictionary is added to that position. For replace error, both the actions are done. For a verb error, a verb is chosen at random from the sentence and is replaced by a random verb form of the same word. Commonly deleted words are also obtained from the parallel corpus.