Autoregressive Diffusion Models

1 INTRODUCTION . Deep generative models have made great progress in modelling different sources of data , such as images , text and audio . These models have a wide variety of applications , such as denoising , inpainting , translating and representation learning . A popular type of likelihood-based models are Autoregressive Models ( ARMs ) . ARMs model a high-dimensional joint distribution as a factorization of conditionals using the probability chain rule . Although very effective , ARMs require a pre-specified order in which to generate data , which may not be an obvious choice for some data modalities , for example images . Further , although the likelihood of ARMs can be retrieved with a single neural network call , sampling from a model requires the same number of network calls as the dimensionality of the data . Recently , modern probabilistic diffusion models have introduced a new training paradigm : Instead of optimizing the entire likelihood of a datapoint , a component of the likelihood bound can be sampled and optimized instead . Works on diffusion on discrete spaces ( Sohl-Dickstein et al. , 2015 ; ∗Work done during as research intern at Google Brain . Hoogeboom et al. , 2021 ; Austin et al. , 2021 ) describe a discrete destruction process for which the inverse generative process is learned with categorical distributions . However , the length of these processes may need to be large to attain good performance , which leads to a large number of network calls to sample from or evaluate the likelihood with discrete diffusion . In this work we introduce Autoregressive Diffusion Models ( ARDMs ) , a variant of autoregressive models that learns to generate in any order . ARDMs generalize order agnostic autoregressive models and discrete diffusion models . We show that ARDMs have several benefits : In contrast to standard ARMs , they impose no architectural constraints on the neural networks used to predict the distribution parameters . Further , ARDMs require significantly fewer steps than absorbing models to attain the same performance . In addition , using dynamic programming approaches developed for diffusion models , ARDMs can be parallelized to generate multiple tokens simultaneously without a substantial reduction in performance . Empirically we demonstrate that ARDMs perform similarly to or better than discrete diffusion models while being more efficient in modelling steps . The main contributions of this paper can be summarized as follows : 1 ) We introduce ARDMs , a variant of order-agnostic ARMs which include the ability to upscale variables . 2 ) We derive an equivalence between ARDMs and absorbing diffusion under a continuous time limit . 3 ) We show that ARDMs can have parallelized inference and generation processes , a property that among other things admits competitive lossless compression with a modest number of network calls . 2 BACKGROUND . ARMs factorize a multivariate distribution into a product of D univariate distributions using the probability chain rule . In this case the log-likelihood of such as model is given by : log p ( x ) = D∑ t=1 log p ( xt|x < t ) , ( 1 ) where x < t is shorthand for x1 , x2 , . . . , xt−1 . ARMs are trained by ensuring that the neural network has a triangular dependency structure , for instance implemented via causal masking . Although this allows parallelized computation of the likelihood for all conditional distributions at once , to sample from the model it requires D iterative sampling steps x1 ∼ p ( x1 ) , x2 ∼ p ( x2|x1 ) towards xD ∼ p ( xD|x1 , x2 , . . . , xD−1 ) . Order Agnostic ARMs Order Agnostic ARMs ( OA-ARMs ) ( Uria et al. , 2014 ) generate variables with a random ordering σ ∈ SD , where SD represents a set of all permutations of the integers 1 , . . . , D. The log-likelihood of this model is given by : log p ( x ) ≥ Eσ∼U ( SD ) D∑ t=1 log p ( xσ ( t ) |xσ ( < t ) ) . ( 2 ) This can be seen as a latent variable model and the log-likelihood is derived via Jensen ’ s inequality : log p ( x ) = logEσ∼U ( SD ) p ( x|σ ) ≥ Eσ∼U ( SD ) log p ( x|σ ) . Mind that hereafter we leave out σ in our notation to avoid clutter . One approach to train OA-ARMs is the procedure as described by Yang et al . ( 2019 ) for XLNet . It takes a permutation equivariant network such as a Transformer that is causally masked . Then , inputs are permuted and outputs are permuted back according to a given order , which models the sequence in that specific order . However , such approaches typically suffer in likelihood score and can not be combined with simple non-equivariant transformations such as convolutional layers . Discrete Diffusion Discrete diffusion models define a destruction process on discrete data . An example is absorbing diffusion ( Austin et al. , 2021 ) , for which each variable has a probability of decaying to an absorbing state . The opposite process to the destruction process is the learned generative process . This generative process models the distribution over variables that are currently absorbed , and generates these with a probability . 3 AUTOREGRESSIVE DIFFUSION MODELS . We introduce Autoregressive Diffusion Models ( ARDMs ) . ARDMs generate variables in an arbitrary order , one by one . Further , ARDMs are able to upscale variables , such as the bit values of a Algorithm 1 Sampling from OA-ARDMs Input : Network f Output : Sample x Initialize x = 0 Sample σ ∼ U ( SD ) for t in { 1 , . . . , D } do m← ( σ < t ) and n← ( σ = t ) x′ ∼ C ( x|f ( m x ) ) x← ( 1− n ) x+ n x′ Algorithm 2 Optimizing OA-ARDMs Input : Datapoint x , Network f Output : ELBO L Sample t ∼ U ( 1 , . . . , D ) Sample σ ∼ U ( SD ) Computem← ( σ < t ) l← ( 1−m ) log C ( x|f ( m x ) ) Lt ← 1D−t+1 sum ( l ) L ← D · Lt pixel . Unlike standard ARMs , ARDMs are trained on a single step in the objective , as in modern diffusion models . In addition , both sampling and inference of ARDMs can be parallelized using dynamic programming with minimal degradation in log-likelihood . Order Agnostic ARDMs The main difficulty of parameterizing an autoregressive model from an engineering perspective , is the need to enforce the triangular or causal dependence . Especially for 2D signals , this triangular dependence is difficult to enforce for arbitrary orders ( Jain et al. , 2020 ) and tedious design is needed for multi-scale architectures ( Salimans et al. , 2017 ) . To relax this requirement , we take inspiration from modern diffusion-based generative models . Using these insights , we derive an objective that is only optimized for a single step at a time . Starting at Equation 2 , a different objective for an order agnostic ARM can be derived , by replacing the summation over t by an expectation that is appropriately re-weighted : log p ( x ) ≥ Eσ∼U ( SD ) ∑D t=1 log p ( xσ ( t ) |xσ ( < t ) ) = Eσ∼U ( SD ) D · Et∼U ( 1 , ... , D ) log p ( xσ ( t ) |xσ ( < t ) ) = D · Et∼U ( 1 , ... , D ) Eσ∼U ( SD ) 1 D − t+ 1 ∑ k∈σ ( ≥t ) log p ( xk|xσ ( < t ) ) Compactly , we can write the expected lower bound as : log p ( x ) ≥ Et∼U ( 1 , ... , D ) [ D · Lt ] , where Lt = 1 D − t+ 1 Eσ∼U ( SD ) ∑ k∈σ ( ≥t ) log p ( xk|xσ ( < t ) ) . Here the term Lt represents the likelihood component for step t. Importantly , we do not need to optimize for all Lt terms of a datapoint simultaneously . Instead , for each datapoint in a minibatch a single Lt term is optimized where t is sampled from a uniform distribution . This objective was originally proposed by Uria et al . ( 2014 ) to train order-agnostic ARMs . We will develop ARDMs starting from this perspective and refer to the special case of an order-agnostic ARM , as an order agnostic ARDM ( OA-ARDM ) . Interestingly , each Lt component can be seen as a BERT-like training objective ( Devlin et al. , 2019 ) , where exactly D − t + 1 tokens are masked and subsequently predicted . Therefore , an OA-ARDM is trained as a collection of D BERTs with loss terms Lt , which contain the reweighting term 1D−t+1 . Another insight is that this generative process is very similar to absorbing diffusion , where the model aims to generate absorbed ( or masked ) variables . In certain situations we might want to refer to loss terms instead of likelihood terms , so we define Lt = −Lt . Parametrization We desire a parametrization for the model distribution log p ( xk|xσ ( < t ) ) for k ∈ σ ( ≥ t ) for all σ and t. For each σ and t it is in principle allowed to have an entirely new neural network . However , this would be very inconvenient as the number of t grows as O ( D ) and the number of σ grows as O ( D ! ) . Instead , a single neural network is utilized and shared for different σ and t. This is implemented by masking variables at the input , and predicting those at the output . To be precise , we let x ∈ X = { 1 , 2 , . . . , K } D represent discrete variables with K classes and a neural network f : X → RD×K that outputs probability vectors for each dimension . Conditioning is done via masking : For a given permutation array σ , we compute the elementwise comparison m = σ < t which produces a Boolean mask . The mask is then used by predicting θ = f ( m x ) , where denotes element-wise multiplication . For each location k ∈ σ ( ≥ t ) , the log probability vectors θk are used . Letting C ( xk|θk ) denote a categorical distribution over xk with class probabilities θk , we choose to model log ( xk|xσ ( < t ) ) = log C ( xk|θk ) . The locations of these relevant indices k ∈ σ ( ≥ t ) are retrieved by using the opposite mask 1 −m . The procedure to sample and optimize an ARDM with this parametrization are given in Algorithms 1 and 2 . A training step is visualized in Figure 2 . Note that opposed to Figure 1 where only a single output was used per step , in the train step all dimensions that were masked are predicted simultaneously . Therefore , there are multiple variables to predict which ensures there is sufficient signal to optimize the model . For clarity some of the implementation details have been left out in the explanation above . However , given that these are important for practical implementations , they are specified in the following . The input to the function f may be different depending on the data modality : For images and audio , the mask is applied to the input so that values are zero after feature normalization . The mask itself is also concatenated to the input as an input representation , which allows the model to identify whether a value is actually zero , or the value is in the absorbing state zero . For language the input representation is augmented and absorbed values are instead set to a new class K +1 , in which case there is no need to provide the mask itself as input to the model . More generally , we can represent the masked state as an absorbing state vector a which has the same shape as x but only contains a pre-specified value . The input to the network is then not the masked m x , but instead the combinationm x+ ( 1−m ) a . In addition , the network f may also take the time component t as input as is typically done in diffusion models ( Ho et al. , 2020 ) . In summary the network takes some additional inputs as θ = f ( i , m , t ) where i =m x + ( 1 −m ) a and the processing of x may be different depending on the type of data .

Summary: This work introduces a new model class combining elements from autoregressive and discrete diffusion models. The new approach is named Autoregressive Diffusion Models (ARDMs). The authors show good performance compared to alternative generative models, in particular in lossless image compression. The authors demonstrate that the new model has particular benefits compared to the autoregressive and discrete diffusion models in terms implementation efficiency, scalability and parallelization.