Paper: A Simple Romance Between Multi-Exit Vision Transformer and Token Reduction (ICLR-2024)

Reviewer comment (Writing):
The phrase "Reduce Ratio" is not a good term to indicate the ratio of how many tokens are kept. Please change to another term like "keep ratio" to make it clear.

Author response:
We have changed the term from "reduce ratio" to "keep ratio" to refer to the proportion of tokens to keep. We have also added an elaboration of this concept in the main text.

Reviewer comment (Presentation):
Table 1 is not well explained. It took me a while to understand the setting of the experiment. The term "Off-the-shelf" is not immediately understandable. It would improve clarity by explicitly explaining the details of the experiments.

Author response:
We have refined the introduction of the off-the-shelf setting in both the main text and the caption of Tab. 1. It should now be clearer.

Reviewer comment (Writing):
It would improve the readability of the paper by changing some words/notations to standard ones, e.g., CSE -> CE for cross-entropy.

Author response:
We have changed the notation for cross entropy to $\mathrm{CE}$.

Reviewer comment (Reproducibility):
There seems no appendix, but at the end of Section 4 first paragraph, it says "See appendix for detailed experiment settings."

Author response:
Sorry for the absence of the appendix in our initial submission. It is now available in the revision.

Reviewer comment (Reproducibility):
Is the loss $L_{me}$ in Eq (8) also added to the $L_{total}$? It is not clearly mentioned in the paper.

Author response:
$\mathcal{L}\_{total}$ does not contain $\mathcal{L}\_{me}$. Specifically, when self-distillation is applied, $\mathcal{L}\_{sd}$ will directly replace $\mathcal{L}\_{me}$, thus they do not coexist.

Reviewer comment (Presentation):
Figure 2 can be improved with better illustration and explanation. Specifically, the arrows from the $A^c$ to the [CLS] are somewhat confusing. And why do the patch tokens have fading colors from the bottom up?

Author response:
Sorry for making you confused. By drawing the blue arrows from $A^{c}$ to the $[\textrm{CLS}]$ token, we hope to convey the idea that information from other tokens flows into the $[\textrm{CLS}]$ token, with $A^{c}$ defining the contribution ratio of each token. In the revision, we have added the explanation to the caption of the figure. The fading colors are used merely to differentiate the patch tokens, enabling the demonstration of the fact that the order of the tokens is changed after token sorting. Based on your comment, we have canceled the fading colors and instead labeled the tokens with id numbers.

Reviewer comment (Evaluation):
The evaluation demonstrates a notable improvement in accuracy compared to the baseline frameworks. It will be helpful to further demonstrate the reduction in latency with fewer FLOPs compared to other baselines.

Author response:
Given that we have already reported the FLOPs in the initial submission, I guess that by mentioning latency, you mean that we should further report the wall-clock latency or actual throughput to demonstrate the effect of METR in practice (If we misunderstood, please correct us).  We would like to first clarify that such indicators were absent in the initial submission because the latency and throughput are METR is totally determined by the base method with which METR is combined (e.g. EViT and DiffRate), rather than METR itself, as METR changes neither the architecture nor the computation at inference time. In other words, while METR can improve the accuracy of existing methods under the same latency, the exact latency is not defined or influenced by METR.  Therefore, we respectfully believe that, unlike most existing works in the token reduction area, latency and throughput are not of great importance in judging the effectiveness of METR. However, we also agree that listing such indicator directly in our paper will make it easier for the readers to form an intuitive understanding over the effect of METR in practice. Therefore, we have supplemented some throughput data in the revision.

Reviewer comment (Theory):
It would be beneficial if the author could offer more insights in the method section, such as explaining how and why this design can enhance performance.

Author response:
We deeply appreciate your suggestions. The fundamental insight and motivation behind METR are primarily detailed in Section 1. Additionally, we have succinctly reiterated the motivation in Section 3.2. As we aim to clarify any ambiguities and enhance our presentation, could you kindly provide more specific details regarding any aspects that may currently seem unclear? Your feedback would be invaluable in guiding our improvements.

Reviewer comment (Theory):
The use of self-distillation loss for the multi-exit training (Eq.9), in place of traditional multi-exit loss of Eq.8, although effective, is not adequately motivated. Self-distillation is typically used to improve the accuracy of the trained exits, which is not a requirement here as these are discarded at inference time. The manuscript would benefit from a more insightful analysis of what motivated this design choice/ why do the authors believe this works better than the traditional approach.

