E di T 5: Semi-Autoregressive Text Editing with T5 Warm-Start (2022)

Abstract: We present EdiT5 - a novel semi-autoregressive text-editing approach designed to combine the strengths of non-autoregressive text-editing and autoregressive decoding. EdiT5 is faster at inference times than conventional sequence-to-sequence (seq2seq) models, while being capable of modeling flexible input-output transformations.This is achieved by decomposing the generation process into three sub-tasks: (1) tagging to decide on the subset of input tokens to be preserved in the output, (2) re-ordering to define their order in the output text, and (3) insertion to infill the missing tokens that are not present in the input. The tagging and re-ordering steps, which are responsible for generating the largest portion of the output, are non-autoregressive, while the insertion uses an autoregressive decoder.Depending on the task, EdiT5 requires significantly fewer autoregressive steps demonstrating speedups of up to 25x when compared to classic seq2seq models. Quality-wise, EdiT5 is initialized with a pre-trained T5 checkpoint yielding comparable performance to T5 in high-resource settings and clearly outperforms it on low-resource settings when evaluated on three NLG tasks: Sentence Fusion, Grammatical Error Correction, and Decontextualization.

## introduction pre-trained seq2seq models such as t5 @xcite , bart @xcite , and mass @xcite have established strong baselines for the majority of text-to-text transduction tasks. a recent trend to massively scale up model sizes, e.g., all the way up to 540b params @xcite , as well as the sizes of pretraining corpora, has further pushed the state-of-the-art without signs of reaching a plateau. from a practical point of view, running inference with such models is prohibitively expensive for most applications, which motivates the work on finding efficient recipes for model distillation, e.g., @xcite and choosing a model architecture that can provide a better trade-off between performance on a given task and inference speed. a typical choice is to distill a large language model into a smaller seq2seq model, e.g., transformer @xcite . in this paper we propose a novel model architecture edit5which blends ideas from a seq2seq t5 @xcite ) and text-editing to provide faster inference without sacrificing on task performance. seq2seq-based models output text token-by-token from scratch, allowing them to model any kind of input-output relationship. however, for many real-world tasks this degree of generality is unnecessary, especially for monolingual tasks where the input and output texts have relatively high degrees of overlap. in such cases a natural approach is to cast conditional text generation as a text-editing task, where the model learns to construct target texts by applying a set of edit operations to the inputs @xcite . typically the set of edit operations is defined ahead of time @xcite @xcite , which on the one hand limits the flexibility of the model to reconstruct arbitrary output texts from the inputs, but on the other, leads to latency improvements as the limited set of allowed operations significantly reduces the output vocabulary of the decoder. in this paper, we propose an approach which is both fast at inference time and flexible, able to model arbitrary rewrites. ## model description the model architecture of edit5 resembles a vanilla transformer @xcite composed of an encoder and a decoder. edit5 decomposes the generation of a text y from an input x into three parts: predicting a sequence of edit tags y t (indicating whether a token from the input should be copied to the output), a permutation of the input tokens π (indicating the order that copied tokens should appear in in the output), and a sequence of tokens y d (indicating additional tokens that should be in the output, and where in the permuted input they should be inserted). y t and π are modeled by the encoder, and y d by the decoder. there are multiple ways to choose the triple (y t , π, y d ) for a given (x, y) pair. during dataset creation we choose a single such triple for each training pair (see section 2.1 for details), in which case the probability of y can be expressed as: during inference, we first greedily set y t to maximize the third term, then π to maximize the second term, and finally y d to maximize the first term. the output text y is realized by applying the tags y t and permutation π to the input sequence x and then inserting the tokens y d . ## experiments we evaluate edit5 on three distinct text-editing tasks: sentence fusion, grammatical error correction, and decontextualization. in addition to reporting previously published results for each task, we also compare to felix @xcite , a recent non-autoregressive text-editing model, and a strong pre-trained t5 baseline implemented in the t5x framework @xcite . modeling. for edit5 we initialize with a t5 base model with a 12-layer transformer encoder, and single-layer transformer decoder. our code is based on the tensorflow model garden's (hongkun yu and li, 2020) tf2 version of t5. after initializing with the t5 checkpoint, we further pre-train on the denoising objective (see section 2.1) using the c4 corpus @xcite , training for 100k steps. for all experiments edit5 is trained using adamw @xcite , additionally the learning rate was decayed using the validation set, and exact match is used for checkpoint selection. tokenization is based on t5's senten-cepiece vocabulary @xcite , with a vocabulary size of 32k. we, however, modify the vocabulary, removing tokens which have punctuation as a suffix, and replacing them with additional span insertion special token, giving edit5 512 span insertion special token. unless otherwise stated, we use an input sequence length of 128. we performed minimal hyper-parameter selection, which is discussed in the appendix. ## latency analysis the tasks on which edit5 outperforms seq2seq models in latency are those that have overlap be-tween sources and targets, but it's unclear how much overlap is required for edit5 to produce latency savings. to answer this question, we split edit5 base, t5 base and t5 slim base into components whose latencies we measure separately and compare. details on how latencies are measured can be found in the appendix c. a seq2seq model decomposes into two parts: the encoder (we include the input embedding here, so we refer to this as encoder* below), and the decoder. edit5 has both of these parts, but also includes a third part (which we call its overhead), comprising of pointer realization and additional transformer layers. to make our analysis simpler and more task-agnostic, we make two simplifying assumptions. first, we assume the worst-case that no tokens are deleted by edit5 and there are no padding tokens in the input foot_1 , in practice this is not the case, and provides significant latency savings for edit5. second, we assume that decoder latency is linear in the number of decoder steps foot_2 . both of these assumptions benefit the latency of seq2seq models more than edit5. results. in from these results we can derive a simple rule for when edit5 will provide a net latency benefit. compared to t5 slim base foot_3 , edit5 base must save on average 4 decoder steps with an input length of 128, and 7 steps with an input length of 512. finally, collating the results in to reduce latency, a number of nonautoregressive (nat) seq2seq methods have been proposed for neural machine translation @xcite @xcite but a quality gap compared to autoregressive methods still exists. to decrease the gap, it is common to run the nat methods iteratively, which, however, limits the inference speed advantage over autoregressive methods @xcite . in contrast, we show that for tasks where inputs and outputs overlap, we can maintain an order-of-magnitude speed-up without compromising on the model quality by treating the problem as a text-editing task and producing the output in a single pass. a number of text-editing models have been proposed as a faster and more sample efficient alternative to seq2seq models like t5 @xcite @xcite . another recently proposed approach to speed up the inference time of transformer models is called aggressive decoding @xcite . closest to our work, @xcite show that adding pointing mechanism for reordering and a separate insertion model allow their text-editing model, felix, to produce an arbitrary output in a flexible manner. felix is a non-autoregressive model which first predicts the tokens to keep, their order, and the locations at which to insert new tokens. then it runs a separate model based on a bert masked language model for inserting new tokens. in contrast, edit5 employs a single, endto-end model which has an autoregressive insertion component. this enables more accurate insertions, while keeping the latency low, given that most of the tokens can be copied from the source non-autoregressively. other text-editing models that employ autoregressive insertion include ed-itnts @xcite , the text-normalization model by @xcite , seq2edits @xcite , esc @xcite and lewis @xcite . however, unlike edit5, these models perform also the edit operation prediction autoregressively, making them potentially slower at inference time. ## conclusions in this paper we have proposed edit5 a low latency solution to text generation, that achieves comparable or betters results, across three distinct tasks, to a strong t5 baseline whilst achieving inference latencies that are up to 25x quicker than the baseline model. in the future we wish to explore the following ideas: 1) the impact of distillation for edit5. distillation has previously been shown to be particularly advantageous to non-autoregressive models.