Autoregressive Distillation of Diffusion Transformers

TL;DR

This work tackles the high compute cost of diffusion transformers by introducing AutoRegressive Distillation (ARD), which conditions a student model on the full historical ODE trajectory rather than only the current denoised sample. ARD employs a modified transformer with time-embedded tokens and a block-wise causal attention scheme, using historical trajectory inputs in the lower layers to capture coarse information and mitigate exposure bias. Empirical results on ImageNet 256p and 1024p text-to-image tasks show ARD achieves up to a 5× reduction in FID degradation compared to baselines, with only a modest increase in FLOPs, and it outperforms public few-step distillation models in text-image alignment. The approach provides a scalable, efficient pathway to high-quality, high-resolution diffusion-generated images and offers robust performance advantages in both class-conditional and text-conditional settings.

Abstract

Diffusion models with transformer architectures have demonstrated promising capabilities in generating high-fidelity images and scalability for high resolution. However, iterative sampling process required for synthesis is very resource-intensive. A line of work has focused on distilling solutions to probability flow ODEs into few-step student models. Nevertheless, existing methods have been limited by their reliance on the most recent denoised samples as input, rendering them susceptible to exposure bias. To address this limitation, we propose AutoRegressive Distillation (ARD), a novel approach that leverages the historical trajectory of the ODE to predict future steps. ARD offers two key benefits: 1) it mitigates exposure bias by utilizing a predicted historical trajectory that is less susceptible to accumulated errors, and 2) it leverages the previous history of the ODE trajectory as a more effective source of coarse-grained information. ARD modifies the teacher transformer architecture by adding token-wise time embedding to mark each input from the trajectory history and employs a block-wise causal attention mask for training. Furthermore, incorporating historical inputs only in lower transformer layers enhances performance and efficiency. We validate the effectiveness of ARD in a class-conditioned generation on ImageNet and T2I synthesis. Our model achieves a $5\times$ reduction in FID degradation compared to the baseline methods while requiring only 1.1\% extra FLOPs on ImageNet-256. Moreover, ARD reaches FID of 1.84 on ImageNet-256 in merely 4 steps and outperforms the publicly available 1024p text-to-image distilled models in prompt adherence score with a minimal drop in FID compared to the teacher. Project page: https://github.com/alsdudrla10/ARD.

Figures (19)

• Figure 1: Samples ($1024\times1024$) generated by our 3-step ARD model, distilled from a 1.7B Emu.
• Figure 3: (a) The proposed transformer architecture for ARD. (b) The visualization of generalized mask options used during training: M1 represents step distillation, while M4 is the default setting of ARD. M2 and M3 are intermediate options between M1 and M4.
• Figure 4: (a) shows an additional inductive bias that we impose by using the historical trajectory in lower layers only. (b, d, f) show the attention scores for each history input (key tokens) during the 2$^{\text{nd}}$, 3$^{\text{rd}}$, 4$^{\text{th}}$ steps when $N=L$. (c, e, g) show the same but with $N=6$. The attention score on input ${\mathbf{x}}_{\tau_{s'}}$ is the sum of attention weights for all key tokens in ${\mathbf{x}}_{\tau_{s'}}$, indicating the portion of ${\mathbf{x}}_{\tau_{s'}}$.
• Figure 5: Generated ImageNet 256p samples from same initial noise ${\mathbf{x}}_{\tau_{S}}$. All distilled models are 4-step models.
• Figure 6: The analysis on design choices (attention mask options & $N$) for 4-step distillation methods.