Author response:
## About Self-Distillation

While the working mechanism of knowledge distillation is not yet fully revealed, recent studies have made some important steps in identifying the aspects from which knowledge distillation may benefit. Among them, the following points could help explain why self-distillation would benefit METR:

1. Label denoising  [1]: existing datasets, though already curated, more or less suffer from inaccurate labeling. Aside from the images whose labels are literally wrong due to mistakes and errors in the data collection process, problems like the co-existence of multiple classes also lead to noises, as the label is always one-hot. Such data noise is adverse to generalization. When it comes to the early-exit heads of METR, this will also mislead the $[\mathrm{CLS}]$ token in determining which patch tokens to attend to. Self-distillation can alleviate this problem to some extent, as the soft predictions from the teacher model (the full-depth model in our case)  make better approximation to the real class distributions of the images. 
2. Another possible advantage of self-distillation is sample reweighting [2]. The inherent property of KL-Divergence causes the fact that the gradient applied to the student is magnified when the teacher model is confident over the sample; when the teacher is uncertain about a sample, the magnitude of student gradient is weakened.  This again plays the role of dataset purification, as the teacher model is more likely to be uncertain about problematic samples than unambiguous ones. When trained with such cleaner supervision, it would be easier for $[\mathrm{CLS}]$ token to learn to differentiate between important and unimportant patch tokens.
3. Additional cues [3]. As pointed out in the original KD paper, probability distribution predicted by the teacher on non-target classes provides additional cues, which could be helpful for the student. The first point listed above (label denoising) can be regarded as the embodiment of this idea on inaccurately-labled samples. On the other hand, even on accurate samples, such implicit cues might also provide some advantage. For example, given two images of dog, suppose that from the appearance perspective, the first one looks more similar to a fox, while the second one looks more similar to a wolf. The soft labels predicted by the final classification head may then reflect such similarity, i.e. allocate more probability to class fox for the first image, and more probability to class wolf for the second image. With self-distillation, such signals will serve as additional cues, which not only enhance the accuracy of early heads, but more importantly motivate the $[\mathrm{CLS}]$ token to allocate its attention in a way through which the $[\mathrm{CLS}]$ token can produce similar representations. For example, for the first image, the $[\mathrm{CLS}]$ token would be driven to pay attention to contents like ears, which grant the dog with fox-like looking. This may then benefit token reduction, as the tokens with high $[\mathrm{CLS}]$ attention scores tend to contain richer information and details.

[1] Feature normalized knowledge distillation for image classification. ECCV, 2020.

[2] Understanding and Improving Knowledge Distillation. arXiv, 2020.

[3] Distilling the knowledge in a neural network. NIPS Deep Learning Workshop, 2014.

Reviewer comment (Evaluation):
Row(2) in Tab.3 seems to be the equivalent of row (4) in Tab.2, where multi-exit and token-reduction fine-tuning are jointly applied (instead of the two-stage ablation in Tab2). If this is the case, it can be deduced that token-aware fine-tuning notably reduces the effectiveness of the proposed approach, leading to significantly smaller gains even when aggressive token reduction takes place. This fact is separate from the commented fading of multi-exit effects after separate fine-tuning and needs to be further investigated/discussed in the manuscript.

Author response:
### Does fine-tuning hurt the effectiveness of METR?

Following your suggestion, we have added a discussion about this problem in Appendix B. For your convenience, we copy the contents here:

