Efficient Visual Transformer by Information Bottleneck Inspired Token Merging

Thank you for your review We appreciate the review and the suggestions in this review. The raised issues are addressed below.

Responses to the Weaknesses

1. “Insufficient and unclear motivation...”

We respectfully point out that this review missed the important insight and motivation which have been acknowledged by other reviewers such as Reviewer m6sM. We have a very important insight as the motivation for token merging with the IB principle. As evidenced in Table 2 and Table 5 in the paper, existing token merging methods, such as ToMe [1] and LTMP [3], already reduce the IB loss of the base models since they reduce redundant tokens with less relevant information for classification, which adheres to the IB principle that aims at learning representations more correlated with class labels while decreasing their correlation with the inputs. By explicitly reducing the IB loss in the token merging process, IBTM-Transformers enjoy lower IB loss compared to existing token merging methods and show better performance with even lower computational costs.

2. “Imprecise argument...”

We respectfully point out that there is a common sense mistake in this review, and there is misunderstanding about the claim “existing token merging methods largely sacrifice the prediction accuracy of the original transformer networks for reduced computation costs.” As a common practice in the token merging and the general model compression literature, we cannot look at the accuracy drop only at a low compression rate, such as a 20% speed-up on ViT-B/AugReg with only a 0.4% accuracy drop mentioned by this reviewer. Instead, we need to examine the performance of the compressed models across various compression rates. Our claim should be understood from a comparative study perspective, which is factually evidenced by the results in Table 1, where IBTM outperforms ToMe with comparable FLOPs/inference time. In addition, Table 8 of the revised paper shows that IBTM, with a 20% speed-up on ViT-B, still outperforms the vanilla ViT-B without compression.

In the revised paper, we clearly indicate that the IBTM causes less drop in prediction accuracy compared to the state-of-the-art token merging methods after this claim. The revised argument reflects this comparative aspect more accurately, acknowledging that while methods like ToMe can achieve efficiency gains with minor accuracy drops, IBTM demonstrates a better trade-off between the computational efficiency and the model performance, benefiting from the IB principle.

3. “Disorganized structure and weak readability of the Introduction section...”

We remark that all the other reviewers appreciate the organization of the instruction section. On the other hand, we respect the opinion of this reviewer and revised the introduction section in the revised paper, focusing on eliminating redundancy and emphasizing key points succinctly. In addition, we integrated the contributions directly into the main text of the Introduction, eliminating the need for a separate subsection.

4. “Unfair comparisons and potential over-claiming...”

We respectfully point out that there is a factual mistake in this review about the unfair comparison and potential over-claiming. As a standard practice in the token merging literature where the token merging models have learnable parameters, these earnable parameters need to be trained in a fine-tuning process. For example, LTMP [3], a token merging method with learnable parameters, loads the pre-trained vision transformers and fine-tunes the additional learnable parameters in the token merging modules. The underlying reason is that these learnable parameters are initialized randomly, so we can only obtain subpar performance without training these parameters. Moreover, it is important to note that the token merging methods without fine-tuning, such as ToMe and ToFu, have a computational step with token matching at all the layers to ensure similar original tokens are merged so that these models still need a computational process before performing token merging.

IBTM requires fine-tuning since it incorporates randomly initialized parameters in the token merging modules that require training. We follow the standard settings of the fine-tuning based token merging in LTMP [3] for our experiments. As evidenced in Table 1, IBTM outperforms the LTMP after fine-tuning for different numbers of epochs. In addition, it is worthwhile to mention that although ToMe [1] and ToFu [2] are free of fine-tuning, they require additional computation costs in performing the token matching at different layers.

Dear Reviewer b1MG,

This is a gentle reminder of your feedback. All the concerns in your original review have been addressed. Since today is the last day for your feedback, we really look forward to your feedback. We will clarify your further doubts/concerns if there are any. Thank you for your time!

Best Regards,

Authors

Thank you for your feedback, and below are our further responses to your remaining concerns.

(1) Further Response for Weakness 1 We will revise the claimed goal, "to merge the output tokens of all the transformer blocks into fewer tokens without largely sacrificing the prediction accuracy of the original vision transformer," to "to merge the output tokens of all the transformer blocks into fewer tokens without largely sacrificing the prediction accuracy of the original vision transformer adhering to a principled Information Bottleneck principle which allows for an informative token merging process". In this way, the unique value of this work is highlighted in the revised goal.

