Beyond In-Place Corruption: Insertion and Deletion in Denoising Probabilistic Models
Abstract
Denoising diffusion probabilistic models (DDPMs) have shown impressive results on sequence generation by iteratively corrupting each example and then learning to map corrupted versions back to the original.
However, previous work has largely focused on in-place corruption , adding noise to each pixel or token individually while keeping their locations the same.
In this work, we consider a broader class of corruption processes and denoising models over sequence data that can insert and delete elements, while still being efﬁcient to train and sample from.
We demonstrate that these models outperform standard in-place models on an arithmetic sequence task, and that when trained on the text dataset they can be used to ﬁx spelling errors without any ﬁne-tuning.
1. Introduction
Although autoregressive models are generally considered state of the art for language modeling, machine translation, and other sequence-generation tasks (Raffel et al., 2020; van den Oord et al., 2016), they must process tokens one at a time, which can make generation slow.
As such, signiﬁcant research effort has been put into non-autoregressive models that allow for parallel generation (Wang & Cho, 2019; Ghazvininejad et al., 2019).
Recently, denoising diffusion probabilistic models (DDPMs) (Sohl-Dickstein et al., 2015) have shown impressive results in a variety of domains (Chen et al., 2020; Ho et al., 2020; Hoogeboom et al., 2021; Austin et al., 2021), in some cases achieving comparable results to autoregressive models with far fewer steps.
In these models, a forward process iteratively corrupts the data towards a noise distribution and a generative model is trained to learn the reverse denoising process. However, these models
Third workshop on Invertible Neural Networks, Normalizing Flows, and Explicit Likelihood Models (ICML 2021).
Copyright 2021 by the author(s). q ( x t | x t − 1 ) p θ ( x t −| x t ) share one limitation: the corruption process always modiﬁes sequence elements in-place.
While convenient, this choice introduces strong constraints that limit the efﬁcacy of the generative denoising process.
For example, if the model makes a mistake and places a word or phrase in the wrong place, it cannot easily compensate.
For sequence-to-sequence tasks, the Levenshtein transformer (Gu et al., 2019) and Insertion-Deletion transformer (Ruis et al., 2020) address this limitation by performing insertion and deletion operations.
However, these models were not designed as purely generative models, and do not in general allow estimation of sample log-likelihoods through both the insertion and deletion phases.
In this work, we integrate insertion and deletion into the
DDPM framework, generalizing multinomial diffusion models (Hoogeboom et al., 2021) and D3PMs (Austin et al., 2021).
We carefully design a forward noising process that allows for tractable sampling of corrupted sequences and computing estimates of the log-likelihood bound.
We show that our models outperform in-place diffusion for modeling arithmetic sequences, and that for text they learn errorcorrection mechanisms that work on misaligned inputs.
2. Background
Here we describe previous work that is needed to introduce our method; see Appendix A for additional related work.
2.1. Denoising diffusion probabilistic models
DDPMs are latent variable generative models deﬁned by a forward Markov process q ( x t | x t − 1 ) which gradually adds noise, and a learned reverse process p θ ( x t − 1 | x t ) that removes noise.
The forward process deﬁnes a joint distribution q ( x 0: T ) =
q ( x 0 ) (cid:81)
Tt =
1 q ( x t | x t − 1 ) where q ( x 0 ) is the data distribution and x 1 , x 2 , ..., x T are increasingly noisy latent variables that converge to a known distribution q ( x T ).
The reverse process p θ ( x t − 1 | x t ) is then trained to match the forward process posteriors q ( x t − 1 | x t , x 0 ), yielding a gradual denoising model with a tractable variational bound on the log-likelihood.
To enable efﬁcient training, q is often chosen such that these posteriors can be computed analytically. For continuous DDPMs, q ( x t − 1 | x t ) is typically a Gaussian. For discrete DDPMs, Hoogeboom et al.
propose setting q to a mixture of a uniform distribution and a point mass at the previous value, and Austin et al.
(2021) consider using a wider class of structured Markov transition matrices.
All recent diffusion models perform corruption in-place: the k th element of x t is a noisier version of the k th element of x t − 1 , with no dependence on other tokens.
2.2. Levenshtein and Insertion-Deletion Transformers
The Levenshtein Transformer (Gu et al., 2019) learns to insert and delete tokens over a series of generation steps.
In each step, it marks tokens in the current sequence x that should be deleted, predicts how many tokens should be inserted at each position, and ﬁnally predicts values for the newly inserted tokens.
It is trained to imitate the optimal sequence of edit actions computed by a dynamic program in order to recover the dataset example x from a noisy proposal x (cid:48) (generated by corrupting x or sampling from the model). The Insertion-Deletion Transformer (Ruis et al., 2020) uses a sequence of insertion steps followed by a single deletion phase.
In each insertion step, it takes a random subsequence x (cid:48) of the original sequence x , and learns to insert at most one token between each element of x (cid:48) according to a random generation order of x from x (cid:48) . In the deletion phase, it takes a (possibly perturbed) proposal from the insertion phase and learns to delete any token that is not part of x .
Both of these approaches have focused on the sequenceto-sequence setting, where there are usually only a small set of possible correct answers. Additionally, neither provide a tractable estimate of the log-likelihood of dataset samples under the model; they are instead trained using hand-designed losses for insertion and deletion phases.
3. Method
Our goal is to design an insertion-deletion-based generative model within the probabilistic framework of diffusion models with a tractable bound on the log-likelihood.
The main considerations are (a) how to deﬁne the forward corruption process so that it leads to a reverse process with insertions, deletions, and replacements, (b) how to parameterize the reverse process, and (c) how to do both tractably within the diffusion process framework.
3.1. Forward Process
The forward corruption process speciﬁes how to gradually convert data x 0 into noise x T by repeatedly applying a singlestep forward process q ( x t | x t − 1 ).
Since the learned reverse process is trained to undo each of these corruption steps, and insertion and deletion are inverses, we can obtain a learned reverse process with deletion, insertion, and replacement operations by including insertion, deletion, and replacement operations in the forward process, respectively.
A challenge is that if a single forward step can apply an arbitrary set of insertions, deletions, and replacements, then there may be many ways to get x t from x t − 1 .
For example, x t can be related to x t − 1 through the minimum edit between the two, or by deleting the full x t − 1 and then inserting the full x t .
In order to compute q ( x t | x t − 1 ), one would need to sum over all these possibilities.
To avoid this, we restrict the forward process so that there is a single way to get each x t from each x t − 1 , by adding two auxiliary symbols into the vocabulary that explicitly track insertion and deletion operations: every insertion operation produces the insertionmarker token INS , and every deletion operation deletes the deletion-marker token DEL .
(We note that, since the reverse process is reversing the forward corruption process, the learned model must instead insert DEL and delete INS .)
We propose the following form for q ( x t | x t − 1 ):
Remove all DEL tokens from x t − 1 . 2.
For each token x in x t − 1 , sample a new value (possibly DEL ) as x (cid:48) ∼ Cat( x (cid:48) ; δ Tx Q t ), where Q t is a Markov transition matrix and δ x is a one-hot vector for x . 3.
Between each pair of tokens in the result, and also at the start and end of the sequence, sample n new i ∼ Geom(1 − α t ) and insert that many INS tokens.
(We explain this choice in Section 3.4.)
We allow Q t to include transitions from INS to any other token, and from any token to DEL , but disallow transitions to INS or from DEL to ensure they only arise from insertions and deletions.
This ensures unique 1-step alignments.
3.2. Parameterization of the reverse process
As an inductive bias, we prefer reverse processes that produce x t − 1 by modifying x t , instead of predicting it from scratch.
As such, the learned reverse process p θ ( x t − 1 | x t ) ﬁrst removes all INS tokens from x t , then predicts two things for each remaining token: the previous value of the token (which might be INS if the token should be removed), and the number of DEL tokens that should be inserted before the token.
(Recall that, since this is the reverse process, the auxiliary tokens have opposite meanings here.)
We also take xa b c d e f xa b
INS c d
DEL f xa
INS DEL h
i j f INS xa g h
i
INS
DEL f
DEL xa b c d e f a
→xa
g h
i
INS
DEL f
DEL inspiration from other work on diffusion models (Ho et al., 2020; Hoogeboom et al., 2021), which ﬁnd improved performance by guessing x 0 and then using knowledge of the forward process to derive p θ ( x t − 1 | x t ), as opposed to specifying p θ ( x t − 1 | x t ) directly. Our full parameterization combines these two ideas: it attempts to infer the edit summary a 0 → t that was applied to x 0 to produce x t (as shown in Fig. 2), then uses the known form of q ( x t − 1 | x t , x 0 , a 0 → t ) to derive p θ ( x t − 1 | x t ).
Speciﬁcally, we compute p θ ( x t − 1 | x t ) ∝ (cid:88) (cid:101) x 0 , (cid:101) a 0 → t (cid:101) p θ (cid:0)(cid:101) x 0 , (cid:101) a 0 → t (cid:12)(cid:12)(cid:12) x t (cid:1) · q (cid:0) x t , x t − 1 , (cid:101) a 0 → t (cid:12)(cid:12)(cid:12)(cid:101) x 0 (cid:1) , (1) where tildes denote predictions that are not directly supervised, and we intentionally use q (cid:0) x t , x t − 1 , (cid:101) a 0 → t (cid:12)(cid:12)(cid:12)(cid:101) x 0 (cid:1) in place of q (cid:0) x t − 1 (cid:12)(cid:12)(cid:12) x t , (cid:101) x 0 , (cid:101) a 0 → t (cid:1) to prevent the model from predicting edits (cid:101) a 0 → t that have zero probability under q ( x t , a 0 → t | x 0 ).
Intuitively, the model predicts a summary of which edits likely happened (at an unknown time s ≤ t ) to produce x t , then q determines the details of which speciﬁc edits appeared in x t − 1 .
This parameterization requires us to be able to compute q (cid:0) x t , x t − 1 , (cid:101) a 0 → t (cid:12)(cid:12)(cid:12)(cid:101) x 0 (cid:1) (discussed in Section 3.4).
3.3. Loss function
We optimize the standard evidence bound on the negative log-likelihood, which can be expressed as
L = E q ( x 0: T ) (cid:34) − log p
[Equation]
(cid:123)(cid:122) (cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)
[Equation]
[Table]
[Equation]
For the L t − 1 terms, we randomly sample t and then compute
[Equation]
inserting
[Equation]
Replace y with z with prob
[Equation]
DEL
[Equation]
Insert INS with prob
[Equation]
Delete DEL if possible
[Equation]
Stop inserting
Replace x with y , or delete x and insert y
Insert any token with some prob.
Delete any token with some prob.
[Equation]
Insert INS
Stop inserting for x t
Insert DEL and delete it immediately
Insert y , then replace y with z
Stop inserting for x t −
Delete
[Figure]
It turns out that we can compute this inner expectation in closed form given ( t , x 0 , x t , a 0 → t ) (see Section 3.4).
For the L T term, we choose q so that q ( x T | x T − 1 ) deterministically replaces every token with DEL and inserts no new tokens; this implies x T will always consist of repetitions of
DEL , so we can simply learn a tabular distribution p θ ( | x T | ) of ﬁnal forward process lengths.
3.4. Computational considerations
While a diffusion model could be trained by simply drawing sequences x 0 , x 1 , . . .
, x T and training the model to undo each step, these models are usually trained by analytically computing the L t − 1 terms for individual timesteps t and samples ( x 0 , x t ), by using closed form representations of q ( x t | x 0 ) and q ( x t − 1 | x t , x 0 ) (Ho et al., 2020).
Unfortunately, doing this for a forward process that inserts and deletes tokens is nontrivial.
Over multiple steps, the INS and DEL markers may be skipped, which means that (as mentioned in Section 3.1) there will likely be many possible sets of insertions and deletions that produce x t from x 0 , with a corresponding wide variety of intermediate sequences ( x 1 , x 2 , . . .
, x t − 1 ).
To address this challenge, we introduce two main ideas: (a) cast the necessary quantities in terms of probabilistic ﬁnitestate transducers (PFSTs), which allow us to marginalize out details about intermediate sequences that do not matter for computing the loss, and (b) choose to condition on the edit summary a 0 → t in addition to ( x 0 , x t ) while analytically computing the loss term L t − 1 in Eq.
(3), which allows us to efﬁciently compute those PFST-based quantities.
A PFST is a probabilistic ﬁnite state machine that has an input tape and one or more output tapes.
It repeatedly makes stochastic transitions based on a set of transition probabili-
[Table]
[Figure]
θ ( x t −| x t )
[Figure]
[Table]
In
[Table]
place
[Table]
ties and the current symbol from the input tape.
As it makes transitions, it consumes input tape symbols and writes to its output tape(s).
In our case, we begin by expressing q ( x t | x t − 1 ) as a PFST, which is possible because geometric random variables can be sampled as a repeated coin ﬂip.
This PFST iteratively consumes the input ( x t − 1 ), transitioning between states and writing to the output ( x t ).
We additionally make use of an algebra over PFSTs that allows composing PFSTs and integrating out output tapes.
By composing PFSTs for q ( x t | x t − 1 ) and q ( x t − 1 | x 0 ), we obtain a two-output tape PFST for q ( x t , x t − 1 | x 0 ), with which we can integrate out x t − 1 to obtain q ( x t | x 0 ).
Fig.
3 shows the high-level structure of each
PFST; full details are in Appendix B.2.
Given a speciﬁc edit summary a 0 → t , we can reconstruct the state transitions in the PFST for q ( x t , x t − 1 | x 0 ), which allows us to compute q ( x t , x t − 1 , a
→ t | x 0 ) and q ( x t − 1 | x t , x 0 , a 0 → t ) in closed form.
Details on how to compute the necessary terms for our loss in Section 3.3 and our model parameterization in
Section 3.2 are given in Appendix B.3 and B.4, respectively.
4. Experiment: Toy sequence datasets
We start by exploring the expressive power of our model on a toy dataset of arithmetic sequences.
We take a 10step multinomial diffusion corruption process (Hoogeboom et al., 2021) and augment it with varying probabilities of insertion and deletion.
As shown in Table 1, moderate insertion/deletion probabilities lead to better log-likelihoods and to generated sequences with fewer deviations from being a valid arithmetic sequence. However, if insertions and deletions are too frequent, the noise overpowers the patterns in the data, leading to lower accuracy.
Figure 1 shows a sequence generated by the 0.6 insert/delete rate model.
See Appendix C.2 for experiment details.
5. Experiment: Text generation
We also investigate training a 32-step multinomial-diffusionbased model augmented with insertion and deletion on the character-level language dataset text8 (Mahoney, 2011).
Although insert/delete models have slightly worse loglikelihood bounds on this dataset (see Table 2 in App. C), the samples are still high quality, and the models show qualitative differences in the generative process: they can correct spelling errors, insert spaces between words, and make other human-like edits.
In Fig.
4 we show a generated sentence from an insert-delete model, and also show that this model can be used to “spellcheck” a badly-human-written sentence without being trained on this task by simply treating the sentence as x 10 and sampling from p θ ( x 0 | x 10 ).
The insert-delete model generates imperfect but intuitive suggestions whereas an in-place model generates nonsense due to misalignment issues.
See Appendix C.3 for experiment details.
6. Discussion
In this work we have opened up the class of denoising-based generative models to more ﬂexible processes that include insertion and deletion in addition to in-place replacements.
While we have motivated these models from the perspective of text generation, this class of models could be useful for several other applications, such as image super-resolution (by inserting and deleting pixel rows and columns), video generation (by inserting and deleting frames), and molecular structure generation (by editing SMILES representations (Weininger, 1988)).
We are also excited about the potential for incorporating other types of non-in-place edits (such as duplication or reordering) into corruption processes as a strategy for improving denoising-based generative models. Acknowledgements We would like to thank Hugo Larochelle, David Bieber and Disha Shrivastava for helpful discussions and feedback, and
William Chan and Mohammad Norouzi for useful context regarding non-autoregressive sequence models.
We would also like to thank Tim Salimans and the anonymous INNF reviewers for reading earlier drafts of this manuscript and giving suggestions for improvement.
A. Other related work
A few other works have studied diffusion-like generative models for structured data, including Seff et al.
(2019), which exploits a structured forward process q ( x t | x t − 1 ) to impose constraints on the generated samples, and Chan et al.
(2020), which iteratively reﬁnes an output sequence jointly with its alignment to an input sequence.
A number of other edit-based generative models have been proposed, including Guu et al.
(2018) which edits prototypical examples in a latent space.
In natural language processing, edit-based models have been proposed for learning to simplify complex sentences into simple ones (Alva-Manchego et al., 2017;
Dong et al., 2019).
In source code applications, it is common to generate edits for bug-ﬁxing (Yin et al., 2019;
Zhao et al., 2019; Dinella et al., 2020; Yao et al., 2021).
There are also models that use edit distances for purposes of supervision (either directly or via imitation learning), but still generate left-to-right (Graves et al., 2006; Bahdanau et al., 2015; Sabour et al., 2018).
B. Computing probabilities with PFSTs
In this section we describe how to compute the necessary probabilities for the forward process and learned reverse process using probabilistic ﬁnite state transducers.
B.1. Notation for PFST representations
We will begin by introducing additional notation which will be useful for representing PFSTs of the multi-step forward process probabilities q ( x t − 1 | x 0 ) and q ( x t , x t − 1 | x 0 ).
For all of our PFSTs, we associate each transition with a label “ p : x
(cid:55)→ y ”, which indicates that, conditioned on x being the next symbol on the input tape, with probability p the PFST consumes x and produces y .
We use ε to denote the empty sequence, and thus p : ε (cid:55)→ y denotes a transition that (with probability p ) inserts y without consuming any input.
Similarly p : x (cid:55)→ ε denotes consuming x without producing any output, which corresponds to a deletion.
For the product transducer q ( x t , x t − 1 | x 0 ), we write p : x (cid:55)→ y
(cid:55)→ z to indicate consuming x from x 0 , writing y to x t − 1 , and writing z to x t .
As stated in Section 3.1, each single step of the forward process is parameterized by a scalar α t and a Markov transition matrix Q t .
To represent the aggregate probabilities over multiple steps, we introduce three parameters α t , β t , and Q t :
• α t is a vector of insertion probabilities, such that [ α t ] i gives the chance of inserting token i when skipping from time 0 to time t .
In particular, [ α t ] (cid:104) INS (cid:105) denotes the probability of inserting INS , and [ α t ] (cid:104) DEL (cid:105) denotes the probability of inserting DEL .
α t is used to summarize inserts at some time s ≤ t followed by a chain of replacements Q s + 1 , . . .
, Q t .
If s < t , we call this a silent insertion .
• Conversely, β t is a vector of deletion probabilities, such that [ β t ] i gives the chance of deleting token i conditional on it appearing in x 0 .
β t is used to summarize a chain of replacements Q 1 , . . .
, Q s that produce DEL at some time s < t .
We call this a silent deletion .
• Finally, Q t is a matrix that speciﬁes how tokens will be replaced over multiple steps, such that [ Q t ] xy denotes the probability of consuming x and producing y conditioned on x appearing in x 0 .
Notably, this encompasses both chains of replacements due to Q t , as well as silent deletion-insertion pairs , where a token is inserted immediately after a deleted token.
(For instance, in Fig. 2, ‘b’ and ‘h’ form a deletion-insertion pair)
Using this
[Figure]
we can fully specify the PFSTs for each process of interest:
[Figure]
B.2. Calculating q ( x t | x 0 ) from q ( x t | x t − 1 )
As discussed in Section 3. [Figure] we can use the transducer representations shown in Fig.
5 to recursively construct probabilities for q ( x t | x 0 ) from the individual step distributions q ( x t | x t − 1 ).
We proceed inductively by constructing a deterministic q ( x 0 | x 0 ) and then repeatedly computing q ( x t | x 0 ) from q ( x t − 1 | x 0 ) and q ( x t | x t − 1 ).
As our base case, observe that q ( x 0 | x 0 ) is the identity transformation, and we can represent it using the following parameters:
[Equation]
[Figure] AA [Figure] BB α t : ε (cid:55)→ INS 1 − α : ε (cid:55)→ ε [ α t − 1 ] (cid:104) DEL (cid:105) : ε (cid:55)→
ε α
Tt − 1 Q t δ z : ε (cid:55)→ z 1 − (cid:80) y [ α t − 1 ] y : ε (cid:55)→ ε [Figure] ] x : x
[Figure]
x [Figure] Now suppose we know α t − 1 , β t − 1 , and Q t − 1 for q ( x t − 1 | x 0 ), and Q t and α t for q ( x t | x t − 1 ), and we wish to compute α t , β t , and Q t for q ( x t | x 0 ).
We start by constructing the product transducer for q ( x t , x t − 1 | x 0 ) by composing the two transducers for q ( x t − 1 | x 0 ) and q ( x t | x t − 1 ), as shown in Fig.
Next, we marginalize out the middle timestep x t − 1 .
This entails removing the middle step from each transition, and instead summing over all possible values for that middle token.
We obtain the two-tape transducer shown in Fig.
Next, we eliminate the middle state S AB , by replacing all paths that pass through it with new transitions that directly connect S AA and S BB .
We note that these paths may enter the loop S AB (cid:55)→ S AB arbitrarily many times without producing any output or consuming any input (this is a silent-insertiondeletion pair ).
The total probability of all paths that take that loop an arbitrary number of times is thus
[Equation]
We obtain the following new transitions.
From S AA to S AA :
[Equation]
[Equation]
α
[Equation]
(cid:55)→ y
[Equation]
From S AA to S BB :
[Equation]
From S BB to S BB :
[Equation]
From S BB to S AA :
[Equation]
Equation (9) is particularly notable, as it corresponds to a silent-deletion-insertion pair , in which q ( x t − 1 | x 0 ) replaces x with DEL and then inserts some other token ( y in Fig.
5, but marginalized out here), after which q ( x t | x t − 1 ) removes DEL and produces z from y .
Combining these new transitions with the old ones between
S AA and S BB gives us the following values for q ( x t | x 0 ):
[Equation]
(Note: Here we assume [ Q
[Equation]
does not allow transitions from DEL or to INS .)
Intuitively,
Eq.
(10) says that inserts occur either as INS -marked inserts at time t or (silent) inserts before time t that are then perturbed; Eq.
(11) says that deletions occur either as silent deletions before time t or as transitions to DEL at time t that are then removed without inserting new tokens; and
Eq.
(12) says that replacements occur either because a token was copied/replaced before time t and then copied/replaced again at t , or because a token x was replaced by DEL at time t − 1, but a new token y was (silently) inserted at or before time t − 1, so that at time t the new token y looks like a replacement for the old token x .
B.3. Closed form of q ( x t − 1 | x t , x 0 , a 0 → t )
We can similarly obtain a closed-form representation of q ( x t − 1 | x t , x 0 , a 0 → t ) by reasoning backwards about the elimination steps in the previous section.
We start by observing that the edit summary a
→ t tells us the sequence of replacements x (cid:55)→ z , insertions ε (cid:55)→ z , and deletions x (cid:55)→ ε executed by the transducer while sampling x t from x 0 .
Suppose we observe a replacement x
(cid:55)→ z (where perhaps x
= z if it was copied unmodiﬁed).
This must have been produced by the Q t edge.
From Eq.
(12) and Fig.
5 we can infer the distribution over the intermediate value x
(cid:55)→ y
(cid:55)→ z , if it exists:
[Equation]
If the event in Eq.
(14) occurs, we can also infer that there was a geometric number n extra
i
[Equation]
) of extra ε (cid:55)→ DEL
(cid:55)→ ε transitions due to the loop in S AB .
Now suppose we observe an insert ε (cid:55)→ z .
If z = INS , we know it was inserted at time t , so it must have been produced by the ε (cid:55)→ ε (cid:55)→ INS transition.
If z is any other token, it must have already existed at time t − 1, with p ( ε (cid:55)→ y (cid:55)→ z | ε (cid:55)→ z ) = [ α
[Equation]
In this second case we also pass through S AB and generate n extra
i
[Equation]
- tions.
Next suppose we observe a deletion x
(cid:55)→ ε (where we know x (cid:44) DEL because there are no deletion markers in the data distribution).
In this case we have
[Equation]
where, like before, the second case passes through S AB and generates n extra
i
[Equation]
transitions.
Finally, we note that every time we move from S A A to S B B (in other words, whenever we stop inserting tokens), there is one more n extra
i
[Equation]
transitions.
Using the above analysis allows us to compute q ( x t − 1 , a 0 → ( t − 1)
| x t , x 0 , a 0 → t ), where the extra information a 0 → ( t − 1) speciﬁes the sequence of x (cid:55)→ ε (cid:55)→ ε , ε (cid:55)→ DEL
(cid:55)→ ε and x (cid:55)→ DEL
(cid:55)→ ε transitions (which are ambiguous from a 0 → t alone).
Since we do not particularly care about this information, we can marginalize it out by noting that the total number n obs of consecutive DEL tokens observed at a particular position in x t − 1 is the sum of the number of explicit deletions x (cid:55)→ DEL (cid:55)→ ε and insertion-deletion pairs ε (cid:55)→ DEL (cid:55)→ ε .
Given a ﬁxed number of explicit deletions, the total number of insertion-deletion pairs is a sum of independent geometric random variables and thus has a negative binomial distribution.
We can thus:
• compute for each deleted token x in x 0 the probability of an explicit x (cid:55)→ DEL (cid:55)→ ε transition using Eq.
(17) • compute for each perturbed x (cid:55)→ z transition the probability of an explicit x (cid:55)→ DEL (cid:55)→ ε transition using Eq.
(14) • compute the distribution of the total number n explicit of x (cid:55)→ DEL (cid:55)→ ε transitions at this location in x t − 1 by noting that it is a sum of independent Bernoulli r.v.s (which can be computed either by taking convolutions of their PMFs, or, if all tokens are deleted with the same probability, by observing that this is a binomial distribution) • use this distribution to compute a mixture of negative binomial distributions: n obs ∼ n explicit + NB( n explicit + 1 , 1 − [ α t − 1 ] (cid:104) DEL (cid:105) ).
B.
Combining (cid:101) p θ ( (cid:101) x 0 , (cid:101) a 0 → t | x t ) with q
The x 0 -predicting parameterization of p θ follows the same general procedure outlined above for inferring x t − 1 from x t , x 0 and a 0 → t .
However, we make a few slight modiﬁcations due to the structure of (cid:101) p θ .
For each token z in x t , the model predicts a modiﬁcation probability (cid:101) p θ ( x (cid:55)→ z ) for each token and an insertion probability (cid:101) p θ ( (cid:15)
(cid:55)→ z ).
We use these as weights to scale the appropriate inference terms in Eqs.
(13) to (15).
Additionally, the model predicts a distribution (cid:101) p θ ( n deli ) of the number x
(cid:55)→ ε transitions that occurred before each position i in x t .
We use this to infer the number n obs i of
DEL placeholders that appear at time t − 1 using the same inference procedure as above, but we now have a mixture of mixtures of negative binomial distributions because we may be uncertain about how many insertions there were.
(Usually, we will have n obs i ≤ n deli , since deletions could have occurred at any time from 0 to t .)
When implementing this parameterization we assume that every token is equally likely to be deleted at each timestep, so that the model only has to predict the number of missing tokens from x 0 ; if this is not the case, it would be possible to predict (cid:101) p θ ( n obsi ) directly instead.
We choose to predict deletion-insertion pairs simply as an insertion preceded by a deletion, instead of reasoning about it as a replacement; this simpliﬁes our computation by avoiding having to separately reason about Eq.
C. Experimental details C.1. Model architecture
For all of our experiments, we use a standard decoder-only transformer following the T5 (Raffel et al., 2020) architecture, with either six or twelve layers depending on the task.
The main modiﬁcation we make is to introduce two output heads instead of one.
The ﬁrst output head, like a standard transformer, predicts a matrix f θ ( x t )
∈ R L × K of unnormalized log-probabilities (logits), where L is the sequence length and K is the vocabulary size.
We interpret f θ ( x t ) iv as the log-probability of the i th token being produced by a replacement edit v (equivalently v (cid:55)→ [ x t ] i in the PFST notation) in the edit summary a 0 → t , and similarly interpret f θ ( x t ) i (cid:104) INS (cid:105) as the log-probability of the i th token of x t being an insertion (or ε (cid:55)→ [ x t ] i ).
We reuse the embeddings for the input vocabulary as the ﬁnal output layer for this head.
The secound output head produces a matrix g θ ( x t ) ∈
R L × L , for which f θ ( x t ) in gives the (unnormalized) log-probability of having n different (or [ x 0 ] j
(cid:55)→ ε ) edges immediately before the i th token of x t .
When running the transformer on an input sequence, we introduce an extra end-of-sequence token EOS that denotes the last position in the input.
The ﬁrst output head f θ is ignored for the EOS token, but we do use the output g θ for the EOS token to determine the number of edges in the
[Figure]
edit summary a 0 → t that occur at the end of the sequence.
As mentioned in Section 3.3, we additionally store a ﬁxedsize table p θ ( | x T | ) ∈ R L , which we ﬁt to the distribution of observed lengths x T .
C.2. Arithmetic sequences
We construct a dataset of arithmetic sequences by randomly sampling a step size s between 1 and 10, a direction (increasing or decreasing), a length (cid:96) between 32 and 64 (with the constraint that s ( (cid:96) − 1) < 509), and ﬁnally a random starting position so that all terms in the sequence are betweenand 511, inclusive.
(0 is used to denote padding in the data loader, and 1 was reserved for preliminary experiments that required additional reserved tokens, but both are treated as ordinary tokens by the model.)
Along with INS , DEL , and an end-of-sequence marker EOS, this yields a total augmented vocabulary of size 515.
We compare four different forward process schedules, each of which is tuned to add less noise for timesteps closer toand more noise as t approaches 10.
We start by choosing an insert/delete rate r ∈ { 0 , 0 . 4 , 0 . 6 , 0 . 8 } .
Next, for 1 ≤ t ≤ 9, we calculate a fraction u t = 0 .
1 t 9 + 0 . 9 (cid:16) t 9 (cid:17) 2 , then choose the insertion probability α t and matrix Q t for each t so that, cumulatively after step t , approximately u t × r of the elements of x 0 have been deleted, u t × r of the elements of x t come from insertions (so that the length of the sequence remains approximately the same), and u t of the remaining elements from x 0 have been replaced by a random integer between 0 and 512.
Finally, at step 10 we append a deterministic step Q 10 that replaces every token with DEL , and set α 10 = 0.
When r = 0 . 0, no insertions or deletions occur until the last step, which is simply used to allow the model to predict the length of the sequence.
We choose [ Q t ] (cid:104) INS (cid:105) n = 1512 for all 0 ≤ n < 512 so that INS is equally likely to transition to any
Bits/char
In place
[Table]
of the 512 numbers in the vocabulary.
The full schedule for r = 0 .
6 is shown in Fig.
For each insert/delete rate r , we train a six-layer transformer model over 100,000 minibatches of 512 random examples, using the Adam optimizer and a learning rate that increases linearly to 2 × 10 − 4 over 5000 steps, then stays constant.
We rerun training with ﬁve random seeds for each schedule.
Since the loss seemed to stabilize at around 90,000 steps, we take averages of the validation metrics computed during the last 10,000 steps of training for each seed, corresponding to ELBO estimates for 46,080 random dataset examples and error rate metrics for 2304 samples drawn from the model.
We then report the average and standard deviation of these per-seed metrics across the ﬁve random seeds for each schedule.
C.3. Text generation on text
For text8, we construct a dataset of training examples by taking randomly-selected 118 character chunks of the full concatenated lower-cased training set.
We use a dataset vocabulary of 28 tokens, including each character ‘a’ through ‘z‘, a space, and an extra token ‘-’ that does not appear in the dataset (again used for preliminary experiments); including
INS , DEL , and EOS gives a vocabulary of size 31.
During training, since we may insert a large number of tokens by chance, we enforce a maximum length of the intermediates x t by rejection sampling until we draw a sample shorter than 128 characters (which we correct for when computing the
ELBO during evaluation).
As in the arithmetic sequence dataset, we compare forward process schedules with four insert/delete rates r ∈ { 0 , 0 . 4 , 0 . 6 , 0 . 8 } , constructed to add less noise near time 0.
In this case, we instead set u t = 0 .
1 t 31 + 0 .
9 (cid:16) t 31 (cid:17) 2 and produce a 32-step corruption process; similarly, when randomizing, we randomly choose from the 28 tokens in the vocabulary instead of the 512 numbers.
For each insert/delete rate r , we train a twelve-layer transformer model over 1,000,000 minibatches of 512 random examples, using the Adam optimizer.
We perform a sweep over four learning rates { 5 × 10 − 5 , 1 × 10 − 4 , 2 × 10 − 4 , 5 × 10 − 4 } and three schedule types: linear increase until 5000 steps followed by constant, linear increase until 5000 steps followed by reciprocal square root decay, and a cyclical cosine schedule with period 100,000.
As a preliminary estimate of performance, and because training seemed to converge before 900,000 steps, we evaluated over a subset of 40,960 length-118 segments sampled from the validation set, averaged over the last 100,000 steps of training.
Table 2 shows preliminary bits/char measurements for the run with the best performance for each value of r .
To produce the typo-repair example on the right side of
Fig.
4, we took the human-written sentence “thisn sentsnetne wasstype vssry babdly”, intended as a typo-ridden version of “this sentence was typed very badly”.
We then padded the sentence out with placeholder text (“lorem ipsum dolor sit amet lorem ipsum dolor sit amet...”) until it had length 119, to be approximately the length of the training examples.
We set this padded sentence as x 10 , then drew ﬁve random samples for both the 0.6 insert/delete rate model and the 0.
insert/delete rate model.
We trimmed off the placeholder text (which the model generally left alone) but did not make any other edits.