Autoregressive Diffusion Models

TL;DR

Autoregressive Diffusion Models (ARDMs) unify and extend existing autoregressive and diffusion models, offering efficient training, parallel generation, and superior performance in data compression tasks, including single data point compression.

Contribution

This paper introduces ARDMs, a new class of models that generalize previous models, enabling scalable training, parallel sampling, and effective lossless compression.

Findings

ARDMs require fewer steps than discrete diffusion models for similar performance.

ARDMs support parallel generation adaptable to various computational budgets.

ARDMs achieve strong results in lossless compression, including single data point scenarios.

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.