💻 computer science

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

CoReDiT is an efficient token pruning framework for Diffusion Transformers that leverages spatial coherence to identify and skip redundant tokens while reconstructing their outputs via neighboring tokens, achieving significant computational speedups and memory savings without compromising visual quality.

Original authors: Zhuojin Li, Hsin-Pai Cheng, Hong Cai, Shizhong Han, Fatih Porikli

Published 2026-05-15

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Zhuojin Li, Hsin-Pai Cheng, Hong Cai, Shizhong Han, Fatih Porikli

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to paint a massive, incredibly detailed mural on a wall. You have a team of artists (the "tokens") standing in front of the wall, and every single artist needs to talk to every other artist to decide exactly what color to put on their specific spot. This ensures the picture looks perfect, but it's exhausting. If you have 10,000 artists, they have to hold 100 million conversations just to agree on the next brushstroke. This is how current AI image generators (called Diffusion Transformers) work: they are brilliant but incredibly slow and hungry for computer power.

The paper introduces a new method called CoReDiT to make this process faster without ruining the painting. Here is how it works, broken down into three simple steps:

1. The "Bored Artist" Detector (Spatial Coherence)

In a typical painting, some areas are very busy and detailed (like a face or a tree), while others are just plain, boring backgrounds (like a blue sky or a smooth wall).

The Problem: The old AI treats every artist the same, forcing the artists painting the boring blue sky to have the same intense conversations as the artists painting the face.
The CoReDiT Solution: CoReDiT looks at the artists and asks, "Who is standing next to someone who looks exactly like them?" If an artist is in a boring, uniform area, their neighbors are practically clones. CoReDiT gives these "bored" artists a high Coherence Score.
The Action: It tells the high-scoring (redundant) artists, "You don't need to talk to everyone right now. Take a break." This skips the expensive conversations for the boring parts of the image, saving a massive amount of time.

2. The "Fill-in-the-Blank" Artist (Reconstruction)

If you just told the "bored" artists to go home and stop working, you'd end up with holes in your mural. The picture would look broken and patchy.

The Problem: Simply removing tokens (artists) creates visual glitches.
The CoReDiT Solution: Instead of leaving a hole, CoReDiT uses the artists who did stay to fill in the blanks. Because the "bored" artists were so similar to their neighbors, the remaining artists can easily guess what the missing ones would have painted.
The Action: It's like a game of "telephone" where the remaining artists whisper the correct color to the empty spots. This ensures the final image looks smooth and complete, with no visible gaps, even though fewer artists were actually doing the heavy lifting.

3. The Smart Manager (Progressive Pruning)

Not all parts of the painting process are the same.

The Problem: At the very beginning of making an image, the picture is just static noise (like TV snow), so almost everything looks the same and is redundant. But at the very end, when the details are being sharpened, every single brushstroke matters. If you cut too many artists too early, the AI gets confused. If you cut too many at the end, the picture gets blurry.
The CoReDiT Solution: CoReDiT acts like a smart manager who learns the perfect schedule. It starts by letting the AI cut out many redundant artists when the image is just noise. As the image gets clearer and more detailed, the manager slowly stops cutting so many artists, ensuring the critical details get full attention.
The Action: It automatically figures out which parts of the AI brain (blocks) and which stages of the process (time steps) can handle the most cuts, adjusting the plan on the fly to keep the quality high.

The Results

By using this method, the authors found that:

Speed: The AI runs 1.33 times faster on powerful cloud computers and 1.72 times faster on mobile phone chips.
Efficiency: It cuts the heavy math work (computations) by up to 55%.
Quality: The images still look just as good as the slow version, with no obvious cracks or weird artifacts.
Bonus: Because it uses less memory, it can even generate higher-resolution images on phones that would normally crash due to running out of memory.

In short, CoReDiT is like hiring a smart foreman for an art project who knows exactly when to let the crew rest and when to keep them working hard, resulting in a masterpiece that gets done much faster.

Technical Summary: CoReDiT

Problem Statement

Diffusion Transformers (DiTs) have established state-of-the-art performance in image and video generation by effectively modeling long-range dependencies through attention mechanisms. However, their computational cost scales quadratically with the number of tokens, making high-resolution synthesis and multi-frame video generation prohibitively expensive for resource-constrained mobile platforms. While a significant portion of tokens in the latent space correspond to visually redundant regions (e.g., uniform backgrounds), existing token pruning methods face three primary challenges when applied to diffusion inference:

Efficient Localization: Identifying low-saliency tokens without incurring the quadratic overhead of computing global attention scores.
Semantic Preservation: Maintaining visual continuity and preventing artifacts (e.g., texture breaks, boundary distortions) when tokens are skipped, which is critical for generative tasks unlike discriminative ones.
Adaptive Scheduling: Determining a pruning strategy that adapts to varying redundancy levels across different transformer blocks and denoising timesteps.

Methodology: CoReDiT

CoReDiT is a structured token pruning framework designed to accelerate pre-trained DiTs while preserving output fidelity. It operates through three core components:

1. Spatial Coherence-Based Token Selection

