FREE: Uncertainty-Aware Autoregression for Parallel Diffusion Transformers

TL;DR

This work tackles the latent latency of Diffusion Transformers (DiTs) by introducing FREE, a feature-level autoregressive drafting framework that leverages the temporal stability and semantic richness of DiT top-block features. A lightweight drafter forecasts multiple future top-block states, which are verified in parallel against the frozen DiT via reflection maximal coupling, ensuring lossless acceleration. To cope with increasing uncertainty in later denoising steps, FREE (relax) adds an uncertainty-guided relaxation that dynamically adjusts acceptance probabilities, further boosting speed. Across ImageNet $256^2$ and $512^2$ settings, FREE achieves significant wall-clock speedups (up to $1.86$–$2.25\times$) with fidelity comparable to the baseline, proving the practicality of staged, parallelized diffusion for large-transformer diffusion models.

Abstract

Diffusion Transformers (DiTs) achieve state-of-the-art generation quality but require long sequential denoising trajectories, leading to high inference latency. Recent speculative inference methods enable lossless parallel sampling in U-Net-based diffusion models via a drafter-verifier scheme, but their acceleration is limited on DiTs due to insufficient draft accuracy during verification. To address this limitation, we analyze the DiTs' feature dynamics and find the features of the final transformer layer (top-block) exhibit strong temporal consistency and rich semantic abstraction. Based on this insight, we propose FREE, a novel framework that employs a lightweight drafter to perform feature-level autoregression with parallel verification, guaranteeing lossless acceleration with theoretical and empirical support. Meanwhile, prediction variance (uncertainty) of DiTs naturally increases in later denoising steps, reducing acceptance rates under speculative sampling. To mitigate this effect, we further introduce an uncertainty-guided relaxation strategy, forming FREE (relax), which dynamically adjusts the acceptance probability in response to uncertainty levels. Experiments on ImageNet-$512^2$ show that FREE achieves up to $1.86 \times$ acceleration, and FREE (relax) further reaches $2.25 \times$ speedup while maintaining high perceptual and quantitative fidelity in generation quality.

Figures (7)

• Figure 1: Overview of the FREE framework. (a) In FREE, a lightweight draft model autoregressively generates six speculative drafts in the top-block feature space. The first five drafts, together with the input, are packed and sent to the target DiT for one pass to obtain the reference outputs. Each draft is sequentially verified against the reference: accepted drafts are kept; Upon rejection, a reflection coupling is triggered and all subsequent drafts are ignored. (b) In FREE (relax), an uncertainty-aware relaxation mechanism is introduced to dynamically adjust the acceptance probability, avoiding over-constraining the draft model under high uncertainty.
• Figure 2: Architecture of the proposed FREE framework. Left: A single forward pass of DiT. Right: The feature-level autoregressive process of the drafter. Red numbers indicate the index of each draft forward pass. "Sampling" denotes drawing samples from the predicted distribution based on the model outputs.
• Figure 3: Temporal stability and in-layer feature change across DiT layers during diffusion sampling. Top: Temporal cosine similarity between features at consecutive steps ($t$). Transformer blocks maintain high similarity ($\ge 0.995$) throughout most denoising steps, especially the top block (block-28), while the final output head (pred-output) shows lower temporal stability. Bottom: In-layer feature change $\|\mathrm{output} - \mathrm{input}\|_2$ at early (step 2) and late (step 998) denoising stages. The top block (layer 28) exhibits the largest change at both stages, consistent with it being the most semantically expressive layer.
• Figure 4: Reflection Maximal Coupling in 1D. $p(x)$ and $q(x)$ denote the draft and target distributions, respectively.
• Figure 5: Prediction uncertainty of DiT under different resolutions and classifier-free guidance (cfg) scales. We run the DiT as both drafter and verifier in speculative sampling. Each subplot shows the mean-squared prediction deviation between the drafts and their reference predicted outputs ($\textcolor{blue}{epsilon\_delta}$), along with the acceptance probability ($\textcolor{red}{accept\_prob}$) across denoising steps.