...why existing methods perform poorly in terms of the IB principle. The existing token merging methods merge the original tokens into target tokens by the weighted average [1, 2, 3]. Since each target token is an aggregation of the original tokens and the number of target tokens is less than that of the original tokens, the token merging process can be viewed as an information compression process where the rich information in the original tokens is compressed into the limited target tokens. As a result, the target tokens are less correlated with the input training features compared to the original tokens, so the mutual information between the target tokens and the input training features is smaller than that between the original tokens and the input training features, leading to a smaller IB loss compared to the baseline model without token merging. As shown in Table 2 in Section 4.2 and Table 7 in Section E.1 of the appendix, the baseline token merging methods, ToMe and LTMP, can already reduce the IB loss compared to the baseline models.

However, the information compression process is not effective enough in preserving the original rich information in the original tokens in the existing token merging methods such as ToMe and LTMP, because there is no information-theoretic measure to ensure that more informative tokens can receive higher importance weights in the weighted average process for token merging so that the target tokens can retain as much information in the original tokens as possible. This drawback of the existing token merging methods is also quantitatively reflected by the IB loss, that is, the IB loss of the existing token merging methods is not small enough. To this end, we propose IBTM, which employs a novel informative token merging process based on a principled information-theoretic measure, the Information Bottleneck, in the token merging process so that more informative and more important tokens contribute more to the merged tokens with a larger importance weight in the weighted average process for token merging. IBTM quantitatively achieves this goal by further reducing the IB loss by a novel and informative token merging mask which further reduces the IB loss compared to the existing token merging methods.

(2) Further Response for Weakness 2

“I was asking for results where ToMe and ToFu are fine-tuned for the same number of epochs as IBTM.”

In order to compare ToMe [1] and ToFu [2] with our IBTM in the fine-tuning setup, we create two baseline models for ToMe and ToFu. For example, there are two baseline models for ToMe so that these ToMe baseline models can be fine-tuned, which are termed ToMe (backbone only) and ToMe (LTMP mask). ToMe (LTMP mask) employs the learnable mask generated by Sigmoid from the LTMP [3], and the final mask for token merging is the product of the binary mask of ToMe and the learnable mask by LTMP. The weights of the neural backbone and the weights for the sigmoid will be trained when fine-tuning the ToMe (LTMP mask) model. In contrast, ToMe (backbone only) does not use the learnable mask generated from the LTMP [3], and only the weights of the neural backbone will be trained when fine-tuning the ToMe (backbone only) model. Similarly, we have designed two baseline models for ToFu, which are ToFu (backbone only) and ToFu (LTMP mask). We remark that ToMe (LTMP mask) and ToFu (LTMP mask) have the same number of learnable parameters for token merging as our IBTM model when the same neural backbone (such as ViT-B) is used.

We have fine-tuned the two baseline models for both ToMe and ToFu for the same number of epochs as IBTM-ViT-B for 1, 5, 10, 25, and 50 epochs following the settings in Section 4.1.1 of our paper. The results are shown in the table below. It is observed that although fine-tuning can improve the performance of the ToMe and the ToFu models, the IBTM-ViT-B still outperforms the ToMe (backbone only) and ToFu (backbone only) models. For example, IBTM-ViT-B outperforms the ToFu (backbone only) by $0.46$% in the top-1 accuracy when fine-tuned for 50 epochs with even faster inference speed. In addition, IBTM-ViT-B also outperforms the ToMe (LTMP mask) and the ToFu (LTMP mask). For instance, the IBTM-ViT-B outperforms the ToFu (LTMP mask) by $0.41$% in the top-1 accuracy when fine-tuned for 50 epochs with even faster inference speed, demonstrating the superiority of the IBTM in token merging under the fine-tuning setup.

All the baseline models for both ToMe and ToFu use the same neural backbone (ViT-B) as IBTM-ViT-B.

(3) Detailed process in token merging in hierarchical ViTs.

To answer all your questions, including the raised questions 3.1 and 3.2, we provide a detailed description about the token merging process in hierarchical ViTs as follows.

The token merging process in hierarchical ViTs. In the token merging process for hierarchical ViTs, the token merging modules are added between the attention modules and the MLP layers. Let the original tokens $X \in R^{N\times D}$ be the features before token merging, where $N=H\times W$ is the number of tokens. Let $\tilde X \in R^{P\times D}$ be the features after the token merging, where $P$ is the number of merged tokens. We explain how to obtain the output features of size $N \times D$ from $\tilde X$ as follows.

