STEP-UNROLLED DENOISING AUTOENCODERS
Year: 2022

Abstract:
In this paper we propose a new generative model of text, Step-unrolled Denoising Autoencoder (SUNDAE), that does not rely on autoregressive models. Similarly to denoising diffusion techniques, SUNDAE is repeatedly applied on a sequence of tokens, starting from random inputs and improving them each time until convergence. We present a simple new improvement operator that converges in fewer iterations than diffusion methods, while qualitatively producing better samples on natural language datasets. SUNDAE achieves state-of-the-art results (among non-autoregressive methods) on the WMT’14 English-to-German translation task and good qualitative results on unconditional language modeling on the Colossal Cleaned Common Crawl dataset and a dataset of Python code from GitHub. The non-autoregressive nature of SUNDAE opens up possibilities beyond left-to-right prompted generation, by filling in arbitrary blank patterns in a template.

heading: 1 INTRODUCTION
text: Autoregressive (AR) models have shown excellent results in generating text (e.g., GPT-3, Brown et al., 2020). However, while their training scales very well, sampling is prohibitively slow for many practical applications. Moreover, there are limitations to the kinds of conditioning AR models can seamlessly handle: the left-to-right restriction makes it hard to “fill in the gaps” in a partially written text draft. Even more importantly, this prohibits iterative refinement of complete text drafts to make them more self-consistent, which is a common task for human writers. Finally, AR models require network architectures to be causal, severely limiting the kinds of neural network architectures that can be used for text-modeling. All of these motivated the machine learning community to make extensive efforts to propose alternatives to AR models.
Machine translation (MT) was perhaps one of the first tasks where non-AR approaches were shown to seriously rival the AR-based state of the art: methods like CMLM (Ghazvininejad et al., 2019) and DisCo (Kasai et al., 2020) show promising results and their decoding speed is excellent compared to AR. However, while their performance is competitive, they are still behind the AR benchmark and actually require distillation of a larger AR model — without which, performance drops considerably.
Non-AR methods have proven hard to apply to the general unconditional language modeling (LM) task. When there is no conditioning, the multi-modality problem becomes paramount, as shown by Gu et al. (2017), which likely makes it problematic to use methods like CMLM and DisCo because their decoding mechanism is deterministic and does not model uncertainty. Yet, recently the community has seen promising results from non-AR models like Multinomial Diffusion (Hoogeboom et al., 2021) and D3PM (Austin et al., 2021). These methods optimize a lower bound (ELBO) on likelihoods and have shown negative log-likelihood (NLL) results approaching AR models on several benchmarks like text8 (Mahoney, 2011) and LM1B (Chelba et al., 2013). However, a major gap in NLL persists, and samples from those models lack coherence.
In this paper we propose a novel non-autoregressive method which shows state-of-the-art results in machine translation on WMT’14 EN→DE raw data (without distillation from AR) amongst non-AR methods and good qualitative results on unconditional language modeling on the Colossal Clean Common Crawl (C4) dataset (Raffel et al., 2019) and a dataset of Python code from GitHub. Our model operates as a time-homogeneous Markov Chain similar to that of Lee et al. (2018): conditioned on the corrupted data, it tries to approximate the original uncorrupted samples by a per-token
∗Shared first authorship.
factorised distribution. During generation, we unroll the chain by sampling from the transition distribution and feeding samples back into the input. We propose a new training mechanism called unrolled denoising, which uses unrolls of the chain during training as well, and which we empirically show to be crucial for practical performance of the algorithm in ablation studies. This feature motivates the name for the model, Step-unrolled Denoising Autoencoder (SUNDAE). The difference between usual denoising and unrolled denoising is illustrated in Figure 1.
To summarize our contributions:
• We present SUNDAE, a new generative model of text that unrolls the denoising process during training.
• SUNDAE achieves state-of-the-art results on WMT’14 English-to-German translation task among non-AR methods.
• We demonstrate good qualitative results for unconditional generation and inpainting on Colossal Clean Common Crawl dataset and a dataset of Python code from GitHub.
• We carefully ablate and analyze the properties of the proposed method, and show that unrolls during training are crucial for the model’s performance.

heading: 2 METHOD
text: We approach the problem of generative modelling of discrete sequences by bringing together the framework of denoising autoencoders and Markov chain models. In this section, we first discuss the definition of the generative model, and then describe the training process, which includes our main contribution, unrolled denoising.
For a fixed prior distribution p0 on some space X consider a process xt ∼ fθ(·|xt−1), where fθ is a parametric transition function. Then {xt}t is a time-homogeneous Markov chain, and transition t-steps ahead has the following form1
pt(xt|x0) = ∑
x1,x2,...,xt−1∈X
t∏ s=1 fθ(xs|xs−1). (1)
For a fixed number of steps T , the prior p0 and decoder pT together determine our model distribution pT (xT ) = pT (xT |x0)p0(x0). We assume X = {1, . . . , v}N to be the space of sequences of length N with values at all positions coming from a vocabulary of size v, and the prior p0 to be uniform overX . Let pdata be a distribution of the data and assume that fθ induces a conditional probability distribution fθ(·|x′) that can be
1for formal derivation, see Appendix A.
factorised into a conditionally-indepedent product as follows:
fθ(x|x′) = f (1)θ (x (1)|x′)f (2)θ (x (2)|x′) · · · f (N)θ (x (N)|x′), for x,x′ ∈ X. (2)
Note that although for t = 1 the distribution p1(·|x0) is limited by product structure (Eq. 2), the subsequent pts are not restricted in the same way, and after each step they belong to potentially more and more expressive family.

