ToMA: Token Merging with Attention For Diffusion Models

1. “ToMA on DiT”:

   Thank you for pointing this out. We’ve added the experimental results for ToMA on FLUX.1 (DiT). Please refer to the general comment for more details. Additionally, we’ve analyzed the locality in the DiT model, and you can find the detailed explanation in Appendix H. We believe the locality of DiT is primarily driven by its positional embedding.

2. “Better Visualization”:

   We appreciate your suggestion, and we’ve improved the visualizations as recommended.

3. “Redundant Paragraphs”:

   Thank you for the careful review. We’ve revised and removed the redundant paragraphs as suggested.

Dear Reviewer,

As the end of the discussion period approaches, we are happy to address any additional questions you may have or further clarify your concerns.

Thank you once again for your time and consideration.

1. “Qualitative Analysis and Concerns About Noise”:

   Thanks for the suggestion. We agree that more visuals are necessary to substantiate our contribution. We have added more images and please check out the general comment.

2. “ToMA on DiT”:

   We have added our experiment result of ToMA on FLUX.1(DiT). You can check the general comment for details.

3. “Sharing in Small Steps”:

   Thanks for pointing this out and we have similar observations. However, from a result-wise analysis, which you can refer to in Table. 7 from Appendix C. 2, we find little difference in image quality between our sharing strategy and the non-sharing strategy.

4. “Scale of Destination Set”:

   In our setting, the size of $D$ is $r\cdot N$ where $r$ is the merge ratio (what percentage of tokens to merge) and $N$ is the original number of tokens. Essentially, the choice of $D$ depends on the level of speed-up & quality trade-off users desire. For theoretical speed-up in terms of $D$, please check the appendix D.

5. “Not Aligned Entries”:

   Thank you for noticing that and we have fixed it.

1. “Novelty of Locality”:

   Thank you for pointing this out. We admit the locality is a common practice in accelerating attention. However, our approach is different from the previous approach. Methods such as Longformer, Sparse Transformers, and Swin Transformer apply locality in attention computation. On the contrary, ToMA applies locality on the merge operation which keeps the computation attention intact. Theoretically, we can apply their method on top of our local merge. Our approach is plug-and-play, requiring no retraining, and maintains computational efficiency while still allowing long-range token interactions. Our experiments demonstrate that even marginal elimination of weak attention scores for distant tokens can result in significant image degradation, highlighting the importance of preserving these interactions.

2. “Comparison to other Baselines”:

   For Flash-attention, there might be a misunderstanding that it is actually orthogonal to our work and we develop ToMA on top of it. It is also a major reason why ToMA performs better compared to ToMeSD or other baselines as their overheads are more notable when utilizing efficient implementation like Flash-attention.

   Thanks for pointing out ToFu, we have added the baseline experiment in Appendix G and please check the general response.

   For Token Pruning, as ToFu is a more advanced version of token pruning, we compare ToMA with ToFu instead.

   For the other two papers, they are not directly related but orthogonal to our work. DiffRate incorporates the compression rate and merges tokens in the vision transformers but in the training stage. FRDiff accelerates diffusion inference by reusing feature maps across time steps but not merging the tokens. But to make the literature review more thorough, we have added them to the related work of the paper.

3. “Qualitative Analysis”:

   Thanks for pointing out the lack of visual analysis. We have added more images and please check out the general comment.

Dear Reviewer,

As the end of the discussion period approaches, we are happy to address any additional questions you may have or further clarify your concerns.

Thank you once again for your time and consideration.

1. “ToMA on DiT”:

   Thank you for your suggestion. We provide our result on the FLUX.1(DiT). Please check our general response.

2. “Inconsistent Comparison with ToMeSD”:

   To our best knowledge, ToMA is the method that has significant speed up with the modern implementation of attention(eg. flash attention/xformers/SDPA).

   The main reason behind the notable inconsistencies is that we utilize a more efficient implementation of attention in the diffusion models of our experiment while the ToMeSD utilizes a much slower version. As a result, ToMeSD’s overhead outweighs its benefits. This issue was also noted by the ToMeSD author in a related GitHub issue https://github.com/dbolya/tomesd/issues/57.

   On the contrary, our method utilize the attention for merging which will continuously benefit from the future more efficient implementation of attention.

   For more details, we use AttentionProcessor2_0, which relies on PyTorch's SDPA kernel with flash_attention2.0. This accelerates attention by approximately 10 times compared to the standard PyTorch implementation (as noted by Dao et al., 2023).

3. “FID Score”:

   In our experiments, the computation of the FID score is implemented using pytorch-fid, with the ImageNet-1k validation statistics provided by heusel2017gans in this https://github.com/bioinf-jku/TTUR.

   For the detailed setup, we generated 3,000 images for each configuration using a classifier-free guidance (CFG) scale of 7.5 with class names as prompts corresponding to ImageNet-1k classes. Each prompt was evaluated with three different random seeds. This setup may explain why our FID scores are much higher compared to classifier-guided generation on ImageNet-1k.

   The scale of our FID scores is also comparable to those reported in the ToMeSD paper, albeit slightly lower. This difference can be attributed to our adherence to the rule of thumb that the number of images generated for evaluation should exceed the hidden dimension of the last layer of the Inception-v3 model (2048). Specifically, we generated 3,000 images per configuration, whereas ToMeSD only generated 2,000 images per configuration. This ensures a fair comparison and yields more robust results in our case.

