PaperlayerView exact revision on arXiv

Source-linked AI summary

Autoregressive Diffusion Models

Emiel Hoogeboom, Alexey A. Gritsenko, Jasmijn Bastings, Ben Poole, Rianne van den Berg, Tim Salimans

arXiv:2110.02037v2cs.LGstat.ML

TL;DR

Autoregressive and discrete diffusion models face constraints in generation order, architectural design, and the number of network calls required for sampling or compression. The paper introduces ARDMs, which combine arbitrary-order generation with diffusion-like training and adaptable parallel generation, achieving competitive generative performance and superior per-image lossless compression. The authors also identify gaps to single-order autoregressive models and scope ARDMs primarily to discrete variables and log-likelihood objectives.

  • Problem

    Autoregressive models require a pre-specified generation order and dimensionality-scaled sampling calls, while discrete diffusion may require long processes and many network calls.

  • Method

    ARDMs learn arbitrary generation orders using a diffusion-like single-step objective, without causal masking, and support parallel generation of multiple variables.

  • Results

    ARDMs perform similarly to or better than discrete diffusion models while using fewer steps, and achieve 2.71 bpd per-image compression versus 3.26 bpd for the next-best model.

  • Takeaways & Limitations

    ARDMs are especially suited to per-image lossless compression while retaining competitive performance across generative-model comparisons.

  • Takeaways & Limitations

    ARDMs still trail single-order autoregressive models on text, are described for discrete variables, and may require different architectures for objectives such as sample quality.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Existing generative models face trade-offs in generation order, sampling cost, and diffusion-step count. ARDMs address these issues by learning arbitrary generation orders while supporting efficient and parallelizable generation.

  • Motivation: Autoregressive models factorize high-dimensional distributions effectively but require a fixed generation order and one sampling call per data dimension.The order may be unclear for modalities such as images.
  • Motivation: Discrete diffusion models learn reverse processes for discrete destruction schemes, but may require long processes and many network calls.Long schedules can affect both sampling and likelihood evaluation.
  • Autoregressive Diffusion Models: ARDMs learn to generate variables in any order and generalize order-agnostic autoregressive and absorbing diffusion models.The paper presents these relationships as part of the ARDM model class.
  • Autoregressive Diffusion Models: ARDMs impose no architectural constraints on neural networks used to predict distribution parameters, unlike standard autoregressive models.This removes the need for the usual causal architectural restriction described for standard ARMs.
  • Autoregressive Diffusion Models: ARDMs can be parallelized to generate multiple tokens simultaneously without substantial performance reduction and require fewer steps than absorbing models for comparable performance.The generation process can therefore be adapted to different computational budgets.

2 BACKGROUND

The background contrasts standard and order-agnostic autoregressive models with discrete diffusion. These approaches differ in factorization, ordering, architectural constraints, and how their generative processes are learned.

  • Autoregressive Models: Autoregressive models factorize a multivariate distribution into D univariate conditionals using the probability chain rule.Their likelihood can be computed in parallel, but sampling proceeds through D iterative conditional steps.
  • Autoregressive Models: Standard autoregressive models enforce triangular dependencies in the network, commonly through causal masking.This architectural structure supports parallel likelihood computation but not parallel ancestral sampling.
  • Order-Agnostic ARMs: Order-agnostic autoregressive models generate variables according to a random permutation σ drawn from the set of all D-dimensional orderings.Their likelihood is treated through an expectation over generation orders.
  • Order-Agnostic ARMs: Permutation-based order-agnostic training can use causally masked, permutation-equivariant networks, but may suffer in likelihood and cannot combine with simple non-equivariant transformations such as convolutions.The limitation is tied to the architectural requirements of that approach.
  • Discrete Diffusion: Absorbing discrete diffusion defines a destruction process in which variables can decay to an absorbing state, then learns the reverse process over absorbed variables.The learned generative process reconstructs variables that are currently absorbed.

3 AUTOREGRESSIVE DIFFUSION MODELS

