ToMA: Token Merge with Attention for Diffusion Models
Authors: Wenbo Lu
Shaoyi Zheng
Yuxuan Xia
Shengjie Wang
Venue: ICML 2025 poster (2025)
Keywords: Diffusion
Token Merge
Attention
Submodular Optimization
Primary area: deep_learning->generative_models_and_autoencoders
TL;DR: We propose an improved token merging algorithm to speed up diffusion for image generation, and is the only working method on diT.

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 **To**ken **M**erge with **A**ttention (ToMA), an off-the-shelf method that redesigns token reduction for GPU-aligned efficiency, with three key contributions: 1) a reformulation of token merging 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% (DINO $\Delta <$ 0.07), outperforming prior methods. This work bridges the gap between theoretical and practical efficiency for transformers in diffusion.

ToMA: Token Merge with Attention for Diffusion Models
Wenbo Lu" Shaoyi Zheng
Yuxuan Xia
Shengjie Wang'
Abstract
gories: attention optimization (e.g.. Flash Attention (Dao, 2023), xformers (Lefaudeux et al., 2022)) and token reduc- tion. While there is extensive research on token reduction for discriminative tasks, including Token Merge (Bolya et al., 2023), Ada (Yin et al., 2022b), and later improve- ments like Diffrate (Chen et al., 2023), these approaches often employ irreversible token pruning or merging, which is unsuitable for generative tasks.
Diffusion models excel in high-fidelity image gen- eration but face scalability limits due to transform- ers' 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 to- ken merge as a submodular optimization problem to select diverse tokens; 2) merge/unmerge as 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% (with DINO A < 0.07), outperforming prior methods. This work bridges the gap between theoretical and practical efficiency for transformers in diffusion.
Generative architectures like diffusion models impose stricter constraints: token counts must be restored after merging (i.c.. "unmerging") to preserve spatial consis- tency for iterative refinement. Earlier attempts, including ToMeSD (Bolya Hoffman, 2023), Token Pruning (Kim et al., 2022), and ToFu (Kim et al., 2023) adapt token reduction to this setting but suffer from a critical flaw: their merge/prune (and unmerge/fill) operations rely on GPU-inefficient primitives (e.g.. sorting, scattered mem- ory writes). This problem intensifies when paired with highly optimized attention implementations like those in the di f fusers framework (von Platen et al., 2022).
In such scenarios, the attention mechanism, once the pri- mary bottleneck, is streamlined to approach hardware effi- ciency limits, leaving only marginal time reductions attain- able. Consequently, these potential gains are dwarfed by the computation costs of unoptimized merging logic, which now dominates the computation time. In other words, the overhead introduced by previous token merging methods becomes negligible only when applied with fast implemen- tations of attention, thereby preventing practical speed-ups.
1. Introduction
Diffusion models (Ho et al., 2020; Song et al., 2021; Dhari- wal & Nichol, 2021) have revolutionized sigh-fidelity im- age generation. Yet, their reliance on transformer architec- tures tures- notably in U-ViT (Bao et al., 2023) and DiT (Pee- bles & Xie, 20ton-insrosuuces bottleneck: the quadratic complexity of self-attention scales prohibitively with token counts, exacerbating latency across denoising steps.
In this paper, propose Token Merge with Attention (ToMA), training-free free framework that bridges this gap by rethinking token merge for GPU-aligned operation. To achieve practical speed-up without image degradation, we reformulate token merge as submodular optimization prob- lem, leveraging its theoretical guarantees and algorithmic toolkit. Specifically, ToMA uses GPU-optimized facility location algorithm select diverse and representative set of "destination" tokens. Merge is then formulated as an attention-lik linear transformation for efficient aggregation of nos-destinttion tokens, with unmerge implemented via its inverse transformation, making full use of hardware-friendee matrix multiplications. This approach down-projects the la- tent space while preserving critical information, accelerating
Efforts to accelerate transformers broadly fall into two cate-
'Equal contribution 'Department of Computer Science, New York University. Correspondence to: Wenbo Lu <wenbo.lu@ @udu..euuu Shaoyi Zheng 1523664@nyy.edde Yux- uan Xia <yx2432@ www.4322@nyuucuu.... Shenji Wan <sw5973 nyu.edu>.
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Token Merge with Attention for Diffusson Models
SDXL
SDXL (Baseline) FlashAttention2
SDXL FlashMitemtion2 ToMe5D
SDXL FlashAttention2 ToMA
1:43x stower
1.25x faster
8.65
6.15
8.7$
4.65
Figure 1. Comparison SDXL-base with four configurations (left to right): Original, +FA2, +ToMeSD, +ToMA (on top of FA2, 0.5).=0.5). While ToMeSD fails to speed up due to overhead, ToMA achieves significant acceleration with negligible loss in image quality.
transformers with negligible overhead or quality loss. To further minimize computational overhead, we exploit two intrinsic properties of diffusion models:
2. Related Work
Efficient Vision Transformer Vision Transform- ers (Dosovitskiy al., 2021) face computational challenges due to quadratic attention complexity. Recent efforts to mitigate this fall into four categories: compact architectures (e.g,, Swin Transformer (Liu et al., 2021a), PVT (Wang et al., 2021)), pruning strategies like X-Pruner (Yu Xiang, 2023), knowledge distillation such DeiT (Touvron et al., 2021), and post-training quantization (Liu al., 2021b). Complementary efforts also explore combined techniques (Papa et al., 2024) to unify these paradigms. While effective, most solutions require retraining and remain inherently token-centric. contrast, ToMA introduces training- free token merging, operating orthogonally, thus enabling seamless integration without conflict.
Latent Space Locality: Tokens exhibit spatial co- herence, allowing parallel merging within non- overlapping local windows (e.g., 8x8 patches).
Sequential Redundancy: Merge patterns persist across I) adjacent denoising timesteps; and 2) consecutive transformer layers. We amortize the overhead by reusing merge patterns over multiple steps and layers.
Theoretical guarantees from submodular optimization ensure that ToMA's token selection approximates optimal coverage. At the same time, its co-design with GPU execution paradigms (e.g., batched matrix operations) eliminates costly operations inherent in prior methods. This synergy translates to real-world speedup, rather than theoretical FL.OP reductions only. For example, ToMA reduces the total generation time for 1DXL-base by 24% and Flux --dev by 23% with negligible degradation of image quality (change in DINO score < 0.07), outperform- ing previous work like ToMeSD and ToFu, which either fails to accelerate modern attention implementations introduce artifacts comparable compression rates. Our contributions are summarized:
Learned Token Reduction Most learned token reduc- tion involves training auxiliary models assess the im- portance of tokens in the input data. Dynamic ViT (Rao et al., 2021) employs a lightweight MLP to generate prun- ing masks based on input token features. A-VIT (Yin et al., 2022a) computes halting probabilities using specific chan- nels of token features to determine the necessity of further processing. Whereas these methods often require additional training for the auxiliary modules, our approach is directly applicable, offering more generalizable solution.
Algorithmic Innovation: A submodular optimization framework for token merging, ensuring provably rep- resentative token selection to enhance quality.
Heuristic Token Reduction Heuristic token reduction strategies can be applied to ViTs without additional train- ing. Our method also falls into this category, making the approaches below natural baselines for comparison. Adap- tive Token Sampling (ATS) (Fayyaz al.. 2022) keeps tokens most similar the class token, which limits its use in pixel-level generation tasks where a class token is absent. Token Downsampling (ToDo) (Smith et al., 2024) down- samples only the key-value in attention, skipping queries, which limits acceleration and causes fine-grained detail loss due to uniform spatial pooling. Token Merging for Sta- ble Diffusion (ToMeSD) (Bolya & Hoffman, 2023) forms
System Co-Design: GPU-aligned implementation strategies leveraging invertible, attention- opera- tions to exploit latent space locality and temporal re- dundancy, minimizing computational overhead.
Empirical Validation: ToMA achieves at least 1.24x practical speedup when paired with FlashAttention2, State-of-thh -rree---tthe--tttte results across different diffusion mod- els (e.g., SDxI.bbase, Flux. 1 1-dev) and GPU ar- chitectures (NVIDIA RTX6000, V100, RTX8000).
2

---