heading: 2.1 TRAINING WITH UNROLLED DENOISING
text: Due to intractability of the likelihood of the model pT as well as the computational requirements of optimising the whole chain {pt}t, we propose an efficient two-step training method which we term unrolled denosing and visualise in Figure 1. We consider a smaller number of steps than T that we intend to use at sampling, but to compensate for that, we unroll the chain starting from corrupted data samples, rather than the prior p0. This way, the model learns to denoise the samples it is likely to encounter during the full unroll used at sample time. While using a single step would resemble the training strategy of BERT (Devlin et al., 2018), using at least two steps is essential for the performance of our model.
Consider a corruption distribution q(xc|x) that replaces a random proportion of tokens in a sentence x with ones randomly sampled from a uniform distribution. Our objective is given by L(1:2) = 1 2 ( L(1) + L(2) ) where
L(t)(θ) := −E x∼pdata x0∼q(·|x)
x1∼fθ(·|x0) ...
xt−1∼fθ(·|xt−2)
[log fθ(x|xt−1)], (3)
i.e. the reconstruction loss of the chain after t steps, starting from a corrupted sample x0. In Appendix A.1 we show that L(t) is an upper bound on the actual negative log-likelihood from the model distribution pt:
L̃(t)(θ) := −E x∼pdata xc∼q(·|x) [log pt(x|xc)]. (4)
Note that L̃(T ) is the actual term we would like to minimise if pT were tractable, and in fact is closely related to evidence lower bound (ELBO) via a fixed encoder q(xc|x) (Appendix A.2). However, due to the inherent non-differentiable nature of discrete Markov chain models, optimisation of L̃(T ) would not only be costly, but also limit the gradient flow to the last step of the chain. We found instead that optimising several steps together without propagating the gradient through sampling is sufficient for obtaining good results. Averaging more L(t) terms can lead to minor improvements in performance (see Figure 4d in Appendix C), but it considerably slows down the training speed. We discuss these and other loss variants further in Section 3.1.
Since function fθ is modelled as a neural network, the logarithm under expectation on the right hand side of Eq. 3 is obtained from predicted logits, thus we term L(1)(θ) and L(2)(θ) logits loss and unrolled logits loss, respectively.
Corruption function. To corrupt texts, we first sample a proportion of tokens uniformly from [0, 1], randomly select positions according to this proportion and then change tokens at those positions to random tokens sampled uniformly from the vocabulary. This way, we want to ensure that samples seen at various steps during a generative unroll (including the first one, sampled from the uniform prior) are well-represented in corruptions applied during training. More details and relations to the forward diffusion process (Hoogeboom et al., 2021) are presented in Appendix A.3.

heading: 2.2 SAMPLING
text: At sampling time we follow the structure of the Markov process and sample sequentially xt ∼ fθ(xt|xt−1) for some fixed number of steps T , beginning from a random sequence x0 (potentially starting with a prompt or a template for conditional generation). To control the speed of convergence, we propose three improved strategies which allow the use of a much smaller number of steps:
• Low-temperature sampling. For temperature τ and trained network fθ we consider modified function fτθ such that log f τ θ (·|x) ∝ 1τ log fθ(·|x). We found that sampling with tempera-
tures lower than 1 leads to much faster convergence than with the original logits, and allowed generation of high-quality samples in a matter of 10-16 steps.
• Argmax-unrolled decoding. At the limit τ → 0 sampling with temperature reduces to deterministic argmax decoding where the most probable token is chosen at each step. Relatedly to our unrolled logits loss, we modify this strategy by resampling the low-certainty tokens in accordance with unrolled logits. This heuristic allowed further improvements to sampling speed while maintaining the high quality of the samples. We discuss it in detail in Section 3.1.
• Updating fewer tokens. In tasks where diversity is paramount (like unconditional text generation), we found that updating a random subset of tokens at each decoding step leads to faster convergence. This likely happens because independent sampling of all tokens might create uncoordinated changes which could take some time to fix in the subsequent steps. We use this strategy in Section 3.2.

heading: 3 EXPERIMENTS

