DADP: Drift-Controlled Denoising-Aware Dynamic Token Pruning for Efficient Diffusion Transformers

Abstract

Diffusion Transformers (DiTs) are powerful backbones for visual pattern synthesis, but iterative denoising repeatedly applies full-token self-attention, limiting high-throughput recognition workflows. We propose drift-controlled Denoising-Aware Dynamic Token Pruning (DADP), a plug-and-play training-free inference framework for efficient DiT sampling. DADP scores tokens using normalized activation energy and centered-token deviation, allocates retention through a layer-timestep scheduler, and stabilizes token support with exponential moving average continuity. Selected tokens are processed as a compact sequence, after which their residual updates are scattered to the original grid while unselected tokens follow an identity bypass. The centered deviation is interpreted as an internal covariance-drift statistic, linking token selection to hidden-distribution preservation without modifying pretrained weights, conditioning paths, or output shapes. On a 28-layer DiT with 50 steps at $256\times256$, DADP has a mean retention ratio of 0.595, reduces attention FLOPs by 58.9\%, and yields $2.43\times$ attention-only and $1.33\times$ measured end-to-end speedups, with FID$_{\mathrm{DiT}}$ 4.5081. Under the conservative $512\times512$ budget, it retains 0.707 of tokens, reduces attention FLOPs by 43.6\%, and yields $1.77\times$ attention-only and $1.27\times$ end-to-end speedups, with FID$_{\mathrm{DiT}}$ 5.7828. At both resolutions, DADP improves dense-reference FID over fixed random and similarity pruning.