Research map / Image generation / Autoregressive and latent generation
ToMA: Token Merge with Attention for 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.
Proposes ToMA, a training-free token merging method for diffusion transformers that selects destination tokens via submodular facility location and performs merge/unmerge as attention-like matrix multiplications, achieving real wall-clock speedups on SDXL and Flux.
Problem
Existing token reduction methods for diffusion models (ToMeSD, ToFu, ToDo) rely on GPU-inefficient operations such as sorting and scattered writes, so their overhead cancels theoretical speedups when combined with optimized attention like FlashAttention.
Approach
Reformulate token merge as submodular facility-location destination selection followed by an attention-style soft merge matrix, unmerged by its transpose; exploit latent locality (tile/stripe regions) and reuse destinations and merge weights across denoising steps and layers.
Key results
- On SDXL-base at 0.5 merge ratio, ToMA reduces latency by 17.0% (5.04 vs 6.07 s/img) with DINO 0.068 and FID 28.9 vs 25.3 baseline.
- On Flux.1-dev, ToMA achieves up to 23.4% speedup on RTX6000 at 0.75 merge ratio (16.12 vs 21.03 s/img) while keeping CLIP-T near baseline.
- ToMA's merge/unmerge kernels are roughly 4-5x faster than ToMeSD's index-scatter implementation at sequence length 1024.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use ToMA instead of ToMe at ratio 0.50 | SDXL-base · SDXL, ratio 0.50, RTX6000 | Sec/img: 8.73 -> 5.04 | Helped |
| Use ToMA instead of ToFu at ratio 0.50 | SDXL-base · SDXL, ratio 0.50, RTX6000 | Sec/img: 6.83 -> 5.04 | Helped |
| Use ToMA instead of ToMe at ratio 0.50 | SDXL-base · SDXL, ratio 0.50 | FID: 26.7 -> 28.9 | Hurt |
| Use ToMA instead of ToFu at ratio 0.50 | SDXL-base · SDXL, ratio 0.50 | FID: 142.1 -> 28.9 | Helped |
| Restrict destination selection to tiles rather than globally (ablation) | ToMA · SDXL, ratio 0.50, destination selection | CLIP: 30.949 -> 31.019 | Helped |
| Use tile rather than stripe destination selection (ablation) | ToMA · SDXL, ratio 0.50, destination selection | DINO: 0.074 -> 0.055 | Helped |
| Increase tile count from 4 to 64 (ablation) | ToMA · SDXL, ratio 0.50, tile granularity | DINO: 0.069 -> 0.055 | Helped |
| Increase tile count from 64 to 256 (ablation) | ToMA · SDXL, ratio 0.50, tile granularity | DINO: 0.055 -> 0.057 | Hurt |
| Replace index-scatter merge with dense matrix multiplication (ablation) | ToMA · Merge kernel, 25% ratio, N=1024 | time (μs): 202.2 -> 39.0 | Helped |
| Replace index-scatter unmerge with dense matrix multiplication (ablation) | ToMA · Unmerge kernel, 25% ratio, N=1024 | time (μs): 160.5 -> 40.2 | Helped |
| Replace transpose unmerge with pseudo-inverse (ablation) | ToMA · SDXL, ratio 0.50, unmerge method | Sec/img: 4.8 -> 10.1 | Hurt |
| Replace transpose unmerge with pseudo-inverse (ablation) | ToMA · SDXL, ratio 0.50, unmerge method | CLIP: 31.027 -> 30.997 | No clear effect |
Surprising findings
- ToMeSD and ToFu add latency on SDXL (ToMe +42.7% slower at 25% ratio), so prior token merging is slower than baseline when paired with optimized attention.
- Transpose unmerge matches pseudo-inverse quality on CLIP/DINO/MSE while being more than twice as fast.
Limitations
- Stripe-shaped and tile-shaped locality incurs reshaping/memory copy overhead on GPU; the tile variant is slower despite better quality.
- Quality degrades notably at 75% merge ratio on SDXL (FID 58.6 vs 25.3).
- Flux requires skipping the first 10 blocks and custom handling of JointTransformer/SingleTransformer blocks; stripe merging is incompatible with RoPE, and V100 runs hit OOM.
Experimental setup
Datasets: GEMRec ImageNet-1K class names
Benchmarks: ImageNet-1K GEMRec
Metrics: FID CLIP-T DINO seconds per image
Baselines: ToMeSD ToDo ToFu Theoretical Lower Bound (TLB)
Builds on: ToMeSD Token Merging (ToMe) FlashAttention-2 ToFu
Closest papers by meaning
- Attend to Not Attended: Structure-then-Detail Token Merging for Post-training DiT Acceleration (CVPR 2025 · similarity 0.90)
- Which Tokens to Merge? Diffusion Dynamics for Efficient Image Generation (NeurIPS 2026 · similarity 0.90)
- Multimodal Promptable Token Merging for Diffusion Models (AAAI 2025 · similarity 0.90)
- SparseDiT: Token Sparsification for Efficient Diffusion Transformer (NeurIPS 2025 · similarity 0.89)
- DiG: Scalable and Efficient Diffusion Models with Gated Linear Attention (CVPR 2025 · similarity 0.89)
- ALTER: All-in-One Layer Pruning and Temporal Expert Routing for Efficient Diffusion Generation (NeurIPS 2025 · similarity 0.89)
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.