Token Merge with Attention for Diffusson Models
Faciity Location
Solve
A'
Astention
Spit
Concat
Spin
Self Attention
In batches
Cross Allentonn or MLP
batches
Same operation for the other chunks
Figure 2. Architectural overview of ToMA. The framework consists of three key stages: (1) Facility Location Algorithm identifies the hest representative token set D C N through submodular optimization to maximize representational diversity; (2) Attention (Merge) constructs an efficient low-rank attention matrix that maps N + via linear transformation for transformer computation (SelfAttn. CrossAttn, MLP) in the reduced space; (3) Inverse (Unmerge) applies the pecuooinveers to recover full-resslution features D + N. The pipeline operates through localized processing latent space regions with parallel batch optimization for efficiency.
source- dentindestination pairs within fixed or randomly tiled re- gions and greedily matches them based on similarity, per- forming unweighted merging followed by simple unmerge that copies the destination embedding back to each source position. ToFu (Kim al., 2023) builds on ToMeSD by dynamically deciding, for each layer, whether merge or prune tokens according a linearity test, thereby combining the benefits of both operations.
The SDPA operation is defined as
QKT Tvd
SDPA(Q, K, T) softmax
V.
(I)
3.2. Submodularity
A submodular function (Fujishige, 2005) is set function ʃ 2 + R defined over subsets of the ground set V. It satisfies the diminishing returns property, which states that the marginal gain of adding a new element e to the set decreases as the context set grows. Mathematically:
Prior token-reduction methods share two flaws: GPU- inefficient matching/merginn operations and heuristic de- signs lacking theoretical guarantees of information preserva- tion. By contrast, ToMA employs a GPU-friendly attention- like linear projection whose destinations are chosen via submodular objective, providing both hardware efficiency and principled foundation. Moreover, ToMA remains compatible with orthogonal blended schemes (alternating between pruning and merging) such as ToFu, allowing addi- trade-offs when combined.
For any subsets A C B V and element v E V B:
f(uAA) 2 f(
where the marginal gain f((l.) is defined as:
tional
f(v)A) = f(A (1)) f(A).
3. Preliminaries
3.1. Attention Notation
Algorithm 1: Greedy Algorithm
The standard Scaled Dot-Product Attention (SDPA) mechanism is widely adopted in mainstream diffusion models. For clarity, we defin the following notations: B for batch size, N for sequence length, for feature dimension, D for the number of all destination tokens, I Nxx for the attention input latent tensor, and Q. Q.K.V V E Bxdx for query, key, and value tensors, respectively, projected from X.
Input: Ground set V. submodular function
f:2% - R. and budget k
Output: Selected subset & of size at most k
Initialize A < 1/:
i 1 to 1: do
Select 8² = arg max ((v')A); {u*}:
Update A
return A;
3

---

Token Merge with Attention for Difflusion Models
This property makes submodular functions well-suited for modeling diversity and coverage in subset selection prob- lems. Naturally, this leads to the canonical problem of submodular maximization under a cardinality constraint:
tokens represents the full token set V by summing, for each token et E V, the maximum similarity Si any destination token Vi E D. Intuitively, this corresponds to asking: for each token in the ground set. how well does the selected subset represent it? A higher value of (ID implies that every token in V is closely matched by a representative in D, making a compact and diverse summary of the input. This naturally aligns with the objective of token merging, where we aim preserve global semantic structure using reduced set of tokens. Notably, our framework is modular- other submodular functions can be substituted for fEL to customize the selection behavior.
max ((A) s.t. SA|
An intuitive approach is the greedy algorithm (Alg. 1), which guarantees a (1 - -/je-aaprrxiimationn of the opti- mal solution (Nemhauser et al., 1978). Starting with A = 0, the algorithm iteratively selects the element with the highest marginal gain until the constraint |A| = l is reached.
The submodular nature of frL provides a theoretical guaran- tee: greedy maximization yields near-optimal subset with provably minimal information loss, as discussed in Sec. 3.2. When optimizing the destination set D using the greedy algorithm (Alg. 1), we iteratively select the token that pro- vides the largest marginal gain [r ID') with respect to the current set D' This marginal gain can be efficiently computed (see Appendix A.1 for derivation) as:
4. Method
Standard token merging reduces the number of tokens pro- cessed in Transformer blocks by identifying and aggregating similar tokens, thereby enabling theoretical speedups pro- portional to the merge ratio and the model's computational complexity (see Appendix for analysis). It works by select- ing destination tokens from the full token set and merging nearby tokens into them based on similarity scores. Dur- ing the unmerge step, the values of the merged tokens are redistributed to their original positions, preserving fidelity.
N max (0, Si, - -m,
arg max D.
Our lightweight and efficient framework, ToMA, improves upon standard token merge at three key stages: 1) Desti- nation Token Selection - efficiently identifying the most representative tokens serve as merge targets; 2) Token Merge performing merge operations as linear transfor- mation, guided by similarity scores computed via attention: 3) Token Unmerge restoring merged tokens after passing through core computational modules (e.g., Attention, MLP) to original positions through reversed linear transformation.
where (I(( = max Sj.k is a cached vector that stores, for each token V, E V, the maximum similarity to any token currently in D'. This caching enables efficient updates: after adding a new token to D' m can be incre- mentally updated in constant time per token. Importantly, all these operations--coppttiin similarities, caching, and evaluating marginal gains can be expressed in matrix form, supported by our derivation in Appendix A. making them highly suitable for parallel execution on GPUs. We include our efficient GPU implementation for the greedy algorithm in Appendix A.2. Also, though the iterative nature of submodular optimization is inherently unavoidable, we manage to parallelize this process by breaking it into smaller ground sets in Sec. 4.3.
To achieve further speedups, ToMA a) exploits the spatial locality of the latent space to parallelize operations within local regions and, b) shares merge-related computations across layers and iterations to reduce runtime overhead.
While there exist more advanced submodular maximiza- tion methods such as the Lazier-than than-Lazy Greedy algo- rithm (Mirzasoleiman et al., 2015), their reliance on opera- tions like random subset sampling introduces irregular mem- ory access patterns that inefficient GPUs. Therefore, the standard greedy approach strikes a practical balance between solution quality and hardware efficiency.
4.1. Jubmodular-Basel Destination Selection
Let S be the cosine similarity matrix between all hid- den states X, where the S.y entry represents the simi- larity between the i th token and the j th token, namely We denote the set of all tokens as
Sy = V (ground set in submodular optimization) and the set of chosen destination tokens as ID.
4.2. (Un)merge with Attention
(D) = Σ max Sij EDED V.E%
(2)
We begin by formulating token merge in its exact form and then show how it naturally generalizes to a linear projection over the input token space, paving the way for our proposed attention-like merging.
The submodular function used for destination token selection is the Facility Location function (FL), as shown in Eq. 2. (r) quantifies how well a subset ID of destination
Let X E RNNxd denote the input matrix of N token em-
4

---

Token Merge with Attention for Diffusson Models
beddings, each of dimension d. Suppose select of ID destination tokens with indices D (J1...--JD) jD) S {I....,N] N.. and partition the input tokens into D disjoint groups:
It is important note that softmax is applied column-wise rather than row-wise, token should not be decomposed into components whose aggregate exceeds 100%. By in- cluding an extra identity matrix I, we obtain an exact form of SDPA (Eq. 1). Note that the included I is functionally redundant and can be omitted in implementation:
G 92, Dt.92.....9m CD.
s.t.
(1.... N). GenG = for k L
A SDPA I.I,T...
1=1
merged E Rd is computed by aggregat- Each merged token z ing the tokens assigned to group GE. In the simplest case, this is done via uniform averaging:
To form a proper merge weight matrix, we normalize the attention matrix row-wise:
A
A,
Σx Ajk
mergod l
[9kl
Ti:
and finally, the merged token representation can be obtained via matrix multiplication:
REGL
or more generally, allow token-specific weights an with normalization:
X merrged = ÁX E Dxd
merged
agiss Zk
Di
where
Σ iEG
aki
Intuitively, this operation softly assigns each token of destination tokens based on similarity. Highly similar source tokens contribute more to the destination representations, while dissimilar ones contribute less.
This formulation can be unified by expressing the merged tokens as a linear projection over the input matrix:
By reducing merge to matrix multiplications and the SDPA kernel, ToMA scales up with token count easily and effec- tively receives free ride from ongoing GPU architecture and ML system improvements for attention, which in- creasingly efficient in both computation and memory usage.
X mergged = WX E Dxd
where W E DxN is non-negative weight matrix with Wa indicating the contribution token to destination This subsumes both hard merging schemes (e.g., ToMeSD, where each row of W one-hot) and soft merging (ToMA, where W contains normalized attention scores).
4.2.2. UNMERGE
Feeding merged tokens into core computation modules,
This linear formulation not only provides principled inter- pretation of token merging but also enables efficient imple- mentation and reuse of merge/unmerge operations (namely the weight matrix across steps and layers.
X' ATTENTION/MI.P(X TR. Pxd
the output we get is still D, and we restore the original token resolution by applying an approximate inverse of the merge projection. Let X' denote the output of core computational modules, with merged tokens as input. The goal is to recon- struct the full-resolution token matrix X' Xmmmrggd E Nxd
4.2. MERGE
Given the formulation above, naturally, one may think of constructing the merge weight matrix W via attention. Specifically, we treat the destination tokens as queries and all input tokens as keys and values, using Scaled Dot- Product Attention (SDPA) to produce similarity scores. These attention scores serve as soft merge assignments.
A principled approach is to apply the Moore-Penrose pseudo-inverse:
X Xunmerged =A''=A''AAT -A
This provides a least-squares reconstruction that minimizes the error between the original and reconstructed tokens, as- suming the merge- trangefrranoormmmmm unnneree process remains ap- proximately linear. However, computing the pseudo-inverse is computationally expensive, requiring matrix decomposi- tions such as SVD QR.
Let X E Nxd denote the input token matrix, and let E IR Dxd be the destination token matrix, formed by selecting subset of D token embeddings from X. We compute an attention matrix A IR DxN between destinations (as queries) and all input tokens (as keys) using SDPA with temperature parameter T:
Fortunately, under certain conditions, much simpler and more efficient alternative is available:
DX
A softmax
X" Yumorged A'' X'
5

---

Token Merge with Attention for Diffussin Models
which happens whenever ÁAT ID. making the pseudo- inverse A+ A (AAT)= collapse the simple trans- pose AT This is equivalent to saying that the rows of A are orthonormal. In ToMA, this condition is approached in practice. The facility-location selection promotes des- tination tokens that cover largely disjoint subsets of the source tokens, making rows of distinct. Moreover, the low attention temperature adopted sharpens the softmax distribution, concentrating each row's mass few source tokens and bringing its f2-norm close to one.
UNet Transformer Blocks
Together, these properties imply:
ÂA ID + ||l|| < 1,
07
21
so that
(AA ID E, and thus Ä+ - x= AT.
Denoising Timestep
Empirically, A remains competitive against the exact pseudo-inverse. Given this high fidelity and the substan- tial computational and memory savings, ToMA adopts the transpose-based unmerge Á X' as the default method.
Figure 3. Re-colored k-means clusters of U-VIT hidden states across transformer blocks and denoising timesteps. A similar visualization on DIT is provided in the Appendix E.1.
4.3. Further Speedups
The primary benefit of this locality constraint is computa- tional. Splitting the sequence into equal-sized tiles yields a dominant 1/1: speed-up for destination selection and even greater 1/k² reduction in computing the attention weight matrices as well as applying (un mungme Detailed complex- ity analysis is provided in Appendix C.
The overhead of ToMA arises from three main sources: 1) selecting destination tokens D through submodular optimization; 2) computing the attention-based (un)merge weight matrices Á and A and 3) applying the merge AX before, and unmerge A+ X' after each module inside the transformer blocks. To further reduce these overheads, we exploit the locality of the feature space, enabling these computations to be performed within localized regions. Additionally, reduce the frequency of steps 1) and 2) by reusing destination selections and (un)merge matrices across multiple iterations and transformer layers.
To leverage these benefits, ToMA limits destination selec- tion and (un) operations to within localized regions using the two partitioning strategies below.
Tile-shaped Regions Tokens are divided into 2-D tiles, preserving both horizontal and vertical proximity. This layout aligns closely with image geometry and gives the best quality, albeit a reshuffling cost on GPUs.
1. LOCALITY-AWARE TOKEN MERGING
A key oversight in prior work is the locality structure of the latent space. As shown in Fig. 3, where we visualize k- means clusters of U-ViT hidden states during the generation of "tommto", the recolored tokens form rough preview of the generated image. For example, in the early denoising steps, the clusters appear as coarse, blocky color regions that gradually refine into a recognizable tomato.
Tokens are grouped by rows di- rectly, maintaining memory contiguity and enabling fast reshaping. Although this ignores vertical proximity, it pro- vides the highest speedup.
Stripe-shaped Regions
Both variants substantially reduce computation by operating on smaller subsets in parallel. Tile-shaped regions offer higher fidelity, while stripe-shaped regions run much faster. Despite substantial acceleration, further acceleration is possible by implementing custom tiled/stripe attention kernels, in which memory copying overhead incurred in either reshape or read as strided no longer exists. However, as the post-literature doesn't 't include such a low-level implementation, we decide to leave as future work for the fairness of comparison. The full locality-aware ToMA algorithm is given in Appendix Alg. 3 in detail.
Because natural images exhibit strong local coherence-each pixel tends to resemble its immediate neighbors their latents inherit this property. As result, merging within a small spatial window aggregates highly similar information, pre- serving global structure while discarding redundancy only, and is therefore as effective as global merging. In terms of selecting destination tokens, restricting the facility location search to local regions ensures diversity within each tile and avoids competition across tiles for the same destination.
6

---

Token Merge with Attention for Diffussion Models
Tokon Intersection Across Timesttos Region Strrs
Inception-V3 V3 feature statistics. For FID CLIP-T and DINO in the main experiments shown in this section, we generate 3,000 images and compute scores against ground-truths from ImageNet-1K. IK. For the ablation experiments listed in Appendix F, we generate images from 50 prompts with three random seeds each and report the average. FID is omitted on GEMRec due to the lack of paired images. Inference latency is reported as the median wall-clock time over 100 runs.
Xerrge Pierncctos
66%
Baselines We compare ToMA with three heuristic token-reduction baselines: Token Merging for Stable Diffusion (ToMe), Token Downsampling (ToDo), and Token-Fusion (ToFu). These approaches were originally developed for UNet-based architectures and perform well in that setting. As a result, they are not compatible with DiT models, since they lack mechanisms to handle positional embeddings in DiTs. Due to this limitation, we evaluate all three baselines only on SDXL- leaving Flux. 1-dev benchmarked solely with ToMA. In terms of implementation, we utilize the official codebases for ToDo and ToMe, and reimplement ToFu based on its paper, as the public code is not available. In our experiment, ToDo uses a fixed merge ratio of 75%, corresponding to a 4-to-1 token downsampling scheme, which represents the lowest merge ratio supported by its implementation.
18
30
50
Figure 4. Average percentage of shared destination tokens at each denoising timestep relative to the first step of its 10-step interval. Each curve represents a different layer in EDXL-base U-ViT model, showing high overlap and gradual divergence over time.
4.3.2. REUSING DESTINATIONS AND MERGE WEIGHTS
Hidden states in diffusion models change gradually, so the destination tokens chosen one step are often similar to those chosen nearby in time. Fig. 4 quantifies this effect: across a 10-step window, more than half of the destinations are reused. Exploiting this redundancy, ToMA reuses the same destination set for several consecutive steps and, because merge and unmerge are linear, also reuses the associated weight matrices across layers. Sharing both the destination selection and the (un)merge matrices sharply reduces the frequency of introducing expensive similarity computation overhead while preserving image quality.
We additionally report TLB (Theoretical Lower Bound), which approximates the maximum attainable speedup by reducing the number of tokens without incurring extra run- time overhead (e.g., gather tokens). To simulate this bound, we do dummy merge-directly drop tokens and proceed with the next module by duplicating retained token features to preserve imput shape, thereby isolating the theoretical ben- efits of token reduction while minimizing implementation- specific costs. Quality metrics are omitted for TLB, as cloned tokens do not yield valid outputs for evaluation.
5. Experiments
5.1. Deseriptions
Setup We evaluate ToMA on two of the most widely used diffusion models: the UNet-based SDXI-base and the DiT-based Flux. 1-dev, to generate 1024 x 1024 images, using the Diffusers framework. Importantly, ToMA is architecture-agnosse and can be readily extended to other diffusion models (e.g., SD2, SD3. 5). Prompts are drawn from the GEMRec dataset (Guo et al., 2024) and ImageNet- IK names of the classes (Deng et al., 2009).
ToMA Variants To analyze the impact of locality on desti- nation selection and (un)merge operations, we evaluate four configurations: 1) ToMA, our default setting, which uses tile-based destination selection and global attention-based merge; 2) ToMAstripe which restricts both destination and merge operations to within stripe regions; 3) ToMAriles which uses tile regions for both destination and merge; and 4) ToMAonccs which improves efficiency by performing mun)merge operations only once per Transformer block- at the beginning and end-rather than around each core com- putation module. As for the hyperparameter, stripe- and tile-based configurations use 64 stripes or tiles, respectively. We reuse the destination for 10 denoising steps and reuse merge weights for 5 steps, with each block of given type sharing one set. No reuse across denoising timesteps in Flux. 1-dev but within blocks of the same kind.
Metrics To assess image quality, we use CLIP-T, DINO, and FID (Radford et al., 2021; Caron al., 2021: Heusel al., 2017). CLIP-T measures semantic alignment between images and prompts via cosine similarity between image-text embeddings. DINO measures perceptual consistency by comparing visual features between the original generated image and its counterpart produced with the merge method applied. FID (Fréchet Inception Distance) quantifies distributional similarity to real images based on

---

Token Merge with Attention for Difflusion Models
5.2. Results
Metries
RTX8000
RTX600
Ratio
Method
FID1
CLIP-T DINOL
Secting
1A
Seefing
Rutio
Method
Metrics
Sec/ing L
Baseline
Flux, 2 1-vvv
31,56
29.03
59.20
0%
21.03
0%
FID.
CLIP-TT
DINOL
KTX6000
V100
RTX8000
0.25
ToMh
30.80
29.07
-4.2%
20.14
-1.2%
Baseline
soxt-bbase]
25.27
29.89
6.1
14.5
16.1
TOMA,
31,49
29.05
0.021
57.47
-7.9%
20.78
-1.2%
ToMA
25.72
29.86
0.048
6.0
14.3
0.50
ToMn
31.70
29.09
0.051
51.M1
-13.1%
18.58
-1166%
ToMee
25.17
29:90
0.054
5.6
12.6
14.5
TOMA,
32,95
29.19
0.032
53.61
-9.4%
19.61
-688%
0.25
ToMAcite
25.45
29.86
0.045
6.2
13.6
15.7
26.31
29.70
0,052
5.5
12.3
13.5
0.75
Tott
33.39
28.98
-15,9%
16.12
-23.%
ToMbusse
33,88
24.34
0,045
49.86
-15.8%
18.30
-12.9%
TLB
5.2
12.1
9.2
TokA
28.88
29.64
0.068
5.0
11.0
12.8
Table 2. Performance comparison of ToMA variants and
ToMAatxipa
29.11
29.51
0.074
4.6
10.1
12.0
0.50
29.19
29.63
0.063
6.3
11.1
13.2
Flux. 1-vvv (Baseline) for 1024x 1024 image generation (35 sampling steps). Best values are highlighted, and relative speed improvements (A) are shown as %. Negative A values indicate faster inference compared to the baseline (lower is hetter).
ToMBence
38.14
29.06
0,080
4.9
9.7
11.5
TLB
4.0
9.9
7.8
ToMA
58.59
27.96
0.098
4.3
8.5
9.8
ToMAscrip
89.93
26.97
0.110
4.5
8.0
9.5
0.75
ToMAcate
58.90
28.17
0.091
6.2
9.1
10.7
ToMPuiee
123.37
24.96
0.106
7.6
TLB
3.1
7.8
6.5
CLIP-T and DINO scores. This srrcciinny-qualiin tradeoff supports our choice of ToMA as the default, while also high- lighting the need for optimized low-level implementations to fully realize the benefits of locality- merging.
Table 1. Performance comparison between ToMA variants and
SDXL-base (Baseline) for generating 1024x1024 images over 50 sampling steps. Best values are highlighted, except for TLB. (1: higher better, 4: lower better).
Benchmark Table 3 presents a comparative evaluation of TOMA (our default method), alongside other token reduction strategies on the SDXL model. Among all methods, ToMA achieves the fastest generation time-up to 28.5% improve- ment-while maintaining competitive perceptual quality, making it the most balanced choice. Although ToMe (Token Merge) delivers slightly better image quality, as reflected in CLIP-T and DINO scores, it suffers from severe latency due to its complex token selection mechanism, which relies on GPU-inefficient operations (c.g., sort). ToFu (Token Fu- sion), while faster than ToMe with lower overhead thanks to its blended strategy that incorporates token pruning, exhibits unstable visual quality some generations appear accept- able, but others are heavily degraded, especially at higher reduction ratios. The high FID values confirm this observa- tion. Lastly, ToDo (Token Downsampling) offers moderate gains in both speed and fidelity but demonstrates notice- able distributional shifts, as indicated by its elevated FID
UNet Results Table shows that ToMAgile consistently performs the best DINO score, indicating strong per- ceptual alignment, but is slowed down by low-level mem- ory copying overhead during token tiling. Our manual in- spection of generated images also confirms that biled merging yields the highest visual quality. In contrast, ToMAstripe benefits from faster runtimes due to its sequen- tial memory access pattern, which enables direct reshaping without copying; however, the absence of a strong locality leads to slightly degraded image quality. ToMA finds a favor- able trade-off, achieving up to % speedup and consistent performance across different GPU architectures. Based on this balance between efficiency and quality, we adopt it as our default method. Finally, the experimental variant ToMA offers the highest acceleration by treating the en- tire Transformer block as single merge unit, significantly reducing overhead. However, it produces the worst qual- ity due to insufficient spatial-contex mixing across layers. Interestingly, in some scenarios, its runtime is even lower than that of the TLB, likely due to less frequent memory copying compared to our dummy merge implementation.
Ratio
Method
FID,
CLIP-TI
DINO!
See/imgl
Buseline
SDXL
25.3
29.89
0
6.07
0%
TOMA
25.7
29.86
0.048
6.03
-0.7%
0.25
ToMe
25.6
29.86
0.054
8.66
+42.7%
ToFu
35.2
29.34
0.072
6.92
+14.0%
ToMA
28.9
29.64
0.068
5.04
-17.0%
DiT Results For the Flux model, we include only TOMA
0.50
ToMe
26.7
29.71
0.071
8.73
+43.8%
and ToMA; le. as stripe-based merging is incompatible with the rotary positional embedding (RoPE) used in Flux. We also exclude results on the V100 GPU due to out-of-memory (OOM) failures. As shown in Tab. 2, both ToMA variants consistently accelerate generation across GPUs, with ToMA achieving up to a 23.4% speedup without compromising im- age quality when merging down to 50% of tokens. Although ToMA,ile is marginally slower due to memory overhead, consistently offers better fidelity, evidenced by stronger
ToFu
142.1
25.04
0.134
6.83
+12.5%
ToMA
58.6
27.96
0.098
4.34
-28.5%
0.75
ToMe
ToFu
161.5
41.2
24.13
29.09
0.084
0.148
8.16
6.76
+11.4%
+34.4%
ToDo
68.6
27.60
0.093
5.67
6.66%
Table 3. Performance comparison of ToMA, SDXL-base (Base- line), and other token reduction methods for generating 1024x 1024 images (50 sampling steps). Best values are high- lighted, and relative speed improvements (A) are shown as %. (-A faster, +A slower; 1: higher better, 4: lower better).
8