> Comparing Tab. 1 with other experiments in the main text, we find that the advantage of METR is numerically more salient in the off-the-shelf setting. We deduce that the reason is that, with finetuning, the model can make some compromises, e.g. redundantly copy and store information among tokens, so that it could gain some robustness even when the token scoring metric is not reliable. In contrast, the off-the-shelf setting is way more challenging and poses stricter demands on the accuracy of token scoring. To empirically validate this deduction, we consider the random token selection baseline. Using the official DeiT-S checkpoint (79.80 top-1 accuracy), we first directly evaluate the model in the off-the-shelf setting, with token reduction applied and keep ratio set to 0.5. When $[\mathrm{CLS}]$ attention is used for token scoring and selection, the accuracy drops to 74.10; in contrast, when tokens are random reduced, the accuracy drops dramatically to 71.46. We then consider the 30-epoch finetuning setting, employing $\mathcal{L}_{final}$ as loss function, with token-reduction operation applied. When $[\mathrm{CLS}]$ attention is used for finetuning and evaluation, the accuracy is 78.47; when random token selection is used for finetuning and evaluation, the accuracy is 77.12 -- the gap is much smaller than the previous setting. This contrast provides an explanation over why the advantage of METR is more pronounced under the off-the-shelf setting.
>
> However, as indicated by our experiment results in main text, even though the improvement in the finetuning setting is not as numerically salient as the off-the-shelf setting, it is still fairly considerable, especially when taking the average developing pace of the token reduction topic into consideration. It indicates that an accurate token scoring metric is still indispensable in fulfilling the potential of token reduction, even with token-reduction aware finetuning applied. An explanation is that, the aforementioned compromises are unlikely to *completely* erase the influence of scoring metric quality. Furthermore, the more unreliable the metric is, the stronger the comprises have to be and the more model capacity would be occupied to achieve it, which would inevitably hurt the representation ability and the accuracy of the model.

Reviewer comment (Reproducibility):
An appendix is mentioned in the manuscript (Sec.4), but was not accessible to the reviewer.

Author response:
### Appendix

Sorry for the absence of the appendix in our initial submission. It is now available in the revision.

Reviewer comment (Theory):
What motivated the use of self-distillation in place of traditional multi-exit training in the proposed setting? What are the authors' insights about the demonstrated effectiveness of this design choice?

Author response:
## About Self-Distillation

While the working mechanism of knowledge distillation is not yet fully revealed, recent studies have made some important steps in identifying the aspects from which knowledge distillation may benefit. Among them, the following points could help explain why self-distillation would benefit METR:

1. Label denoising  [1]: existing datasets, though already curated, more or less suffer from inaccurate labeling. Aside from the images whose labels are literally wrong due to mistakes and errors in the data collection process, problems like the co-existence of multiple classes also lead to noises, as the label is always one-hot. Such data noise is adverse to generalization. When it comes to the early-exit heads of METR, this will also mislead the $[\mathrm{CLS}]$ token in determining which patch tokens to attend to. Self-distillation can alleviate this problem to some extent, as the soft predictions from the teacher model (the full-depth model in our case)  make better approximation to the real class distributions of the images. 
2. Another possible advantage of self-distillation is sample reweighting [2]. The inherent property of KL-Divergence causes the fact that the gradient applied to the student is magnified when the teacher model is confident over the sample; when the teacher is uncertain about a sample, the magnitude of student gradient is weakened.  This again plays the role of dataset purification, as the teacher model is more likely to be uncertain about problematic samples than unambiguous ones. When trained with such cleaner supervision, it would be easier for $[\mathrm{CLS}]$ token to learn to differentiate between important and unimportant patch tokens.
3. Additional cues [3]. As pointed out in the original KD paper, probability distribution predicted by the teacher on non-target classes provides additional cues, which could be helpful for the student. The first point listed above (label denoising) can be regarded as the embodiment of this idea on inaccurately-labled samples. On the other hand, even on accurate samples, such implicit cues might also provide some advantage. For example, given two images of dog, suppose that from the appearance perspective, the first one looks more similar to a fox, while the second one looks more similar to a wolf. The soft labels predicted by the final classification head may then reflect such similarity, i.e. allocate more probability to class fox for the first image, and more probability to class wolf for the second image. With self-distillation, such signals will serve as additional cues, which not only enhance the accuracy of early heads, but more importantly motivate the $[\mathrm{CLS}]$ token to allocate its attention in a way through which the $[\mathrm{CLS}]$ token can produce similar representations. For example, for the first image, the $[\mathrm{CLS}]$ token would be driven to pay attention to contents like ears, which grant the dog with fox-like looking. This may then benefit token reduction, as the tokens with high $[\mathrm{CLS}]$ attention scores tend to contain richer information and details.

[1] Feature normalized knowledge distillation for image classification. ECCV, 2020.

[2] Understanding and Improving Knowledge Distillation. arXiv, 2020.

[3] Distilling the knowledge in a neural network. NIPS Deep Learning Workshop, 2014.

