# ToMA: Token Merge with Attention for Diffusion Models

arXiv: https://arxiv.org/abs/2509.10918 · HTML: https://mlsanity.com/paper/2509.10918/

## Key facts

- Cited by 18 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 24 of its 32 references that are in the map.
- First posted on arXiv 13 Sep 2025.

## Placed in

- Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/)
- Efficient LLM inference and serving / [Visual token pruning in multimodal models](https://mlsanity.com/inference/visual-token-pruning-in-multimodal-models/)

## Builds on (most cited first, 24 of 24)

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Emerging Properties in Self-Supervised Vision Transformers](https://mlsanity.com/paper/2104.14294/) (2021)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Training data-efficient image transformers & distillation through attention](https://mlsanity.com/paper/2012.12877/) (2020)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions](https://mlsanity.com/paper/2102.12122/) (2021)
- [Token Merging: Your ViT But Faster](https://mlsanity.com/paper/2210.09461/) (2022)
- [DynamicViT: Efficient Vision Transformers with Dynamic Token Sparsification](https://mlsanity.com/paper/2106.02034/) (2021)
- [All are Worth Words: A ViT Backbone for Diffusion Models](https://mlsanity.com/paper/2209.12152/) (2022)
- [AdaViT: Adaptive Tokens for Efficient Vision Transformer](https://mlsanity.com/paper/2112.07658/) (2021)
- [Post-Training Quantization for Vision Transformer](https://mlsanity.com/paper/2106.14156/) (2021)
- [Token Merging for Fast Stable Diffusion](https://mlsanity.com/paper/2303.17604/) (2023)
- [Adaptive Token Sampling For Efficient Vision Transformers](https://mlsanity.com/paper/2111.15667/) (2021)
- [Token Fusion: Bridging the Gap between Token Pruning and Token Merging](https://mlsanity.com/paper/2312.01026/) (2023)
- [Lazier Than Lazy Greedy](https://mlsanity.com/paper/1409.7938/) (2014)
- [DiffRate : Differentiable Compression Rate for Efficient Vision Transformers](https://mlsanity.com/paper/2305.17997/) (2023)
- [A survey on efficient vision transformers: algorithms, techniques, and performance benchmarking](https://mlsanity.com/paper/2309.02031/) (2023)
- [Upsample Guidance: Scale Up Diffusion Models without Training](https://mlsanity.com/paper/2404.01709/) (2024)
- [GEMRec: Towards Generative Model Recommendation](https://mlsanity.com/paper/2308.02205/) (2023)

## Built on by (most cited first, 18 of 18)

- [A Survey on Cache Methods in Diffusion Models: Toward Efficient Multi-Modal Generation](https://mlsanity.com/paper/2510.19755/) (2025)
- [TC-Pad'e: Trajectory-Consistent Pad'e Approximation for Diffusion Acceleration](https://mlsanity.com/paper/2603.02943/) (2026)
- [DDiT: Dynamic Patch Scheduling for Efficient Diffusion Transformers](https://mlsanity.com/paper/2602.16968/) (2026)
- [One Model, Many Budgets: Elastic Latent Interfaces for Diffusion Transformers](https://mlsanity.com/paper/2603.12245/) (2026)
- [Foveated Diffusion: Efficient Spatially Adaptive Image and Video Generation](https://mlsanity.com/paper/2603.23491/) (2026)
- [Token Pruning for In-Context Generation in Diffusion Transformers](https://mlsanity.com/paper/2602.01609/) (2026)
- [Shiva-DiT: Residual-Based Differentiable Top-k Selection for Efficient Diffusion Transformers](https://mlsanity.com/paper/2602.05605/) (2026)
- [FastSTAR: Spatiotemporal Token Pruning for Efficient Autoregressive Video Synthesis](https://mlsanity.com/paper/2603.07192/) (2026)
- [CoReDiT: Spatial Coherence-Guided Token Pruning and Reconstruction for Efficient Diffusion Transformers](https://mlsanity.com/paper/2605.14191/) (2026)
- [Multi-Resolution Flow Matching: Training-Free Diffusion Acceleration via Staged Sampling](https://mlsanity.com/paper/2607.01642/) (2026)
- [NanoFLUX: Distillation-Driven Compression of Large Text-to-Image Generation Models for Mobile Devices](https://mlsanity.com/paper/2602.06879/) (2026)
- [SemanticDialect: Semantic-Aware Mixed-Format Quantization for Video Diffusion Transformers](https://mlsanity.com/paper/2603.02883/) (2026)
- [Dynamic Video Generation: Shaping Video Generation Across Time and Space](https://mlsanity.com/paper/2605.21042/) (2026)
- [Rethinking Token Reduction for Diffusion Models via Output-Similarity-Awareness](https://mlsanity.com/paper/2605.22011/) (2026)
- [Keep The Essentials: Efficient Reference Conditioned Generation via Token Dropping](https://mlsanity.com/paper/2606.23682/) (2026)
- [TORINO: Token Reduction via Interpretable Concept Overlap in Vision-Language Models](https://mlsanity.com/paper/2607.04593/) (2026)
- [Importance-Aware OBS Pruning for Diffusion Models](https://mlsanity.com/paper/2607.20048/) (2026)
- [Level-of-Token Diffusion](https://mlsanity.com/paper/2610.05816/) (2026)

## Abstract (arXiv)

Diffusion models excel in high-fidelity image generation but face scalability limits due to transformers' quadratic attention complexity. Plug-and-play token reduction methods like ToMeSD and ToFu reduce FLOPs by merging redundant tokens in generated images but rely on GPU-inefficient operations (e.g., sorting, scattered writes), introducing overheads that negate theoretical speedups when paired with optimized attention implementations (e.g., FlashAttention). To bridge this gap, we propose Token Merge with Attention (ToMA), an off-the-shelf method that redesigns token reduction for GPU-aligned efficiency, with three key contributions: 1) a reformulation of token merge as a submodular optimization problem to select diverse tokens; 2) merge/unmerge as an attention-like linear transformation via GPU-friendly matrix operations; and 3) exploiting latent locality and sequential redundancy (pattern reuse) to minimize overhead. ToMA reduces SDXL/Flux generation latency by 24%/23%, respectively (with DINO Delta < 0.07), outperforming prior methods. This work bridges the gap between theoretical and practical efficiency for transformers in diffusion. Code available at https://github.com/WenboLuu/ToMA.