heading: 3.1 MACHINE TRANSLATION
text: We first evaluate SUNDAE on Machine Translation (MT) benchmarks. We compare SUNDAE against AR and non-AR models in terms of the translation quality using BLEU (Papineni et al., 2002) as the metric. We demonstrate that without using techniques like sequence-level knowledge distillation (Kim & Rush, 2016), SUNDAE performs almost as well as the AR model and outperforms all other methods that do not rely on AR models.
Target Length Prediction. Unlike their AR counterparts, non-AR models do not explicitly learn to predict an end of sequence, but instead have been shown to benefit from auxiliary target length
prediction (Lee et al., 2018; Ghazvininejad et al., 2019; Kasai et al., 2020). A common approach is to treat it as a classification task, predicting either the exact length or the difference between the source and target lengths. The decoding process has to commit to the predicted length, and often multiple length beams are used to get the best translation results. In our models, we use a separate network that receives source encodings as input and predicts corresponding target length. Target length embedding vector is prepended to source embeddings so that decoder can attend to it. During training, we use the reference target length, while at sampling, a predicted one, although the model does not need to fully commit to the predicted length like some previous non-AR models (Saharia et al., 2020). Note that the length classification loss does not affect the encoder parameters. Our models are also trained to predict the padding tokens at the end of the text. For more details, see Appendix B.
Experiment Settings. We conduct experiments on WMT’14 parallel corpora using EN↔DE (4.5M pairs) and EN→FR (36M pairs) translation tasks. The raw texts are encoded using BPE (Sennrich et al., 2015) as the subword units, and we use the same preprocessed data as in Vaswani et al. (2017) for fair comparisons. We evaluate the performance by measuring BLEU (Papineni et al., 2002; Post, 2018) on the test split of each translation task2. We use the encoder-decoder Transformer architecture for MT (Vaswani et al., 2017), but remove the causality masking in the decoder. There are 6 attention layers for both encoder and decoder, 8 attention heads, 512 model dimension and 2048 feedforward hidden dimension. The total number of parameters is 63M including the target length prediction module described in 3.1. We use dropout (p= 0.1) and label smoothing ( =0.1) during training for all tasks, except for AR-distilled EN→DE, where we found lower dropout (p=0.05) and no label smoothing produces better results on validation. The training batch size is 4096, and we use Adam (Kingma & Ba, 2014) with β1=0.9, β2=0.999, =10−6 and weight decay of 0.1 (Loshchilov & Hutter, 2017). We warm up the learning rate from 10−7 to 10−4 in the first 5K steps and decay it to 10−5 using cosine annealing (Loshchilov & Hutter, 2016). Our models were trained for 106 steps using 16 TPU accelerators using bfloat16 precision. We crop sequences that have more than 128 tokens (this occurs less than 0.2% in training data). Finally, we average the last 10 checkpoints to obtain a single model for the evaluation. All hyperparameter tuning is performed on a held-out validation set.
Decoding. We decode translations from SUNDAE using two different types of heuristic decoding methods in MT experiments. The decoding process always begins with an array of random integers sampled from the discrete uniform prior y0 ∼ p0, while the encoder input - the source sentence - remains the same throughout the process. The first decoding method is low-temperature sampling, where we iteratively sample yt ∼ fτθ (·|yt−1) for T iterations (with the decoder output logits divided by τ , see Section 2.2), with T ≤ 16 and τ ∈ [0.1, 0.6] determined based on the validation performance. The second method is argmax-unrolled decoding, which requires a smaller number of iterations compared to its stochastic counterpart. In argmax-unrolled decoding, we first compute the logits λ1 = fθ(·|y0) of the initial sample array y0 and obtain the samples y1, pass a tuple (y1, λ1) to the next iteration. At each step t≥ 2, the method finds top-ρ share of the tokens that are sorted by the log-probability in descending order (i.e. uncertain tokens) from λt−1, where ρ ∈ [0.1, 0.6] is another hyperparameter searched for in the validation stage. We compute unrolled logits for those top-ρ tokens and then apply argmax to obtain unrolled tokens. For the rest of the input tokens of yt−1 we compute logits once and obtain λt = fθ(·|yt−1) and then apply argmax to obtain predicted tokens. Finally, we combine unrolled tokens for uncertain input positions with predicted tokens for remaining positions to obtain yt, and the tuple (yt, λt) is provided as the input to the next iteration. This procedure is repeated over a fixed number of iterations T . We always decode n samples in parallel and rerank them based on the model score. We show the relative speed gain of SUNDAE in Table 2, the Transformer base is used as the AR baseline, for which we perform incremental sampling by caching the previous attention states.
Baselines. We compare SUNDAE with Transformer base (Vaswani et al., 2017) as the AR baseline and several non-AR models including CMLM (Ghazvininejad et al., 2019) and DisCo (Kasai et al., 2020). The last two are both iterative non-AR models that share a substantial amount of common ground with SUNDAE. However, their best results come from the variants distilled from Transformer large (Vaswani et al., 2017), with only few test results available without distillation. As we aim to remove any kinds of dependencies on AR approaches, we would like to make a distinction
2Due to limited space, we report BLEU using SacreBLEU in Appendix G.
Steps (T ) Relative Speed Improvement 4 4.7x 8 2.6x 10 2.2x 16 1.4x
Table 2: Relative speed gain of SUNDAE over AR Transformer base (greedy decoding) on WMT’14 EN→DE validation set.
between raw, fully-non-AR models, and AR-distilled ones, i.e. those trained via distillation from large AR teachers (which typically perform better than the student models), or via AR reranking.
Results and Ablation Studies. In Table 1, we show test BLEU scores from SUNDAE and other baselines. On EN→DE, SUNDAE achieved 26.25 BLEU, and it performed the best among the raw non-AR models. To the best of our knowledge, SUNDAE achieved the closest result (4 = 1.05) to the score of Transformer base (Vaswani et al., 2017) without the aid of AR models.3 In the raw setting, argmax-unrolled decoding outperforms mask-predict (Ghazvininejad et al., 2019) and easyfirst (Kasai et al., 2020) by 0.74 BLEU when using the same amount of the worst-case compute budget. At T =8, SUNDAE performs better than Imputer (Saharia et al., 2020), which is a strong baseline that can generate high-quality translation within few (4-8) iterations. We discovered that as T increases, low-temperature sampling outperforms argmax-unrolled decoding, and at T =10, lowtemperature sampling outperforms Mask-Predict (Ghazvininejad et al., 2019) and Easy-First (Kasai et al., 2020) by 1.5 BLEU score. SUNDAE achieved 30.80 BLEU on DE→EN, where there are not that many raw baselines to compare with. The difference between SUNDAE and the AR baseline on DE→EN is 0.98. Finally, SUNDAE achieved 37.53 BLEU on EN→FR at T = 10, while the Transformer base (Vaswani et al., 2017) reported 38.1 BLEU on this task. Thus, SUNDAE without distillation is only 0.57 BLEU behind the standard AR baseline on EN→FR. We present more details on EN→FR scores in Table 7 in Appendix G. While knowledge distillation was outside our main focus, we also report scores for AR-distilled SUNDAE.4 It again performed well, establishing new SotA on EN→DE translation among AR– distilled models, and being second only to JM-NAT (Guo et al., 2020) in DE→EN task. We ablate the number of steps of unrolled denoising on EN→DE test split. We observe that having at least one unrolled denoising step is crucial to obtain good performance in translation. With only the L(1) loss, SUNDAE achieves 11.19 BLEU, whereas using L(1:2), the score improves to 26.57 BLEU. Using L(1:3), i.e. averaging more unrolled denoising loss terms, did not improve the performance further, scoring 26.25 BLEU. The target length prediction also turned out to be an important component of our model: without it, the translation performance degrades on average by 2 BLEU score points on EN→DE (see Appendix C). Finally, we qualitatively show how translation improves along the sampling steps by elimination of incoherent/repeated tokens in Table 5 in Appendix E.

