Accelerating Transformers with Spectrum-Preserving Token Merging

Summary

This paper proposes Protect Informative Tokens Merging (PIToMe), a token merging method that safeguards crucial information-bearing tokens prior to the merging step. PIToMe focuses on the token-splitting step before token merging. Specifically, PIToMe defines and calculates an energy score for tokens to be merged, and marks the highest scored tokens as suitable candidates for merging. Experiments show that PIToMe consistently outperforms previous token merging methods such as ToMe, DiffRate, ToFu, etc. Theoretical analysis reveals that PIToMe preserves the intrinsic spectral properties to the original token space.

Strengths

1. Being the first method to consider the graph relations of tokens during token merging, the design of PIToMe is novel. Thorough theoretical analysis demonstrates the design insights of PIToMe. 2. Although some improvements are not significant enough, PIToMe successfully outperforms the previous token merging methods. Considering that many investigations towards token merging have been conducted by many works, I consider the improvements, albeit minor, as strengths. 3. Visualizations demonstrate the design insights of PIToMe. From those demonstrations, I could see that PIToMe merges the correct tokens in the image.

Weaknesses

1. ToMe is actually capable of conducting token merging at a very high compression rate (see Table 10 of ToMe paper, ViT-L and ViT-H can achieve about 2.5x speedup). Another recent work on arXiv [1] also demonstrate that this compression rate can even increase to >3x. However, in the experiments reported by PIToMe, the biggest speedup I could find is about 2x. I wonder whether PIToMe still works out for higher compression rates, since the metric vital for finding token clusters gradually defines more tokens as informative tokens (Eq. 4 in this PIToMe article). [1] PYRA: Parallel Yielding Re-Activation for Training-Inference Efficient Task Adaptation (arXiv 2403.09192)

Questions

1. Also on the topic of compression rate. In ToMe, r can be randomly set to achieve any compression rates. Can I randomly set the r values for each layer in PIToMe?

Rating

7

Confidence

3

Soundness

4

Presentation

4

Contribution

4

Limitations

See weaknesses & questions. I may change the rating after carefully checking the replies and reviews from other reviewers.

Summary

This paper introduces an energy-based approach to the token merging process, utilizing energy scores to avoid erroneous merges by distinguishing between informative or isolated tokens and redundant tokens. This method enhances the efficiency of heuristic merging approaches and preserves the spectral properties between the original and merged graph, thus accelerating token-based architectures such as ViT. Specifically, the technique employs cosine similarity with neighboring tokens to evaluate the energy score. Tokens with a high energy score are considered redundant and can be merged, indicating they belong to large clusters (e.g., background). Conversely, tokens with low energy scores are treated as foreground tokens. Furthermore, the authors also present theoretical findings demonstrating that the proposed method more effectively approximates the spectral distance between the initial token spaces and the merged token set compared to existing approaches. The method has shown competitive results across four tasks: image-text retrieval, visual question answering with LLMs, image classification, and text classification, underscoring its effectiveness and significance in the field.

Strengths

* This paper is well-written and clearly presents both the proposed solution and the experimental results. * In addition to thorough experiments demonstrating the effectiveness of the proposed method, the paper provides a theoretical derivation to explain why PiToMe outperforms another related approach (e.g., ToMe).

Weaknesses

* Some formulas and inferences are ambiguous, leading to misunderstandings that undermine the correctness of the theoretical explanations. * The lack of results for higher merging ratios weakens the practicality of the proposed approach. * There is a lack of discussion on related works (e.g., CrossGET [1], TRIPS[2]). [1] CrossGET: Cross-Guided Ensemble of Tokens for Accelerating Vision-Language Transformers, ICML 2024 [2] TRIPS: Efficient Vision-and-Language Pre-training with Text-Relevant Image Patch Selection, EMNLP 2022

Questions

