Research map / Image generation / Fast diffusion sampling
Attention-Driven Training-Free Efficiency Enhancement of Diffusion Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Training-free run-time token pruning for SD-XL guided by attention maps (G-WPR ranking, similarity-based token copy for ResNet compatibility) plus a denoising-steps-aware pruning schedule, cutting FLOPs ~39% with near-full FID.
Problem
Diffusion models are expensive at inference, mainly due to attention blocks, and existing efficiency methods require costly retraining; training-free architectural compression of diffusion models is largely unexplored.
Approach
Uses attention maps as directed graphs to rank tokens with a generalized Weighted PageRank (G-WPR) that works for both self- and cross-attention, recovers pruned tokens via a similarity-based copy before ResNet blocks, and prunes less in early denoising steps (DSAP) where attention maps have low variance.
Key results
- At 4.1 TFLOPs (38.8% saving vs SD-XL 6.7 TFLOPs), AT-EDM reaches FID 28.04 vs 35.27 for ToMe and 27.30 for full SD-XL at its sweet-spot CFG.
- AT-EDM achieves 52.7% latency speed-up without FO and 15.4%/17.6% with FO under CI/DI, respectively.
- Under 30-40% FLOPs saving, AT-EDM gets both lower FID and higher CLIP score than ToMe.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| replace ToMe token merging with AT-EDM attention-driven pruning at 4.1 TFLOPs | SD-XL · MS-COCO 2017 validation, 5k images | FID: 35.27 -> 28.04 | Helped |
| replace ToMe token merging with AT-EDM attention-driven pruning at 4.1 TFLOPs | SD-XL · MS-COCO 2017 validation, 5k images | CLIP: 0.3198 -> 0.3209 | Helped |
| replace ToMe token merging with AT-EDM attention-driven pruning at 2.9 TFLOPs | SD-XL · MS-COCO 2017 validation, 5k images | CLIP: 0.2954 -> 0.2784 | Hurt |
| use self-attention rather than cross-attention maps for G-WPR (ablation) | AT-EDM · SD-XL image generation | image quality | Helped |
| use similarity-based copy rather than zero padding to recover pruned tokens (ablation) | AT-EDM · SD-XL image generation | image quality | Helped |
| use similarity-based copy rather than bicubic interpolation to recover pruned tokens (ablation) | AT-EDM · SD-XL image generation | image quality | Helped |
| use similarity-based copy rather than direct copy to recover pruned tokens (ablation) | AT-EDM · SD-XL image generation | image quality | Helped |
| add the DSAP prune-less schedule for the first 15 denoising steps (ablation) | AT-EDM · SD-XL image generation | image quality | Helped |
| apply DSAP to ToMe by leaving corresponding attention blocks unmerged in early steps (ablation) | ToMe on SD-XL · SD-XL image generation | image quality | Helped |
Surprising findings
- Cross-attention-based WPR prunes too many background tokens and performs worse than self-attention-based WPR.
- Naive padding, bicubic interpolation and direct copy of pruned tokens severely degrade image quality, so token recovery is critical.
Limitations
- Latency gains shrink substantially when fused attention (xformers) is used, and the current xformers implementation does not expose attention maps.
- Pruning at the second feature level incurs overhead exceeding its latency benefit under fused operations.
Experimental setup
Datasets: MS-COCO 2017 validation
Benchmarks: MS-COCO 2017 validation FID/CLIP
Metrics: FID CLIP score FLOPs latency
Baselines: SD-XL ToMe
Builds on: WPR (ZeroTPrune) Token Merging (ToMe) SD-XL
Closest papers by meaning
- ALTER: All-in-One Layer Pruning and Temporal Expert Routing for Efficient Diffusion Generation (NeurIPS 2025 · similarity 0.91)
- DiP-GO: A Diffusion Pruner via Few-step Gradient Optimization (NeurIPS 2024 · similarity 0.90)
- Diffusion Models Without Attention (CVPR 2024 · similarity 0.90)
- SaRA: High-Efficient Diffusion Model Fine-tuning with Progressive Sparse Low-Rank Adaptation (ICLR 2025 · similarity 0.89)
- Dual-Rate Diffusion: Accelerating diffusion models with an interleaved heavy-light network (NeurIPS 2026 · similarity 0.89)
- EfficientDM: Efficient Quantization-Aware Fine-Tuning of Low-Bit Diffusion Models (ICLR 2024 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.