heading: 3.2 TEXT GENERATION
text: Unconditional Text Generation (Qualitative). We train our method on a large high-quality publicly available Colossal Clean Common Crawl (C4) dataset (Raffel et al., 2019) to demonstrate samples. We tokenize the data using SentencePiece (Kudo & Richardson, 2018) 5 with vocabulary size 32K and train on randomly cropped sequences of length 32. The network is the same as the Transformer decoder part for MT in Section 3.1 but is much larger — 335M parameters: 24 layers,
3As shown in Table 7 in Appendix G, allowing larger compute reduces the gap even further, to 4=0.73. 4For distillation, we follow the protocol used by Saharia et al. (2020). Scores in Table 1 were obtained with distillation from Transformer–Big to allow direct comparisons with CMLM (Ghazvininejad et al., 2019) and SMART Ghazvininejad et al. (2020); for Transformer–Base–distilled scores see Table 6 in Appendix F.
5https://github.com/google/sentencepiece
1024 embedding size, 4096 hidden size, 16 attention heads. Same as in the MT setup, we remove the causal mask since our model is non-autoregressive. We train up to 400K steps with Adam optimizer with batch size 4096 and use cosine annealing for scheduling the learning rate with minimum 10−5, maximum 2 ∗ 10−3 and linear warm-up of 10K steps from starting value of 10−7. The trained model is then used for unconditional sampling. Starting with 32 random tokens, we iteratively perform up to 1K steps from our model, stopping earlier if samples do not change anymore. To make convergence faster, we use temperature 0.8 and update only 30% of all tokens at each step, chosen randomly — as explained in Section 2.2. We show 10 non-cherry-picked samples from our model in Table 8 of Appendix H. Of those samples, all but one resemble reasonable internet texts.
Unconditional Text Generation (Quantitative). While C4 is one of the largest high-quality datasets publicly available, it does not have an established benchmark among non-AR methods. To quantitatively evaluate our algorithm, we train it on the EMNLP2017 News6 dataset, which has numerous text GAN baselines, including a strong ScratchGAN method (d’Autume et al., 2019). We also compare to several other baselines reported in (d’Autume et al., 2019; Caccia et al., 2020): SeqGAN (Yu et al., 2017), LeakGAN (Guo et al., 2017), RankGAN (Lin et al., 2017), MaliGAN (Che et al., 2017) and autoregressive (AR) language models.
We use the same architecture and optimizer as in qualitative experiments and train with batch size 1024 for 8K steps (chosen to achieve the lowest validation loss). We follow the same tokenization strategy as d’Autume et al. (2019), with vocabulary size of 5.7K and maximum length 52 tokens, padding shorter sequences to this maximum length. During sampling, we perform 1K steps and update all the tokens at each step (which we found beneficial for the chosen metrics). The quality/variation trade-off is controlled by changing the sampling temperature in the range [0.1, 2.0], using 60 values overall. For every temperature, we sample 10K texts for evaluating quality and another 10K texts for evaluating diversity, as done by d’Autume et al. (2019).
The results are shown in Figure 2. SUNDAE demonstrates good quality/variation trade-off, as measured by BLEU/self-BLEU metrics, outperforming 4 out of 5 GAN baselines, and competing with ScratchGAN and AR model. In the higher quality region, SUNDAE outperforms ScratchGAN, while in higher diversity region it slightly underperforms. The advantage of ScratchGAN in higher diversity region could be explained by its autoregressive decoding (it is a hybrid between AR and GAN). Interestingly, the curve for SUNDAE has a more complicated shape than the baselines — possibly because of its Markov Chain sampling, which can empirically exhibit something akin to “phase transitions”. We show samples from SUNDAE in Table 9 of Appendix I.
Text In-painting. The non-autoregressive nature of our model opens up possibilities beyond leftto-right prompted generation, such as filling in arbitrary blank patterns in a template. In this section, we qualitatively investigate this new capability, using the C4-trained model described in previous sections, and another one, trained on a dataset of Python code from GitHub. All the architecture/training settings remain the same, except we train twice longer: up to 800K steps. Sampling settings are also the same, except we update all the tokens at each step as diversity is less of an issue with stronger conditioning. We construct the code dataset by extracting files ending in .py from open source GitHub repositories with licenses that are one of apache-2.0, mit, bsd-3-clause, bsd-2-clause, unlicense, cc0-1.0, isc, artistic-2.0. We perform document level de-duplication checking for exact matches.
For the C4 dataset, we present a model with a particular type of prompt [CTX1][MSK1][CTX2][MSK2], where [CTX] stands for “context” and [MSK] for “masked”, in Table 3, which shows how information can be “teleported” from arbitrary location to another arbitrary location, while taking into account bidirectional context. We change a species of the animal in [CTX2] and observe that the model actually uses the right word while in-painting
6http://www.statmt.org/wmt17/
[MSK1] and plausibly continues [MSK2] according to the species. This kind of task would be impossible to perform for an AR model: with left-to-right causality it would be unable to correctly guess the name of the animal, while with a right-to-left one it would have to start filling [MSK2] tokens first, so it would just ignore the conditioning. By contrast, our model succeeds in this task.
For GitHub Python data, we use two different kinds of prompts for our model in Table 4: [CTX1][MSK1][CTX2] and [CTX1][MSK1][CTX2][MSK2][CTX3]. In the first example, our model guesses that n should be the length of the input sequence and correctly in-paints [MSK1] region. This would again be impossible to do for AR models. Looking from left-to-right, it is impossible to guess that n will be the name of the length variable and t is meant to be a threshold. Looking from right-to-left, it is unclear what variables n and t represent and what is the function being computed. Only a bidirectional model can perform well in this task. In the second example, our model correctly follows the suggestion to reuse the function which it also has to in-paint.

