Autoregressive Distillation of Diffusion Transformers

TL;DR

ARD (AutoRegressive Distillation) is a few-step diffusion-distillation recipe that conditions the student on the history of the ODE trajectory rather than only the most recent denoised sample. The motivation is exposure bias: standard step-distilled students take their own last output as input and accumulate error along the few-step trajectory. By exposing the student to the entire sequence of previous trajectory states (token-wise time embedding to mark each, block-wise causal mask, history injected only in lower transformer layers), ARD cuts FID degradation 5× vs prior baselines on ImageNet-256 at 1.1% extra FLOPs, reaches FID 1.84 at 4 steps on ImageNet-256, and matches teacher prompt adherence on 1024p T2I distillation.

Key claims

Existing few-step distillation methods condition the student only on the most recent denoised sample, which makes them susceptible to exposure bias as errors accumulate over the trajectory [§1, §3].
Conditioning on the predicted historical trajectory (multiple prior states, not just one) supplies coarse-grained information that is less corrupted by accumulated single-step error than the latest sample alone [§3].
Architectural delta from the teacher DiT: token-wise time embedding to tag each history input with its trajectory index, plus a block-wise causal attention mask so each step attends to itself and earlier steps [§3, Fig. 3].
Injecting history only in the lower transformer layers (not all layers) improves both performance and efficiency — the deeper layers stay focused on refining the current step [§3, ablations].
ImageNet-256 class-conditioned: 5× reduction in FID degradation vs baseline few-step distillation methods at 1.1% extra FLOPs [Abstract, §4].
ImageNet-256 4-step generation: FID 1.84 [Abstract, §4].
1024p T2I distillation: outperforms publicly available 1024p T2I distilled models on prompt-adherence score with a minimal FID drop relative to the teacher [Abstract, §4].

Method

ARD takes a pretrained diffusion transformer teacher and trains a few-step student whose input at each step is not the single denoised latent $x_{t_i}$ but the concatenated history $(x_{t_0}, x_{t_1}, \ldots, x_{t_i})$ along the ODE trajectory. Each history entry gets a token-wise time embedding identifying which step it came from, and the transformer applies a block-wise causal mask so step $i$ attends to steps $0..i$ but not future ones. To keep extra compute small, the history channels are only fed into the lower transformer layers; upper layers operate on a single-step representation as in the teacher. Training mirrors standard few-step distillation but uses these history-augmented inputs as the student’s view of the trajectory, so the student learns to predict the next ODE step from a trajectory prefix rather than from its own latest noisy guess.

Results

ImageNet-256, class-conditioned, 4 NFE: FID 1.84 [Abstract] — strong few-step result on a standard benchmark.
5× reduction in FID degradation vs baseline few-step distillation at matched step budgets, at only 1.1% extra FLOPs [Abstract].
1024p T2I: ARD beats publicly available 1024p T2I distilled models on prompt-adherence score with minimal FID drop relative to the (many-step) teacher [Abstract].
The history-injection-in-lower-layers-only ablation shows this is both a performance and an efficiency win versus full-layer history injection [§3].

Why it’s interesting

ARD (April 2025, ICLR 2025) is the conceptual antecedent of the AR-video distillation lineage on the wiki — Causal Forcing (Causal Forcing: Autoregressive Diffusion Distillation Done Right for High-Quality Real-Time Interactive Video Generation) and Causal Forcing++ (Causal Forcing++: Scalable Few-Step Autoregressive Diffusion Distillation for Real-Time Interactive Video Generation) both treat distillation as an AR-over-trajectory problem and address the same exposure-bias failure mode, but in the temporal-chunk dimension of video rather than the denoising-step dimension of a single image. ARD’s “history in lower layers only” is a clean parameter-efficient way to widen the conditioning context that complements Streaming Autoregressive Video Generation via Diagonal Distillation‘s chunk-axis asymmetric schedule and Cross-Resolution Distribution Matching for Diffusion Distillation‘s resolution-axis cascade — three orthogonal ways to give a few-step student more useful conditioning than the bare last denoised sample. It also contrasts cleanly with Diversity-Preserved Distribution Matching Distillation for Fast Visual Synthesis (DP-DMD), which fixes a different failure mode (reverse-KL mode-seeking) by role-separating which objective each step uses while still conditioning each step on just the previous output.