Attend to Not Attended: Structure-then-Detail Token Merging for Post-training DiT Acceleration

TL;DR

This work tackles the high computational cost of diffusion transformers (DiTs) by introducing SDTM, a post-training, finetuning-free token merging framework that leverages structure-then-detail denoising priors. By applying two stage-specific visual token mergers—similarity-prioritized structure merging (SSM) for early structure-focused steps and inattentive-prioritized detail merging (IDM) for later detail-focused steps—coupled with compression ratio adjusting and prompt reweighting, SDTM dynamically reduces redundancies during generation. Extensive experiments across backbones, schedulers, and datasets demonstrate that SDTM achieves around $1.55\times$ acceleration with negligible image quality loss and improved CLIP/FID metrics in many configurations. The method integrates seamlessly into existing DiT architectures and maintains performance without requiring fine-tuning, offering a practical pathway to faster diffusion-based generation in real-world deployments.

Abstract

Diffusion transformers have shown exceptional performance in visual generation but incur high computational costs. Token reduction techniques that compress models by sharing the denoising process among similar tokens have been introduced. However, existing approaches neglect the denoising priors of the diffusion models, leading to suboptimal acceleration and diminished image quality. This study proposes a novel concept: attend to prune feature redundancies in areas not attended by the diffusion process. We analyze the location and degree of feature redundancies based on the structure-then-detail denoising priors. Subsequently, we introduce SDTM, a structure-then-detail token merging approach that dynamically compresses feature redundancies. Specifically, we design dynamic visual token merging, compression ratio adjusting, and prompt reweighting for different stages. Served in a post-training way, the proposed method can be integrated seamlessly into any DiT architecture. Extensive experiments across various backbones, schedulers, and datasets showcase the superiority of our method, for example, it achieves 1.55 times acceleration with negligible impact on image quality. Project page: https://github.com/ICTMCG/SDTM.

Figures (14)

• Figure 1: Illustration. Upper: Our SDTM represents a dynamic multi-resolution generation process by reducing feature redundancies in areas not unattended by the denoising process. Lower: Compared to the baseline method, our approach achieves 1.55$\times$ acceleration with negligible impact on generation quality.
• Figure 2: Preliminaries. Left: Evolution of Denoising Process: The L2 norm of (a) LL, (b) LH, HL, and HH subbands of estimated noise during the DiT denoising process. Right: Evolution of Feature Redundancies: (c) Location and (d) Degree evolution of token redundancies across DiT different steps and layers. Experiments were conducted on MMDiT mmdit with 50-step RF schedule rectifiedflow based on 10k samples.
• Figure 3: Overview. Grey: Our SDTM compresses weak-structure redundancies in the early stage and weak-detail redundancies in the later stage. Blue: Compression ratio adjusting (CRA) dynamically adjusts the ratio or threshold to control the pruning degree. Yellow: Prompt token reweighting (PTR) categorizes each prompt token into structure or detail groups, optimizing the denoising direction by reweighting attention map. Here, increase and decrease in red and green, while bold implies intensity.
• Figure 4: Visual Token Merging. By measuring structure similarity, unmerged frequency, and detail inattentiveness, SSM and IDM target different types of feature redundancies for reduction.
• Figure 5: Qualitative comparison on COCO2017 and PartiPrompts under varying data complexities. For ToMeSD and AT-EDM, we use versions with approximately 1.3$\times$ acceleration, while others use approximately 1.5$\times$ versions. Best viewed when zoomed in.