heading: 4 RELATED WORK
text: In this section we provide an overview of generative models, focusing on applications in text domain.

heading: 4.1 AUTOREGRESSIVE MODELS
text: Early attempts to perform autoregressive modelling of text using neural networks began with Bengio et al. (2003); later Sutskever et al. (2011) extended this approach with RNNs. Since then, improvements in neural network architecture, such as the Transformer (Vaswani et al., 2017), and scaling up as with GPT-3 (Brown et al., 2020) dramatically improved the performance of AR models. While the architectures evolved, the loss function remained the same with minor modifications.
Insertion (Stern et al., 2019) and Levenshtein (Gu et al., 2019) Transformers are deviations from the standard paradigm of training and could potentially overcome some of AR problems. However, in the parallel decoding regime, they have to rely on AR distillation for obtaining competitive results.

heading: 4.2 NON-AUTOREGRESSIVE MODELS IN GENERAL
text: Non-autoregressive models evolved in parallel with autoregressive ones but have so far struggled to address the same spectrum of tasks with the same quality as AR models.
Diffusion for generative models was originally introduced by Sohl-Dickstein et al. (2015). Recently Ho et al. (2020) proposed a different parametrization: instead of modeling one-step transitions, one could instead model p(x0|xt) and later convert those to p(xt−1|xt) by probabilistic reasoning. This was recently up-scaled for image modeling by Nichol & Dhariwal (2021) and extended to continuous time by Song et al. (2020). In terms of applications to text generation, two recent works addressed this: Hoogeboom et al. (2021) and Austin et al. (2021). Likelihoods obtained in those works are promising, but still behind AR and lacking sample quality as well.
Some diffusion works like Mittal et al. (2021) also approached modeling discrete sequences in another way: first encode sequences with a VAE into a continuous space and then apply continuous diffusion approaches commonly used for image generation.
Variational Autoencoders (VAEs) were proposed by Rezende et al. (2014) and Kingma & Welling (2013) as a likelihood-based extension of autoencoders. Application of such methods to text has been problematic empirically: as it was first noted in Bowman et al. (2015), good results come from using a strong AR decoder instead of a simple factorised one, but the latents are of-
ten ignored — which was named “posterior collapse problem”. This problem was later tackled by Delta-VAEs (Razavi et al., 2019), with limited success in the text domain (Bosc & Vincent, 2020).
Normalising Flows (Dinh et al., 2014; Rezende & Mohamed, 2015; Dinh et al., 2016) are another prominent family of generative models, originally proposed for real-valued data and recently revisited for text generation by Hoogeboom et al. (2021). While conceptually interesting, the text samples of such methods lack in fidelity.
Generative Adversarial Networks (GANs) were originally introduced by Goodfellow et al. (2014) and remain one of the dominant methods for image generation. They were later adapted to text generation in Yu et al. (2017); Guo et al. (2017); Lin et al. (2017); Che et al. (2017). However, text generation with such methods is still a challenge, with a more recent ScratchGAN (d’Autume et al., 2019) showing relatively low-quality samples. At least part of the problem with applying GANs to text generation comes from non-differentiability of discrete text samples, which requires usage of zeroth order optimization methods.
Energy-Based models have a long history dating back to Hopfield (1982); Hinton (2002); LeCun et al. (2006); Ranzato et al. (2007). Recently, Deng et al. (2020) shows promising results for text.

heading: 4.3 NON-AUTOREGRESSIVE MODELS FOR MACHINE TRANSLATION
text: There have been many attempts to apply non-autoregressive methods to machine translation. Latent transformer (Kaiser et al., 2018) generated latent variables autoregressively and then decoded feedforwardly. NAT (Gu et al., 2017) was the first to characterize a problem with non-autoregressive generation — “multi-modality”. FlowSeq (Ma et al., 2019) applied Normalising Flows to machine translation. LVM-DAEs (Lee et al., 2018) are perhaps most related to our method, but these models do not have unrolled denoising and their decoding method is deterministic. The leading methods in this area are CMLM (Ghazvininejad et al., 2019) and DisCo (Kasai et al., 2020). While conceptually similar to LVM-DAEs (Lee et al., 2018), these works have introduced significant improvements to the training and decoding procedure, and, as a result, shown strong competition to AR methods. Both of them also do not have unrolled denoising and sampling like we do, but they are still powerful baselines. Later, there were attempts to combine CMLM with local AR (Kong et al., 2020). Perhaps the most related to our work is the SMART (Ghazvininejad et al., 2020) follow-up of the CMLM. However, this method is not probabilistic and it does not show unconditional generation capabilities. Finally, a recent line of work called Imputer (Chan et al., 2020; Saharia et al., 2020) achieves strong results in machine translation by aligning target to source via dynamic programming.

