One-sentence Summary

Researchers from multiple institutions propose ToMA, an off-the-shelf token reduction method for diffusion models that reformulates token merge as a submodular optimization, implements merge/unmerge via GPU-friendly attention-like matrix operations, and leverages latent locality and pattern reuse, achieving $24\%$24% and $23\%$23% latency reductions for SDXL and Flux respectively ($\text{DINO } \Delta < 0.07$DINO Δ<0.07), outperforming prior methods like ToMeSD and ToFu.

Key Contributions

• ToMA is introduced as a training-free token reduction method that reframes token merging as a submodular optimization problem, offering provable guarantees for selecting diverse and representative tokens to preserve generation quality.

• The method redesigns merge and unmerge operations as attention-like linear transformations computed through GPU-friendly matrix operations, sidestepping the sorting and scattered-write overheads of prior approaches and enabling seamless integration with optimized attention implementations such as FlashAttention.

• ToMA exploits latent-space locality and temporal pattern reuse across denoising steps to minimize runtime overhead, and experiments show it cuts SDXL and Flux generation latency by 24% and 23% respectively with DINO Δ below 0.07, achieving at least 1.24x practical speedup with FlashAttention2 across different models and GPU architectures.

Introduction

Diffusion models deliver high-fidelity image generation, but their transformer backbones, such as U-ViT and DiT, face a scalability bottleneck from the quadratic complexity of self-attention. Prior token reduction techniques (ToMeSD, ToFu) adapted merging to generative tasks, yet they rely on GPU-inefficient primitives like sorting and scattered writes, which become the dominant cost when paired with fast attention implementations, leaving only marginal speedups. To address this, the authors propose Token Merge with Attention (ToMA), a training-free framework that rethinks token merge as a submodular optimization problem. ToMA uses a GPU-optimized facility location algorithm to select representative destination tokens, performs merge and unmerge via attention-like linear transformations, and exploits latent space locality and sequential redundancy to amortize overhead, achieving practical end-to-end speedups (1.24x or more) with negligible quality loss.

Method

The authors propose ToMA, a lightweight framework designed to accelerate diffusion models by optimizing the token merging process. Standard token merging reduces computational load by aggregating similar tokens, but ToMA refines this pipeline at three critical stages: destination token selection, the merge operation itself, and the subsequent unmerge step.

To select the most representative tokens to serve as merge targets, the authors utilize a submodular optimization approach. Specifically, they employ the Facility Location (FL) function to quantify how well a subset of destination tokens represents the entire token set. The objective is to maximize the sum of maximum similarities between each token and the selected destinations. This is solved using a greedy algorithm that iteratively selects the token providing the largest marginal gain. To ensure efficiency on GPUs, the marginal gain computation is optimized using cached vectors, allowing for parallel execution via matrix operations.

Once the destination tokens are selected, the merge operation is formulated as a linear projection guided by attention scores. The authors treat the destination tokens as queries and all input tokens as keys and values in a Scaled Dot-Product Attention (SDPA) mechanism. By applying softmax column-wise and normalizing the resulting attention matrix row-wise, they obtain a merge weight matrix $\tilde{A}$A~. The merged token representations are then computed as $\boldsymbol{X}_{\text{merged}} = \tilde{A}\boldsymbol{X}$Xmerged​=A~X.

After the merged tokens pass through core transformer modules (such as Attention or MLP layers), they must be restored to their original resolution. The authors approximate the inverse of the merge projection using the transpose of the weight matrix, $\tilde{A}^{\top}$A~⊤. This simplification is justified because the low attention temperature and the disjoint nature of the destination subsets make the rows of $\tilde{A}$A~ nearly orthonormal, effectively collapsing the Moore-Penrose pseudo-inverse into a simple transpose operation.

To further enhance efficiency, the framework exploits the spatial locality inherent in the latent space. As illustrated in the figure below, the hidden states in the U-ViT model exhibit strong local coherence, where tokens within small spatial windows share similar features.