4. “Comparison with ToDo”:

   We have tested “Token Downsampling for Efficient Generation of High-Resolution Images”. In general, ToDo is slower than ToMA and gives lower-quality images. Please check Appendix G for details.

5. “Figure Change”:

   Thanks for the suggestion, we have changed the X-axis. Please check the paper for a modified version

Dear Reviewer,

As the end of the discussion period approaches, we are happy to address any additional questions you may have or further clarify your concerns.

Thank you once again for your time and consideration.

1. “Title and Scope Revision”:

   We have taken your suggestion into account and revised the title to more accurately reflect the focus of our work. Specifically, the revised title highlights the emphasis on generative tasks within diffusion models, ensuring that the scope of our contributions is clear. We sincerely thank you for this valuable suggestion.

2. “Evaluation on Discriminative Tasks”:

   Our current work primarily aims to enhance the efficiency of token merging specifically within the context of text-to-image (T2I) generation using diffusion models. We acknowledge that token merging is a versatile approach that could be evaluated in a broader set of tasks, including discriminative tasks and integration with baseline models. However, evaluating such applications falls outside the scope of this paper. We consider this as an exciting direction for future work, and we are actively exploring these applications as part of ongoing research.

3. “Related Work on CrossGET and TRIPS”:

   Thank you for pointing out these two papers for us and we have added these two papers to the related work. CrossGET combines token but on vision-language models with tasks like image captioning and image-text retrieval. TRIPS proposes text-relevant image patch selection but it accelerates the image-language model pertaining. The major difference between these two papers and ToMA is the task. ToMA focuses on image generation of diffusion models and they focus on other image-language tasks like image captioning and pretraining.

4. “Implementation of Tile-Shaped Regions”:

   We extract the tile-shaped region like a sliding window without overlapping area. Implementation-wise, it relies on the permute operation or the as_strided operation provided by Pytorch. For more details, please check the code in the supplementary material at the tile_wise_batched_facility function in the facility_location.py.

5. “Comparison with Farthest Point Sampling (FPS)”:

   Below, we address the distinctions between our submodular-based destination selection approach and Farthest Point Sampling (FPS) [Eldar et al., 1997] in terms of representativeness, theoretical foundation, and computational efficiency:

   • Representativeness and Diversity:

     The submodular-based approach we adopt is fundamentally different from FPS. In FPS, the next point $v^*$ is selected to maximize the minimum distance from all previously selected points:

     $$v^* = \arg max_{v \in (V \setminus A)} \min_{v' \in A} \text{dist}(v, v'),$$

     where $\text{dist}(v, v')$ measures the distance (e.g., Euclidean or cosine similarity) between points $v$ and $v'$. While FPS selects points to maximize spatial separation, our method ensures both representativeness and diversity in the selection of merge destinations by optimizing a gain function derived from the facility location problem. Specifically, the greedy algorithm selects the token $v^*$ from $V \setminus A$ that maximizes:

     $$v^* = \arg max_{v \in (V \setminus A)} f(v | A),$$

     where $f(v | A)$ measures the marginal contribution of a new token $v$ to the representative set $A$. This ensures that the selected tokens not only capture the most diverse features but also provide a comprehensive representation of the input, outperforming FPS in scenarios with complex feature distributions.

   • Theoretical Foundations:

     Unlike FPS, which primarily serves as a heuristic without strong theoretical guarantees, our method is rooted in the theory of submodular optimization, particularly the facility location problem. The gain function:

     $$f(v | A) = \sum_{u \in V} \max(0, \text{sim}(u, v) - \max_{v' \in A} \text{sim}(u, v')),$$

     provides a principled framework for selecting tokens.

   • Efficient GPU Implementation:

     We recognize that submodular optimization could introduce computational overhead. To address this, we have implemented the gain function computation on GPUs using a highly parallelized design. By vectorizing similarity calculations and marginal gain updates, our method achieves comparable running time to FPS on GPUs while maintaining superior token selection quality. Thus, the computational cost is significantly reduced, making the method practical for large-scale applications.

6. “Discrepancy in Inference Time with ToMeSD”:

   To our best knowledge, ToMA is the only method that has significant speed-up with the modern implementation of attention(eg. flash attention/xformers/SDPA).

   The main reason behind the notable inconsistencies is that we utilize a more efficient implementation of attention in the diffusion models of our experiment while the ToMeSD utilizes a much slower version. As a result, ToMeSD’s overhead outweighs its benefits. This issue was also noted by the ToMeSD author in a related GitHub issue https://github.com/dbolya/tomesd/issues/57.

   On the contrary, as ToMA utilizes attention for merging, its speed-up will continue to benefit from any future improvement with the implementation of attention, from which other methods suffer.

   For more details, we use AttentionProcessor2_0, which relies on PyTorch's SDPA kernel with flash_attention2.0. This accelerates attention by approximately 10 times compared to the standard PyTorch implementation (as noted by Dao et al., 2023).