* There are some ambiguities that need to be clarified in the methodology section. The authors claim that $m$ is a fixed constant margin; however, according to the context, $m$ seems to be the threshold of the margin, and $1-m$ is the margin. If I am wrong, the following description is confusing: In lines 162-163, "$m = 0.9 − 0.9 × l_i/l$, ... indicating an increasing margin as tokens move to deeper layers." This sentence seems to be the complete opposite of the conclusion. Also, does $\alpha = 1.0$ mean that when token $j$ is outside the margins, $f_m(x)$ is a negative constant? Furthermore, there are no related ablation studies or discussions of the impact of the $\alpha $. * In most of the experiments, the ratio $r$ is 0.95, 0.975, 0.9, and 0.925. Only the image classification results include a ratio $r = 0.7$. Existing approaches (e.g., ToMe) can achieve ratios of $r = 0.5 \sim 0.6$. I am curious about the consistency of the proposed method at higher merging ratios. * Some existing work, such as TRIPS and CrossGET, has also investigated token merging. The lack of thorough discussion or comparison with previous work weakens PiToMe's claimed novelty. I suggest the authors more clearly present their merits and differences in the paper.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

- The authors discuss some limitations in the last section.

Summary

This paper proposes a new strategy for reducing the number of tokens used by a vision transformer by merging similar tokens without. Such methods are very sensitive to the type of clustering/partitioning used on the tokens. The proposed approach, `PITOME` explicitly identify highly informative tokens, which should be kept out of the merging process. More specifically, the tokens are place in a graph where each node corresponds to a key embedding and the edges are the cosine distance between said keys. Based on this graph, an "energy" is computed for each token, capturing how mergeable they are. Only the top $2k$ are considered mergeable, while the rest is preserved. The actual following merging step follows a similar procedure as previous token merging methods: In summary, an "energy" score is computed on each token, using their key embeddings as base features. Using these energy scores, `EPITOME` draws itself apart from other merging techniques such as `ToME` in two ways; - explicitly exclude some tokens from the merging process - split tokens across the two sets $A$ and $B$ for bipartite matching as to maximise the change of max

Strengths

- The proposed method is thoroughly evaluated on image, text and image+text tasks - Generally, `PITOME` achieves higher accuracy while preserving the expected number of FLOPs and memory usage - the issues addressed by `EPITOME` are well explained and motivated.

Weaknesses

- Generally, the **experiments section** is really hard to parse: * the figure placement is a bit all over the place * some results would benefit to be moved to the appendix; for instance, in my opinion, the trade-off curves are much more informative than the tables * some figures are hard to read (Figure 3) (e.g. small captions, overlapping curves due to the y/x axes scales) * ablation experiments are much too condensed - In terms of ablation, there is little discussion about the term $m$ when computing the energy score: It seems to be an important hyperparameter as it impacts the "merge ability" of the tokens as the layer depth increase. Since `ToMe` and other related work are presented as having a lack of robustness to the token splitting strategy, it would be interesting to also show how robust `PITOME` is to this hyper parameter. - Some claims didn't seem substantiated enough to me - e.g. line 120 *"`PITOME` remains robust to token partitioning strategies"*: However `PITOME` uses a specific token partitioning algorithm. And the ablation on it (random split in Table 1) show that it actually suffers from a different portioning - line 276. *"We contend that this improvement stems from the merging of less significant tokens in PITOME potentially enhancing the robustness of the language mode"* -> this seems like a very vague hypothesis which in my opinion does really need a justification to begin with, since some other merging methods also outperforms the baseline. It may as well be a consequence of the noisiness of question answering datasets.

Questions

- I didn't fully understand the role of **Table 2**: isn't `PITOME` agnostic to the choice of the architecture ? it would send a stronger message to combine `PITOME` with various backbones, rather than compare it to different architectures.

Rating

6

Confidence

3

Soundness

3

Presentation

2

Contribution

3

Limitations

there is a limitation section discussing some current drawbacks of token merging methods in general

Summary

This paper proposes a novel method called PITOME, which enhances the efficiency of Vision Transformers (ViTs) by merging tokens in a way that preserves crucial information. Unlike previous methods, PITOME uses an energy score to prioritize and protect informative tokens while merging redundant ones. This approach reduces the computational and memory requirements by 40-60% with minimal impact on accuracy. PITOME achieves superior performance in tasks such as image classification, image-text retrieval, and visual question answering, demonstrating its effectiveness and robustness compared to existing token merging techniques. Furthermore, it theoretically preserves the spectral properties of the original tokens, contributing to its efficiency and performance stability.

