Research map / Image generation / Fast diffusion sampling
Autoregressive Distillation of Diffusion Transformers
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes AutoRegressive Distillation (ARD), which distills a diffusion transformer's PF-ODE into a few-step student that conditions on the full historical trajectory, reducing exposure bias on ImageNet-256 and 1024p text-to-image (Emu).
Problem
Few-step distillation of diffusion transformers suffers from exposure bias because the student conditions only on the most recent (possibly off-trajectory) sample, and errors accumulate across sampling steps.
Approach
Generalizes step distillation to an autoregressive formulation where each student step attends to the entire historical ODE trajectory (including the ground-truth initial noise) via token-wise time embeddings and block-wise causal attention, with history used only in lower transformer layers to save compute.
Key results
- 4-step ARD on ImageNet-256 with discriminator loss reaches FID 1.84, outperforming the 25-step teacher (FID 2.89) and the step-distillation baseline (FID 3.84 at 4 steps with R+D).
- ARD's FID degradation from teacher at 4 steps (1.43) is about 5x lower than step distillation's (7.36) on ImageNet-256 with regression loss only.
- On CompBench, 3-step ARD distilled from Emu (CFG 7.5) achieves average 59.9 vs 59.5 for DMD2 (4 steps, 2.6B), and ARD (N=6) adds only 1.1% FLOPs over step distillation.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| add historical ODE trajectory inputs with block-wise causal attention | Step Distill. on ImageNet 256p, 4 steps · ImageNet 256p | FID: 10.25 -> 4.32 | Helped |
| add discriminator loss (ablation) | ARD on ImageNet 256p, 4 steps · ImageNet 256p | FID: 4.32 -> 1.84 | Helped |
| increase student steps from 2 to 4 (ablation) | ARD on ImageNet 256p · ImageNet 256p | FID: 6.29 -> 4.32 | Helped |
| increase student steps from 2 to 4 (ablation) | Step Distill. on ImageNet 256p · ImageNet 256p | FID: 10.92 -> 10.25 | Helped |
| replace M4 with window attention mask M2 (ablation) | ARD on ImageNet 256p, 4 steps, N=6 · ImageNet 256p | FID: 4.32 -> 4.75 | Hurt |
| replace M4 with initial-noise attention mask M3 (ablation) | ARD on ImageNet 256p, 4 steps, N=6 · ImageNet 256p | FID: 4.32 -> 4.45 | Hurt |
| use historical inputs in 28 layers instead of 6 (ablation) | ARD on ImageNet 256p, 4 steps · ImageNet 256p | FID: 4.32 -> 4.80 | Hurt |
| use initial noise as an additional input via M3 instead of recent-history window M2 (ablation) | ARD on ImageNet 256p, 4 steps, N=6 · ImageNet 256p | FID: 4.75 -> 4.45 | Helped |
| use all trajectory history rather than step distillation's current sample only (ablation) | Step Distill. on ImageNet 256p, 2 steps · ImageNet 256p | FID: 10.92 -> 6.29 | Helped |
| distill with ARD rather than Step Distil. from the same Emu teacher | 3-step text-to-image generation · MS-COCO 2017 (zero-shot FID 5k) | FID-S: 30.51 -> 30.03 | Helped |
| use ARD rather than public 3-step distillation models | text-to-image generation, CFG 7.5 · CompBench | AVG: 52.6 -> 59.9 | Helped |
| add historical trajectory to the student (ablation) | 4-step distillation on ImageNet 256p · ImageNet 256p | FID: 10.25 -> 4.32 | Helped |
Surprising findings
- ARD with the discriminator loss outperforms its own teacher on ImageNet-256 (FID 1.84 vs 2.89).
- Step distillation's FID barely improves from 2 to 4 steps (10.92 to 10.25) and its recall drops, whereas ARD keeps recall stable.
Limitations
- Text-to-image experiments use an internal pre-training dataset and internal Emu teacher, limiting reproducibility.
- Requires precomputing or online computing of teacher ODE trajectories, adding training cost.
- Gains over step distillation shrink when the discriminator loss is used in some settings; the N and mask choices are tuned on ImageNet.
Experimental setup
Datasets: ImageNet MS-COCO 2017 CompBench internal text-to-image pre-training dataset
Benchmarks: ImageNet 256p MS-COCO 2017 (zero-shot FID 5k) CompBench
Metrics: FID IS Precision Recall FLOPs latency CompBench alignment score FID drop from teacher
Baselines: Knowledge Distillation (KD) Step Distillation DMD2 Lightning LCM-LoRA ADD Imagine Flash Pixart-delta BigGAN StyleGAN-XL GigaGAN MAR
Builds on: DiT step distillation (progressive distillation) DMD2 Emu VAR
Closest papers by meaning
- Diffusion Model is Effectively Its Own Teacher (CVPR 2025 · similarity 0.90)
- Revisiting Non-Autoregressive Transformers for Efficient Image Synthesis (CVPR 2024 · similarity 0.90)
- Distilled Decoding 1: One-step Sampling of Image Auto-regressive Models with Flow Matching (ICLR 2025 · similarity 0.89)
- Faster Diffusion: Rethinking the Role of the Encoder for Diffusion Model Inference (NeurIPS 2024 · similarity 0.89)
- One-Step Diffusion Distillation through Score Implicit Matching (NeurIPS 2024 · similarity 0.89)
- EM Distillation for One-step Diffusion Models (NeurIPS 2024 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.