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: Autoregressive Diffusion Models (ARDM) are discrete diffusion models that extend D3PMs and generalize order agnostic auto-regressive models (OA-ARM) . Similarly to other diffusion models, training OA-ARDMs only require evaluating a single transition step (i.e $p(x_{\sigma(t:t+k)} | x_{\sigma(<t}))$, where $\sigma$ is a random permutation). In contrast to traditional ARMs, OA-ARMs are trained to generate k steps $p(x_{\sigma(t:t+k)})$ at each transition.  Furthermore, the authors supplement the OA-ARM with two additional contributions: - factoring the generative model across the variable *depth* (named *depth upscaling*), which allows unfolding the generative process into an even greater number of steps ($S \times T$ steps instead of $T$) - an algorithm to parallelize generation using dynamic programming (allows trading quality for generation speed)  The authors demonstrate the effectiveness of OA-ARM on character-level text modelling (text8), image modelling/compression (CIFAR-10).