Because natural images possess local coherence, merging tokens within restricted spatial regions is as effective as global merging while significantly reducing computational overhead. The authors implement two partitioning strategies to leverage this property: tile-shaped regions, which preserve 2D proximity for higher fidelity, and stripe-shaped regions, which group tokens by rows to maintain memory contiguity and maximize speed.

Finally, the method reduces overhead by reusing computations across time and layers. The analysis below demonstrates that destination tokens selected at one denoising step remain highly relevant to subsequent steps.

Given that hidden states in diffusion models change gradually, the destination sets and associated merge weights exhibit high overlap across consecutive timesteps. ToMA capitalizes on this redundancy by reusing the same destination set for several consecutive steps and sharing the merge and unmerge matrices across transformer layers. This strategy sharply reduces the frequency of expensive similarity computations without compromising image quality.

Experiment

ToMA is evaluated on UNet-based SDXL and DiT-based Flux.1-dev across 1024x1024 images, comparing its tile, stripe, and once variants with heuristic baselines (ToMe, ToDo, ToFu) using CLIP-T, DINO, and FID. The method consistently achieves speedups up to 28.5% while maintaining perceptual quality, with the tile variant offering the best fidelity but slower runtime and the once variant being fastest yet worst in quality. Baselines suffer from high latency, unstable outputs, or distributional shifts, whereas qualitative results show ToMA outputs remain nearly indistinguishable from originals even at 75% token merging. Ablations confirm design choices such as tile-based merging and transpose unmerge for an optimal speed-quality tradeoff.

ToMA variants achieve faster generation than the SDXL baseline while maintaining competitive image quality. Among the variants, ToMA provides the best balance of speed and quality, whereas tile-based merging offers higher fidelity but slower runtimes due to memory overhead. Stripe-based merging is faster but slightly less effective in preserving visual quality. Tile-based merging achieves the strongest perceptual alignment but incurs memory copying overhead, making it slower. Stripe-based merging benefits from faster memory access patterns yet yields slightly degraded image quality. ToMA delivers up to 24% speedup while maintaining consistent performance across GPU architectures. Compared to other token reduction methods, ToMA is the fastest while keeping competitive perceptual quality.

ToMA variants accelerate Flux.1-dev image generation while preserving or improving quality. The default ToMA achieves substantial speedups, especially at higher token reduction ratios, while the tile-based variant offers better perceptual fidelity but runs slower due to memory copying overhead. This trade-off supports ToMA as the default choice. ToMA delivers up to about 23% faster inference on RTX6000 at 75% token merging, with minimal quality impact. ToMA tile consistently achieves stronger CLIP-T and DINO scores but is slower, particularly at lower merge ratios. At 25% token merging, both variants speed up by around 4% or less, with ToMA showing slight FID improvements.

ToMA achieves the fastest generation time among all compared token reduction methods on SDXL, with up to 28.5% speed improvement, while maintaining competitive perceptual quality. ToMe offers slightly better image quality but suffers from severe latency due to GPU-inefficient operations, whereas ToFu shows unstable visual quality and high FID, especially at higher reduction ratios. Overall, ToMA strikes the best balance between efficiency and quality, making it the recommended default method. ToMA achieves the largest speedup (up to 28.5% faster) while keeping FID and CLIP-T scores close to the baseline and ToMe. ToMe delivers marginally better CLIP-T and DINO scores but incurs over 40% slowdown due to sorting-based token selection. ToFu exhibits unstable visual quality with FID rising sharply at higher reduction ratios (e.g., 142.1 at 50% ratio). At both 25% and 50% token reduction, ToMA consistently outperforms ToFu in all metrics and remains faster than baseline. ToMA's speed advantage grows with higher reduction ratios, achieving -17.0% time at 50% merging while keeping quality competitive.

ToMA consistently achieves the fastest inference among token reduction methods on SDXL and Flux.1-dev, with speedups up to 28.5% at higher merge ratios while preserving competitive perceptual quality. Its tile-based variant improves fidelity but incurs memory copying overhead, making it slower, whereas stripe merging is faster but slightly degrades image quality. Compared to ToMe, which is slower due to sorting operations, and ToFu, which shows unstable quality at high reduction, ToMA offers the best speed-quality trade-off as the default choice.