ToMA: Token Merge with Attention for Diffusion Models

Dear Reviewer,

Thank 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.

Regarding your other concerns about efficiency and submodularity, please find our detailed responses below.

Response:

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.

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.

At 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.

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.

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.

Dear Reviewer,

Thank you for taking the time to review our work. We have addressed your concerns in detail in the responses below.

Response:

Lack of Comparison on DiT-Based Diffusion Model:

Thank 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.

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.

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.

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.

We 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.

Dear Reviewer,

We appreciate your insightful comments and the time you've taken to evaluate our work. Below are our responses to your concerns:

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.

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.

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.

FLOP Comparison Table

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.

Combined Memory Usage Table (FLUX vs ToMA vs ToMA_tile)

Combined Memory Usage Table (SDXL vs ToMA vs ToMA_stripe vs ToMA_tile)

The 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.

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.

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.