Accelerating Transformers with Spectrum-Preserving Token Merging

Thank you for your feedback and high score, and for recognizing the main contribution of our work. We would now address your concerns regarding the compression rate of PiToMe.

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

Our PITOME can indeed achieve up to a $3\times$ speedup, as evidenced in Figures 5 and 6 (in the submission). In these experiments, we benchmarked PiToMe against various baselines using different ratios of $r$, where the compression rate ranged from approximately $70\%$ down to $30\%$. Generally, performance dropped by approximately $0\%$ to $3\%$, keeping more than $95\%$ of the baseline model's performance. Remarkably, on text-classification tasks (figure 10 in the submission), PiToMe can even compress down to $20\%$ of the number of tokens. This resulted in training and inference speeds approximately $5\times$ faster while maintaining accuracy above $85\%$ (about a $7\%$ accuracy drop compared to baseline ). Please refer to figures 2, 3, and 4 in the global rebuttal, where we conduct experiments on a wider range of ratio $r$.

Thank you for recommending arXiv [1], which was accepted by ECCV24 on July 01, 2024. Indeed, modulating token features during token merging can mitigate the performance drops under a low compression rate and achieve up to $\times3$ speed. However, we do not recommend compression rates higher than $60\%$ because compressing beyond this rate leads to significant performance degradation in many tasks. As shown in the arXiv [1] (PYRA) and ToMe, achieving a $3\times$ faster inference speed typically results in a $4\%$ to $7\%$ drop in accuracy compared to the baseline in image classification tasks, which is undesirable. Also, please note that factors other than the compression rate, such as hardware and batch size, can affect the model's speed. Therefore, we are unable to fully compare PyRA with PiToMe as the author hasn't published their code.

2. 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)

First, we want to clarify that Equation 4 defines the token energy score. In our approach, informative tokens (low energy) are clustered into a preserved group, which will not be merged, and we only merge the less informative tokens (top 2k tokens with highest energy). Like other compression methods, PITOME can handle higher compression rates; however, there is a trade-off between performance and compression rate.

It is important to note that the compression process primarily merges less informative tokens while preserving the informative ones, which is true as long as a reasonable ratio $( r )$ is maintained. This is demonstrated in all the trade-off figures, where PiToMe can maintain model performance before a significant drop occurs if $( r )$ is too low. In the final version, we will include more ablation studies on higher compression rates with various $( r^i )$ ratios for each layer $( l^i )$, as suggested by the reviewer.

3. 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?

In PITOME, the compression rate can be flexibly controlled based on the ratio r, representing the fraction of remaining tokens. Specifically, each layer $l^i$ is controlled by a ratio $r^i$, meaning that $1- r^i$ of the tokens will be merged. The overall percentage of remaining tokens after compression is given by:

For example, if we consider a ViT with 12 layers $(N = 12)$ and $r = 0.9$, the percentage of remaining tokens after compression is $r_{\text{remain}} = 0.538$. Similarly, for a BLIP2 model with 48 layers, $r_{\text{remain}} = 0.362$ with $r = 0.9$.

Thus, it is feasible to set the value $r^i$ for each layer randomly. However, it is important to consider the trade-off between performance and compression. As shown in our results, performance dramatically drops when $r < 0.85$ for all layers.

Thank you for your thoughtful feedback and for pointing out substantiated claims. We will now address your concerns one by one.

1. The experiments section is hard to parse

We will address these concerns in the final revision using the additional page. This will allow us to provide more details about the experimental results. Furthermore, we will move some less important parts to the appendix. For example, in Figure 3, we will remove redundant space at the bottom of the top-row figure and the top of the second-row figure to increase the size of the curves. Regarding the ablation study, we will include more detailed descriptions for each setting and provide an analysis of their impacts to improve clarity.

2. Discussion about the term $m$ in the ablation when computing the energy score

We acknowledge that the ablation study for $m$ is a valuable addition to the paper. To support readability, we include Figure 1 in the global rebuttal to visualize the impact of $m$ when experimenting with multiple margins $m$ on different tasks using various thresholds. From figure 1, it can be seen that, in both tasks, the adaptive version $m=0.9 - 0.9 \cdot l_i/l$ achieves the best results. While models with fixed $m$ tend to have the accuracy drop sharply when $r$ is lower than some threshold. The reason might be as the token space becomes more spare in deeper layers, PiToMe with fixed $m$ will likely assign the same energy score to every token which is undesirable since we want to identify isolated tokens to protect them while merging the others.

3. Some claims didn’t seem substantiated enough

• Line 120, which is “PITOME remains robust to token partitioning strategies.”