Strengths

1. The introduction of the energy score metric for token merging is a novel idea that addresses the limitations of previous methods. This innovative approach effectively distinguishes between informative and redundant tokens, ensuring the preservation of critical information during the merging process. 2. The theoretical foundation for preserving the spectral properties of the original token space is a unique contribution. This aspect of the work sets it apart from existing methods, which often do not consider the impact on the spectral properties of the token space. 3. The paper provides extensive empirical validation across multiple tasks, including image classification, image-text retrieval, and visual question answering. 4. The theoretical analysis that supports the empirical findings enhances the overall quality of the paper. By demonstrating that PITOME preserves spectral properties, the authors provide a solid theoretical justification for the observed performance improvements.

Weaknesses

1. The proposed method, PITOME, does not show significant improvement over the baseline, particularly ToMe, in many tasks. The empirical results indicate that in some cases, PITOME's performance is only marginally better or even comparable to existing methods, which raises questions about its practical benefits. 2. As shown in Figure 11, PITOME fails to outperform ToMe in visualizations, and in some instances, the token merging appears unreasonable. For example, in Figure 11(d), the merging around "a tennis racquet" is poorly executed, leading to doubts about the method's effectiveness in preserving critical details. 3. The paper includes extensive mathematical derivations and proofs, which significantly impact its readability. The dense mathematical content makes it difficult for readers to quickly grasp the core contributions and the practical implementation of the proposed method. 4. The paper primarily focuses on improving ToMe by addressing its token merging strategy. While this is a valid research direction, the scope of the problem tackled is relatively narrow. The significance of the improvement may not be high enough to warrant extensive interest or application, as it addresses a specific aspect of an existing method rather than introducing a fundamentally new concept or problem.

Questions

1. In Figure 11, PITOME's token merging appears unreasonable in some cases (e.g., the tennis racquet example). Could you provide a deeper explanation or justification for these visualizations and any potential improvements to address these issues? 2. The mathematical derivations and proofs are quite complex and impact readability. Are there ways to simplify these explanations or provide more intuitive summaries of the key points to improve accessibility for readers?

Rating

5

Confidence

5

Soundness

2

Presentation

2

Contribution

2

Limitations

The authors have acknowledged some limitations of their work and provided suggestions for future research directions. The paper does not explicitly address any potential negative societal impacts. Given the focus on efficiency improvements in Transformers, the primary societal impact would likely relate to the broader implications of making these powerful models more accessible and efficient.

Follow Up Our Rebuttal

Dear Reviewer zfJf, We sincerely appreciate the time you have taken to provide feedback on our work, which will help us to improve its clarity, among other attributes. We want to follow up to check if our response fully addressed your concerns/questions before ending the discussion period. Sincerely, The Authors

Follow up our rebuttal

Dear Reviewer ZkPm, Because we have around two days left before ending the discussion phase, we wonder if our rebuttal above answers your concerns. Your major questions about discussing **related works**, **compressing rates**, and **ablation studies** on margin $m$ were provided. Please let us know if there is any points you are still unclear. Thank you very much for taking the time to give our feedback. We look forward to hearing your response soon. Regards The Authors

Follow up on our rebuttal

Dear Reviewer 4Z7M, We thank you again for your positive feedback on our work. Regarding your remaining questions about the **merging capacity of PiToMe**, we wonder if our explanation is clear enough. If you have other questions, please feel free to raise it. We are happy to discuss it before the **discussion period ends in two days**. Thank you very much Regards The Authors

Follow up on our rebuttal

Dear Reviewer jrpm, We would like to thank you very much for your feedback, and we hope that our response addresses your previous concerns, for e.g., about the experimental section in Table 2. If there exist some points we have not responded to so far, please feel free to let us know soon, as the discussion period is expected to end soon. We would be more than happy to address any additional concerns from you. Thank you again for spending time on the paper; we really appreciate it! Sincerely Authors.