$\tilde X$ will be fed into the MLP layers in the transformer block. Let $\tilde Z = MLP(\tilde X)\in R^{P\times D}$ be the output features of the MLP layers. Before feeding $\tilde Z$ to the window attention of the next layer, we need to pad the features $\tilde Z$ so that the spatial size of $\tilde Z$ is still $N = H \times W$. Let $Z$ denote the padded version of $\tilde Z$. This padding process is achieved by repeating the merged tokens at the positions where the original tokens are merged into them. For example, if two original tokens, $X_j$ and $X_k$, are merged into a target token $\tilde X_i$ in the token merging process, both $Z_j$ and $Z_k$ co-located at the positions of $X_j$ and $X_k$ in the feature map will be filled by copying the target token $\tilde Z_i$. In this way, the features $\tilde Z$ are padded to the feature map $Z$ with a shape of $N \times D$ again, which can be reshaped to $H\times W \times D$. Therefore, the tokens in the padded features $Z$ as the input to the window attention in the next transformer block have the same spatial relationship as the original tokens in $X$. In addition, the computational cost is reduced since the number of tokens processed in the MLP layers is reduced from $N$ to $P$.

To answer your question 3.1, the “same features” are the copied target tokens in the description titled The token merging process in hierarchical ViTs.

To answer your question 3.2, please refer to the token merging process in the same description above. Even though the tokens in a window are merged into only a few tokens, the input features to the next window attention will still be in the same shape as the vanilla swin model due to the padding strategy. Therefore, the features $Z\in R^{N\times D}$ are of the same shape as the vanilla swin model, and the tokens can be placed into the window in the same manner as in the vanilla swin model. Again, as mentioned in The token merging process in hierarchical ViTs, the tokens in the padded features $Z$ as the input to the window attention in the next transformer block have the same spatial relationship as the original tokens as the input to the current transformer block.

References

[1] Bolya, Daniel, et al. "Token merging: Your vit but faster." ICLR, 2023.

[2] Kim, Minchul, et al. "Token fusion: Bridging the gap between token pruning and token merging." WACV, 2024.

[3] Bonnaerens, Maxim and Dambre, Joni. "Learned thresholds token merging and pruning for vision transformers." TMLR, 2023.

We appreciate the review and the suggestions in this review. The raised issues are addressed below.

Responses to the Weaknesses

1. “The related work on efficient vision transformers could be more comprehensive...”

In Section 2.1 of the revised paper, we have conducted more comprehensive discussions on the related works on efficient vision transformers, including [1].

2. “The description of the method could be more concise and clearer...”

We appreciate the suggestion, and we promise to revise the method section to be more concise and clearer in the final revision.

Responses to the Questions

1. “I wonder if incorporating token pruning techniques could lead to better results?”

Token merging methods can be easily combined with token pruning methods. For instance, token pruning modules can be applied right after token merging modules in transformers to further compress the vision transformers. We have performed a comparison between the token pruning methods and the IBTM in Section B.4 of the appendix of the revised paper. We will further investigate the effectiveness of combining the token merging methods with the IBTM.

References

[1] Wu X, Zeng F, Wang X, et al. PPT: Token Pruning and Pooling for Efficient Vision Transformers. arXiv preprint arXiv:2310.01812, 2023.

We appreciate the review and the suggestions in this review. The raised issues are addressed below.

Responses to the Weaknesses

1. “The major concern of this method is the requirement of training...”

As a standard practice in the token merging literature where the token merging models have learnable parameters, these learnable parameters need to be trained in a fine-tuning process. For example, LTMP [1], a token merging method with learnable parameters, loads the pre-trained vision transformers and fine-tunes the additional learnable parameters in the token merging modules. The underlying reason is that these learnable parameters are initialized randomly, so we can only obtain subpar performance without training these parameters. Moreover, it is important to note that the token merging methods without fine-tuning, such as ToMe and ToFu, have a computational step with token matching at all the layers to ensure similar original tokens are merged so that these models still need a computational process before performing token merging.

2. “Lack of experiments to validate the transferrability...”

We respectfully point out that we have adopted the IBTM models pre-trained on ImageNet as the feature backbone for other tasks. Our work focuses on performing token merging in vision transformers for computer vision tasks, and we have assessed the transferability of IBTM models on object detection and semantic segmentation in Section A.2 and Section A.3 of the appendix of our paper. We adopt the IBTM models pre-trained on ImageNet as the feature backbone for the object detection and semantic segmentation networks. It is observed in Table 4 and Table 5 that the IBTM models show consistent improvements over different vision tasks.

