CoReDiT: Spatial Coherence-Guided Token Pruning and Reconstruction for Efficient Diffusion Transformers

TL;DR

CoReDiT introduces a spatial coherence-guided token pruning method for diffusion transformers, significantly reducing computational costs while preserving image quality and enabling efficient on-device image and video generation.

Contribution

It proposes a novel structured token pruning framework with coherence-guided reconstruction and a progressive pruning schedule for diffusion transformers.

Findings

Achieves up to 55% reduction in self-attention FLOPs.

Speeds up inference by 1.33x on cloud GPUs and 1.72x on mobile NPUs.

Maintains high visual quality despite pruning.

Abstract

Diffusion Transformers (DiTs) deliver remarkable image and video generation quality but incur high computational cost, limiting scalability and on-device deployment. We introduce CoReDiT, a structured token pruning framework for DiTs across vision tasks. CoReDiT uses a linear-time spatial coherence score to estimate local redundancy in the latent token lattice and skips high coherence (redundant) tokens in self-attention. To maintain a dense representation and avoid visual discontinuities, we reconstruct skipped attention outputs via coherence-guided aggregation of spatially neighboring retained tokens. We further introduce a progressive, block-adaptive pruning schedule that increases pruning gradually and allocates larger budgets to blocks and denoising steps with higher redundancy. Across state-of-the-art diffusion backbones including PixArt-{\alpha} and MagicDrive-V2, CoReDiT achieves up to 55% self-attention FLOPs reduction and inference speedups of 1.33x on cloud GPUs and 1.72x on mobile NPUs, while maintaining high visual quality. Notably, CoReDiT also increases on-device memory head-room, enabling higher-resolution generation.