Autoregressive Diffusion Models
Authors: Emiel Hoogeboom
Alexey A. Gritsenko
Jasmijn Bastings
Ben Poole
Rianne van den Berg
Tim Salimans
Venue: ICLR 2022 Poster (2022)
Keywords: diffusion
autoregressive models
lossless compression

Abstract: We introduce Autoregressive Diffusion Models (ARDMs), a model class encompassing and generalizing order-agnostic autoregressive models (Uria et al., 2014) and absorbing discrete diffusion (Austin et al., 2021), which we show are special cases of ARDMs under mild assumptions. ARDMs are simple to implement and easy to train. Unlike standard ARMs, they do not require causal masking of model representations, and can be trained using an efficient objective similar to modern probabilistic diffusion models that scales favourably to highly-dimensional data. At test time, ARDMs support parallel generation which can be adapted to fit any given generation budget. We find that ARDMs require significantly fewer steps than discrete diffusion models to attain the same performance. Finally, we apply ARDMs to lossless compression, and show that they are uniquely suited to this task. Contrary to existing approaches based on bits-back coding, ARDMs obtain compelling results not only on complete datasets, but also on compressing single data points. Moreover, this can be done using a modest number of network calls for (de)compression due to the model's adaptable parallel generation.

Published as a conference paper at ICLR 2022
AUTOREGRESSIVE DIFFUSION MODELS
Emiel Hoogeboom; Alexey A. Gritsenko, Jasmijn Bastings, Ben Poole, Rianne den Berg, Tim Salimans
Google Research
wwo.geboom@uva.... nl, fagritsennie pooleb, bastings, sal saliimans)@goggle.com ri annevdberg@@gai...... com
ABSTRACT
We introduce Autoregressive Diffusion Models (ARDMs), model class encom- passing and generalizing order-agnostic autoregressive models (Uria al., 2014) and absorbing discrete diffusion (Austin et al., 2021), which we show are special cases of ARDMs under mild assumptions. ARDMs are simple to implement and easy to train. Unlike standard ARMs, they do not require causal masking of model representations, and can be trained using an efficient objective similar to modern probabilistic diffusion models that scales favourably to highly-dimensional data. At test time, ARDMs support parallel generation which can be adapted to fit any given generation budget. We find that ARDMs require significantly fewer steps than discrete diffusion models to attain the same performance. Finally, we apply ARDMs to lossless compression, and show that they uniquely suited to this task. Contrary to existing approaches based on bits-back coding, ARDMs obtain compelling results not only on complete datasets, but also compressing single data points. Moreover, this can be done using a modest number of network calls for (de compresson due the model's adaptable parallel generation.
INTRODUCTION
Deep generative models have made great progress in modelling different sources of data, such as images, text and audio. These models have wide variety of applications, such as denoising, in- painting, translating and representation learning. A popular type of likelihood-based models are Autoregressive Models (ARMs). ARMs model a high-dimensional joint distribution as a factor- ization of conditionals using the probability chain rule. Although very effective, ARMs require pre-specified order in which generate data, which may be an obvious choice for some data modalities, for example images. Further, although the likelihood of ARMs can be retrieved with single neural network call, sampling from 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 datapoint, component of the likelihood bound can be sampled and optimized instead. Works on diffusion on discrete spaces (Sohl-Dickstecr et al. 2015
"Work done during research intern Google Brain.
86
p(tan)
Figure 1: Generation of Autoregressive Diffusion Models for the generation order = (3, 2, 4). Filled circles the first and third layers represent respectively the input and output variables, and the middle layer represents internal activations of the network.

---

Published as a conference paper at ICLR 2022
Hoogeboom al. 2021 Austin al. 2021 describe 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 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 introduce Autoregressive Diffusion Models (ARDMs), 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 standard ARMs, they impose no architectural constraints on the neural networks used to predict the dis- tribution 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 substantial reduction in performance. Empirically we demonstrate that ARDMs perform similarly to or better than discrete diffusion models while being more efficient 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 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 modest number of network calls.
2
BACKGROUND
ARMs factorize multivariate distribution into product of D univariate distributions using the probability chain rule. In this case the higeliill of such model given by:
log p(x)) Σ log pa
where Ø<t is shorthand for F2,422.... 21-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 once, to sample from the model it requires D iterative sampling steps x1 p(x1), p(x2/z1) towards uD ~ p(xp|=1, ,2,,.,,,--... ID-)).
Order Agnostic ARMs (OA-ARMs) (Uria et al. 2014) generate variables with a random ordering a E Sp, where Sn represents a set of all permutations of the integers 1,..., D. The likeiiildhood of this model is given by:
Order Agnostic ARMs
log p(z) 2 Ea-u(Sn)
Σ log Pogoooocoaaeee
This can be seen as latent variable model and the log likeliholo is derived via Jensen's 's inequality:
logp(z>) = log E. (ptp:/ssssppppp/// p(z))) 2 Eo- log p(z)()).
Mind that hereafter we leave out o in our notation to avoid clutter. One approach to train OA-ARMs the procedure described by Yang al. (2019) for XLNet. It takes permutation equivariant network such as a Transformer that is causally masked. Then, inputs are permuted and outputs are permuted back according to given order, which models the sequence in that specific order. However, such approaches typically suffer in likelihood score and cannot be combined with simple non-equivariant transformations such convolutional layers.
Discrete Diffusion Discrete diffusion models define destruction process on discrete data. An example 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 gen- erutive process. This generative process models the distribution over variables that are currently absorbed, and generates these with probability.
3 AUTOREGRESSIVE DIFFUSION MODELS
We introduce Autoregressive Diffusson Models (ARDMs). ARDMs generate variables in an arbi- trary order, one by one. Further, ARDMs are able upscale variables, such the bit values of
2

---

Published as a conference paper at ICLR 2022
Algorithm 1 Sampling from OA-ARDMs
Algorithm 2 Optimizing OA-ARDMs
Input: Network
Input: Datapoint æ, Network f
Output: Sample x.:
Output: ELBO L
Initialize x = 0
Sample t ~ ...,...,
Sample σ ~ (Sp)
Sample o ~ (((p))
for L in {I...., D] do
Compute m (o < <1)
m 5 (a < 1) and n t (a=
1 + m) @ log C(x: [(m 0 x))
x' C( @ x))
L.
sum(l)
x - (1
+ L.
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 lielliieldoood.
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 21) 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- insights, we derive an objective that is only optimized for single step at a time. Starting at Equa- tion 2 a different objective for an order agnostic ARM can be derived, by replacing the summation over by expectation that is appropriately re- weighte::::
generative models. Using these
log p(x:) 2 F(o Ei log (0)
= ( D. Brade D) log ((Fo(t)) (o(()))
D (L,, Bar D - + ExEr(20) log logp(ex/=a(e))
Compactly, we can write the expected lower bound as:
log p(x) 2 EE Ci], where L; D-1 E3 ((Sp)
log p(æx (o(()))
kEa
Here the term C. represents the likelihood component for step t. Importantly, we do not need to optimize for all L. terms of a datapoint simultaneously. Instead, for each datapoint in minibatch a single L, term is optimized where 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 ag- nostic ARDM (OA-ARDM). Interestingly, each L, component can be seen BERT-like training objective (Devlin et al. 2019h, where exactly D t + 1 tokens are masked and subsequently pre- dicted. Therefore, OA-ARDM is trained as collection of D BERTs with loss terms L,, which contain the reweighting term D- Another insight is that this generative process is very similar to absorbing diffusion, where the model aims to generate absorbed (or masked) variables. certain situations might want to refer to loss terms instead of likelihood terms, define I Lt.
Parametrization We desire a parametrization for the model distribution log p(z (<t)) for k (2 t) for all a and f. For each a and é it is in principle allowed to have an entirely new neural network. However, this would be very inconvenient as the number of L grows as O(D) and the number of grows as OD!). Instead, a single neural network is utilized and shared for dif- ferent a and t. This is implemented by masking vari- ables at the input, and predicting those at the output. To be precise, we let æ E X = 21, KJD represent discrete variables with K classes and neural network f X + R D that outputs probability vectors for each dimension. Conditioning is done via masking: For a given permutation array o, we compute the element- wise comparison m a which produces Boolean
J
122
J3
Figure 2: ARDM training step. This step optimizes for step t = 2 for all possible permutations o simultaneously which satisfy a(1) 3.
3

---

Published as a conference paper at ICLR 2022
mask. The mask is then used by predicting 0 = f(m @ x). where denotes element wise multiplication. For each location &o E o2 t), the log probability vectors 0s are used. Letting C(z)(0g) denote categorical distribution over 24 with class probabilities og. choose to model log(2z/ (T)(<))) = log Car (x)) The locations of these relevant indices A E a( 1) are retrieved sample and optimize parametrization are given in Algorithms Maaanss[aman and 2 A training step is visualized in Figure Note that opposed Figure I where only 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.
by using the opposite mask 1 m. The procedure
ARDM with this
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 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 input representation, which allows the model to identify
is also concatenated to the input 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 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 a pre-specified value. The input to the network is then not the masked m () æ, but instead the combination m @ x + (1 - m) @ a. In addition, the network 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 0- f(i m, where i m may be different depending on the type of data.
absorbing state vector a which has the same shape ær but only contains
m) @ a and the processing of
3.1 PARALLELIZED ARDMS
An important property of our parametrization is that the distribution over multiple variables is pre- dicted at the same time. In this section, we will leverage this parameterization to allow parallel independent generation of variables. Essentially, we desire distributions over Zo(t+R) for positive k while conditioning only on Fo(((). First we make an observation regarding a connection between predicting future variables and our likelihood terms: For k- = 1,2, D l:
B,, log ((()=(+))) = E., log p(#ac)) [o(ct))] = C.,
(3)
due to the uniform expectation over permutations. In other words, it does not matter which step t + k the model predicts. in expectation these all have the same associated likelihood. As result, order agnostic generation of & tokens independently, starting from the f-th variable will result in a log-probability contribution of k: L, in single step, whereas the traditional approach would take
k steps the cost of E Li+i- This knowledge is sufficient to construct dynamic programming algorithm as described by Watson et al. (2021) to compute how many parallel steps to take at which moment, given budget. Since dynamic programming is typically described from minimization perspective we define the loss component L, - Lt, which is measured in bits. In terms of loss, generating i variables at timestep & will cost k La bits. Further, we define the transition cost matrix L1,4+ = k- Le for positive integers k and Li+ktt 0 otherwise. So Lt,++ exactly describes how much costs to model the next & variables in parallel starting at the t-th position for all relevant t and k. Using this transition cost matrix, the dynamic programming algorithm can be utilized to find which steps should be parallelized. For instance, in the example in Figure B| hypothetical 20-step problem is given budget of steps. Typically, the algorithm will spend more steps on regions with large differences between L, components and fewer steps on regions where the Le components are
Figure 3: Loss components for Parallelized ARDMs using a budget of 5 steps for problem of 20 steps. Left: individual loss component for every step. Right: parallelized policy extracted from the dynamic programming algorithm. Components of the same height are modelled simultaneously, so they are inferred and generated in parallel.
4

---

Published as a conference paper at ICLR 2022
approximately equal. Parallelizing an ARDM may ineur some cost, for well-calibrated model:
(4) but can be traded off for faster generation because fewer steps are used. In other words, the loss com- ponents L, are monotonically decreasing over t and parallelizing a model incurs a cost, which the algorithm aims to minimize. Recall that this is under the assumption that model is well-calibrated, which is observed in practice. See Figure 3 for an example of a parallelized schedule.
L. = E. (espp((() & E, [log p(To(e+1)) [ = Li+1,
3.2 DEPTH UPSCALING ARDMS
Order agnostic ARDMs learn to generate variables in random order. As a result, decisions on very detailed information (such as the least significant bit in an image) are modelled relatively early in the generative process. Instead, we can structure the process into stages, where for each stage a refinement of the variable is generated. We refer to this process as upscaling. For example, instead of generating an entire 256-categoorial variables at once, we can first generate the most significant bit, and then the subsequent bits in order of significance. To define the process, it is helpful to first imagine the opposite process to upscaling, which is the destructive process downscaling. Formally, we can define maps via transition matrices p(0) that define how a data variable downscales from its data value towards a common absorbing state. For simplicity assume single dimensional variables at this moment. Denote the absorbing state as a one-hot vector (0)),, where all values are zero except at a prespecified index so that xa = 1. From a diffusion perspective, upscaling is complementary downscaling destruction process where each variable decays by zeroing its least significant bit.
Let P(1), ((S) define sequence of downscaling maps that for any categorical one-hot data variable (9)) E {0, 1)K it holds that p(1) .............................. p(S) (8) xx))) In other words, any category K decays to the common absorbing state after S downscaling maps. We now define the upscal- ing generative process by learning the reverse of the downscaling map, specifically by modelling p(xx ((x (S)1) . ((x (x) (((z(()) The transition matrices allow easy transitions between
the different stage variable x via the following rules:
x p((+1)) = (++1)) (SS),
where P p(() p(++1) .....())
The matrices P (1+1)
are computed as cumulative matrix multiplications, and allow a transition di- rectly from a datapoint (SS) to the corresponding downscaled variable ( This is particularly useful during training, where the model will only be optimized for single specific stage per data-
point. For implementations is generally useful to define p(5+1) = I as an identity matrix so that the above equation also holds when s = S. To train Upscale ARDMs, we can extend Algorithm [ In addition to sampling a timestep f, stage i ~ (111...., S) to optimize sampled. For this par- ticular stage, the ARDM models ((x)(*) |x: [=))) by sampling a permutation a within the stage and represents a stage that modelled with an order agnostic ARDM. This highlights an interesting property of ARDMs: Although sampling from a model may take up to D S steps, the training complexity has not changed by modelling multiple stages. As a result, one can experiment with adding arbitrary number of stages with- out an increase in computational complexity during training. Depth upscaling is reminiscent of the upscaling networks proposed in (Kalchbrenner et al. 2018 Menick & Kalchbrenner 2019). with the important differences that Upscale ARDMs model the variables order-agnostic and only utilize single neural network to parametrize all stages. For a more detailed explanation that includes the dimensionality of the variables (x)*)) } see Appendix
a timestep within the stage. Every term ((xx())
Bit Upscaling Depth upscaling is easiest demonstrated with an example, bir-upscaling. Consider the task of generating standard image with pixel values (0.... 255} so that an image with D dimensions can be represented by (8) E {0, 1} Dx256 in onehot notation. Imagine downscaling process defined by the function that removes the least significant bits: lsb, ) = [k/2*] - via this function can define our transition matrices:
p(8+1->) = if = lsb, (k) and & Im(lsb,-1)
otherwise p|8+1 1 0,
PLA
where (P())} are indexed starting at zero Notice that in this case 8 stages map every value to the absorbing state O. because (k) = 0 for any l E {O...., 255}. See Figure 4 for a visualization of such matrices for problem with less categories.
5

---

Two methods can be used to parametrize the distri- Published as a conference paper at ICLR 2022
P()
(
P P
P
P
Figure 4: Bit upscaling matrices for data with eight categories and hence three stages, meaning = 3. Entries that are white represent zeros, coloured entries represent ones.
Depth upscaling is not confined to bits, and indeed more general formulation is given by the downscaling map = [k/b*] 8f, for branching factor b. When b is set to 2, the bit upscaling transitions are retrieved as special case. When b is set to higher values, then variables can be generated in fewer stages, 5 = [log, [ to be exact. This allows for a unique trade-off between the number of steps the model takes and the complexity that each modelling step inhibits. Other hand-crafted transitions are also imaginable, not excluding transitions that augment the space to new categories, but these not considered in this paper.
Parametrization of the Upscaling Distributions Although is now defined how a datapoint (S) downscales to (S--1), 1), ................................ (1) and to its absorbing state (0)), is not immediately clear to parametrize the distributions [xx=))//((000)))))))))))))))))
bution. The first is direct parametrization. In the example of the bit-upscaling model above, one models the s-th significant bits given the (s - 1) significant bits. The direct parametrization generally more computationally efficient, as it requires only distribution parameter outputs that are relevant for the current stage. This is especially useful when the number of classes is large (such as with audio, which has 2 classes). However, it can be somewhat tedious figure out exactly which classes are relevant and should be modelled.
data parametrization which is similar to the parametrization in Austin al. 2021) An important difference with their work is that the downscaling matrices p(+) represent deterministic maps while theirs represent stochastic process. For this parametrization, the network outputs a probability vector 0 that matches the shape of the data x((5), which transformed and converted to the relevant probabilities in stage via:
Alternatively
((+) P(s) x
(1+1)
where p(zr(0)/z)0-11) = C( (0(+))
The advantage this parametrization that one only has to define the transition matrices (P(.) ). As result, the appropriate probabilities can be automatically computed which is ideal for experi- mentation with new downscaling processes. The disadvantage may be that modelling full probabil- ity vectors for problems with high number of classes may be expensive and not even fit in memory. Empirically in our experiments on image data we find that there is no meaningful performance dif- ference between the parametrizations.
4 RELATED WORK
Autoregressive Models Autoregressive Models (ARMs) factorize joint distribution into a prod- of conditional distributions (Bengio & Bengio 2000 Larochelle & Murray/ (011)) Advances in deep learning have allowed tremendous progress on various modalities, such as images (van den Oord et al. 2016b] Child al. 2019. i.a.), audio (van den Oord et al., 2016a| Kalchbrenner al.] 2018 i.a.), and text (Bengio et al. 20033 Graves, 2013 Melis et al. 2018 Merity et al. 2018 Brown ct al. 2020 i.a.,, where for the latter they are referred to as language models.
Although evaluating the likelihood of datapoint is generally efficient with ARMs, sampling re- quires iterative process with as many network calls as the dimensionality of the data. Parallelized
ARM approaches often rely either cutting many dependencies in the conditioning (Reed et al. 2017) which tend to suffer in log-likee bigeliiel.oodd. Alternatively, ARMs can be solved using ixxed-point
iteration algorithms in fewer steps without sacrificing likelihooddd (Wiggers & Hoogeboom 2020
Song et al. 2021). but these methods typically still require a large number of steps to converge.
Order agnostic sequence modelling was introduced in (Uria et al. 2014) and utilizes the same objec- tive as AO-ARDMs to optimize the model, operating by masking and predicting variables. Differ- from their method, ARDMs have more choices in absorbing states, parallelization support and

---

Published as a conference paper at ICLR 2022
Table 1: Order Agnostic model performance (in bpc) on the text8 dataset. The OA-Transformer learns arbitrary orders by permuting inputs and outputs described in XLNet. A Transformer learning only a single order achieves 1.35 bpc.
Order Agnostic modelling perfor- mance (in bpd) on the CIFAR-IO dataset. The upscaling model generates groups of four most significant categories, equivalent to 2 bits at a time.
Table 2:
Model
Steps
NLL
Model
Steps
NLL
OA- DA-Transformer
250
1.64
ARDM-OA
3072
2.69 0.uss
D3PM-uniform
1000
1.61 10.20
Parallel ARDM-OA
50
2,74
JsPM-absorhing
1000
1.45 #0,00
DSPM-absorbing
256
1.47
Parallel ARDM-Upscale 4
ARDM-Upscoie 4
4 x 3072
4 x 50
2.64
2.68
0.002
OA-ARDM (ours)
250
1.43
#00001
D3PM Absorbing
1000
4.40
D3PM-absorhing
20
1,56
#00000
D3PM Gaussian
1000
3.44
Parallelized OA-ARDM (ours)
20
1.51
#:0007
depth upscaling techniques, in addition to modern advances to fit larger scale data. An alternative approach for order agnostic modelling is via causally masked permutation equivariant models such as Transformers (Yang et al.) [219 Allorm & Nguyen| (021), but these have had limited success in likelihood-basee tasks. In (Ghazvininejad et al. 2019) a mask predict method is proposed, although it does not contain a likelihood analysis. In other work, mixtures of ARMs over certain orders trained by overriding convolutional routines for masking (Jain et al. 2020). In a different context in (Liu et al., 2018)) graph edges connected node are modelled without order. However, the model is not entirely order agnostic because it models edges centered around focus nodes.
Diffusion Models Diffusion models learn to denoise a Gaussian base distribution into the dis- tribution of the data via a chain of latent variables (Song & Ermon 2019 Sohh-Dickktein et al.)
2015 Ho et al. 2020). Diffusion and score-matching methods have shown large improvements in image (Dhariwal & Nichol (021) and audio sample quality (Chen et al. 2220 Kong et al. (021), well as Tikelihood improvements with variational interpretations of diffusion models (Kingma al.)
2021 Huang et al. [202I). Although faster sampling schedules for continuous diffusion models have been explored (Jolicouur Martteuur-aaannnauu et al. 202T; Kong & Ping 202T). little is known about shorter generative processes for discrete diffusion.
Discrete diffusion models operate directly on discrete spaces. In Sohl-Dickstein et al. (2015) dif- fusion for binary data was proposed which was extended for categorical data in Hoogeboom et al. (20T)) Whereas these approaches uniformly resample categories, in Austin et al. (2021) a wide variety of transition distributions was proposed. This work finds that absorbing diffusion produces the best performing models in log-likelihood for text data, but these models still demand a large number of steps. OA-ARDMs are equivalent to the infinite time limit of absorbing diffusion, which makes them maximally expressive. Simultaneously, ARDMs upper bound the number of steps to the dimensionality of the data. More details on the connections between these model types are in Appendix C Other discrete diffusion processes have been explored in (Johnson et al. 2021).
RESULTS
Order Agnostic Modelling To better understand how ARDMs compare to other order agnostic generative models, we study their performance on a character modelling task using the text8 dataset (Mahoney] 2011). ARDMs are compared to D3PMs that model the inverse absorbing diffusion process (Austin al. 2021), and causally masked Transformers that are directly optimized on ran- domly permuted sequences as done in XLNet (Yang et al. 2019). The different methods all use the same underlying neural network architecture which the Transformer used in (Austin al.) 202T), which has 12 layers, 786 hidden dimensions and 12 heads. For the DA-Transformer baseline the ar- chitecture is causally masked, and inputs are permuted to model the sequence in a specific order. In addition to the standard positional embeddings for the input, the embeddings for the output are also concatenated to the token embedding. This can be seen as an implicit method to condition on the permutation that is currently generated. The specific hyperparameters of the optimization procedure are specified in Appendix and are the same as reported in (Austin et al., 2021), with the exception of different learning rate schedule and further train steps.
Performance of these methods is presented in Table perform very well compared to the other models. This result matches the behaviour that was found
I Firstly, the OA- Transformer baseline does not
7

---

Published as a conference paper at ICLR 2022
Figure 5: Visualization of æ through the generative process for an ARDM Upscale 4 model.
by Yang et al. 2019), who observed underfitting behaviour and limited the task complexity by only predicting a subset of the permuted tokens. Further, as expected the performance of our OA-ARDM with 1.43 bpe is very close to the performance of D3PM 33PM-absorbing at 1000 steps with 1.45 bpc. This is expected, since OA-ARDMs are equivalent to the continuous time limit of D3PM-absorbing models. For sequences containing only 250 dimensions, the D3PM schedule with 1000 steps starts to approximate the jump process where generally only a single variable is absorbed at a time. The important takeaway from this comparison is that OA-ARDMs perform similar to large-steps D3PM absorbing models while only requiring quarter of the steps. When the D3PM model is forced to take 256 steps which comparable to our OA-ARDM model, then its performance degrades further towards 1.47 bpd. In addition, a Parallelized ARDM with only 20 steps has performance of 1.51 bpd over similar D3PM which has 1.56 bpd. This pattern translates to CIFAR- 10 (Krizhevsky) et al. 2009) where ARDMs also outperform D3PMs and degrade more gracefully under fewer steps. This comparison to D3PM is however less direct, as the underlying architectures differ.
Lossless Compression To validate that ARDMs can form a viable basis for practical neural network-based compressors, we study their performance when compressing CIFAR-IO images and comparing them to existing methods. Since ARDMs provide probabilities for a sequence of sym- bols, they can be directly used together with an off-the-shelf entropy coder for lossless compression. In this experiment we use the range- based entropy coder rANS (Duda) 2009). To use ARDMs the order of the coding process needs to be fixed for all images. To avoid an unlucky sample, before coding we evaluate the log-likelihood of few random permutations on the train and pick the best performing one. Empirically, there is very little difference in performance (< 0.02 bpd) between different permutations.
Several deep learning based lossless compression methods in literature rely on bits-back coding (Townsend et al. 2019), such as LBB (Ho et al. 2019), HiLLoC (Townsend al. 2020) and VDM (Kingma et al. 2021). Although bits- back coding methods can perform well on large datasets, they have a large overhead when used as per-image compressors. This is caused by the large number of initial bits that are required. Further, the dataset is often interlinked, meaning that if an image in the middle of the dataset needs to be accessed, requires all images earlier the bitstream to also be decompressed. Therefore per-image compression is important for practical applications, because it is desirable to be able to send a specific image without sending an entire datasct. On the other hand, direct compressors such as L3C (Mentzer al.) 2019), IDF (Hoogeboom et al., 2019) and IDF++ (van den Berg et al.; 2021) do not incur an intial message overhead and their dataset performance translates directly per-image compression. A more conventional codec is FLIF (Sneyers & Wuille 2016), which is a recent lossless compression codec with machine learning components that outperforms traditional codecs such PNG.
Performance of ARDMS and related methods in literature is presented in Table 3 ARDMs signif- icantly outperform all methods on compression per image, requiring only 2.71 bpd versus 3.26 for the next best performing model, IDF++. In addition, even compared to setting where an entire dataset needs to be compressed, ARDMs perform competitively to VDM, which attain 2.72 bpd. Moreover, ARDMs degrade more gracefully when fewer steps are used to encode the data.
Note that the lossless compressor based on VDM was trained on non-augmented data, whereas the best-performing likelihood model of Kingma et al. (2021) was trained with data augmentation. As a result, it is likely that their dataset compression results could be somewhat improved when trained on augmented CIFAR-10 10. Also, it is not a coincedence that HiLLoC and FLIF have the exact same compression per image performance. HiLLoC compresses the first few images using the FLIF format to fill the initial bitstream, and compresses the remaining images in the dataset with back- coding (Townsend et al. 2020). As a result, on a per-image compression benchmark the method is equivalent to FLIR.
Effects of Depth-Upscaling A natural question that might arise is how standard order-agnostic modelling performs compared to order agnostic bit-upscaling, and how bit-upscaling compares to
8

---

Published as a conference paper at ICLR 2022
Table 3: CIFAR-10 lossless compression performance (in bpd).
Model
Steps
Compression per image
Dataset compression
VDM (Kingma et al.
2021
1000
2.72
VDM (Kingma et al.
2021
500
2.72
OA-ARDM (ours)
500
2.73
2.73
ARDM-Upscale (ours)
500
2.71
2.71
VDM (Kingma et
2021
100
2 8
2.91
OA-ARDM (ours)
100
2,75
2.75
ARDM- -eDDMle 4 (ours)
100
2.76
2.76
LBB (Ho al. 2019
28
3.12
IDF HHogebooom
2019
3.34
3.34
IDF4+ Berg aL. 2021
3.26
3.26
HiL.LoC Townsend et
4.19
3.56
FI.IF (Sneyers &
4.19
4.19
Table 4: Audio (SC09) depth upscaling test set performance (in bpd). A WaveNet baseline learning only single order achieves 7.77 bpd.
Table 5: Image (CIFAR-10) depth upscaling performance (in bpd).
Model
Steps
Performance
Model
Steps
Performance
OA-ARDM
D) 16000
7.93
OA-ARDM
D 3072
2.69
ARDM Upscale 256
X D
6.36
ARDM Upscale 16
X D
2.67
ARDM Upscale
6.30
ARDM Upscale
D
2.64
ARDM Upscale
8 x D
6.29
ARDM Upscale 2
3xD
2.67
ARDM Upscalc 2
16 D
6.29
the upscaling with larger values. Due to the constant training complexity of ARDMs, one can easily train models that have generative processes of arbitrary length. To test this, we train ARDMs on image data from CIFAR-10 and audio data from SC09 (Warden] 2018). For the audio data, the total number of categories is 26 which is typically too large in terms of memory to model as a single softmax distribution. For that reason, the single stage OA-ARDM is trained using a discretized logistic distribution because is computationally cheaper for high number of categories. For the same reason, the Upscale ARDMs for audio can only trained using the direct parametrization, whereas for images they trained with the data parametrization.
For images, the best performing model has an upscaling factor of with 2.64 bpd (see Table 5) and for audio the best performing model upscales by factor of 2 with 6.29 bpd (see Table The hypothesis is that as the upscale factor becomes smaller, the generative process generally be- comes more structured and easier model. However, although for audio this pattern is consistently observed, for images an upscale factor of 4 has better performance than an upscale factor of 2. It possible that for certain data, some point smaller upscale factors give diminishing returns for performance. We hypothesize that by prolonging the generative process, the model may get less gra- dient signal per optimization step, leading to the decreased performance smaller upscale factors in some situations.
6 LIMITATIONS AND CONCLUSION
Notwithstanding the good results in this paper. there are some limitations to ARDMs. 1) Even though ARDMs outperform all other order-agnoste approaches on text, there still gap to the performance of single-order autoregressive models. In preliminary experiments, upscale variants for language did not perform better than the order-agnostic versions. 2) In the current description, ARDMs model discrete variables. In principle tinuous distributions. 3) Finally, in this work we have focused on optimizing for log-likelihood because it directly corresponds to coding length in lossless compression. However when optimizing for other objectives such as sample quality, different architectural choices may give better results.
could also define absorbing processes for
In conclusion, introduced ARDMs, models and discrete diffusion models. ARDMs perform competitively with existing generative mod- els, and outperform competing approaches per-image lossless compression.
class of models at the intersection of autoregressive

---

Published as a conference paper at ICLR 2022
REPRODUCIBILITY AND ETHICS STATEMENT
To ensure the work is as reproducible as possible, in this paper we have described in detail both the training algorithms and the sampling algorithms. The main ideas are presented in Section 3 and further clarifications that may be important for re-implementation are given in Appendix [ The hyperparameter settings to run experiments are presented Section [ and further clarified in
Appendix D In addition, we plan to release the code that can be used to reproduce the experimental results in this paper.
In terms of ethics, we do not see immediate concerns for the models we introduce. However, deep generative models have a wide variety of applications such as representation learning, image inpaint- ing, special effects in video, outlier detection and drug design. On the other hand, the generation of images, text and video may have negative downstream applications such as making false media seem realistic. Further, to the best of our knowledge datasets were used that have known ethical issues.
REFERENCES
TensorFlow Datasets, collection of ready-to-use datasets.
https://ww.tennorfflo...
datasets
Michael A. Alcorn and Anh Nguyen. The DEformer: An order-agnostic distribution estimating transformer. In ICMI. Workshop on Invertible Neural Networks, Normalizing Flows, and Explicit Likelihood Models, 2021.
Jacob Austin, Daniel D. Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. Structured denoising diffusion models in discrete state-spaces. CoRR, abs/2107.03000 2021.
Samy Bengio and Yoshua Bengio. Taking on the curse of dimensionality in joint distributions using neural networks. IEEE Trans. Neural Networks Learn. Syst., 2000.
Yoshua Bengio, Réjean Ducharme, Pascal Vincent, and Christian Janvin. A neural probabilistic language model. The journal of machine learning research, 3:1137-1155, 2003.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. Advances in Neu- ral Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS, 2020.
Nanxin Chen, Yu Zhang, Heiga Zen, Ron J Weiss, Mohammad Norouzi, and William Chan. Wave- Grad: Estimating gradients for waveform generation. arXiv preprint WXXiv:2009.007333 2020.
Rewon Child. Very deep vaes generalize autoregressive models and can outperform them on images. In 9th International Conference on Learning Representations, ICLR. OpenReview.net, 2021.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. CoRR, abs/1904.10509, 2019. URL
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HIIT, pp. 4171-4188 Association for Computational Linguistics, 2019.
Prafulla Dhariwal and Alex Nichol. Diffusion models beat gans on image synthesis. CORR, abs/2105.05233, 2021.
Jarek Duda. Asymmetric numeral systems. arXiv preprint arXiv:0902.0271, 2009.
10

---

Published as a conference paper at ICLR 2022
Stefan Elfwing, Eiji Uchibe, and Kenji Doya. Sigmoid weighted linear units for neural network function approximation in reinforcement learning. Neural Networks, 107:3-11, 2018.
Marjan Ghazvininejad, Omer Levy, Yinhan Liu, and Luke Zettlemoyer. Mask-predict Parallel de- coding of conditional masked language models. In Kentaro Inui, Jing Jiang, Vincent Ng, and Xiaojun Wan (eds.). Proceedings of the 2019 Conference on Empirical Methods in Natural Lan- guage Processing and the 9th International Joint Conference on Natural Language Processing, FMNNP-IJCNILP, 2019.
Alex Graves. Generating sequences with recurrent neural networks. CoRR, URL http://arxiv. wwwp///g/////...//////////////./ 0850
2013.
Jonathan Ho, Evan Lohn, and Pieter Abbeel. Compression with flows via local bits-back coding. In Advances in Neural Information Processing Systems 32: Annual Conference on Neural Informa- fion Processing Systems 2019, NeurlPS, 2019.
Jonathan Ho. Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS, 2020.
Emiel Hoogeboom, Jorn W. T. Peters, Rianne van den Berg, and Max Welling. Integer discrete flows and lossless compression. In Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, 2019.
Emiel Hoogeboom, Didrik Nielsen, Priyank Jaini, Patrick Forré, and Max Welling. Argmax flows and multinomial diffusion: Learning categorical distributions. CoRR, abs/2102.05379, 2021.
Chin-We Huang. Jae Hyun Lim, and Aaron C. Courville. A variational perspective on diffusion- based generative models and score matching, CoRR, abs/2106.02808, 2021. URL https://
wwx.v.org/aaie..ee............at owg. abgg/bb//2106... 02808
Ajay Jain, Pieter Abbeel, and Deepak Pathak. Locally masked convolution for autoregressive mod- els. In Ryan P. Adams and Vibhav Gogate (eds.), Proceedings of the Thirty-Sixth Conference on Uncertainty in Artificial Intelligence, UAI, 2020.
Daniel D. Johnson, Jacob Austin, Rianne van den Berg, and Daniel Tarlow. Beyond in-place cor- ruption: Insertion and deletion in denoising probabilistic models. CoRR, abs/2107. 07665,, 2021.
Alexia JolicoeurMaartincau, Ke Li, Rémi Piché-Taillefer, Tal Kachman, and Ioannis Mitliagkas. Gotta go fast when generating data with score-based models. CoRR, abs/2105.1400 14080, 2021.
Heewoo Jun, Rewon Child, Mark Chen, John Schulman, Aditya Ramesh, Alec Radford, and Ilya Sutskever. Distribution augmentation for generative modeling. In Proceedings of the 37th Inter- national Conference on Machine Learning, ICML, 2020.
Nal Kalchbrenner, Erich Elsen, Karen Simonyan, Seb Noury, Norman Casagrande, Edward Lock- hart, Florian Stimberg, Aäron van den Oord, Sander Dieleman, and Koray Kavukcuoglu. Efficient neural audio synthesis. In Proceedings of the 35th International Conference on Machine Learn- ing, ICML, 2018
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.00 1412.6980, 2014.
Diederik P. Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. CoRR, abs/2107.006333 2021.
Zhifeng Kong and Wei Ping. On fast sampling of diffusion probabilistic models. abs/2106.00112,, 2021. URL
CoRR,
Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. DiffWave: A versatile dif- fusion model for audio synthesis. 9th International Conference on Learning Representations, ICLR, 2021.
11

---

Published as a conference paper at ICLR 2022
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Hugo Larochelle and lain Murray. The neural autoregressive distribution estimator. In Proceed- ings of the Fourteenth International Conference on Artificial Intelligence and Statistics, AISTATS, 2011.
Qi Liu, Miltiadis Allamanis, Mare Brockschmidt, and Alexander L. Gaunt. Constrained graph vari- ational autoencoders for molecule design. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa-Bianchi, and Roman Garnett (eds.), Advances in Neural Informa- tion Processing Systems 31: Anmual Conference on Neural Information Processing Systems 2018, NeurlPS, pp. 7806-7815, 2018.
Matt Mahoney, Large text compression benchmark, 2011.
Gábor Melis, Chris Dyer, and Phil Blunsom. On the state of the art of evaluation in neu- ral language models. In International Conference on Learning Representations, 2018. URL
https::/oppnreviiee net/ By
Jacob Menick and Nal Kalchbrenner. Generating high fidelity images with subscale pixel networks and multidimensional upscaling. In 7th International Conference on Learning Representations, ICLR, 2019.
Fabian Mentzer, Eirikur Agustsson, Michael Tschannen, Radu Timofte, and Luc Van Gool. Practical full resolution learned lossless image compression. In IEEE Conference Computer Vision and Puttern Recognition, CVPR, pp. 10629-10638. Computer Vision Foundation IEEE, 2019.
Stephen Merity, Nitish Shirish Keskar, and Richard Socher. Regularizing and optimizing LSTM
language models. International Conference on Learning Representations, 2018. URL https:
//openreview net wwt httt//frummii-SyyyppPP
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, ICML, 2021.
Scott Reed, Aäron van den Oord, Nal Kalchbrenner, Sergio Gomez Colmenarejo, Ziyu Wang, Yutian Chen, Dan Belov, and Nando de Freitas. Parallel multiscale autoregressive density estima- tion. In Proceedings of the 34th International Conference on Machine Learning, ICML, 2017.
Tim Salimans, Andrej Karpathy, Xi Chen, and Diederik P. Kingma. PixeICNN+:: Improving the PixeICNN with discretized logistic mixture likelihood and other modifications. In 5th Interna- tional Conference on Learning Representations, ICLR 2017. OpenReview.net, 2017.
Samarth Sinha and Adji B. Dieng, Consistency regularization for variational auto-encoders. CoRR, abs/2105. 14859, 2021.
Jon Sneyers and Pieter Wuille. FLIF: free lossless image format based on MANIAC compression. In 2016 IEEE International Conference on Image Processing. ICIP, pp. 66-70. IEFE, 2016
Jascha Sohl-Dickstein, Eric A. Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamies. In Francis R. Bach and David M. Blei (eds.), Proceedings of the 32nd International Conference on Machine Learning, ICML, 2015.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribu- tion. In Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS, 2019.
Yang Song, Chenlin Meng, Renjie Liao, and Stefano Ermon. Accelerating feedforward computation via parallel nonlinear equation solving. In Proceedings of the 38th International Conference on Machine Learning, ICMI., 2021.
James Townsend, Tom Bird, and David Barber. Practical lossless compression with latent variables using bits back coding. CoRR, abs/1900 048666 2019.
12

---

Published as a conference paper at ICLR 2022
James Townsend, Thomas Bird, Julius Kunze, and David Barber. Hilloc: lossless image com- pression with hierarchical latent variable models. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, 2020.
Benigno Uria, lain Murray, and Hugo Larochelle. A deep and tractable density estimator. In Pro- ceedings of the 31th International Conference on Machine Learming, ICML 2014, Beijing, China, 21-26 June 2014, volume 32 of JMLR Workshop and Conference Proceedings, pp. 467-475. JMLR.org, 2014.
Arash Vahdat and Jan Kautz. NVAE: A deep hierarchical variational autoencoder. In Advances in Neurul Information Processing Systems 33: Annual Conference on Neural Information Process- ing Systems 2020, NeurlPS, 2020.
Rianne van den Berg, Alexey A. Gritsenko, Mostafa Dehghani, Casper Kaae Sønderby, and Tim Salimans. IDF++ analyzing and improving integer discrete flows for lossless compression. In 9th International Conference on Learning Representations, ICLR. OpenReview.net, 2021.
Aäron van den Oord, Sander Dieleman, Heiga Zen, Karen Simonyan, Oriol Vinyals, Alex Graves, Nal Kalchbrenner, Andrew W. Senior, and Koray Kavukcuoglu. Wavenet: A generative model for raw audio. In The 9th ISCA Speech Synthesis Workshop, 2016a.
Aäron van den Oord, Nal Kalchbrenner, and Koray Kavukcuoglu. Pixel recurrent neural networks. In Maria-Florina Balcan and Kilian Q Weinberger (eds.), Proceedings of the 33nd International Conference on Machine Learning, JCML, 2016b.
P. Warden. Speech Commands; A Dataset for Limited-Vocabulee Vocabulary Speech Recognition. ArXiv e-prints, 2018. URL http::tp::/// htt///rxivv..oo wttp://oggabbbb/booeee 03209
Daniel Watson, Jonathan Ho. Mohammad Norouzi, and William Chan. Learning to efficiently sample from diffusion probabilistic models. CoRR, abs/2106.03802, 2021. URL https:
//arxiv. org/ abs/21b6..1.../ 03802
Auke J. Wiggers and Emiel Hoogeboom. Predictive sampling with forecasting autoregressive mod-
els. In Proceedings of the 37th International Conference on Machine Learning, ICMI., 2020.
Zhilin Yang, Zihang Dai, Yiming Yang, Jaime G. Carbonell, Ruslan Salakhutdinov, and Quoc V. Le. XLNet: Generalized autoregressive pretraining for language understanding. In Advances in Neu- ral Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS, 2019.
13

---

Published as a conference paper at ICLR 2022
A FURTHER DETAILS OF AUTOREGRESSIVE DIFFUSION
Next to given descriptions, the implementation has been open-sourced at https://githubb... com/
diffssion
A.1 DEPTH UPSCALING
This section explains further details that are important to optimize and sample from Depth Upscaling ARDMs, which are summarized in Algorithm and Recall that for depth- upscaling models, the variables are modelled in stages .....
3..., (S) and the model learns
p((=1)))a(S-11).. .......... p( (12x(1)ar (xx())). Here xx()) is a constant absorbing state and x(S) represents the data. The transition matrices (P(.)) describe the destructive maps which end up in the absorbing state. They form the destructive counterpart of the generative process.
Instead of optimizing for all stages simultaneously, we sample stage uniformly U(1.....
and optimize for that stage. Here the cumulative matrix products p( allow us to directly transition
to a specific stage, since (+() Sppppttt (++1) (S), To be precise, for a single dimension i the variable (e)
a is represented as onehot vector and then transformed using the matrix multiplication m!!)
P (*+1)) (5) For multiple dimensions this matrix multiplication is applied individually, meaning
that P (()) [P (++1) (5) P 22 ................................ P ((++1) (1)) (z) (4)
(s+1) (8)
................................ .......................... ( x+))
For a optimization step, stage 8 U(1,..., 5) and a step t ~ U(1...., D) are sampled, in addition to a permutation a ((p)). Then using the cumulative matrices, from a datapoint æ = (5) the variables (x(*) and (= are computed. As before, the mask m
I gives the locations of the variables that are conditioned on. For those locations the values in (*)) may already m, instead the values from (=5-1) are accessed. This
be accessed. For the opposite locations leads to the expression for the input i ( + (1-m) @x((=1) The target of the network will be to predict distribution for (*) at the locations at m. The network will take in the computed input é together with variables to clarify in which stage of the generative process the model is, m, and t. In case of the data parametrization, the probabilities O are appropriately normalized and reweighted to (*)) using transitions (P(*)). Then, the log probabilities log C(xx!*) (((+)) are computed clementwise over dimensions and subsequently masked with 1 - m. These quantities are then summed and reweighted to get stochastic estimate for the ELBO.
For the sampling, the model traverses through each stage, and for each stage through every dimen- sion in a different order. In each step the network together with the transition matrices produces probability vector ((*) from which elementwise samples are taken x' C(x(()))(())) but only the values at locations n + (a = l) are filled in, corresponding to the current generation step. By traversing through all steps and stages, the variable (S) is generated.
Algorithm 3 Sampling from Upscale-ARDMs
Algorithm 4 Optimizing Upscale-ARDMSs
Input: Network
Input: Datapoint æ. Network
Output: Sample x
Output: ELBO C
Initialize æ (0)
Sample ~ U(1 ..,
for in (1....,S} S do
Sample ~ (I....., D)
Sample 0 (SD) for in (1...., D} do
Sample ~ U(SD)
(+) + P (s+1)
x and (=--1) f- (1)
m x o <
Compute m +- o <t
+ (o = t)
i <- m ④ (+) (1 - m) 2 (0
0 < f(x m, 8,
0 -- f(i, m. 8,
((5) x p(s) x P (s+1)
(s)) x p(s) T
x (s-1) + p ((+1)) 0
x' ((x()) ((+))
+ (1- m) @ log ((x()) (0()))
x + (1 - n
x
L D- sum(l:)
14

---

Published as a conference paper at ICLR 2022
A.2 DETAILS ON PARALLELIZED ARDMS
This section discusses further details Parallelized ARDMs, and provides JAX version of the dynamic programming algorithm from (Watson et al.) 2021) that was written in NumPy. Since the algorithm scales with O(D3) this implementation important scale to larger dimensional problems. To clarify, the upscale ARDMs can be seen as a S sequential OA-ARDMs that model p(x(() (=(+=1)), and when parallel schedule is computed. it computed for each stage separately. It is also possible to the dynamic programming algorithm for all S D steps simultaneously, which could even choose to distribute steps unevenly over stages, but that is not done in this paper.
the algorithm matrix L is needed which gives the cost of travelling from generation step to another. It is constructed so that Lt,+* = k- L, for positive and otherwise, which represents the cost of generating k variables in parallel where the loss component. In practice this is implemented via a cumulative sum of triangular mask. This part is relatively computationally cheap.
Recall that
import jax
from jax import numpy jnp import numpy as np
def get nel_o matrimat (loss componentt: ndarray): num_timesteps len closss_componetts)
# Creates multiplicative mask. E.g. if num_timesteps then:
#
[1 2 3]
# triu = [0 1 2].
#
0
triu np tritriucCn (np ones (inum_timesteps, num_timesteps))
triu = np www.uuuntriuu::::....
(xis=0)(::-1)
# Compute nelbos t] which contains -logp (x_s x_t) nelbos_ coss_components(; None] * triu
# Pad last row first column.
nelbos = np. zeros ((num_tinmnnn timestemmeeesss 1, num_timesteps + 1))
nelbos[:-1. 1:] nelbos_
nelbos
The most expensive part of the algorithm is the loop which has computational complexity ( This is the most important extension of the NumPy version and reduces runtime from 5 minutes to about seconds for D 3072, which would be very impractical where = 16000, which now take less than half a minute to run. Through JAX this loop is XLA- compiled with the operation, limiting overhead when running the algorithm.
for audio experiments
@jax @jax.jit
def nelbos: jnp. ndarray, first_cost: jnp ndarray): ... Inner eax-loop that computes the cost dimension matrices. ... num_timesteps first_
def ctteeeeeex__ccttpprrreeee bpds prev_cost [:, None] + nelbos new_dimension = jnp argmin axis=0) new_cost = jnp min (bpds, axis=0) return new_cost, (new_cost, new_dimension)
jnp ndaaray, jnp. ndarray)
(costs, dimensions) = jax lax.scant penpute_nex_ccot.. mnit=first_cost. xs= jn=jnn. arange (1, num num_ tum_timesteps+))
costs, dimensions
15

---

Published as a conference paper at ICLR 2022
The inner algorithm logic is then called via the function below. It first builds the loss transition matrix L which is referred to as nelbos and then calls the inner loop. As an output it gives the cost and dimension matrices that can be used to 1) find optimal path and 2) describe how expensive such paths are. As can be seen in Figure 6 the running average of the loss components {T4) might be somewhat noisy, which can negatively influence the algorithm. As straightforward method to reduce variance of the values {Le). they are sorted before they are given to the algorithm. This uniquely possible for ARDMs, as we expect L, to be monotonically decreasing over (see also Equation 4. For Upscale ARDMs that have multiple stages, the loss components are seperately sorted per stage.
wwweaddddddiinnnaeeeeeeeeeeoooee coon_mmtricesse ces es [loss_components: np ndarray):
Compute cost and assignment matrices, , in JAX. num_timesteps len (loss_components
# First row of the costs matrix. first_cost np. full www.nuulnnnnlmmesteeee [irst_costt
1,), np.inf)
first_cost jnp array (first_cost)
# First row of the dimensions matrix. The first row just contains -1 and is never used, but this way it aligns with the matrix. first_dimension jnp. full ((num_timesteps
1), -1, dtype=np.int32
# nelbos[s, t] is going to contain the value 1xgp x_t) array (get_nelbo matrii components))) costs, dimensions inner cone
nelbos array
(nelbos, first_cost)
# Concatenate first rows to the matrices. costs = jnp concatenate ([first_cost [None, :] costs], axis=0)
dimensions jnp. concatenate axis=0)
[None, :], dimensions],
costs = array (costs) dimensions np. array (dimensions)
return costs, dimensions
The final part of this algorithm is used to retrieve the path that needs to be taken to attain a certain cost. This algorithm takes as input budget and the cost & dimension matrices, and returns the corresponding path to traverse.
def get_optimal__azt paallllpphh_weee wwt (budget: int, costs: np. ndarray,,
dimensions: np ndarray):
num_timesteps len (costs) t num_timesteps path np zeros (budget, dtype=np int32) cost costs (budget, num_ting tum_ttmmsteps] for k in reversed (range (1, b)dget+1)): dimensions Ik, t] path(k-1] return path, cost
16

---

Published as a conference paper at ICLR 2022
B ADDITIONAL RESULTS
B.I RELATION TO OTHER LIKELIHOOD-BASED GENERATIVE MODELS
In this section we show how ARDMS perform compared to existing likelihood based generative models in literature. These results are presented in Table 6 The best performing model is the Variational Diffusion Model (VDM) (Kingma et al.) (021). ARDMs perform competitively with a best score of 2.64 bpd, and are the best performing model among discrete diffusion approaches.
Table 6: CIFAR-10 generative modelling.
Model
Type
NLL
ARDM-AO (ours)
Discrete Diffusion U ARM
2.69
ARDM-Upscuae (ours)
Discrete Diffusion U ARM
2,64
D3PM Gaussian Austin et al. 2021
Discrete Diffusion
3.44
DDPM (Ho ct al.
2020
Diffusion
3.69
Improved DDPM Nichol & Dhariwal 2021
Diffusion
2,94
VDM (Kingma et al. 2021]
Diffusion
2.49
PixeICNN++ (Salimans et al.
2017
ARM
2.92
SPN (Menick & Kalchbrenner
2019
ARM
2,90
Sparse Transtormer (Jun al.
2020
ARM
2.52
NVAE (Vahdat Kautz 2020
VAE
2.91
Very Deep VAE Child 2021
VAE
2.87
CR-VAE (Sinha & Dieng 2021
VAE
2.52
B.2 ADDITIONAL AUDIO EXPERIMENTS
In Table 7 we present additional experimental results from our best Upscale ARDM model for the SC09 dataset (branching factor 4), which we consider smaller computational budgets. Recall that dimensionality D = 16000 for SCO9 data.
B.3 LOSS COMPONENTS OVER TIME
Since the training algorithm estimates the NLL by sampling a step t for each input in the batch, we can collect and keep track of the loss components {I.) and plot them as a function of (see Figure [6. These are collected by updating an exponential moving average during training, and are used in the dynamic programming routine. As expected by Equation 4 the components L, are monotonically decreasing over the step & within stage. The height is re-normalized so that the average height represents the total bits per dimension. As a result, in the upscale model the value divided by number of stages S represents the actual uncertainty of generating that token.
The loss plot of the upscale model allows for interesting observations: For instance, After an initially high uncertainty (Rx 8/S 2 bits) the most significant bits become increasingly easier model very fast (< 2/5 = 0.5 bits). In contrast, for each stage that follows, the average height of that stages increases. This indicates that the model is uncertain for the less significant bits. This can be combination of two things: The model may have more difficulty in modelling these less significant
Table 7: Audio (SC09) depth upscaling test set performance (in bpd) for various computational budgets.
Model
Steps
Performance
ARDM Upscale 4
8 16000
6.29
ARDM Upscale 4
8 1000
6.30
x 500
6.30
x 100
6.32
8 x 50
6.32
17

---

Published as a conference paper at ICLR 2022
bits (i.e. high KL between data and model distribution), and the data distribution may be more uncertain those regions (i.e. high entropy of the data distribution).
8-
6-
5-
2-
500
1000
1500
2000
2500
3000
2000
4000
6000
8000
10000
12000
Figure 6: Loss components over model step on CIFAR-10. The height is normalized so that the average represents the total bits per dimension. Left: loss terms for the OA-ARDM. Right: loss terms for the ARDM-Upscale 4, which comprises four stages.
B.4 SAMPLES FROM ARDMS
Language Sampling from ARDMs can be visualized at different steps to highlight the genera- tive process. Recall that for models trained on language. the absorbing state augments the space, meaning that an additional index that is added for the absorbing state. We visualize this token by the underscore character The process at four selected steps the process are presented in Figure where the last sentence represents the resulting sample.
egy
is is
ot
ce pr
er onal W
to_r__li_e s_rategy for _
hort s_riee.... b.r.....brr_.ee whethh se 1s. _l fo a. ot us
h nex__rri ce
i erational
po_s_in_t go_ g_ne_cc u
m. u. e....
fy
to r__li_e s_rategy for autom__ics o. the c_n ex short b.rikeebombers_wheett i r_use it_is deli_erate ga_auls ha_e _ot_ntially pro i_ject crational M_a in te_game_car _s_i_divid_ mod_1 urro........d0..ffff...
to realize strategy for automatics it the experience for shert strike bonbers whether i n reuse it is deliberate for gasauls have potentially unsucked proef or inject operational wea pons the game individunl madel urrealise identifying
_e_5_. DM5
he_
bl.
t.a.a... i.land omp
0. d
1e_s_re_shows_..
1_0 _ta _ta___ fta.a_f__yyy___ya by_sa_el_ _ard_ ro__5 he_
ita_e
y_mi fis u.r bl. e
i_lddd_ itmp..ittooo
____-_rgrr__ d_tun mokg___nggs_____.....___
_o__
leisure shows... t itage of t... ata_. main fish b uer george_ clea_son a.. diltun peak amongsthe g _onn_
bl_ el erpa daa tin_a._s island compo_it_on o.. sta...
by ellitte _oadd a.d troo
leisure shows it the pratection of the stamp falls by satellite board and troops lace tenua he ritage the catalay nain fish b fuerta robla el and georges clearson and diltun peak amongsthe general commu
eden tingalas island composition
Figure 7: Two generative processes of an OA-ARDM trained on text8. The resulting sample from the model is at the bottom of each frame.
18

---

Published as a conference paper at ICLR 2022
The generative processes for images visualized in Figure g In constrast with the language model, here the absorbing state takes on a specific value in the domain of the image itself. In the case of OA-ARDMs, the absorbing state is 128 that it is 0 when normalized in the network architecture. constrast, the absorbing state of the Upscale ARDM is because zeroing least significant bits until everything is zero. The right-most grid represents the resulting samples from the model. The generative processes are very different: whereas the upscale ARDM first generates a coarses version of the images with fewer bits, the order agnostic ARDM generates each value at once.
Images
defined by
(a) Generative process of an Upseale ARDM. This model was trained without data augmentation and with dropout, which explains that all images are generated upright. The performance this model is approximately 2.71 bpd whereas the same model trained with data augmentation has 2.64 bpd.
(b) Generative process of OA-ARDM, starting the absorbing state This model trained with data augmentation, which is known to somewhat degrade sample quality and naturally sometimes samples rotated or reflected images.
Figure 8: Visualization of the generative process for æ, ending with the resulting samples at the right-most grid.
19

---

Published as a conference paper at ICLR 2022
C EQUIVALENCE OF AO-ARDMS AND ABSORBING DIFFUSION IN CONTINUOUS TIME
In this section we examine the connection between absorbing diffusion and AO-ARDMs more closely. We start with description of the independent process of absorbing diffusion as used in (Austin et al. 202I) and highlight potential complications of this process. Then, we will show that version of absorbing diffusion models.
AO-ARDMs are equivalent to a continuous-
The Independent Absorbing Process from Austin et al.
In absorbing diffusion as described by (Austin et al, (021), each dimension can independently get absorbed with some small probability for each time step. Specifically, letting vector x(t) represent a Markov process as a collection of random variables index by integers i, where x=0) is the data distribution. Each dimension (t) as equal and independent chance of getting absorbed according to rate >t) at index t to the absorbing state a, Define the cumulative chance of retaining
the state as att) = = (1 - (T)). This allows the direct expression for the distribution over x, (1) as categorical on data and the absorbing state (0. a; with probabilities {a((), 1 - a())) Typically. the decay rate is chosen so that (x() = 0 for some large integer T. For example in (Austin et al.] (02T) it is set T 1000 for most experiments. We refer to the absorbing process from (Austin et al.) 2021) as independent absorbing process, due to its independent absorbing probabilities between dimensions.
The reverse of this absorbing process is the generative process. As described above, the chance of a dimension absorbing is independent. As a result when T is small, it is inevitable that multi- ple dimensions decay at once. This has direct consequence for the generative process, which is parametrized to model dimensions independently. The generative process will have to model the variables of these multiple absorbed dimensions
independent factorized distribution, which causes a loss in modelling performance. This problem can be overcome by setting I' to a larger value. Indeed, when T' is larger the chance of multiple dimensions decaying once decreases. Dur- ing training, ''' can be set arbitrarily high without incurring costs. However, to sample or evaluate the likelihood of specific datapoint, the computational cost scales directly with T so it is desired to keep T as low as possible.
As an example, consider the experiment from (Austin et al.] 2021) where text sequences of length 256 are modelled using 1000 timestep absorbing diffusion model. When sampling from this model, at least 744 of the neural network forward passes do nothing to the latent variable and are needless compute. When T' is reduced, performance degrades. In addition, it can be difficuult to determine beforehand how high a T should be sufficient, and it depends on the data and the decay rate.
ARDMs model the reverse of a Fixed Absorbing Process
Our ARDMS can be viewed as learning the generative process of a slightly modified absorbing process. Instead of independent absorbing costs, exactly one dimension decays a time step until all dimensions have been absorbed. Since only one variable decays at a time, we refer to this process as a fixed absorbing process. This ensures that T = D exactly, where D is the dimensionality of the data.
An equivalent way describe this process is by sampling permutation of the indices 1...., D and decaying in that order towards the absorbing state. The corresponding generative process is then modelling the variables exact opposite order of the permutation: an AO-ARDM. As a result the generative process with a fixed absorbing process process only requires at most D steps.
ARDMs are equivalent to Continuous Time Absorbing Models
The absorbing diffusion model from (Austin et al. 202T) can be relaxed to continuous time. Define a continuous-time discrete-space absorbing state Markov process as a collection of Markov random variables (x(t)} in dimension D, parameterized by E I Starting in state (0), each of the ele- ments (4) of the state vector independently decays towards an absorbing state a at rate (t), such that at time the distribution on the vector elements is categorical on (2)( a:). with probabilities (a<t), 1 (f)), with at)) exp(- So (s))ss). This last equivalence is obtained via the first order logarithmic Taylor expansion log(1 - x) = -2 which holds for the small (s)dss
An equivalent way of describing this stochastic process is as a finite set of D random transition times (7)) for i {1,..., N) describing the time where the element z; has transitioned into the absorbing
20

---

Published as a conference paper at ICLR 2022
state. Specifically, we say that T; is the latest time for which z; is non-yet absorbed, so 2t ag for > T.. From this perspective, x is only changing at the transition times T₃ and remains the same at other times. Then, the reverse process to model {x(t)} for all is equivalent to only modelling the finite dimensional (T)( Ta). T;) other words, to model the reverse process, we only need to model the transition times {r:} and the variable right before it absorbed.
To show an equivalence between the continuous time absorbing process and ARDMs, we will show that we can model the reverse process given by (7;) Ta by sampling the transition times inde- pendently, and by using an AO-ARDM for the transitions in æ.
The distributions {Ta} already known, they are given by the distribution with the cumulative distribution function - a(t). That leaves the modelling of {az (T:)) to model the reverse process.
An important question that now arises whether the transition times {T,} provide any additional information in modelling the variables (a;( Apart from knowing that the variable will be
un-absorbed, the answer no. This is because the values are distributed as the data distribution so v, (.(0))a (()) ~ (x(1)))=(() and the actual continuous value of does not matter, specifically (o)( 1 (T.))a(().
As consequence, the model for the reverse process does not need to be conditioned on the precise values (Ti) and can instead be solely conditioned on f4(T++1) to model x, (T, for dimensions Recall that this process is equivalent to the generative process of our AO-ARDM: Each new timestep, a new dimension of the variable is modelled. The order in which the variables are modelled depends on the decay times } and since these are all identically distributed, the order is uniform over all possible permutations.
We claim that we can write the VLB follows:
log ((x(()())) 2 By(m(>0))a(()) logp(x(>0)) + logp(=(0) log -og
logp(a(T)) "o
= -Esa)((S))
Σ log ( (a(e)))
where <<> 0) denotes all values of æ(t) for t > O. The first equivalence follows from the above described equivalent representation of the continuous process. And indeed when the transition times and the values of æ at the transition times are given, the remaining variables of the continuous process can be reconstructed so it can be ensured that:
L(TD), TI. (Ilg(ala(nn......(())))
(ttp:///)ppppp)o.)))_aaaa) a(DD)
0.
In addition, recall that any transition variable is distributed according any chosen cumulative distribution att). Therefore, we can simply set our generative process to the same distribution, which ensures that:
(alg(T,...,TT)llle .T)) 0.
At this point observe that the sampling times only determine the order in which the dimensions are modelled. In fact when modelling xTTx) (=(+11) only one dimension dimension changes from Ti+1 Tr. Since all ] are independently and equally distributed, the distribution over the order of (T.) is uniform. A subtle detail is that the reverse of the order of 7; describes the generative order since we model timestep t; given T-+1- Nevertheless, since the distribution over orders is uniform, the distribution over reverse orders is also uniform. Therefore:
Ey((E. pl (or (T)))(() (Eg
Σ (()x())
where the latter equation contains the same omitted notation as the main paper: The model is aware which dimensions are conditioned on and which are not. In practical terms, this means that (o(<s) should be viewed as a masked vector and not as an order-less set. For the curious reader, technically the ability move from order-less to the structured masked vector is enabled by condi- tioning on σ. In summary, modelling the generative process of a continuous absorbing jump process equivalent to AO-ARDM. This is beneficial, viewing the model AO-ARDM gives simple bound to the number of required steps and allows an easier perspective on the model and its implementation.
21

---

Published as a conference paper at ICLR 2022
D EXPERIMENTAL DETAILS
In this section further details are given on the experimental setup,
Images For CIFARIO (Krizhevsky et al., 2009) we train the model using a fixed number of steps using the typical splits and evaluate the test log-likelihood after 3000 epochs of training. The results that are reported with standard deviations results are based on runs with three different initial seeds. The other results are based on single-run experiments. The runs take approximately 2 weeks to complete training on TPUv4 devices, although good performance (+% 2.8 bits per dimension) is already achieved after couple of days.
As base architecture for the OA-ARDM and the Upscale ARDMs, the exact same U-Net architec- ture in (Kingma al., 202 I) are used. This architecture has 32 ResBlocks at resolution 32 x 32, a single middle attention layer and then another 32 ResBlocks at resolution 32 x 32. Throughout the network feature maps have 256 channels. This architecture is typical for NLL optimization and lossless compression, which typically require many high-resolution feature maps (Mentzer et al. 2019). The models are trained for 3000 epochs with Adam using learning rate of 0.0001 and beta parameters (0.9/ / 0.999). The input processing uses a combination of the floating point represen- tation and an embedding layer. The integer-valued input is put through a simple normalization by dividing by the total number of classes, and subtracting 0.5. To this normalized input the mask m then concatenated. In the case of upscale ARDMs, the current stage in one-hot representation is also converted to mask with S channels, and is also part of m. So in that case m has 5 + channels. Then a 3 x convolutional layer maps the concatenated inputs to 3/4 of the channels. In addition, the integers are also fed through an embedding layer to 1/4 of the channels. These two outputs are then combined which produces the feature maps with 256 channels. This given as an input to the U-Net architecture described above. Following Austin et al.] (2021) we include the Cross-Entropy (CE) objective Lee = EE_u((......) D(D - L + 1)4x (i.e. the unnormalized likelihood components), with a small factor of 0.001. However, in an experiment without the Le loss included, no substantial differences in performance were found for our ARDMs. Since the like- lihood of the dataset is estimated with the ARDM objective, the results are computed over multiple dataset passes to reduce the stochastic effects from sampling l and a. For evaluation the exponential moving average of the parameters is used with a momentum coefficient of 0.9999. The models are trained with batch size of 128. The gradient is clipped at 100.
Language For the text8 dataset (Mahoney 2011 I we train using the typical 90- 10 /5- 10 10%/5- 10 splits in characters. Because the text8 dataset is a long string of characters, there is predictive infor- mation between segments when chunked. For this reason there is a big difference between model performance in literature in the reported scores on the text8 benchmark. Some methods consider a larger context before the current sequence, which greatly improves the information available to the model and gives better log-likelihoods. The runs take approximately a week to complete on 4 TPUv4 devices.
Since interested in the pure modelling capacity of models, follow (Hoogeboom et al., 2021]| et adl]] and consider chunked text segments without any additional context. How- ever, since the text& splits are not evenly divisible by 256, we slightly adjust the chunk size to 250 characters, to avoid dropping the last batch. We validated empirically with a baseline Transformer that this small change does not meaningfully change the performance with the 256 version. For reference, a baseline 12 layer Transformer attains 1.35 bpc on this problem.
As base architecture, we use a 12 layer Transformer as used in (Austin al. 2021 It has 768 dimensions, 12 heads, 3072 MI.P dimensions. For the ARDM architectures we followed (Austin) ct al., (02T) and used batch size of 512 with no dropout. For standard language model baseline, since we observed overfitting the batch size was lowered and dropout of 0.1 was added. The mod- els are trained for 3- 10 training steps. ARDMs are optimized with Adam with a learning rate of 0.0005 which has a linear warm-up for the first 5000 steps. The additional Leg loss was included with a factor 0.0001. The gradient clipped at 0.25. For evaluation the exponential moving aver- age of the parameters is used with a momentum coefficient of 0.995. All models use sinusoidal
22

---

Published as a conference paper at ICLR 2022
positional embedding. However, the DA-Transformer based on the XLNet approach (Yang al. 2019) requires both the imput and target positional embeddings to infer which permutation currently needs to be generated. In (Alcorn & Nguyen 202T), this is handled by interleaving input and target nodes in the sequence. A downside to this approach is that is increases the sequence length by two, which increases the quadratic computational complexity of the attention layers by four. In contrast, we concatenate the input and target embeddings, which does not alter the sequence length.
Audio For audio experiments we used subset of the SC09 dataset (Warden 2018) obtained by filtering out all non-digit commands from the dataset without changing the structure of the train/- validation/test splits. The resulting dataset contains 31 (1158/3666 /3////444/1000 training/validationnnes audio clips that are 1 second long and sampled at 16 kHz. In few rare cases when the audio clips were shorter than second, we right-padded them with zeros; and all considered models were trained and evaluated on padded data. A Tensorflow Datasets TFD version of this dataset is provided with the open-source code. Training takes approximately 4 days.
For both, the AO-ARDM well the Upscale ARDM experiments, closely followed the Dif- fWave setup (Kong et al. 2021) in terms of the network architecture and size, bit adapted the input embedding layer to take input masks into account. Specifically, we used non-causal WaveNet van den Oord et al. 201 architecture with 36 blocks, 256 channels and a dilation cycle of 11 (i.e. maximum dilation of 2048); input embedding were obtained by concatenating 1) 64-channel em- beddings of integer input values; with 2) 192-channel mask and continuous input value embeddings output by a width 3 convolution; the shared time embedding was obtained by mapping the standard 256-channel sine and cosine representation through two dense layers with 1024 features and Swish nonlinearities (Elfwing et al. 2018).
Audio AQ-ARDM and Upscale ARDM models were trained using the Adam optimizer (Kinmma & Ba 2014) with beta paramters 0.9 / 0.999 for 10 steps with batch size of 256 and linear learning rate warm-up over the first 15000 steps followed by constant learning rate of 10 10-4, During training we tracked an exponential moving average (EMA) of the model parameters using a momentum of 0.995, and employed the EMA parameters during evaluation. As in the case of image ARDMs, the models were optimized using a combination of the ELBO and CE objectives the latter taken with tiny weight of 10-4. 4. No further regularisation
used.
Due to the large output space (216 classes) audio AO-ARDM modelled the output distribution using a mixture of discretized logistics (DMoL.) with 30 components, although experimentally found the number of components to not make a big difference. To aid with training. the DMoL was initialized as an approximately uniform distribution with different mixtures responsible for the different parts of this distribution; and gradients with an Le norm larger than 1000 were re-normalized. Owing to the smaller per-stage output space, we were able to utilize the categorical softmax parameterization for the Upscale ARDMs. Empirically we observed this model class to demonstrate more stable training behaviour (in addition to significantly improved likelihoods), which we (partially) attribute to the choice of parametrization.
For our autoregressive single-order baseline we sought to deviate from the above AO-ARDM setup as little as possible, and used a causal version of the WaveNet architecture above. However, we observed that the single-order baseline overfits quickly on the training data. To overcome this, we found it necessary to use weight decay (0.01), smaller batch size (64) and fewer channels (128) for the baseline model.
23


Reviews:
[{"review_id": "ew7Hb9AG7-", "reviewer": "Reviewer_y93W", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strength\n- The paper is well-written.\n- Various experiments are conducted (across diverse domains), and hence, it shows the extensibility of their work.\n- The experimental results are very promising in two ways: performance measure, and efficiency on the number of steps.\n- (Minor) The authors explicitly provide the limitation of the their work.\n\nWeakness\n- The implementation is not provided (even though the authors provide algorithms and pseudo-codes, and promised to release the code).\n- No error bars are provided in the experimental result.\n- Question: does each step takes similar time? Else, time consumption might be provided.", "summary_of_the_paper": "The paper proposes Autoregressive Diffusion Models (ARDMs), which is a combination of two concepts: autoregressive models and (discrete) diffusion models. The proposed ARDMs generalize order-agnostic ARMs [1] and discrete diffusion models [2]. ARDMs are efficient in several ways: they do not require the causal dependency so that the can be trained in efficient way, and they provide the parallel generation. The authors provide proper experiments that back up their work.\n\n[1] Uria et al., A deep and tractable density estimator, 2014.\n[2] Austin et al., Structured denoising diffusion models in discrete state-spaces, 2021.", "summary_of_the_review": "I lean to positive on this work. The work generalizes two frameworks: OA-ARM and discrete diffusion model. The provided method is simple to implement and efficient is generating mode. Also, the authors conduct experiments across various domains, and the result are very promising. Therefore, I recommend this paper to be accepted in ICLR 2022.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "de1Ir-T9WjF", "reviewer": "Reviewer_g6Qq", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "## 1. Strengths\n\n- a. Strong contributions:\n     - ARDM generalize a wide class of models (autoregressive models, order agnostic autoregressive models (XL-Net))\n     - ARDM allow unfolding of the generative process with *depth upscaling* (translates in using more denoising steps)\n     - ARDM allow parallelizing the generative process for fast generation  \n- b. even the authors introduce multiple contributions, the paper is well written, easy to read and well-structured. \n- c. Equations stated in the main text are correct\n- d. Related work is well described. This is very clear how the method relates to D3PM.\n- e. Ideas are supported by a complete set of experiments\n    - The idea translates into impressive empirical results \n    - The experiments are well designed and interpreted, this gives a good understanding of the functioning of the OA-ARDM\n    - Additional results in the appendix give a good picture of the underlying generative process \n\n## 2. Weaknesses\n- a. Training time and required resources are not mentioned. \n- b. Experiments do not describe how models trained with a single seed are selected (does the method consistently converge to the reported results, or have you selected the best performing seed?). \n- c. The code needs to be released for full reproducibility (which is planned according to the authors) \n\n## 3. Minor comments / questions to the authors\n- a. Could you please elaborate on the limitations of the dynamic programming approach for parallelisation?\n- b. Is the method sensitive to the choice of hyperparameters?\n\n## 4.Disclaimer \n- I am confident in my evaluation of the definition of the OA-ARDM, overall terms of parallelisation, depth upscaling and modelling \n- I only have a general understanding of the described dynamic program and limited experience with the compression experiments", "summary_of_the_paper": "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.\n\nFurthermore, the authors supplement the OA-ARM with two additional contributions:\n- 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$)\n- an algorithm to parallelize generation using dynamic programming (allows trading quality for generation speed)\n\nThe authors demonstrate the effectiveness of OA-ARM on character-level text modelling (text8), image modelling/compression (CIFAR-10).", "summary_of_the_review": "This paper is well-written and introduces significant contributions. The contributions are theoretically significant (generalization of ARMs) and are of practical use (parallelization of the OA-ARDMs, depth upscaling). The method is well related to the literature. The theoretical findings are supported by a complete set of experiments.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "cyoOyBkJos", "reviewer": "Reviewer_Y5ay", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Let me start by confessing that I really struggled with this paper, because (in my humble opinion) it is written with the assumption for the reader to be familiar with not yet published (or recently, where recently means days, accepted) work. In principle, I would be fine with this, if the submitted paper was self-contained such that it would be possible to have a grasp on the presented ideas, to be refined upon reading the literature.\nUnfortunately, at least for me, this was not the case: I imperatively had to gain familiarity with unpublished work to fully appreciate the technical merit of this paper. Specifically, this is true for discrete diffusion processes (heavily based on a paper recently accepted at NeurIPS 2021, and a paper updated on Arxiv as recently as few days ago) and for the dynamic programming part (an arxiv paper appeared a few months ago).\n\nAs another disclaimer (which instead is not affecting my positive judgment for the paper), I found the experimental section to require extensive knowledge in \"application\" domains where I have passing familiarity (language modeling, neural compression) but by no means I consider myself as an expert.\n\nWith all that being said, here is my detailed review:\n\n* Section 1:\nOne suggestion I have to improve the intro is to focus more on the key issues (especially computational ones) that are addressed by this work. Also, upscaling seems to be the most original contribution of this work, and it is only mentioned in passing. I would instead avoid to go into technicalities that can be known only by substantial familiarity with unpublished or recently accepted work, such as Austin et al. 2021.\n(By the way, I really liked the paper Austin et al. 2021; this is a necessary read prior to working on this paper!)\n\n* Section 2:\nThe background is not enough, and all the accent is put on order agnostic methods. This is, in my humble opinion, problematic because there is no mention to discrete diffusion models, which are necessary to understand this work.\nAnother note: with reference to Yang et al. 2019, it is brought to the reader attention that if you shuffle the order of the r.v., you can still have a reasonable bound to the log likelihood, but the training procedure requires model architectures that are invariant to the permutation order, and this looks like a limitation you want to avoid. This is important in my opinion, and not sufficiently exploited or emphasized, e.g. in the experimental section.\n\n* Section 3:\nI think here I need more references and a more thorough discussion about the requirement for triangular or causal dependence to parametrize an autoregressive model, which seems to be the one pain point you aim to address. Also, I think that the goal of upscaling variables comes out of the blue, and it is not until later in the paper that we learn about the modeling approach to see it as a diffusion process, that we learn the connection with the presented work.\n\nFig. 1 and 2 are very useful to understand the key idea of the proposed parametrization! One thing that induced me in error the first (of the many) time I read this paper is the use of $\\theta$: I generally consider these are parameters, whereas this is the output of a \"neural network\".\n\nOverall, I really miss a clear connection to discrete diffusion processes in the first part of Section 3. Instead, the reader is assumed to be familiar with the details, and the final discussion before subsection 3.1 does not help.\n\n* Section 3.1:\nI consider this as a key contribution, even though it is based on an existing dynamic programming approach. The observation in eq (3) is key.\nI was wondering, what would happen if the \"shape\" of $L_t$ wouldn't be so well behaved as you depict in Figure 3. Although $L_t$ is monotonically decreasing, it could be so that with a small budget you would incur in a larger error than we may think: for example, consider a linearly decreasing $L_t$.\n\n* Section 3.2:\nThis is, in my opinion, the main contribution of this work, in that as I see it, it is the most original. However, it again assumes complete familiarity with an unpublished work (Austin et al., 2021), and little context and background information is given. As a consequence, it is difficult to appreciate, for example, the parametrization presented in this work (the transition matrices here it are deterministic, theirs are stochastic).\n\nSection 5 is very well executed, I have no major comments (nor all the required expertise to nitpick on experimental setup choices).\n\nSo to conclude:\n\n== Strength:\n* methods to improve the efficiency of autoregressive models\n* upscaling as a discrete autoregressive diffusion model\n* compelling experimental results\n\n== Weaknesses:\n* the paper is largely not self-contained, making it hard to fully grasp its merits\n* the discrete probabilistic diffusion nature of the proposed models is not put forward enough\n* main methodological contribution drowned into a series of careful engineering contributions\n\n", "summary_of_the_paper": "This is a very dense paper that proposes a new autoregressive modeling (ARM) approach, which enjoys several properties. The proposed approach builds on two main ideas: 1) order-agnostic ARM which they exploit to randomize over the generation of variables according to a random order picked uniformly from all possible permutations of orders, 2) denoising diffusion probabilistic models extended to work on discrete variables.\nThen, one contribution of this paper is to improve the computational cost of order-agnostic ARM. This stems from rewriting the log likelihood component at a given timestep $t$, and noticing that this single term can be optimized for all datapoints in a minibatch simultaneously. The parametrization of the neural network proposed in this work is also an improvement over a traditional approach to generation in an autoregressive diffusion model.\nFurthermore, eq (3) is key for the realization that the un-ordered generation of $k$ tokens can happen independently, because this would contribute to $k$ times the log likelihood at a given timestep $t$. Using ideas from a related work, in this paper the authors show how to exploit this fact and a dynamic programming approach to parallelize autoregressive diffusion models for a given budget.\n\nAn additional contribution is to consider the application of the proposed model to the problem of upscaling. This is modeled as a (discrete) diffusion process where the forward process is destructive (downscaling) and the backward process is constructive (upscaling).\nAs discussed originally in the work used by the authors as a basis for discrete diffusion processes, upscaling can be defined using simple transition matrices, which define the transitions in the Markov chain supporting the process from a downscaled version of the data to an upscaled version. The authors present two variants of a parametrization approach to learn the backward distributions.\n\nAn extensive experimental campaign complements the ideas presented in this paper, with application ranging from the image to the text domains, including upscaling and lossless compression.", "summary_of_the_review": "It is not an easy task (for me) to come up with a recommendation for this paper. I didn't find any technical flaws, I think the topic is very interesting and important, and the empirical validation of the proposed method(s) is extensive.\n\nNevertheless, I find this paper has problems with the positioning, and with the choice of what the authors consider their main contribution. Also, a lot of the proposed methods are based on other (recent) work, which are assumed to be perfectly known by the reader: this makes this paper not self-contained, and to some extent reduces the surface of the contributions, which seem additional engineering prowesses on top of existing methods.\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "25MKzkgK9Y3", "reviewer": "Reviewer_Pog3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n+ Intuitive and novel approach that aims to combine benefits from two different modeling techniques\n+ Extension of a widely used model class (autoregressive models)\n+ Vast practical applicability\n\nWeaknesses:\n- Performance gains in benchmarking tests appear to be rather moderate\n- The number of benchmarking tests is limited; robustness of the result remains unclear\n\n", "summary_of_the_paper": "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.  ", "summary_of_the_review": "Overall, this is a well reported work that presents useful new modeling ideas that bring together previously separate methodology.  The work could, however, benefit from a more comprehensive set of benchmarking with different data sets / case studies to establish how robust the claimed performance gains are.\n\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper introduces Autoregressive Diffusion Models (ARDMs), which generalises order-agnostic autoregressive models and absorbing discrete diffusion.

All reviewers appreciated the paper with a few also finding it very dense. The experimental section is a bit lacking in detail. This has to some degree been answered in the discussion and should also be included in the final version of the paper. 

Acceptance is recommended.
Decision: Accept (Poster)