Reviewer comment (Evaluation):
Is token-reduction aware fine-tuning indeed limiting the effectiveness of the proposed approach? If yes, this should be commented in the manuscript.

Author response:
### Does fine-tuning hurt the effectiveness of METR?

Following your suggestion, we have added a discussion about this problem in Appendix B. For your convenience, we copy the contents here:

> Comparing Tab. 1 with other experiments in the main text, we find that the advantage of METR is numerically more salient in the off-the-shelf setting. We deduce that the reason is that, with finetuning, the model can make some compromises, e.g. redundantly copy and store information among tokens, so that it could gain some robustness even when the token scoring metric is not reliable. In contrast, the off-the-shelf setting is way more challenging and poses stricter demands on the accuracy of token scoring. To empirically validate this deduction, we consider the random token selection baseline. Using the official DeiT-S checkpoint (79.80 top-1 accuracy), we first directly evaluate the model in the off-the-shelf setting, with token reduction applied and keep ratio set to 0.5. When $[\mathrm{CLS}]$ attention is used for token scoring and selection, the accuracy drops to 74.10; in contrast, when tokens are random reduced, the accuracy drops dramatically to 71.46. We then consider the 30-epoch finetuning setting, employing $\mathcal{L}_{final}$ as loss function, with token-reduction operation applied. When $[\mathrm{CLS}]$ attention is used for finetuning and evaluation, the accuracy is 78.47; when random token selection is used for finetuning and evaluation, the accuracy is 77.12 -- the gap is much smaller than the previous setting. This contrast provides an explanation over why the advantage of METR is more pronounced under the off-the-shelf setting.
>
> However, as indicated by our experiment results in main text, even though the improvement in the finetuning setting is not as numerically salient as the off-the-shelf setting, it is still fairly considerable, especially when taking the average developing pace of the token reduction topic into consideration. It indicates that an accurate token scoring metric is still indispensable in fulfilling the potential of token reduction, even with token-reduction aware finetuning applied. An explanation is that, the aforementioned compromises are unlikely to *completely* erase the influence of scoring metric quality. Furthermore, the more unreliable the metric is, the stronger the comprises have to be and the more model capacity would be occupied to achieve it, which would inevitably hurt the representation ability and the accuracy of the model.

Reviewer comment (Evaluation):
In Tab.1,2,3 does “reduce ratio” refer to number of tokens or GFLOPs? Both should be reported to get the full picture.

Author response:
### About 'Reduce Ratio'

In the initial submission, we used the term 'Reduce Ratio' to refer to the ratio of tokens to $keep$ at each token-reduction block. As pointed out by you and Reviewer G39w, we agree that the phrase is not proper here, and we have changed it to 'Keep Ratio' in the revision. Furthermore, we have also added an elaboration of this concept in the main text.

Reviewer comment (Writing):
Notation in Sec.3 is quite loose. Consider defining the dimensionality of each introduced symbol (X,x,A,...).

Author response:
### Typos and notations

Thank you for pointing out these problems! We have corrected the typos, and refined the mathematical notations based on your suggestions.

Reviewer comment (Writing):
In Sec.4.1.1, the phrase "without incorporate" should be corrected to "without incorporating", and "Subsequently, We (...)" should be corrected to "Subsequently, we (...)".

Author response:
### Typos and notations

Thank you for pointing out these problems! We have corrected the typos, and refined the mathematical notations based on your suggestions.

Reviewer comment (Evaluation):
It would be better to shows the GPU-throughput and compare it with those of SOTA.

Author response:
Thank you so much for your recognition of our work!
## About GPU throughput

We would like to first clarify that we did not report throughput in the initial submission because the latency and throughput of METR are *totally* determined by the base method with which METR is combined (e.g. EViT and DiffRate), rather than METR itself, as METR changes neither the architecture nor the computation at inference time. In other words, while METR can improve the accuracy of existing methods under the same throughput, the exact throughput is not defined or influenced by METR at all. Therefore, we respectfully believe that, unlike most existing works in the token reduction area, latency and throughput are not of great importance in judging the effectiveness of METR. However, we also agree that listing such indicator directly in our paper will make it easier for the readers to form an intuitive understanding over the effect of METR in practice. Therefore, we have supplemented some throughput data in the revision.