Reproduce “ToMA: Token Merge with Attention for Diffusion Models” on OpenResearch
Turn your favorite coding agents into research agents that review literature, develop hypotheses, and run experiments.
Universal app for Apple silicon and Intel. Requires macOS 11 or later.
About this paper
Wenbo Lu, Shaoyi Zheng, Yuxuan Xia, Shengjie Wang
ToMA introduces a training-free framework for token reduction in diffusion models, achieving practical speedups of up to 24% for models like SDXL-base and Flux.1-dev with negligible image quality degradation by using GPU-aligned attention-like merge/unmerge operations and submodular optimization for token selection.
Key insights
- The critical bottleneck for achieving practical speedups in prior token reduction methods was not the attention computation itself, but the GPU-inefficient merge/unmerge operations when paired with optimized attention kernels.
- Token selection can be framed as a submodular optimization problem, allowing for theoretically robust and GPU-efficient identification of representative tokens to minimize information loss during reduction.
- Leveraging the inherent spatial coherence and temporal redundancy within diffusion model latent spaces enables significant overhead reduction through locality-aware processing and reuse of merge weights, without compromising image fidelity.
Summary by alphaXiv.
Tip: change arxiv.org to autoarxiv.org in any arXiv abstract link to open the paper here.