ToMA: Token Merge with Attention for Image Generation with Diffusion Models

2025-09-13

AI Summary

Diffusion models face generation efficiency bottlenecks due to the quadratic computational complexity of Transformer self-attention. Existing token compression methods (e.g., ToMeSD) rely on GPU-inefficient operations—such as sorting and scatter writes—that hinder compatibility with highly optimized attention kernels like FlashAttention, yielding limited practical speedup. This work proposes a GPU-friendly, efficient token merging framework: it formulates token selection as a submodular optimization problem; designs an attention-inspired linear transformation enabling purely matrix-based computation; and integrates DINO-guided feature selection with latent locality modeling to reduce redundancy. The method is fully compatible with FlashAttention. On SDXL and Flux, it reduces generation latency by 24% and 23%, respectively, while maintaining DINO feature distance increments below 0.07—significantly outperforming baseline approaches in both efficiency and fidelity.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Natural Language Processing: Generation
• Machine Learning: Large Multimodal Models (LMMs)

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization

Abstract

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 $Δ< 0.07$), outperforming prior methods. This work bridges the gap between theoretical and practical efficiency for transformers in diffusion.

Problem

• Reducing quadratic attention complexity in diffusion models
• Eliminating GPU-inefficient operations in token reduction methods
• Bridging theoretical-practical efficiency gap in transformer optimization

Innovation

• Submodular optimization for diverse token selection
• GPU-friendly linear transformation via matrix operations
• Exploiting latent locality and sequential redundancy