---

Token Merge with Attention for Diffusion Models
Baseline 0%
To M (ours)
Baseline
To M A (ours)
10%
20%
30%
40%
50%
0%
25%
50%
75%
SDXL-base vs. TOMA
Flux. 1-dev
Figure 5. Visual comparison of baselines (SDXL & Flux) versus (ToMA) under different higher token-merge ratios. Despite merging up 50% tokens, TOMA preserves sharp details and overall scene coherence. Additional qualitative examples are provided Appendix
score. Hwang et al. discuss the possible cause that the noise reduction due downsampling creates a difference in the signal-to-noise ratio that leads to suboptimal performance if applied directly (2024). Overall, these results highlight ToMA as the most robust and well-rounded token merging strategy for efficient high-resolution image generation.
choices, such as selecting alle-based merge with 256 tiles and transpose-basee baasdd unmerge for optimal speed-qualite quality bal- ance. Appendix G provides memory profiling across multi- ple models and sparsity levels, showing that ToMA variants incur negligible memory overhead compared to dense base- lines. Appendix H complements this by providing a detailed FLOP analysis, saving computations up to 3.4x.
Figure provides a side-by-side visual comparison of images generated by the original models ver- sus our accellrace variants at several token-merge ratios. On the left, six different prompts SDXL and ToMA 10% 50% token reduction. Even the merge ratio increases, ToMA's outputs (second through sixth columns) remain nearly indistinguishable from the originals. Similarly, on the right, we present six different prompts for F Flux vs. ToMA at 25%, 50%, and 75% token reduction. Even at 75% merging, ToMA's results (fourth column) faithfully reproduce key details.
Qualitative Result
6. Conclusions
shown for
In this work, we introduce Token Merge with Attention for Diffusion Models (ToMA), co-designed merging frame- work that advances prior methods in three key aspects: 1) more representative token selection via submodular opti- mization with theoretical guarantee; 2) a flexible and ef- ficient merge-unmerge mechanism implemented through attention based operations; and 3) the incorporation of locality- and shared merging strategies maximize runtime gains. Together, these improvements yield sub- stantial speedups in practice while preserving high image fidelity across different GPU and model architectures (U- VIT & DiT), establishing ToMA as a robust and deployable solution for efficient high-resolution generation.
Ablation & Otheres. ToMA's robustness is further vali- dated by comprehensive ablations (Appendix F), which in- vestigate the impact of merge frequency, tile/stripe granular- ity, unmerge strategies (c.g., transpose vs. pseudo-inverse), and sharing schedules. These studies confirm our design