Responses to the Questions

3. “...compared to the existing upper bounds such as CLUB [1]?”

Although CLUB [2] also proposes an upper bound for the Information Bottleneck (IB), the derivation of the upper bound in CLUB assumes that the $p(\tilde X|X)$ is known, where $\tilde X$ and $X$ are the random variables representing the learned feature and the input feature. When $p(\tilde X|X)$ is unknown, they adopt the Gaussian Mixture Model (GMM) parameterized by an external neural network to approximate the upper bound. In contrast, the variational upper bound for the IB derived in our paper does not have such an assumption. As shown in Lemma C.1 in Section C.2 of the appendix of our paper, $p(\tilde X|X)$ can be directly computed from the training data without the need for training another neural network.

References

[1] Bonnaerens, Maxim and Dambre, Joni. "Learned thresholds token merging and pruning for vision transformers." TMLR, 2023.

[2] Cheng, Pengyu, et al. "Club: A contrastive log-ratio upper bound of mutual information." International conference on machine learning. PMLR, 2020.

We appreciate the review and the suggestions in this review. The raised issues are addressed below.

Responses to the Weaknesses

1. “...additional finetuning epochs...extra training time to train from scratch...”

As a standard practice in the token merging literature where the token merging models have learnable parameters, these learnable parameters need to be trained in a fine-tuning process. For example, LTMP [3], a token merging method with learnable parameters, loads the pre-trained vision transformers and fine-tunes the additional learnable parameters in the token merging modules. The underlying reason is that these learnable parameters are initialized randomly, so we can only obtain subpar performance without training these parameters. Moreover, it is important to note that the token merging methods without fine-tuning, such as ToMe and ToFu, have a computational step with token matching at all the layers to ensure similar original tokens are merged so that these models still need a computational process before performing token merging.

2. “...whether soft masks will also work for the same purpose...IBTM (thus fewer tokens) perform even better than the original models...”

IBTM first generates a token binarized mask $M^{(l)}$ for the $l$-th IBTM block. This mask is then elementwise multiplied by $\tilde G^{(l)}$, which is derived from the gradient on the variational upper bound of the IB loss. Since $\tilde G^{l}$ is a softmask, the token merging mask used by the IBTM block is indeed a softmask. In addition, it has been observed in the model compression literature [1,2] that compressing parameter-heavy models, such as vision transformers, can even improve the accuracy rather than compromising it due to the reduced over-fitting effect in the compressed models. In addition, features learned by the IBTM models exhibit lower IB loss compared to the vanilla models, as evidenced in Table 2 of the paper. Lower IB loss indicates a more effective representation of information, which could further explain the superior performance of the IBTM models.

Responses to the Questions

1. “Table 1 suggests that using IBTM will not increase the number of parameters compared to the original models yet Table 3 and 4 report different scenarios... ”

IBTM first generates a token binarized mask $M^{(l)}$ for the $l$-th IBTM block, employing the same module as used in the LTMP [3]. This mask is then multiplied by $\tilde G^{(l)}$, which is derived from the gradient on the variational upper bound of the IB loss. Therefore, IBTM incorporates the same number of additional parameters as LTMP. Nevertheless, the increase in parameters is marginal and does not alter the first digit after the decimal point of the number of parameters. The parameter sizes reported in Tables 3 and 4 were incorrect due to typographical errors and they have been corrected in the revised version of the paper.

*2.“Why IBTM is faster than LTMP in inference...“

IBTM demonstrates faster inference compared to LTMP because we configure the token compression ratio of IBTM to be lower than that of LTMP to assess the performance of IBTM with lower FLOPs/ inference time. For our experiments in Table 1, we set the compression ratio for IBTM models at $0.7$ and for all other competing token merging methods at $0.75$. It is observed that our IBTM models not only achieve superior top-1 accuracy but also do so with fewer FLOPs and enhanced inference speed, outperforming state-of-the-art token merging methods.

References

[1] Chen, Tianlong, et al. "Chasing sparsity in vision transformers: An end-to-end exploration." Advances in Neural Information Processing Systems 34 (2021): 19974-19988.

[2] Kong, Zhenglun, et al. "Spvit: Enabling faster vision transformers via latency-aware soft token pruning." European conference on computer vision. Cham: Springer Nature Switzerland, 2022.

[3] Bonnaerens, Maxim and Dambre, Joni. "Learned thresholds token merging and pruning for vision transformers." TMLR, 2023.