We apologize for the confusion the statement has caused. Indeed, this was a typo mistake. What we meant to convey is that "PITOME remains robust compared to other token partitioning strategies." In addition, we want to highlight that PITOME offers a better trade-off between speed and accuracy compared to previous BSM-based algorithms (i.e, ToMe, Diffrate, etc), while still being able to preserve spectral properties of the token space, similar to loop-based algorithms (i.e k-mean, spectral clustering, graph coarsening, etc).

• In 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”:

We agree that the dataset may contain noise. However, this noisiness can be interpreted as either biased or diverse perspectives. A possible explanation for the performance improvement in our token merging method is that the model treats all objects in the image equally, regardless of their size and the number of tokens they have. Large objects have their tokens merged, while smaller objects' tokens are preserved. Therefore, the number of tokens representing an object is not much different. This approach allows the model to avoid bias caused by object size and focus on the objects themselves rather than being biased by larger objects with a greater number of tokens. Additionally, by paying equal attention to all scene objects, the model aims to capture diverse perspectives.

4. The role of Table 2

In Table 2, we aim to demonstrate the following points:

• PITOME can be combined with various backbones (as mentioned by the reviewer). In particular, we test PITOME with 2 backbones in Table 2 (BLIP, BLIP2) and two additional backbones in Figure 3 (CLIP and ALBEF).
• We show that PiToMe, when combined with various backbones still outperforms state-of-the-art architecture (CLIP, ALBEF, UNITER, VilT, etc...).
• In addition to the image-text retrieval task, we evaluated PITOME on various other tasks using multiple backbones. For the image classification task, we tested our algorithm on six backbones with two different pre-training styles (DeiT and MAE). We used two backbones (LLaVA-7B and LLaVA-13B) for the visual question-answering task. We tested on two backbones (BERT and DistilBERT) in the text classification task. These experiments confirm the versatility of our algorithms.

Thank you for your helpful feedback. In what follows, we will address your concerns individually.

1. The authors claim that 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.

We're sorry for the confusion. As we use the cosine similarity to compute the energy score, the value of $x$ in Eq. 4 will be in the range $(-1,1)$. A value close to 1 indicates that the two vectors point in the same direction (identical tokens). Thus, since $m=0.9-0.9\cdot l_i/l$ , we have that for the first layer, the margin will have a value of 0.9, and only neighbors that have a cosine similarity larger than 0.9 will be considered as neighbors. Tokens outside this margin will have the value $x$ replaced by a negative constant $\alpha$. The margin becomes adaptively "wider" in deeper layers as $m$ decreases.

2. Does mean that when the token is outside the margins, is a negative constant? Furthermore, there are no related ablation studies or discussions of the impact of the $\alpha$

Yes, due to the strict page limits, we could not include the ablation regarding these parameters, the margin $m$, and the lower bound $\alpha$ . Recall that we used $\alpha(\exp(x-m) - 1)$ to smoothen the $f_m(\cdot)$ function to consider neighbor tokens that are outside but are close to the margin $m$. The larger $\alpha$, the more consideration for these neighbors, and $\alpha=0$ means we completely ignore them. The table below shows experiment results on the image-text retrieval task using different values of $\alpha$ and $m$.

3. In most of the experiments, the ratio is 0.95, 0.975, 0.9, and 0.925. Question about the consistency of the proposed method at higher merging ratios.

Thank you for your comment. We define the ratio $r$ differently from the ratio $r$ as mentioned in your question. We define the ratio $r$ as the number of tokens remaining in each layer after merging, not as the percentage of total remaining tokens after compression. Since PITOME is applied to all layers, $(1-r)$ percent of tokens will be merged after each layer. For example, let's consider a ViT model with 12 layers and a ratio $r=0.9$. This means that 10% of the tokens will be merged after each layer. The percentage of remaining tokens after compression is thus given by:

$r_{\text{remain}}$ is the ratio $r$ mentioned in your question, which means reducing nearly half the number of tokens the model needs to process and leading to about two times the throughput. This percentage can further decrease for larger models with more layers. For instance, in the BLIP2 model with 48 layers, we can achieve $r_{\text{remain}} = 0.362$  if $r = 0.9$ . For large language models like LLaVA, $r_{\text{remain}}$ can be even lower (in our paper we tested with $r_{\text{remain}}$ ranging from 0.732 down to 0.277).

4. PiToMe compared to TRIPS and CrossGET

For better readability we added a table to compare these works:

Since these papers do not provide their source code, it is difficult to reproduce their results. Also, we would like to explain more clearly the weakness of the CSGM algorithm used in CrossGET:

• CSGM's implementation divides all tokens into two disjoint sets, $T_s$ , and $T_d$, allowing only one token from $T_s$ to merge with one token in $T_d$. This is suboptimal for imbalanced clusters, such as large clusters with many tokens (e.g., backgrounds) and small clusters (e.g., small objects), which are common in practice.
• For example, let's assume we have four tokens where three are close to each other and one is isolated. CSGM would merge two of the three close tokens and then merge the remaining close token with the isolated one, causing significant information distortion. The correct behavior should be to merge the three close tokens and leave the isolated one untouched.
• PiToMe, on the other hand, assigns high energy scores to tokens in large clusters, prioritizing their merging while protecting isolated tokens. In the example above, the three close tokens are identified as mergeable, while the isolated token remains protected. Moreover, PiToMe can also benefit from cross-modality guidance techniques as add-ons to enhance performance.

Thank you for your valuable feedback and for recognizing the strengths of our work. We appreciate your thoughtful comments. We will now address each of your concerns individually.

1. Performance comparison with Tome

We would like to emphasize that PiToMe demonstrates significant improvements over the baselines, particularly ToMe. We acknowledge that the condensed presentation of experiments in the paper might have made it difficult to discern the performance gains of our method. We highlight the performance gains of PiToMe over ToMe as follows:

• Image-text retrieval: At a 60% compression rate, PiToMe outperforms ToMe by a margin of approx 11% for recall score at top 1.
• Visual Question Answering: With 50% compression rate, PiToMe maintains a more stable performance drop below 0.05% while ToMe's drop exceeds 1%.
• Text classification: PiToMe achieves a better trade-off, compressing over 60% of tokens with accuracy above 90% (about a 2% drop) and over 80% of tokens with accuracy above 85%.
• Image classification: PiToMe saves 40-60% of FLOPs, with an average performance drop of 1.4% compared to 2.6% for ToMe.

2. Figure 11

We kindly argue that a misunderstanding might cause the observation that the merging is unreasonable. For clarity, we discuss how we generate these visualizations and then provide a definition of good token merging:

• Implementation: Tokens are represented as grid boxes with black borders, colored by the mean color of the pixels they cover. Tokens with higher attention scores will have a bolder cyan border. During the merging process, result tokens will have their colors replaced by the mean color of all the merged tokens. If merged tokens are adjacent, the black borders will be removed.
• Good merging criteria: Objects represented by a small number of tokens should be preserved, while larger objects with repeated textures, such as background areas, can be considered for merging.

In Figure 11d, we observe that PITOME effectively minimizes information distortion by preserving important tokens representing "the tennis racquet," "the man," and the "Adidas logo" while merging tokens that contain redundant information representing the background. Despite some mis-merged tokens in the background regions, PITOME largely preserves the distribution of all tokens. This is illustrated by the attention map, which closely resembles that of the original model. On the other hand, In ToMe, there are several wrongly merged tokens; for example, all tokens representing the “tennis racquet" are merged into a single token. Similarly, for the “Adidas logo”. Also, ToMe merges tokens representing the man's arm and the white skirt with background tokens, which can be seen from the resulting false color.

3. Mathematical derivations and proof

The theory section, with mathematical derivations and proofs, is intended to show the spectrum preservation property of a graph built based on merged tokens using our method compared to the original token graphs in ViT. To improve readability, we will emphasize the key steps in deriving our theoretical results and provide high-level summaries in the final version.

4. Scope and significance

We focus on ToMe because it pioneered this line of research, making it a prime candidate to demonstrate our theoretical results. To address the reviewer's concern, we first provide a comparison table, below

We summarize our contribution:

• Conceptual novelty: We introduce a new concept called energy score. When applying this into other applications like LLM, 3D point clouds, video processing, or PDEs, it could open new research directions for designing an optimal energy score that captures specific underlying data distributions. Furthermore, energy minimization theory has been well-studied in fields such as chemistry and physics. Thus, we believe this will pave the way for a new research direction in these fields.
• Theoretical results: Our theoretical results provide a solid foundation for understanding and applying the energy score concept, offering valuable insights for further research and development.
• Extensive applications: To demonstrate the versatility of PITOME, we provided extensive experiments for multiple scenarios (image image-text retrieval, image classification, text classification, and VQA using LLAVA-7B and 13B). Furthermore, we compared our method with SOTA compression models in each specific downstream task, showcasing the broad applicability and effectiveness of our approach. Moreover, PITOME’s accuracy can also benefit from existing techniques (differential merging schedule, cross-modal guidance, etc.) to improve its accuracy.

Most importantly, due to the widespread use of Transformers and the associated high computational resources required, significant efficiency improvements can lead to substantial resource savings and open up new applications.

Dear Reviewer zfJf,

We have around five hours from this post before ending the discussion phase. Could you please let us know whether our responses answer your main concerns?

Regards

Authors