---

Token Merge with Attention for Diffusion Models
Acknowledgements
Dhariwal, and Nichol, A. Diffusion models beat gans on image synthesis. Advances in neural information process- ing systems (NIPS/NewrIPS), 34:8780-8794, 2021.
We gratefully acknowledge the New York University High Performance Computing (NYU HPC) facility for providing the GPU clusters and technical support that enabled the experiments in this work.
Dosovitskiy, A., Beyer, L.,, Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer,
M., Heigold, G., Gelly, S., Uszkoreit, J., and Houlsby, N. An image worth 16x16 words: Transformers for image recognition at scale, 2021. URL https:
Impact Statement
Our work on Token Merge with Attention (ToMA) improves the efficiency of diffusion models for image generation. While this advancement democratizes access to high-quality Al art creation, it's important to acknowledge that such tech- nologies can be misused to generate misleading content deepfakes. Additionally, as these models are trained on large internet-scraped datasets, there's s a risk of perpetuating soci- etal biases in the training data. We recognize these ethical considerations and emphasize the importance of responsible development and use of such technologies.
/ //arxiv.com wwwooo///aaao 2010
Fayyaz, M., Koohpayegani, A., Jafari, F. R., Sengupta, S., Joze, R. V., Sommerlade, E., Pirsiavash, H., and Gall, J. Adaptive token sampling for efficient vision transformers. In ECCV, pp- 396 414, 2022.
Fujishige, S. Suumoddlarr functions and optimization. Else- vier, 2005.
Guo, Y., Liu, H., and Wen, H. Gemrec: Towards gen- erative model recommendation. In Proceedings of the 17th ACM International Conference on Web Search and
References
Bao, F., Nie, S., Xue, K., Cao, Li, C., Su, H., and Zhu, All are worth words: A vit backbone for diffusion models. In Proceedings of the LEEE/CVE conference computer vision and pattern recognition, pp. 22669-22679 2023.
Data Mining, volume of WSDM '24, pp. 1054-1057. ACM, March 2024. doi: 10.1145/33118855.333377777
URL http:///a dww..oi...ooooooooooooo.o..oo... http://cc......orr.11111////66.. 3635700.
Bolya, D. and Hoffman, J. Token merging for fast stable diffusion, 2023. URL pestps://arxiie org/ abs/ 1303.17604
Heusel, M., Ramsauer, H., Unterthiner, T., Nessler, B., and Hochreiter, S. Gans trained by a two time-scale update rule converge to local nash equilibrium. NeurlPS, 30, 2017.
Bolya, D., Fu, C.-Y., Dai, X., Zhang, P., Feichtenhofer, C., and Hoffman, J. Token merging: Your ViT but faster. ICLR, 2023.
Ho, J., Jain, A., and Abbeel, P. Denoising diffusion proba- bilistic models. Advances in neural information process- ing systems (NIPS/NeurIPS), 33:6840-6851 2020.
Caron, M., Touvron, H., Misra, I., Jégou, H., Mairal, J. Bojanowski, P. and Joulin, A. Emerging prop- erties in sulf-super vision transformers. CoRR, abs/2104.14220 2021. URL
Hwang. J., Park, Y.-H., and Jo, J. Upsample guidance: Scale up diffusion models without training. arXiv preprint arXiv:240..01709, 2024.
Kim, M., Gao, S., Hsu, Y.-.,, Shen, Y., and Jin, H. Token fusion: Bridging the gap between token pruning and token merging, 2023. URL ptttps://axxiv....
Chen, M., Shao, W. Xu, P., Lin, M., Zhang, K., Chao, F,, Ji, R., Qiao, Y., and Luo, P. Diffrate: Differentiable compression rate for efficient vision transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 17164-17174, 2023.
abs abs/2312.. 01026.
Kim, S., Shen, S., Thorsley, D.,, Gholami, A., Kwon, W., Hassoun, J., and Keutzer, K. Learned token pruning for transformers, 2022. URL abs 0t1//2110.0000
Dao, T. Flashattention-2: Faster attention with bet- ter parallelism and work partitioning. arXiv preprint arXiv:2307.08691, 2023.
Lefaudeux, B., Massa, F., Liskovich, D.. Xiong, W., Caggiano, V., Naren, S., Xu, M., Hu, J., Tintore, M., Zhang, S., Labatut, P., Haziza, D., Wehrstedt, L., Reizen- stein, J., and Sizov, G. xformers: A modular and hack- able transformer modelling library. https: pestps://gitan //githbb. conm/facebookreseae www.arrch//foororo.//eoo 2022.
Deng, J., Dong, W., Socher, Li, L.-J,, Li, K., and Fei-Fei, L. Imagenet: A large-scale hierarchical image database. In 2009 IEEE Conference on Computer Vi- sion and Pattern Recognition, pp. 248-255, 2009. doi: 10.1 109/CVPR22009 52060448888888..
10

---

Token Merge with Attention for Diffusion Models
Liu, Z., Lin, Y., Cao, Y., Hu, H., Wei, Y., Zhang, Z., Lin, S., and Guo, B. Swin transformer: Hierarchical vision transformer using shifted windows. In ICCV, pp. 10012- 10022, 2021a.
Berman, W., Xu, Y., Liu, S., and Wolf, T. Diffusers:
State-of-the-art diffusion models. https: /tpss///tt /tee://gittbb.
com/hhigsingfacc//cfffeeess 2022.
Wang, W., Xie, E., Li, X., Fan, D.-P., Song, K., Liang. D., Lu, T., Luo, and Shao, Pyramid vision transformer: A versatile backbone for dense prediction without convo- lutions. In ICCV, pp. 568-578, 2021.
Liu, Z., Wang, Y., Han, K., Ma, S., and Gao, W. Post- training quantization for vision transformer. arXiv preprint arX1v:2106.11156 2021b
Mirzasoleiman, B., Badanidiyuru, A., Karbasi, A., Vondrák, J., and Krause, Lazier than lazy greedy. In Proceed- ings of the AAAI Conference on Artificial Intelligence, volume 29, 2015.
Yin, H., Vahdat, A., Alvarez, J. M., Mallya, A., Kautz, J., and Molchanov, P. A-ViT: Adaptive tokens for efficient vision transformer. In CVPR, pp. 10809-10818, June 2022a.
Nemhauser, G. L., Wolsey, L. A., and Fisher, M. L. An analysis of approximations for maximizing submodular functions-i. Mathematical programming, 14:265- 294, 1978.
Yin, H., Vahdat, A., Alvarez, J. M., Mallya, A., Kautz, J., and Molchanov, P. A-vit: Adaptive tokens for efficient vision transformer. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pp. 10809-10818, 2022b.
Papa, L., Russo, P., Amerini, I., and Zhou, A survey on efficient vision transformers: Algorithms, techniques, and performance benchmarking. IEEE Transactions on Pattern Analysis and Machine Intelligence, 46(12):7682- 7700, 2024. doi: 10. 10.1100 109/TPAMM..2024.39999
Yu, L. and Xiang, W. X-pruner: Explainable pruning for vision transformers. In CVPR, pp. 24355-24363, 2023.
Peebles, W. and Xie, S. Scalable diffusion models with transformers. Proceedings of the IEEE/CVF Interna- tional Conference Computer Vision, pp. 4195-4205, 2023.
Radford, A., Kim, W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., Krueger, G., and Sutskever, Learning transferable visual models from natural language supervision, 2021. URL hwt.h:tpp:/..mm...tmm / aawiiv.iivv org/ abs/ 00020.
Rao, Y., Zhao, W., Liu, B., Lu, J., Zhou, J., and Hsieh, C.-J. Dynamie VIT: Efficient vision transformers with dynamic token sparsification. In NeurlPS, November 2021.
Smith, E., Saxena, N., and Saha, A. Todo: Token downsam- pling for efficient generation of nigh-resolution images. arXiv preprint arXiv:2402.13577 2024.
Song. Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon. S., and Poole, B. Score-based generative model- ing through stochastic differential equations. In Interna- tional Conference on Learning Representations (ICLR), /
2021. URL https:tt/////tt http///ppeeeeeeeeeeeeeeeeeeeeeee wtwprrewwwwwweeeeeeeeeeeeeeeeeee www./eernuuttt 1d-PxTI612RRIS
Touvron, H., Cord, M., Douze, M., Massa, F., Sablayrolles, A., and Jégou, H. Training data-efficient image trans- formers & distillation through attention. In ICML, pp.
10347-10357, 2021.
von Platen, P. Patil, S., Lozhkov, A.. Cuenca, P., Lam- bert, N., Rasul, K., Davaadorj, M., Nair, D., Paul, S.,