heading: 4.4 DENOISING OBJECTIVE IN GENERAL
text: With the advent of BERT (Devlin et al., 2018), the denoising objective became very popular for text representation learning. The original idea was later simplified in RoBERTa (Liu et al., 2019) by removing the unnecessary next-sentence-prediction task and only retaining the denoising task (the cloze task). This is similar to our work, however, we do not apply masking to corrupt the input tokens, and instead only switch them to random ones. More importantly, RoBERTa does not use unrolled denoising. Another work, Electra (Clark et al., 2020a), deals with the same kind of corruption that we use, but instead of predicting the original token before the corruption, Electra predicts whether it was corrupted or not. Electric (Clark et al., 2020b) does not perform masking, but instead uses noise-contrastive loss to learn a representation of text. BART (Lewis et al., 2019) uses both BERT-style denoising and autoregressive modeling for learning representations. All of these methods were originally motivated by improving representation learning of text. There were a few attempts to heuristically decode models like BERT and turn them into text generative models such as in Wang & Cho (2019), however, samples from these models lack coherence.

heading: 5 CONCLUSION
text: In this paper we have proposed a novel non-autoregressive method that operates within the framework of denoising autoencoders and, crucially, unrolls the denoising process during training. Unrolled denoising allows us to achieve state-of-the art results in WMT’14 English-to-German translation task amongst non-AR methods (without distillation from large AR models) and good qualitative results in unconditional text modeling on C4 dataset. We have also qualitatively demonstrated new inpainting capabilities on C4 and GitHub Python data, which opens up new avenues for creative text editing where a human could more naturally collaborate with a machine on writing and even programming.
AUTHOR CONTRIBUTIONS
Nikolay Savinov came up with the idea of unrolled denoising, wrote the first prototype of unconditional generation, contributed to the machine translation codebase and co-led the project. Junyoung Chung wrote most of the machine translation codebase, came up with the length prediction conditioning, suggested argmax-unrolled decoding and co-led the project. Mikołaj Bińkowski came up with theoretical insights about our model, wrote most of the evaluation codebase, implemented model improvements and co-led the project. All three first authors contributed equally to writing the paper. Erich Elsen gave high-level scientific guidance, suggested text in-painting experiments and made substantial edits to the paper draft. Aäron van den Oord originally suggested to look into denoising generative models for text, gave high-level scientific guidance, proposed machine translation experiments and made substantial edits to the paper draft.

heading: ACKNOWLEDGMENTS
text: We would like to thank Jean-Baptiste Alayrac and Miruna Pı̂slar for valuable discussions on the machine translation experiments, William Chan and Chitwan Saharia for sharing distilled translation datasets, Mihaela Rosca and Cyprien de Masson d’Autume for sharing text GAN evaluation code, Yujia Li and Jack Rae for sharing a dataset of Python code from GitHub, Sander Dieleman and Oriol Vinyals for in-depth feedback about our work, Jungo Kasai for responses to our questions about DisCo, Dani Yogatama for sharing knowledge on machine translation evaluation, and Jeff Donahue for feedback and discussions.