ARDMs generate variables in arbitrary orders while training with a single-step objective, avoiding causal architectural constraints. Their parametrization supports parallel generation and depth upscaling through staged refinement.

  • ARDMs generate variables in an arbitrary order, extending order-agnostic autoregressive modeling with a diffusion-like training perspective.
  • Training samples one likelihood component L_t uniformly for each datapoint instead of optimizing all components simultaneously.The objective reweights the sampled component by D.
  • A shared masked network predicts distributions for variables conditioned on the currently revealed subset, without requiring a separate network for every permutation and step.The input mask identifies observed variables, while the complementary mask selects variables whose categorical distributions are predicted.
  • 3.1 PARALLELIZED ARDMS: Parallelized ARDMs use dynamic programming to schedule multiple variables at once, trading a calibration-dependent likelihood cost for fewer generation steps.Variables assigned the same schedule height are inferred and generated in parallel.
  • 3.2 DEPTH UPSCALING ARDMS: Depth upscaling generates increasingly refined representations in stages, such as successive significant bits, by learning reverse transitions of a downscaling process.Each stage models p(x^(s)|x^(s−1)); stage-specific permutation and timestep sampling leaves training complexity unchanged despite up to D · S sampling steps.

4 RELATED WORK

ARDMs relate order-agnostic autoregressive models and discrete diffusion while addressing limitations in arbitrary-order parameterization and iterative generation. Related approaches include masked permutation-equivariant models, conventional diffusion, and other parallelized autoregressive methods.

  • Autoregressive Models: Autoregressive models factorize joint distributions into conditional distributions, but sampling conventionally requires as many iterative network calls as data dimensions.
  • Order Agnostic Models: Order-agnostic sequence models train by masking and predicting variables, whereas ARDMs add absorbing-state choices, parallelization, and depth-upscaling techniques.
  • Order Agnostic Models: On text8, Table 1 reports order-agnostic model performance in bits per character, with a single-order Transformer achieving 1.35 bpc.
  • Order Agnostic Models: On CIFAR-10, Table 2 evaluates order-agnostic modeling in bits per dimension, including an upscaling model that generates four categories, equivalent to two bits, at a time.
  • Diffusion Models: Discrete diffusion operates directly on discrete spaces through destruction and learned reverse processes; absorbing diffusion was reported as the best-performing transition choice for text log-likelihood among the cited approaches.

5 RESULTS

ARDMs perform competitively across order-agnostic modelling and depth-upscaling experiments, while showing especially strong results for per-image lossless compression. Their performance remains comparatively robust when generation or encoding uses fewer steps.

  • Order Agnostic Modelling: 1.43 bpc for OA-ARDM is close to D3PM-absorbing at 1000 steps with 1.45 bpc, while requiring one quarter as many steps.At 256 steps, D3PM performance degrades to 1.47 bpd; a 20-step Parallelized ARDM achieves 1.51 bpd versus 1.56 bpd for a similar D3PM.
  • Lossless Compression: The compression experiment uses ARDM probabilities with the range-based entropy coder rANS for lossless CIFAR-10 image compression.A fixed coding order is selected for all images after evaluating log-likelihood before coding.
  • Lossless Compression: ARDMs significantly outperform competing methods on CIFAR-10 per-image compression, achieving 2.71 bpd versus 3.26 for the next-best IDF++ model.They also perform competitively with VDM's 2.72 bpd when compressing an entire dataset.
  • Lossless Compression: ARDMs degrade more gracefully than related methods when fewer steps are used to encode data.This robustness supports adapting the generative process to different step budgets.
  • Effects of Depth-Upscaling: For images, the best depth-upscaling model uses factor 4 and achieves 2.64 bpd, while the best audio models use factors 2 or 4 and achieve 6.29 bpd.Smaller upscale factors may provide diminishing returns, and can sometimes reduce performance by prolonging the generative process.

6 LIMITATIONS AND CONCLUSION

