Step-unrolled Denoising Autoencoders for Text Generation

1 INTRODUCTION . 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 . 2 METHOD . 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 . 2.1 TRAINING WITH UNROLLED DENOISING . 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 . 2.2 SAMPLING . 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 . 3 EXPERIMENTS . 3.1 MACHINE TRANSLATION . 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 .

Summary: 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.