heading: A METHOD DETAILS
text: Using Markov property and the form of pt = fθ(·|xt−1), for any t > 0 we obtain pt(xt|x0) = ∑
xt−1∈X pt−1(xt−1|x0)fθ(xt|xt−1) (5)
= Ext−1∼pt−1(·|x0)fθ(xt|xt−1). (6)
By induction, we can keep unrolling the chain for all steps s = t − 1, t − 2, . . . , 1, eventually obtaining
pt(xt|x0) = ∑
x1,...,xt−1 ∈X
t∏ s=1 fθ(xs|xs−1)
= E x1∼p1(·|x0) ···
xt−1∼pt−1(·|xt−2)
[fθ(xt|xt−1)] , (7)
which proves the from of the model distribution (Eq. 1).
A.1 UPPER BOUND ON THE LOSS.
Using Jensen’s inequality we can swap log and expectation in Eq. 7 to obtain
− log pt(x|x0) ≤ −E x1∼p1(·|x0) ···
xt−1∼p(·|xt−2)
[log fθ(x|xt−1)] , (8)
hence
L̃(t)(θ) = −E x∼pdata x0∼q(·|x) [log pt(x|x0)] ≤ −E x∼pdata x0∼q(·|x)
x1∼fθ(·|x0) ...
xt−1∼fθ(·|xt−2)
[log fθ(x|xt−1)] = L(t)(θ). (9)
A.2 RELATION TO VAE
Variational Autoencoders (Kingma & Welling, 2013; Rezende et al., 2014) optimise the expected evidence lower bound (ELBO) on the marginal likelihood
log p(x) ≥ −DKL(q(z|x)‖p0(z)) + Ez∼q(·|x) [log p(x|z)] , (10)
with respect to parameters of both the variational distribution q and generative distribution p, where p0 remains the prior for for the latent variables. Denosing autoencoders limit the optimisation to the second term, assuming a fixed encoder q, which is also the case with SUNDAE: identifying encoder with our corruption distribution and generative distribution p with the distribution of our chain at the last step pT , expected ELBO becomes
Ex∼pdata [log pT (x)] ≥ −Ex∼pdata [DKL(q(z|x)‖p0(z))] + E x∼pdata xc∼q(·|x) [log pt(x|xc)]
≥ −Ex∼pdata [DKL(q(z|x)‖p0(z))]− L(T )(θ). (11)
where the second inequality is a consequence of the upper bound 3. Hence, minimizing L(T ) would lead to maximisation of the lower bound on the model likelihood.
A.3 CORRUPTION FUNCTION
Our corruption function is obtained in stages:
• sample uniform expected corruption proportion α ∼ U([0, 1]), • sample Bernoulli mask, independently for each token m ∼ (Bernoulli(α))N , • sample noise η ∼ (U({1, 2, . . . , v})N ,
• compute corruption xc = cα,m,η(x) = (1 −m) · x +m · η (with multiplications taken element-wise).
The distribution pc(·|x)pdata(x) obtained in this process covers the entire input space X , with probability skewed towards samples whose parts come from pdata. Note that except for corruption proportion, all steps are independent for all individual tokens. Given α, pcα(y|x) can be factorised as ∏
n pcα(y(n)|x(n)) = ∏ n h(y(n))TQαh(x (n)),
where Qα = αI + (1 − α)V is a transition matrix, V is a v × v matrix with all entries equal to 1 v , and h(a) denotes one-hot vector representation of a ∈ {1, 2, . . . , v}. In this context, corruption with rate α corresponds to a multinomial-diffusion forward process with probability α of remaining in the same state, as proposed by (Hoogeboom et al., 2021).

heading: B DETAILS OF TARGET LENGTH PREDICTION
text: We provide a simplified high-level overview of the target length predictor in Figure 3. In the rest of this section we focus on implementation details.
The target length prediction module consists of 6 residual blocks. The length prediction module takes the source encodings hx= encoder(x) as the input, where x is a sequence of source tokens. We first project hx into a vector vx ∈ RdLP , where dLP = 128 is the hidden size of the target length prediction module. We add a source length embedding vector vlen = S(Nsource×dLP )(lx) to vx and obtain the final encoder state representation v′x = vx + vlen, where lx is the number of the source tokens, and Nsource = 128 is the maximum number of the source tokens. The target length prediction module then takes vx as the input and predicts the (downsampled) target length l̃d. We found beneficial downsampling the target lengths for the length prediction by a factor of 2, i.e. ld = dl/2e as compared to exact prediction of l; since the maximum number of target tokens is N=128, the maximum size of the outcome of length prediction is Nd=64.
We prepend an embedding hlen ∈ Rdenc of target length to the source encodings hx ∈ RNsource×denc for each source example x and allow the decoder to attend to it (denc here denotes the dimensionality of the encoder). hlen is obtained by applying an embedding matrix V(Nd×denc) to ground truth (downsampled) length ld at training time, or to predicted one l̃d during sampling. Note that optimisation of the length classification loss does not affect the encoder parameters.

heading: C EFFECTS OF TARGET LENGTH PREDICTION
text: It is shown empirically that the length of the translated results can affect BLEU score. Approximate search methods such as beam search usually incorporate a length normalization term in order to prevent the search algorithm to prefer shorter sentences. This is due to the fact that the scoring function
is usually chosen as the likelihood modelled by the translation system, and beams containing shorter sequences tend to score higher likelihoods. In non-AR models, the end of sequence is not explicitly learned by the models, thus they can benefit by knowing the target sequence length in advance to generating the translation. However, the length of the target sequence cannot be known at inference time, therefore, a model should predict it instead. We show how BLEU scores can differ in WMT’14 translation tasks with and without the target length prediction in Figure 4a-4c.

heading: D EFFECTS OF UNROLLED LOGITS LOSSES
text: We show how varying the number of unrolled logits can affect the MT performance in Figure 4d.

heading: E TRANSLATION ANALYSIS
text: We show how our algorithm works step-by-step for translation in Table 5. It is interesting to observe how the multimodality problem (Gu et al., 2017) gets resolved with more steps. Multimodality usually manifests in repeated tokens in translation, which intrinsically comes from inability to coordinate decisions when independently sampling multiple tokens at once. However, after a few steps all repetitions are corrected, because the model conditions on the previous step and hence can better coordinate further changes.
Figure 5 shows how the model and the BLEU scores evolve during sampling for various temperatures. The model score here refers to the cross-entropy between translation and model logits produced by giving this translation as the input. While for very low temperatures the scores initially improve faster, they are eventually outperformed by scores for higher temperatures, possibly due to low diversity induced by near-deterministic sampling. For very high temperatures, on the other hand, the scores improve slowly, and hence using such temperatures is rather impractical.

heading: F DISTILLATION SCORES OBTAINED WITH TRANSFORMER–BASE
text: In Table 6 we also present results for SUNDAE trained on data distilled from autoregressive Transformer-Base model, which allows direct comparison with similarly obtained Imputer scores (Saharia et al., 2020). The results show that SUNDAE noticeably outperforms Imputer at 8 steps for both EN→DE and DE→EN language pairs.

heading: G SACREBLEU SCORES FOR WMT’14 EXPERIMENTS
text: SacreBLEU (Post, 2018) is a library for computing BLEU score7 that does not require the user to manually tokenize the reference and candidate translations. We report BLEU scores measured with SacreBLEU library in order to provide a reference for future studies. We show all the BLEU scores in Table 7.

heading: H UNCONDITIONAL SAMPLES FOR C4
text: We provide samples for our model trained on C4 dataset as described in the main text (Section 3.2). All of them resemble reasonable-quality internet texts, except perhaps sample #4.
7https://github.com/mjpost/sacrebleu

heading: I UNCONDITIONAL SAMPLES FOR EMNLP2017 NEWS
text: We provide samples from our model trained on EMNLP2017 News dataset in Table 9, accompanying the quantitative results in the main text (Section 3.2). These samples are obtained at temperature 0.8. None of them appear in the training set, so the model does not merely memorize the data it is given. To provide a point of reference, we also include samples from ScratchGAN (d’Autume et al., 2019).

heading: J PSEUDOCODE OF SUNDAE
text: We provide Python-like pseudocode for our method in Listing 1. The main functions are build loss fn, which should be used for training, and sampling fn which should be used for sampling. For simplicity, we consider the decoder-only setup as in unconditional generation experiments (Section 3.2).
1 def get_random_text(shape, vocab_size): 2 random_text = rand_int(shape, minval=0, maxval=vocab_size) 3 return random_text 4
5 6 def corrupt_text(batched_text, vocab_size): 7 corruption_prob_per_sequence = rand_uniform([batched_text.shape[0], 1]) 8 rand = rand_uniform(batched_text.shape) 9 mask = rand < corruption_prob_per_sequence
10 random_text = get_random_text(batched_text.shape, vocab_size) 11 corrupted = mask * random_text + (1 - mask) * batched_text 12 return corrupted 13
14 15 def build_logits_fn(vocab_size, n_unrolled_steps, enable_sampling): 16 17 def logits_fn(batched_text): 18 model = Transformer(vocab_size, use_causal_mask=False) 19 20 def fn(input_batched_text): 21 logits = model(input_batched_text) 22 return logits 23 24 def unrolled_fn(input_batched_text): 25 samples = corrupt_text(input_batched_text, vocab_size) 26 all_logits = [] 27 for _ in range(n_unrolled_steps):
28 logits = fn(samples) 29 samples = stop_grad(rand_categorical(logits)) 30 all_logits += [logits] 31 final_logits = concatenate(all_logits, axis=0) 32 return final_logits 33 34 if enable_sampling: 35 return fn(batched_text) 36 else: 37 return unrolled_fn(batched_text) 38 39 return logits_fn 40
41 42 def build_loss_fn(vocab_size, n_unrolled_steps=2): 43 logits_fn = build_logits_fn( 44 vocab_size, n_unrolled_steps, enable_sampling=False) 45 46 def loss_fn(batched_text): 47 logits = logits_fn(batched_text) 48 # repeat batched_text to fit unrolled logits 49 targets = concatenate([batched_text] * n_unrolled_steps, axis=0) 50 one_hot_targets = one_hot(targets, vocab_size) 51 loss_per_token = -sum( 52 one_hot_targets * log_softmax(logits), axis=-1) 53 loss = mean(loss_per_token) 54 return loss 55 56 return loss_fn 57
58 59 def sampling_fn(logits_fn, steps, temperature, batch_size, 60 sequence_length, vocab_size): 61 batched_text = get_random_text( 62 shape=[batch_size, sequence_length], vocab_size) 63 for _ in range(steps): 64 logits = logits_fn(batched_text) 65 samples = rand_categorical(logits / temperature) 66 batched_text = samples 67 return batched_text
Listing 1: Pseudocode of SUNDAE.

Review summaries:
This paper presents a training method of non-autoregressive text generation that is trained from scratch without distillation. The method is motivated as an unrolled denoising approach, which starts with a sequence of random tokens and transforms it in multiple stages to a valid sequence. To make training feasible, the proposed approach starts with a corrupted target sentence instead of a totally random sequence and does only two steps of unrolling. The training loss is an approximation of this process and is computed as a sum of loss on direct denoising and two steps of denoising. During inference, it starts with a random sequence and applies several steps of denoising to arrive at the desired sequence. Several heuristics are also applied to make generation feasible. The main result of the paper on non-autoregressive machine translation on which the authors show decent improvements over baselines on non-distilled settings however slightly underperforms AR distilled models. The authors also show examples for the case of unconditional generation and text in-painting

This paper proposes a new training technique for an iterative non-autoregressive (NAR) model and achieves significant improvements over the previous model. Inspired by the success of the diffusion model, the author formalizes the iterative NAR decoding process with a denoising autoencoder framework, then unrolls the denoising process during training. The proposed method (SUNDAE) achieves new state-of-art performances in machine translation by deep iteration decoding and parallel reranking. More experiments on unconditional text generation tasks further show its promising capability.

The paper proposes an unrolled denoising method for non-autoregressive text generation. Specifically, the method first corrupts the original input and then re-generate uncertain tokens in it, and the result is taken as the output of the denoising autoencoder to compute loss with the target. Experiments are conducted on machine translation and unconditional text generation. The method achieves good results on MT without utilizing knowledge distillation.

The present paper proposes a Step-untolled denoising autoencoder (SUNDAE), a NAR text generation model which marries de-noising auto-encoder and Markov chain models. The major contribution of this paper is the unrolled denoising training scheme, which is a training step method. In the first step, the original text is randomly corrupted in the same way as what the denoising autoencoder did. Subsequently, in the second step, the noisy text from the first step is first denoised by sampling from the model, and the denoised text is then used as a sample for training the denoising autoencoder.   The proposed model has been tested on two tasks that involve text generation: machine translation as well as language modeling.