---

Token Merge with Attention for Diffusion Models
A. Facility Location for Selecting Destination Tokens
A.I. Mathematical Foundations for Efficient Greedy Maximization
We begin by defining the notations used in the greedy selection procedure. Let V denote the full set tokens, and let C V be the current set of selected representative tokens. each step, the greedy algorithm selects the next token v" E V that maximizes the marginal gain:
F² arg max A), "E(V(A)
where the gain function f(v|A) is defined as the increase in a coverage objective when v is added to A. Formally,
f(x|A) = f(A') f(A)
Σ
max sim(us (z, Ua)
Σ
max sim(u, up), where A' = (VJUA.
u.EA'
EA
WEY
wEV
Here, sim( denotes the similarity between two tokens (we use cosine similarity in practice for computational effipiiitooaalfiiiieeeeee) The first term in the equation measures for each token v E V, how well represented in the updated set A' = {v) U A. The second term asks the same question with the current set A. Their difference quantifies the marginal gain including in the representative set.
We simplify the first term in the gain function by observing that:
Σ max sim(u, un)) = Σ uEV
max max sim(u, un), sim(u, v)
u.
u.EA
wEY
in which the identity holds because for each u E the maximum similarity with the updated set A is either its similarity with u, or its previous maximum over A. So, we take the maximum of the two.
Substituting this into the definition of [()]A), we obtain:
f(o|A)=
Σ wEY
max
max sim ue ua), sim v) U.EA
max sim(u, us)
EA
Σ
0, sim(u, v) max sim(u, un)
max
useh
This simplification shows that the marginal gain of adding v depends only on the tokens V for which provides new maximum similarity beyond what is already achieved by the current set A.
As result, the greedy selection objective becomes:
max O0 sim(m, v) max sim u. us)
arg max
(E(PTA) WEY
M,EA
This formulation enables efficient computation of the marginal gain for each candidate token and facilitates the selection of the token that most improves the representational coverage of the current set.
12

---

Token Merge with Attention for Diffusion Models
A.2. Facility Location Algorithm
This algorithm implements greedy approach select D tokens based on the Facility Location objective. It begins by computing the sum of similarities for each token. The first token chosen is the with the highest sum of similarities. Then, iteratively selects the remaining D tokens.
In each iteration, the algorithm computes its gain for every unselected token, representing the overall similarity improvement. The token with the highest gain is chosen and added to the selection. After each selection, the algorithm updates the maximum similarities achieved so far. This process continues until tokens are selected.
By selecting tokens that maximize the marginal gain in similarity at each step, this approach effectively covers the input space while avoiding redundancy, ensuring diverse and representative set of tokens.
Algorithm 2: Greedy Algorithm for Token Selection
Input:
Similarity matrix RRN number of tokens select Output: Selected destination token indices d d < ft:
Initialize:
Sum over each row: s = Σ S.y
Select the first token index: t <- arg max, 8; Add the greedy choice to destination tokens: d + d {ti):
Create the cache vector m,(II))) from the corresponding in the S: m,()))) Sei Set
avoid re-selection: St, 0;
for k 2 to D do
for each token index i not in d do Compute the marginal gain efficiently with cache: g Egn max 10, S.j m, mI(II)) )):
end
Select the next token greedily: & + arg maxied 9s:
Add the newly selected token index: d d U {tx}: Update largest row: my (TV/) max (TT'), St,) Set to zero to avoid re-selectoin: St, 0;
end
return d
13

---

Token Merge with Attention for Diffusion Models
B. Overall Detailed Algorithm of ToMA
Algorithm 3: ToMA with Local Regions
Input: Tensor X e Bxxxxx (input sequence), D (number of destination tokens), 7 (attention temperature), F(-) (core
computational module (e.g., MLP, Attention)
Split into local regions
Partition the sequence dimension into P blocks, X +- (XI...., (p)) with Xp BBxNNxx and Noc N;
Dioe t D/P:
X + X.reshape(BB P. Moce d);
Step 1: Facitity-location token selection
"i....,T..P). + Greedy (FFL: Dise: X);
Xr 4- (X XB.P.TB p) E BP Dus xd,
Step 2: Merge
- SDPA(XT: (XT. X, 1, 7) B
A t A/ A/
X merged f- AX E Bdux
Computational layer
X' < (-((Xmrrgon Cenergod-reshie reshape B, D, d)):
Step 3: Unmerge
X unmerged + Á X':
Reassemble X unmerggd to reverse the local-region split;
return X unmerged
Description, Algorithm 3 details a single ToMA layer equipped with local-region processing. Given an input tensor X R. Bxdxd, we first shard the sequence into equally sized local regions of length Nise (lines 1-3), then assign budget of Dioe destination tokens each region (line 2).
In Step we invoke GPU-friendly greedy facility-location algorithm to pick, for every mini-batch x region pair, the set of tokens whose neighbourhoods best cover the local region (lines 5-6). These destinations are gathered into Xy.
Step 2 forms a scaled-dot-produce attention map from each destination to all tokens in its region, row-normalises it to Å, and left-multiplies to obtain the merged representation Xmerged whose sequence length is reduced from N to D (lines 8-10). Any computational block F(.)-e.g. e.g., a transformer layer or UNet block can now process the shorter sequence (line 12), yielding a computed reduction factor of N/D without altering the block's internal parameters.
Now, in Step 3, we distribute the updated destination embeddings back to their original token positions via A and undoes the initial region partitioning (lines 14-15), so that the next layer receives a full-resolution tensor. Because the (un)merge operations are purely linear and share weights across the batch, their overhead is negligible, and they can be fused with existing attention kernels. Repeating this procedure layer-by-layer yields significant wall-clock speed-ups while exhibiting scarcely any perceptible degradation in image quality.
14

---

Token Merge with Attention for Diffusion Models
C. Computational-comleiitg analysis
retain explicit constant factors in the flop count because they translate directly to empirical speedups on modern GPUs. Throughout, N is the original sequence length, D the length after token merging, d the embedding dimension, and
D N
(merge ratio, i.c. the fraction of tokens kept).
Baseline self-attention block. Treating each matrix multiplication as collection of dot products, the total number of scalar multiplications for a standard self-attention block is
Case 4 + 2d
The first term (Q, K, V projections and the output projection) scales linearly in N; the second term (QKT and attention-value
product) scales quadratically.
Token-merged self-attention.
After reducing the token count from N to D the attention cost becomes
Cama = 4dD + 2d = 4d + 2dr
Hence the ideal speedup (ignoring merge overhead) is
base
4d 2N
SpeeduPiteal
Canu ())
4dr + 2N ,
Overheads introduced by ToMA. Token merging incurs several additional costs:
Submodular-selection overhead (computing marginal gains for all pairs): Csab 22d.
Merge-attention projection (computing pairwise weights): Cpre = Nd.
Merge operation (applying the computed weights to produce merged tokens); Cmerge = Dd.
Unmerge operation with transpose-style redistribution: Cunmerge NDd.
Summing the three linear-in-D overhead terms yields Cin = 3NDd = Nr.
Total cost with ToMA. The overall computational cost after merging is therefore
(r)
dd + N2
22d submodular
Nr
atlention black
merge unmergge
Realistic speedup.
The practical speedup of ToMA relative to the baseline is
Chase (r) 4dr 2 ( 3r 2r
4dN + 2N
Speedup practical
Discussion. For practical transformer settings we have N d, the quadratic terms 2d 22r2N2 (remaining attention cost) and 2d (one-shot sub-muuuue modular selection) dominate. Consequently, the empirical speedup is well approximated by Speedup 4d/N))2((2+ + 1 3r). which approaches the analytic bound Ч, whenever A 0.5 and d/N is small. At moderate merge ratios E [0.25, 0.5) this yields the 24 - 28% latency drop we observe on SDXL/Flux. If we push below ~ 0.1, the linear-in-r merge overhead 3r 2dd and the fixed 2 selection cost start to dominate, so further merging brings diminishing returns-undersoomm underscoring why our locality-aware pattern reuse amortises these costs across multiple layers.
15

---

Token Merge with Attention for Diffussion Models
D. More Qualitative Results
D.I. Images Generated with SDXL
We present additional visual comparisons between SDXL, ToMeSD, and ToMA below in Fig 6. The prompts are sampled from the GemRec and ImageNet- 1K datasets. As the visuals reveal, ToMA maintains image quality more faithfully compared other methods, especially in preserving fine-grained details and spatial coherence. This evident in both synthetic scenes (e.g., fantasy landscapes, bowls of fire) and natural subjects (e.g., animals, boats, and portraits). ToMA Additional images generated with ToMA on SDXL are provided on the next page in Fig. 7.
SDXL
SDXL (Baseline)
ToMeSD
ToMA
ToMeSD
ToMA
(Baseline)
Figure 6. Qualitative comparison between Baseline SDXL-base, ToMeSD, and ToMA.
16

---

Token Merge with Attention for Diffusson Models
SDXL (Baseline)
ToMA
ToMA
ToMA 30%
TOMA
ToMA 50%
SDXL (Baseline)
ToMA
ToMMA
ToRMA 30%
ToMA
ToMA 50%
Figure 7. Visual examples of ToMA. Even with half of the tokens merged, ToMA consistently preserves image quality and often demonstrates greater robustness compared to other methods (ToDo, ToFu, and ToMeSD).
17

---

Token Merge with Attention for Diffusion Models
D.2. Images Generated with Flux Please refer Fig. below for more images generated with ToMA on Fluxl. O O-dev.
Flux. 1-dev
ToMA 25%
TOMA 50%
Flux.1-dev (Baseline)
ToMA 25%
TOMA 50%
ToMA
(8aseline)
75%
75%
Figure Qualitative comparison between Baseline Flux. 1 1-dev and ToMA.
18

---

Token Merge with Attention for Diffusion Models
E. Diffusion Transformers (DiT)
E.1. Locality in DiT
We inspected the hidden states of lux. Using simple visualizations (K-means coloring) at the start of each block and across denoising timesteps, we observed that even without convolutions the hidden states already resemble the target image (Fig. 9). This locality introduced mainly by the rotary embeddings in Flux. Empirically, when we apply submodular token selection within local windows, the model still produces high-quality images.
FluxJointZrann weantrranatta former8io
1239
Denosing Timestep
Figure Re-colored K-means clusters hidden states in F Flux. 1-dev across blocks and denoising steps.
E.2. Transformers and Positional Embeddings in DIT
DiT blocks differ from the usual "self-attention - cross-attention - MLP" pattern, so off-the-shelf merging methods such ToMeSD, ToFu, or ToDo break the model (all-black/ / noisy and nonsense outputs). Two issues arise:
1. The DIT block order (attention+MLP fusion) is not aligned with the assumptions in those methods.
2. Positional embeddings (c.g. RoPE) are mixed with both image and text tokens; careless merging discards useful tokens.
We therefore add two simple rules:
Skip the first 10 blocks. Early blocks fuse text and image features; skipping them avoids over-merging.
Handle the two DiT block types separately.
Joint Transforrrr Text and image tokens are projected separately, then concatenated before RoPE. We merge text and image tokens independently, then concatenate; RoPE indices are gathered accordingly. Single Transorrrnssormmr Tokens are already concatenated. We first split the hidden state back into text and image parts, merge each part, then re-concatenate; RoPE is gathered in the same way.
These lightweight changes respect both modality boundaries and positional embeddings, allowing our token- -merging variant to run on Flux with no perceptible loss in image quality while still delivering the intended speedup.
19

---

Token Merge with Attention for Diffusion Models
F. Ablation Study
All the ablations reported this appendix are conducted on SDXL with a default merge ratio of = 0.5 unless otherwise noted. Each experiment isolates one design choice that its direct impact on quality measured by CLIP, DINO, and pixel MSE, and on runtime-measurge in seconds per generated image, all can be clearly assessed.
F.1. Destination-selectio strategy
The first experiment evaluates four types of selection windows for choosing destination tokens: a global window that considers every pair of tokens, local tile facility that performs the same selection inside non-overlapping windows, horizontal stripe window, and a random baseline the full window. Tab. 4 shows clear pattern. Restricting the search to tiles yields the best CLIP and DINO scores and the lowest MSE, while running more than six times faster than the exhaustive global search. The result confirms our locality hypothesis from Fig. 9: most informative tokens lie close to one another, so global scan is unnecessary and wasteful.
Type
CLIP f
DINO +
MSE +
Sec/img +
Global
30.949
0.069
1,637
33.2
Tile
31.019
0.055
1,274
5.1
Stripe
30.986
0.074
1,730
5.2
Random
30.553
0.090
2,029
4.5
Table Comparison of destination-sellctioon rules. Tile-based selection delivers the best quality and the lowest latency. Values are averaged over the SDXL validation prompt set; CLIP/DINO rounded to three decimals, MSE to the nearest integer.
F.2. Tile granularity
Once the tile-based facility was chosen, we next varied the number of tiles in order control the spatial extent of each selection window. Intuitively, fewer tiles correspond larger windows and therefore allow tokens to compete across a wider context, whereas many small tiles enforce highly local competition. Tab. 5 shows four granularities ranging from 4 large tiles to 256 small tiles. Moving from 4 to 16 tiles yields large quality jump: DINO improves by 17% and MSE by 14% while latency is cut almost in half. The improvement continues when the window count rises to 64, which records the best DINO and MSE and also the lowest runtime. At 256 tiles, DINO and MSE drift upward again, indicating that extremely small windows over-constrain the matching pool; nevertheless, CLIP remains tied with the best value, and the latency does not increase further because the GPU remains compute- bound.
Because the numerical differences between 64 and 256 tiles are modest, 64 tiles strike cleaner balance: delivers the strongest quality metries while preserving the same throughput and avoiding the bookkeeping overhead that arises when thousands of tiny windows must be indexed. For these reasons, we adopt 64 tiles as the default granularity in the main paper.
# Tiles
CLIP
DINO 4
MSE 4
Sec/img 4
4
30.775
0.069
1,564
11.4
16
30.991
0.057
1,345
6.4
64
31.019
0.055
1,274
5.0
256
31.027
0.057
1,296
5.0
Table Influence of tile granularity at 50% merge ratio. Using tiles achieves the best DINO and MSE while matching the runtime 256 tiles.
E.3. Merge and unmerge latency
The third experiment benchmarks the merge and unmerge kernels at fixed sequence length of N = 1024 tokens. We compare the dense linear formulation used in ToMA with the index-scatter implementation in ToMeSD, In ToMeSD the algorithm first builds a cestinationninder array, then gathers features with torch. index_seles and finally scatters them back with index_add_. Because both gather and seatter operate on the full index list, their cost grows linearly with the merge ratio r. Moreover, the discontinuous memory accesses inherent these calls leave many GPU warps idle.
20

---

Token Merge with Attention for Diffusson Models
ToMA eliminates the two passes by replacing them with a single dense matrix multiplication AX, where À E BDxN and D (1 r)N. The operation therefore depends only on the output length D; its cost is constant with respect the number of removed tokens and maps efficiently to single GEMM that fully utilizes GPU compute units. As reported in Table 6, ToMA is consistently four to five times faster than ToMeSD for both merge and unmerge across all tested values of r.
Time ((8)
Speedupt
Operation
Method
25%
50%
75%
25%
50%
75%
ToMe ToMA
202.2
202.1
193.2
Merge
39.0
38.8
38.8
5.2x
5.2x
5.0x
ToMc ToMA
160.5
160.1
144.0
Unmerge
40.2
40.5
39.6
4.0%
3.9x
3.6x
Table 6. berchbbbnchmrrks sequence length 1024 (median over 1,000 rums an NVIDIA RTX6000) across different merge ratios. Shaded cells indicate the best result per column. ToMA roughly four five times faster than ToMeSD over merge ratios r.
F.4. Transpose versus pseudo-inverse for unmerge
We experiment whether a mathematically exact unmerge, obtained via the Penrose pseudo- inverse of the merge matrix, offers any quality advantage over the much cheaper transpose. In theory, the pseudo-inverse should restore the pre-merge feature space more faithfully, because it inverts the least-squares projection implicit in the merge. In practice, the merge matrix used by ToMA is highly sparse and close to orthogonal, so the transpose already provides an excellent approximation. Computing the peeuooinveery requires QR or SVD decomposition of the D x N merge matrix, followed by two matrix multiplications to apply the result. These decompositions are memory-bandwidth bound and cannot be fused with the surrounding transformer layers, so the cost paid in every unmerge step.
Table confirms that the extra work wasted. Across 300 generated images, the pseudo innuree gains no measurable improvement: CLIP, DINO, and MSE differ by less than 1%. Meanwhile, latency more than doubles because the decomposition incurs additional global synchronizations the GPU. Given the negligible benefit and the clear timing penalty, we adopt the simple transpose as the default unmerge method.
Unmerge Method
CLIP t
DINO
MSE
Sec/img
Transpose
31.027
0,057
1,296
4.8
Pseudo- inverse
30.997
0.057
1,288
10.1
Table Transpose versus pseudo-inverse unmerge 50% merge. Quality metrics are identical, but transpose is more than twice as fast.
E.5. Recompute schedule
The final ablation varies how often destination indices and attention weights recomputed during denoising. Tab. 8 indicates that refreshing attention every step gives the best overall accuracy, whereas destination indices can be updated ten times less frequently with minimal loss. A schedule of "destination every 10 steps, attention every 5 steps" preserves 99% of the peak quality while roughly halving recompute FLOPs, and is therefore adopted in the main experiments.
Recompute D
Recompute Á
CLIP
DINO 1
MSE
See/img |
Every 50 steps
Every 50 steps
30.043
0.077
2,489
4.84
Every 10 steps
Every 10 steps
30.817
0.073
1,735
4,97
Every steps
Every steps
30.865
0.070
1,632
5.00
Every 10 steps
Every step
30.997
0.067
1,525
5.06
Every 5 steps
Every steps
30.892
0.069
1.609
4.92
Every 1 step
Every step
30.920
0.067
1.552
5.05
Table Effeet recomputation frequency at 50% merge. The shaded cells denote the best metric each column.
21

---

Token Merge with Attention for Diffusion Models
G. Memory analysis
Table 9 provides a peak-memory audit on our methods. For each model, we record both the maximum allocated memory, which reflects the live tensor footprint, and the maximum reserved memory, which includes CUDA's internal caching. Across all three merge ratios the numbers remain tightly clustered: on fluux, the largest deviation from the dense baseline 0.3% increase in allocated memory for plain ToMA 25% merging; SDXL-base, the worst case is 1.9% rise in reserved memory, again for plain ToMA at 25%. The tile variant is even closer and occasionally dips below the baseline because smaller activation tensors leave more room for the allocator to reuse blocks.
Max Memory (MB).
Model
Metric
Method
25%
50%
75%
Baseline
34,640
34,640
34,640
Alloc
TOMA
34,744
34,710
34,675
ToMAA
34,647
34,647
34,642
Flux.1-ca
Baseline
37,002
37,002
37,002
Resv.
37,050
36,976
36,954
ToMAche
37,054
37,006
36,950
Baseline
10.721
10,721
10,721
ToMA
10,931
10,857
10,797
Alloc
10,722
10,719
10,718
l'oMAc11e
10,725
10,720
10,719
SDXL-base
Baseline
14,150
14,150
14,150
TOMA
14,460
14,260
14,130
Resv.
ToMAstripe
14,114
14,188
14,222
ToMAclle
14,158
14,158
14,182
Table Peak GPU memory (MB) for different ToMA variants on P1 ux -dev Sixi,-base across three merge ratios. We report both maximum allocated and reserved memory. (J: lower is hetter.)
H. FLOP Analysis
Table 10 reports layer-wise
(FLOP) breakdown, restricted to the dominant computational modules inside each transformer block: the QKV output projections and the attention matrix products. Results are shown for the largest block Flux lux and for the two block types that occur in SDXL. Applying ToMA with a 50% merge ratio yields a 2.3x reduction on Flux and up to 3.4x reduction on SDXL. The additional FLOPs introduced by ToMA- JoMA-namely submodular token selection, the merge weight computation, and the linear (un)merge kernels amount to less than 1 of the new total and therefore negligible the scale of the overall savings.
Layer Size
FLOPs (G)
Model
Reduction
(Seq Dim)
Original
ToMA (50%)
Overhead
Flux
4608 x 3072
520
225
1.01
~2.3
SDXL
4096 640
106
32
0.42
~3.4x
SDXL
1024 1280
30
13
0.06
~2.4x
Table 10. Layer-level FLOP counts before and after applying ToMA at 50% merge ratio. The "Overhead" column includes sub-modular selection, merge weight computation, and the linear (un)merge operations.
22


Reviews:
[{"review_id": "lerhpAgPy3", "reviewer": "Reviewer_fssN", "summary": "This paper introduces ToMA (Token Merge with Attention), a GPU-optimized token merging method for transformer-based diffusion models, addressing inefficiencies in existing token merging techniques such as ToMeSD and ToFu. ToMA achieves its efficiency improvements through submodular optimization for token selection and attention-based linear projections for merging and unmerging. Experiments demonstrate that ToMA achieves 24% speedup on SDXL and 23% on Flux.1-dev while maintaining image quality, outperforming previous work like ToMeSD and ToFu.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "*Dear Reviewer*,\n\nWe appreciate your insightful comments and the time you've taken to evaluate our work. Below are our responses to your concerns:\n\n***Concerns about High Image Quality and Little Speed Up:*** At a ratio of 0.25, the acceleration is limited primarily because the reduction in sequence length is not substantial, which naturally results in limited speedup. Additionally, some unavoidable overhead remains, including operations within the transformer block that are not accelerated, as well as the overhead introduced by our own method (though relatively small, it still contributes). That said, further speedup could be achieved through engineering optimizations, such as low-level implementations. For instance, fusing the kernel to merge directly into the attention computation could significantly improve efficiency. However, in order to maintain a fair comparison with the baseline and other existing methods (which do not achieve the same level of speedup as ToMA), we defer such engineering enhancements to future work.\n\n***FLOP, Memory, and Efficiency:*** The efficiency primarily stems from the reduced sequence length in attention inputs achieved by merging, which significantly lowers the computational load of the attention mechanism. \n\n**FLOP**: Our FLOP analysis focuses on the transformer blocks, specifically the projection and attention computations. In both FLUX and SDXL, ToMA achieves a substantial FLOP reduction, with a maximum improvement of approximately 3.4×. The overhead introduced by ToMA operations—including submodular token selection, merge token computation, and (un)merge—is negligible compared to the overall FLOP savings.\n\n### FLOP Comparison Table\n\n| Model | Example Layer (Seq Len × Dim) | Original FLOPs(GFLOPs) | ToMA (0.5) FLOPs(GFLOPs) | ToMA Operation FLOP(GFLOPs) | Reduction |\n| --- | --- | --- | --- | --- | --- |\n| **FLUX** | 4608 × 3072  | 520  | 225  | 1.01  | ~2.3× |\n| **SDXL** | 4096 × 640 | 106 | 31.54 | 0.42 | ~3.4× |\n| **SDXL** | 1024 × 1280 | 30 | 12.76 | 0.06 | ~2.4× |\n\n---\n\n**Memory**: Regarding memory usage, please check the table below showing the max memory allocated and reserved during the inference process. Experiments in both FLUX and SDXL settings indicate minimal impact on memory allocated or reserved, suggesting our method incurs only minor additional memory allocation. \n\n### Combined Memory Usage Table (FLUX vs ToMA vs ToMA_tile)\n\n| Metric | FLUX | **ToMA** |  |  | **ToMA_tile** |  |  |\n| --- | --- | --- | --- | --- | --- | --- | --- |\n|  |  | 0.25 | 0.5 | 0.75 | 0.25 | 0.5 | 0.75 |\n| **Max Allocated (MB)** | 34640.07 | 34743.86 | 34709.85 | 34675.21 | 34647.18 | 34646.66 | 34642.15 |\n| **Max Reserved (MB)** | 37002.00 | 37050.00 | 36976.00 | 36954.00 | 37054.00 | 37006.00 | 36950.00 |\n\n### Combined Memory Usage Table (SDXL vs ToMA vs ToMA_stripe vs ToMA_tile)\n\n| Metric | SDXL | **ToMA** |  |  | **ToMA_stripe** |  |  | **ToMA_tile** |  |  |\n| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |\n|  |  | 0.25 | 0.5 | 0.75 | 0.25 | 0.5 | 0.75 | 0.25 | 0.5 | 0.75 |\n| **Max Allocated (MB)** | 10720.92 | 10930.77 | 10856.75 | 10796.86 | 10722.08 | 10718.62 | 10718.23 | 10725.21 | 10720.06 | 10718.67 |\n| **Max Reserved (MB)** | 14150.00 | 14460.00 | 14260.00 | 14130.00 | 14114.00 | 14188.00 | 14222.00 | 14158.00 | 14158.00 | 14182.00 |\n\nThe efficiency comes from the shortened attention sequence length enabled by token merging, allowing ToMA to significantly reduce FLOPs while introducing negligible additional memory overhead compared to the baseline.\n\n***Sparse Attention:*** The sparse attention mechanism naturally complements token merging. The underlying strategies of sparse attention and token merging are fundamentally orthogonal: token merging explicitly reduces the input length of the attention computation, whereas sparse attention selectively attends to subsets of tokens. As these methods operate independently, they can easily be combined—sparse attention can directly replace flash attention within ToMA to further improve efficiency. Moreover, many sparse attention methods like Longformer typically employ a sliding window that aligns with the local structure exploited by ToMA and can be directly utilized to restrict both merging and attention operations to sliding window regions, further enhancing computational efficiency.\n\n***Overhead and Image Quality at Extreme Merge Ratio:*** In general, we can achieve speedup by merging more tokens. Specifically, with a merge ratio above 90%, ToMA achieves approximately 4.1 seconds generation time on SDXL and 15.05 seconds on FLUX. However, at such high ratios, we observe a noticeable degradation in image quality as the number of tokens processed by the transformer becomes very limited, directly impacting the output. That said, even at these extreme ratios, we did not observe any inefficiencies significant enough to offset the overall speedup in generation time.", "claims_and_evidence": "Yes, most of the claims are supported by experiments, such as ToMA achieves speedups by leveraging GPU-friendly operations and submodular token selection works reasonably.", "methods_and_evaluation_criteria": "Yes. The paper evaluates the proposed method on both SDXL and Flux against ToFu and ToMeSD.", "theoretical_claims": "N/A.", "experimental_designs_or_analyses": "Yes. Overall the experimental designs are sound. However, it lacks detailed computational complexity breakdown, such as memory usage, FLOP reductions, or additional latency savings beyond just inference time.", "relation_to_broader_scientific_literature": "The paper follows the idea of previous token reduction methods such as ToMeSD and ToFU but optimizes for GPU execution. It also extends submodular optimization into token selection for generative models.", "essential_references_not_discussed": "The paper misses the discussion of sparse transformer [1], which also leverages locality for token reduction.\n\n[1] Child et al. Generating Long Sequences with Sparse Transformers. Arxiv 2019.", "other_strengths_and_weaknesses": "Strengths:\n\n- ToMA is shown to work efficiently with modern GPU optimizations like FlashAttention2, avoiding memory bottlenecks.\n\n- Empirical results confirm the speed improvements of ToMA over prior methods on SDXL and Flux.\n\nWeaknesses:\n\n- At 0.25 merge ratio, the speedup is minimal, raising concerns about whether ToMA is beneficial in scenarios requiring high image quality.\n\n- The paper reports inference speed but lacks FLOP breakdowns and memory usage analysis. It remains unclear where efficiency gains come from.", "other_comments_or_suggestions": "N/A.", "questions_for_authors": "- Would ToMA benefit from sparse attention techniques like Longformer?\n\n- What is the computational overhead of ToMA at extreme merge ratios (>90%)? Does merging too aggressively introduce additional inefficiencies or degrade quality significantly?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hv6Nht2Rip", "reviewer": "Reviewer_J1my", "summary": "The authors present a token merge algorithm for image diffusion models, based on the theory of submodular optimization, operations that are more friendly to GPU, and some extra tricks to further enhance the speedup. The proposed method first selects destination tokens that are most representative and then merges these tokens considering the attention and locality. The experimental results show speedup over the unoptimized baseline and several other token merging baselines. The absolute speedup on image generation is around 30% when reducing 75% of the token, which is a notable improvement over the baselines compared.\n\n## update after rebuttal\n\nThe rebuttal solved all my concerns, therefore, I raised the score to 4.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "*Dear Reviewer*,\n\nThank you for taking the time to review our work. We have addressed your concerns in detail in the responses below.\n\n***Response***:\n\n***Lack of Comparison on DiT-Based Diffusion Model:*** \n\nThank you for raising this point. Table 2 includes fewer comparisons with other methods in the DiT-based model setting because existing approaches often produce mosaics or noise—they are not compatible with both RoPE positional embeddings and the specialized DiT architecture (Joint and Single Transformers). As a result, we focus on Table 2 for analyzing our ToMA variants. To address your concern, we include comparisons with the ToMA_tile implementation. We omit the ToMA_stripe, as stripe is incompatible with the 2D RoPE. Additionally, we report results on the RTX8000, excluding V100 due to its insufficient memory for the FLUX model.\n\n| Ratio | Method | FID↓ | CLIP↑ | DINO↓ | Sec/img (RTX 8000)↓ | Sec/img (RTX 6000)↓ |\n| --- | --- | --- | --- | --- | --- | --- |\n| 0.00 | Baseline | 31.56 | 29.03 | **0** | 59.20 (0%) | 21.03 (0%) |\n| 0.25 | ToMA | **30.80** | **29.07** | 0.043 | **56.70 (-4.2%)** | **20.14 (-4.2%)** |\n|  | ToMA_tile | 31.49 | 29.05 | **0.021** | 57.47 (-2.9%) | 20.78 (-1.2%) |\n| 0.50 | ToMA | **31.70** | 29.09 | 0.051 | **51.44 (-13.1%)** | **18.58 (-11.6%)** |\n|  | ToMA_tile | 32.95 | **29.19** | **0.032** | 53.61 (-9.4%) | 19.61 (-6.8%) |\n| 0.75 | ToMA | **33.39** | 28.98 | 0.064 | **49.83 (-15.9%)** | **16.12 (-23.4%)** |\n|  | ToMA_tile | 33.88 | **29.34** | **0.045** | 49.86 (-15.8%) | 18.30 (-12.9%) |\n\n***ToMA Stripe, ToMA Tile, and ToMA*:*** Thank you for pointing this out. We ultimately selected the tile method as it aligns well with the inherent hidden states locality, and in FLUX scenarios, it complements the 2D RoPE structure. We believe ToMA_tile can be further optimized—e.g., via contiguous tile-shaped memory access and kernel fusion for (un)merging operation with the attention computation. We believed that through optimization, the ToMA tile could achieve acceleration comparable to the ToMA_stripe. However, such optimization poses a significant engineering challenge, and to keep a fair comparison with baselines, we defer this to future work. Thus, we introduced the ToMA_stripe as a practical fallback due to its lower overhead and immediate applicability, despite slightly weaker performance in quality metrics. As for ToMA*, it is an exploratory variant aimed at reducing overhead by applying (un)merging once at the model level rather than per transformer block. While this significantly improves runtime, it currently poses challenges in preserving generation quality, suggesting a promising direction for future architectural or training improvements.\n\n***ToDo:*** The reason is that the implementation of ToDo only supports merging exactly 4 tokens into 1, inherently limiting the merge ratio to 75%. Therefore, we report results for ToDo only at the 75% merge ratio.\n\n***Runtime Breakdown:*** Before presenting our result, we would like to clarify that the reported runtime does not fully reflect the actual time consumed, as we explored several timing measurement methods, each with its limitations. Torch Profiler introduces overhead, while using cuda.Event to measure individual components adds extra synchronization, potentially disrupting concurrency. Both inflate the runtime measured. To address this, we adopted a method where we selectively disabled specific components via commenting and measured the time difference. This approach yields results that are more consistent with theoretical expectations and offers a more reasonable estimation of the runtime impact. \n\nWe have included a comparison of merge and unmerge operation runtimes for both ToMeSD and ToMA in Section F of the Appendix. Additionally, the table below provides a detailed runtime breakdown. As shown, ToMA introduces significantly less operational overhead compared to ToMeSD. However, when the computation time is significantly reduced, the relative impact of overhead becomes more pronounced, resulting in a speedup that is smaller than the theoretical gain suggested by FLOPs. The runtime breakdown for FLUX closely aligns with that of SDXL.\n\n| Time Category | Subcomponent | SDXL + ToMA (s, %) | SDXL + ToMe (s, %) |\n| --- | --- | --- | --- |\n| **Computation Time** |  | **0.550 (45.32%)** | **0.550 (24.54%)** |\n|  | SelfAttention | 0.151 (12.45%) | 0.151 (6.74%) |\n|  | CrossAttention | 0.107 (8.80%) | 0.107 (4.77%) |\n|  | FeedForward | 0.292 (24.07%) | 0.292 (13.03%) |\n| **Token Merge Operation** |  | **0.191 (15.73%)** | **1.168 (52.11%)** |\n|  | ComputeMerge | 0.015 (1.21%) | 0.419 (18.67%) |\n|  | Merge | 0.090 (7.44%) | 0.431 (19.25%) |\n|  | Unmerge | 0.086 (7.10%) | 0.318 (14.19%) |\n| **Other Overhead** |  | **0.472 (38.91%)** | **0.523 (23.35%)** |\n| **Total Block Time** |  | **1.214 (100%)** | **2.241 (100%)** |", "claims_and_evidence": "Most of the claims are supported. For details, see weakness.", "methods_and_evaluation_criteria": "The method and benchmark make sense.", "theoretical_claims": "Yes, they are correct.", "experimental_designs_or_analyses": "Yes.", "relation_to_broader_scientific_literature": "N/A", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "The overall quality of the paper is good with noticeable improvement over the baseline together with sufficient qualitative comparison. However, the paper can be polished by addressing the following weaknesses:\n1. Table 2 is apparently included in a rush with significantly less content compared to Table 1. Thorough valuation and comparison on DiT-based diffusion model is crucial for readers to see the merit of this paper.\n2. ToMA stripe usually performs worse in terms of the metrics despite the improvements in runtime, and ToMA tile provides comparable quality but at the cost of overhead, which makes the reader wondering what the real benefit of using these two variants is. And there is the ToMA* which leads to significant degradation in quality.\n3. In table 3, the proposed method is only compared with ToDo at 75% of token reduction.\n4. The authors are suggested to include runtime breakdown and comparison to properly prove that the designed method is GPU-friendly and superior over the prior works.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "Please refer to weaknesses.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VXwQUObGwx", "reviewer": "Reviewer_6HTM", "summary": "The paper investigates token reduction via submodular optimization.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer,\n\nThank you for your review. We appreciate your feedback and pointing out several typos that need polishing. Due to ICML 2025 regulations, we are currently unable to modify the submitted PDF. However, we will certainly address and correct these issues in the updated version.\n\nRegarding your other concerns about efficiency and submodularity, please find our detailed responses below.\n\n***Response***:\n\n***Parallel of Submodular Optimization:*** Thank you for pointing this out. You're correct—the iterative nature of submodular optimization is inherently unavoidable, as each selection step depends explicitly on the previously selected items. However, to maximize parallelism, we indeed parallelize the computation by performing token selection independently across multiple local regions, each selecting tokens with maximal marginal gains. Since the submodular optimization in each local region is independent, this approach significantly reduces iteration time, as submodular optimization is now executed over smaller subsets with fewer tokens to select at each step.\n\n***Values of N and d:***  In our formulation, *N* denotes the sequence length and *d* the channel dimension. In SDXL, as the UNet goes deeper, *N* decreases while *d* increases—for 1024×1024 image generation, typical configurations include [4096, 640] and [1024, 1280]. In FLUX, *N* is composed of both text (512) and image (4096) tokens, totaling 4608, with a fixed channel dimension of 3072. \n\nAt 75%, ToMA achieves 1.4x speed-up instead of a higher reduction because of its own overhead. Though ToMA introduces significantly less overhead than other token merging methods—as demonstrated in Appendix Section F—its operations still produce a non-negligible runtime cost. Yet, this overhead can be further reduced through engineering efforts such as kernel fusion and custom CUDA optimization. It's also worth noting that the speed-up in FLUX is slightly lower than in SDXL, mainly due to the additional computation required for RoPE reconstruction. This involves gathering original RoPE positions and recomputing positional encodings based on the merged token set to maintain image quality.\n\n***Orthogonal***: Thank you for your question. To clarify, ${\\tilde{A}}$ refers to the **merge matrix** computed from attention scores, not the token feature matrix. Since ${\\tilde{A}}$ is a non-negative matrix, a strict orthogonality would mean that each source token attends to exactly one destination token, with no overlap. In the approximate case, the overlap across different rows is small—i.e., each source token mainly attends to a unique destination token. This naturally requires the destination tokens to be diverse, or otherwise, a source token may attend to two non-diverse destination tokens and result in overlap. Therefore, **diversity is a necessary condition** for orthogonality or approximate orthogonality in this setting. Moreover, this connects closely to the **facility location objective** used in our submodular selection: it ensures that every source token has at least one highly similar (i.e., representative) token in the selected subset, which means that every source token mostly attends to its representing destination token.\n\n***Submodular Optimization Related**:* Thank you for your question. Indeed, we select a subset of up to D tokens. However, linear dependency issues are unlikely in our setting. Specifically, our pseudo-inverse operation is applied to the merge matrix ${\\tilde{A}}$, which is obtained by multiplying the selected token subset with the entire token set. This merge matrix is non-negative and naturally has the largest values along its diagonal (or, more precisely, every row has a distinct index of the largest value), ensuring that the selected tokens remain linearly independent and, thus, preventing numerical instability during the pseudo-inverse computation.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "The proposed methods are valid, however, having $75\\%$ token reduction leads only to $\\approx 1.4$ faster inference than the baseline. Please check the question section.", "theoretical_claims": "The paper relies on the known plain greedy algorithm for maximizing a submodular function. However, I find some information concerning the unmerging part lacking; see the questions section.", "experimental_designs_or_analyses": "The experimental designs and corresponding analyses are sound, however, as mentioned earlier, when reducing a large amount of tokens, the model does not benefit much in terms of decreasing inference time; see the questions section.", "relation_to_broader_scientific_literature": "The idea of token reduction via submodular is quite innovative and fitting. Such methods could alleviate standard techniques in Deep learning, becoming inseparable tools that allow faster performances with a theoretical backbone (that would work on vast amounts of models and not tailored to work with a specific set of diffusion models).", "essential_references_not_discussed": "The paper references the right papers that are essential to understand the contribution of the paper.", "other_strengths_and_weaknesses": "The strengths of the paper are:\n1) Formulating the token reduction as a submodular function maximization problem is fitting.\n2) Utilizing the fact that submodular optimization aims to ensure a diverse subset of tokens into the token merging and unmerging is quite fascinating, however, I do have some questions on this matter. \n3) The architecture changes done to ensure the full potential of utilizing the token reduction phase, merging and unmerging are necessary and innovative, and efficient from a practical point of view.\n\n\nThe weaknesses, however, are:\n1) The submodular phase takes $O(N^2d)$ time (as mentioned in section C in the appendix). This time might explain why the method while being faster than the competitors, is not much faster than the baseline model; See questions section.\n2) The paper needs polishing; see the comments section.", "other_comments_or_suggestions": "In what follows, some typos will be listed:\n1) In the appendix, please replace line 1071 \"Tab ??\" with Table 4.\n2) In the appendix, please replace line 1076 \"Tab ??\" with Table 5.\n3) Remove line 452.\n4) Line 370, second column, I think the authors meant to mention Table 1.\n5) Table 2 was never referenced in the paper.", "questions_for_authors": "1) The authors claim to implement the naive greedy submodular optimization maximization algorithm concerning the facility location problem applicable to GPU. Since the algorithm in itself is iterative, it can not be parallelized. However, did the authors mean to parallel the similarity matrix computation? In addition, did the authors parallel the finding of the item with the maximal marginal gain? \n\n2) What are the values of $N$ and $d$ throughout the experiments? This aims to clarify why $75\\%$ does not lead to a higher reduction in inference time.\n\n3) The authors hint that due to diversity in the sampled set (a given byproduct of the submodular optimization process), the resulting $\\tilde{A}$ is somewhat orthogonal. Is this claim true? If so, why would it be orthogonal? \n\n4) A follow-up question on the previous: Usually with submodular optimization, we would have a set that contains up to $D$ items. Is this the case also here? This would explain whether the set is forced to have exactly ensure the token reduction ratio -- in such a case, the items in the set may not be independent (from a linear point of view) which then would lead to numerical issues with the psuedo inverse", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper proposes a method for token merging in diffusion models. This is an area of significant interest at this time. The method attacks the problem both from the optimization and the GPU utilization viewpoint achieving significant improvements. All reviewers recommend acceptance and the AC concurs
Decision: Accept (poster)