Instead of relying on global attention scores (which are $O(N^2)$ ), CoReDiT utilizes a Spatial Coherence (SC) score to estimate local redundancy in linear time ( $O(N)$ ).

Mechanism: The token lattice is partitioned into non-overlapping grids. For each token, the SC score is calculated as the mean cosine similarity between the token's normalized embedding and the mean normalized embedding of its grid.
Logic: Tokens with high SC scores are deemed redundant because they are well-explained by their spatial neighbors.
Selection: The top- $K$ tokens with the highest SC scores are selected for skipping ( $X_S$ ), while the remaining tokens ( $X_R$ ) proceed to the Multi-Head Self-Attention (MHSA) layer.

2. Coherence-Guided Token Reconstruction

To prevent the spatial discontinuities caused by skipping tokens, CoReDiT reconstructs the attention outputs for skipped tokens ( $Y_S$ ) using the outputs of retained neighbors ( $Y_R$ ).

Approximation: Rather than computing pairwise similarities for reconstruction (which adds overhead), the method approximates the similarity weights using the pre-computed SC scores of the neighbors.
Local Sub-grids: To avoid over-smoothing, the reconstruction is performed within smaller sub-grids.
Micro-Designs:
- Alternating Grid Sizes: Grid dimensions alternate between consecutive blocks to shift boundary artifacts and allow information propagation across grid borders.
- m-Stride Retention: A deterministic stride pattern ensures at least one token is retained in every sub-grid, providing necessary anchors for reconstruction.

3. Progressive, Block-Adaptive Pruning Schedule

Recognizing that redundancy varies across transformer blocks and denoising steps, CoReDiT employs a progressive pruning strategy during lightweight fine-tuning.

Block Adaptivity: The pruning budget is allocated to blocks with the highest measured redundancy (estimated by the sum of SC scores of candidate tokens).
Timestep Decay: Pruning is more aggressive in early denoising steps (where noise dominates and structure is coarse) and reduced in later steps (where fine details are refined). A two-phase decay schedule is used to balance efficiency with the complexity of managing step-specific computation graphs.
Training: The model is fine-tuned using a distillation loss that combines the original model output loss and a pruned-block loss to recover performance.

Key Contributions

Linear-Time Redundancy Estimation: The introduction of a spatial coherence score that identifies redundant tokens in $O(N)$ time, avoiding the computational bottleneck of global attention analysis.
Dense Lattice Preservation: A reconstruction mechanism that synthesizes skipped attention outputs from neighbors, ensuring the token lattice remains dense for subsequent layers and preventing visual artifacts.
Adaptive Pruning Strategy: A progressive, block-adaptive schedule that dynamically allocates pruning budgets based on empirical redundancy statistics across depth and time, facilitating stable adaptation.

Experimental Results

CoReDiT was evaluated on state-of-the-art DiT backbones, including PixArt-α, PixArt-Σ, and MagicDrive-V2.

Efficiency Gains:
- Achieved up to 55% reduction in self-attention FLOPs.
- Inference Speedups: 1.33× on cloud GPUs (Nvidia H100) and 1.72× on mobile NPUs (Qualcomm Snapdragon 8 Elite).
- Memory Headroom: Enabled generation at resolutions up to 1920 on mobile devices, where the baseline model encountered Out-of-Memory (OOM) errors at 1600.
Quality Preservation:
- On PixArt-α-1024 with 40% pruning, CoReDiT maintained an FID of 28.7 (vs. baseline 27.3) and competitive CLIP/IS scores. With distillation, the FID gap narrowed to 27.4.
- On PixArt-Σ-2048, a 23% average pruning ratio reduced self-attention latency by 22% with only a modest drop in quality metrics.
- On MagicDrive-V2 (video), a 39% reduction in self-attention FLOPs preserved per-frame quality (PSNR, LPIPS, SSIM) and conditional alignment (mAP, mIoU), though a slight drop in temporal consistency (FVD) was noted.
Ablation Studies:
- Random token selection resulted in catastrophic quality degradation (FID 644.5), validating the necessity of coherence-based selection.
- Disabling reconstruction led to significant quality loss (FID 30.4 vs. 28.5), confirming the importance of the reconstruction step.
- Uniform pruning followed by fine-tuning performed worse than the proposed progressive adaptive schedule.

Significance and Claims

The paper claims that CoReDiT offers a practical solution for deploying high-fidelity diffusion models on edge devices by addressing the specific challenges of generative token pruning. Its significance lies in:

Scalability: Enabling higher-resolution synthesis and video generation on mobile hardware by reducing both compute and memory footprints.
Architectural Compatibility: Maintaining compatibility with existing DiT architectures without requiring structural changes or complex routing modules.
Efficiency-Quality Trade-off: Demonstrating that structured, coherence-guided skipping can drastically reduce inference costs (up to 55% self-attention FLOPs) while maintaining perceptual quality comparable to the full model.

The authors note that while the current approach focuses on spatial redundancy, future work aims to extend efficiency to FFN/MLP components and incorporate cross-frame coherence objectives to further improve video temporal consistency.