The paper concludes that ARDMs connect autoregressive and discrete diffusion models and perform competitively with existing generative models. Their strongest reported advantage is per-image lossless compression, although limitations remain for language, continuous variables, and objectives beyond likelihood.

  • LIMITATIONS: ARDMs outperform all other order-agnostic approaches on text but remain behind single-order autoregressive models.Preliminary language experiments found that upscale variants did not outperform order-agnostic versions.
  • LIMITATIONS: ARDMs are described for discrete variables, while analogous absorbing processes for continuous distributions are left as a possible extension.This marks the current scope of the model formulation.
  • LIMITATIONS: The paper focuses on log-likelihood because it directly corresponds to coding length in lossless compression, while other objectives may favor different architectural choices.The authors specifically identify sample quality as an objective where alternative architectures may perform better.
  • CONCLUSION: ARDMs are introduced as a model class at the intersection of autoregressive and discrete diffusion models and outperform competing approaches on per-image lossless compression.The conclusion also characterizes their performance as competitive with existing generative models.

REPRODUCIBILITY AND ETHICS STATEMENT

The paper details its training, sampling, and parallelization procedures and describes plans to release code for reproducing the experiments. It also reports no immediate ethical concerns, while acknowledging possible negative downstream uses of generative models.

  • REPRODUCIBILITY: The paper describes its training and sampling algorithms in detail to support re-implementation and reproducibility.Main ideas appear in Section 3, with further clarifications and hyperparameter settings in the appendices.
  • ETHICS: The authors report no immediate ethical concerns, but note that generative models can support harmful applications such as making false media seem realistic.They also state that, to their knowledge, no datasets with known ethical issues were used.
  • REPRODUCIBILITY: Code for reproducing the experimental results is planned for release, and implementation details are open-sourced.The open-source implementation is identified by its repository URL.
  • PARALLELIZATION: The parallelized ARDM routine computes cost and dimension matrices for finding generation paths under a chosen budget.The transition costs encode generating multiple variables in parallel, while the matrices support optimal-path retrieval.
  • PARALLELIZATION: The JAX implementation reduces runtime from 5 minutes to about 2 seconds for D = 3072 and enables audio experiments with D = 16000 to run in under half a minute.The underlying loop has computational complexity O(D^3) and is XLA-compiled through scan.

B ADDITIONAL RESULTS

On CIFAR-10, ARDMs are competitive with likelihood-based generative models and achieve the best result among discrete diffusion approaches, although VDM performs best overall.

  • ADDITIONAL RESULTS: 2.64 bpd is the best reported ARDM score on CIFAR-10 and the strongest result among discrete diffusion approaches.The Variational Diffusion Model is the best-performing model overall.

B.2 ADDITIONAL AUDIO EXPERIMENTS

Additional SC09 experiments evaluate an Upscale ARDM with branching factor 4 under smaller computational budgets for data with dimensionality D = 16000.

  • B.2 ADDITIONAL AUDIO EXPERIMENTS: The SC09 audio experiments test the best Upscale ARDM with branching factor 4 across smaller computational budgets.SC09 data has dimensionality D = 16000.

B.3 LOSS COMPONENTS OVER TIME

Loss components are tracked over generation steps and used in dynamic programming to select paths under computational budgets. In depth-upscaling models, earlier bits become easier to model while later-stage bits remain more uncertain.

  • LOSS COMPONENTS OVER TIME: Loss components are estimated by sampling steps during training, tracked with an exponential moving average, and supplied to the dynamic programming routine.Within each stage, the components decrease monotonically over step t.
  • LOSS COMPONENTS OVER TIME: After an initial uncertainty of approximately 2 bits, the most significant bits become easier to model quickly, reaching below 0.5 bits.The stated values are approximately 8/S = 2 bits initially and below 2/S = 0.5 bits afterward.
  • LOSS COMPONENTS OVER TIME: Later stages have higher average loss, indicating greater uncertainty for less significant bits from model difficulty, data entropy, or both.Figure 6 plots normalized loss components for OA-ARDM and four-stage ARDM-Upscale 4 on CIFAR-10.

B.4 SAMPLES FROM ARDMS

