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 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