ARDM sampling visualizations show how language and image generation proceed through partially specified states before producing final samples. For images, order-agnostic and upscale ARDMs follow distinct generative processes.

  • Language: Language samples can be inspected at selected steps, with the underscore representing the absorbing state and the final frame showing the generated sentence.The absorbing state augments the language space with an additional index.
  • Images: Image visualizations end with generated samples in the right-most grid.
  • Images: The OA-ARDM assigns the image absorbing state to value 128, normalized to 0 in the network architecture.
  • Images: The Upscale ARDM uses 0 as its absorbing state by zeroing least significant bits until all values are zero.
  • Images: The Upscale ARDM first generates a coarse low-bit image, whereas the order-agnostic ARDM generates each image value at once.

C EQUIVALENCE OF AO-ARDMS AND ABSORBING DIFFUSION IN CONTINUOUS TIME

The section establishes that AO-ARDMs are equivalent to continuous-time absorbing models by representing the process through independently sampled transition times and values immediately before absorption. It also explains why fixed absorption avoids the inefficiency and factorization loss of independent discrete absorption.

  • Independent Absorbing Process: Independent absorbing diffusion can absorb multiple dimensions simultaneously, forcing its reverse process to model those variables with an independent factorized distribution.This factorization causes information loss, especially when the number of timesteps is small.
  • Independent Absorbing Process: At least 744 of 1000 forward passes do nothing for length-256 text sequences, while reducing T degrades performance.The sufficient value of T is difficult to determine in advance and depends on the data and decay rate.
  • Fixed Absorbing Process: In a fixed absorbing process, exactly one dimension decays per timestep, giving T = D and requiring at most D generative steps.The process can equivalently be described by sampling a permutation and absorbing dimensions in that order.
  • Continuous-Time Equivalence: The continuous absorbing process is represented by transition times and the variable values immediately before absorption, which suffice to reconstruct the remaining process.The reverse process therefore models the finite-dimensional collection {xi(τi), τi}.
  • Continuous-Time Equivalence: The transition-time distributions are known from the cumulative distribution 1 − α(t), and matching the generative process makes the transition-time KL term zero.The derivation uses the small-step approximation log(1 − x) ≈ −x.
  • Continuous-Time Equivalence: Independently and identically distributed transition times induce a uniform modelling order, matching the generative order of an AO-ARDM.The reverse process models one changing dimension at a time and need not condition on the precise transition times.

D EXPERIMENTAL DETAILS

The experiments use established architectures with input-mask adaptations for ARDMs, dataset-specific training and evaluation procedures, and adjusted baselines where overfitting was observed. Setups cover images, language, and audio.

  • Images: CIFAR10 models use typical splits and are evaluated by test log-likelihood after 3000 training epochs, with some results averaged over three seeds.Runs take approximately two weeks on eight TPUv4 devices.
  • Images: Image ARDMs use the same U-Net architecture as Kingma et al., with 32 ResBlocks before and after a middle attention layer and 256 feature channels.The architecture is designed for NLL optimization and lossless compression.
  • Language: Text8 experiments use chunked segments without additional context, with chunk size adjusted from 256 to 250 characters because the dataset splits are not evenly divisible.A baseline Transformer showed that this adjustment does not meaningfully change performance.
  • Language: Language models use a 12-layer Transformer, batch size 512 for ARDMs, no dropout for ARDMs, and 3 · 10^6 training steps.The standard language-model baseline instead uses a smaller batch size and dropout because of observed overfitting.
  • Audio: The SC09 audio subset contains 31158/3643/4107 training/validation/test clips, each one second long at 16 kHz.Shorter clips are right-padded with zeros.
  • Audio: Audio ARDMs use a non-causal WaveNet with 36 blocks, 256 channels, and maximum dilation 2048, while the single-order baseline uses a causal version.The baseline additionally requires weight decay, a smaller batch size, and fewer channels to address rapid overfitting.
  • Audio: Audio models are trained with Adam for 10^6 steps using batch size 256, learning-rate warm-up, EMA evaluation, and a 30-component discretized-logistic mixture output.The mixture output accommodates the large 2^16-class audio output space.
Loading 2110.02037v2…