Attention Is All You Need For Mixture-of-Depths Routing

Authors: Advait Gadhikar, Souptik Kumar Majumdar, Niclas Popp, Piyapat Saranrittichai, Martin Rapp, Lukas Schott
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
Keywords: Mixture of Depths, attention, parameter free routing.

Abstract:
Advancements in deep learning are driven by training models with increasingly larger numbers of parameters, which in turn heightens the computational demands. To address this issue, Mixture-of-Depths (MoD) models have been proposed to dynamically assign computations only to the most relevant parts of the inputs, thereby enabling the deployment of large-parameter models with high efficiency during inference and training. These MoD models utilize a routing mechanism to determine which tokens should be processed by a layer, or skipped. However, conventional MoD models employ additional network layers specifically for the routing which are difficult to train, and add complexity and deployment overhead to the model. In this paper, we introduce a novel attention-based routing mechanism *A-MoD* that leverages the existing attention map of the preceding layer for routing decisions within the current layer. Compared to standard routing, *A-MoD* allows for more efficient training as it introduces no additional trainable parameters and can be easily adapted from pretrained transformer models. Furthermore, it can increase the performance of the MoD model. For instance, we observe up to $2$\% higher accuracy on ImageNet compared to standard routing and isoFLOP ViT baselines. Furthermore,  *A-MoD* improves the MoD training convergence, leading to up to $2\times$ faster transfer learning.

Decision: reject
Reviewer scores: 3, 5, 3, 5

Reviews and discussion:
<thread>
creation date:
1730059846416

thread_type:
official_review

<thread_object>
creator: reviewer_je9b

<content>
summary:
Mixture of Depth methods have recently been proposed as an alternative to Mixture of Expert methods. Instead of deciding to which expert to route a token, MoD takes a decision whether or not to skip the current expert. The main purpose shared by these methods is that of finding a better compute/performance tradeoff, which is currently of particular interest in the context of LLMs and large generative models. One of the main challenges for building such methods is the design of stable routing mechanisms. The authors propose A-MoD, a routing method which performs pruning of tokens traversing the current layer, based on the respective attention scores in the previous layer. The authors evaluate their method on classification tasks on a set of DeiT and DiT models of varying size.

strengths:
- The method is simple and clearly explained throughout the paper. It seems possible to implement it easily for small models operating on short token sequences.
- The proposed method is simple and the idea of eliminating learnable routing, a major source of training instabilities, is an interesting research direction.
- Some performance improvements are shown under some settings of a classification task.

weaknesses:
- The method requires access to the attention matrices. Modern attention implementations (e.g. "FlashAttention") do not materialize attention matrices due to their large materialization cost. An implementation of A-MoD viable for an LLM or large generative model would require non-trivial implementation effort to match the performance of modern attention implementations.
- The chosen evaluation setting is composed only of classification tasks. Token pruning is particularly suited for this task, as all tokens with information not relevant to the depicted object can be safely discarded and the model output consists in only an aggregation of the tokens. A more significant evaluation would have taken segmentation, language modeling, or image/video generation problems into account. In these tasks token pruning can be problematic, as every token will be returned as an output. Pruned tokens are at risk of generating artifacts in such context. Lack of evaluation in these context is a significant limitation.
- The method presents a limitation in the maximum achievable speedup with respect to a corresponding dense architecture. In particular, as each A-MoD block bases its pruning on the previous block, it can only be applied at alternating blocks, posing a maximum achievable speedup to 2X over a corresponding non A-MoD model. In contrast, a learnable routing mechanism would allow application of MoD to each layer, allowing for higher theoretical speedups, making the pursuit of this paradigm more promising than A-MoD.
- Some recent MoE architectures (Snowflake Arctic LLM) show that it is beneficial to apply MoE at every layer. A-MoD can apply it only every other layer
- Table 1 proposed evaluation with an isoFLOP baseline. Under 3/8 configurations, the isoFLOP baseline has <90% of the flops of the A-MoD model. Such comparison does not appear fair
- Table 1 would benefit from being accompanied by a Pareto curve visualization of performance vs compute, which could highlight better the Pareto front for the C=12.5% case, where compute of the baseline and A-MoD method is mismatched.
- The paper does not take inference throughput into consideration. In this paper, FLOPs are an incomplete indicator of model cost in that they do not highlight performance losses caused by the need to materialize the attention matrix. It is likely that the isoFLOP baseline when implemented carefully with an efficient attention implementation would outperform A-MoD in terms of inference throughput and memory utilization, especially if evaluated in practically-relevant language modeling or generative tasks with large amounts of tokens.
- Discussion in 4.5 and Fig. 7 are not convincing. In Fig 7 (a) performance is aligned to the isoFLOP baseline, in Fig 7 (b) compute is too mismatched between baselines to draw conclusions.

In summary, the pursuit of A-MoD seems less promising than traditional MoD or MoE paradigms due to its limited theoretical maximum achievable speedup of 2X. Evaluation is performed on classification only, missing some important performance metrics (throughput/latency) and in some circumstances under unfair evaluation settings for the baseline. Evaluation does not comprise important dense prediction tasks such as language modeling or image/video generation. No discussion or remedy is presented for the need of materializing full attention matrices, limiting applicability of the method for large models operating on long sequences, which could benefit most from MoD/MoE.

For these reasons, I do not believe the method implementation and evaluation are solid enough for the paper to be accepted.

questions:
- Can the authors provide an implementation of an efficient attention mechanism such as "Flash Attention" with support for A-MoD? If yes, how would the performance of this implementation compare to the original implementation?
- Can the authors show performance of the method on a task such as semantic segmentation, language modeling or image/video generation that involves producing an output token for each input token to confirm that token dropping would not affect output quality?
- Could the authors revise Table 1 to perform a fairer comparison with isoFLOP baselines and show results using a Pareto curve?
- Could the authors revise qualitative evaluations to include throughput measures rather than FLOPS only, ensuring baseline methods make use of efficient attention implementations?
- Could the authors revise quantitative evaluation to include a task involving large numbers of tokens (>=4k tokens) to show that the attention operation implementation can remain efficient as the input size grows?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730505937004

thread_type:
official_review

<thread_object>
creator: reviewer_n1y6

<content>
summary:
In this work, the authors aims to deal with an interesting problem: How to reduce the token number need to be computed in a layer, and proposed A-MOD. A-MOD is simply based on attention score to select tokens need to be calculated, and the rest will be directly skiped. Experimental results show promising results.

Such a topic is always interesting and important, especially in the era of large models. However, almost none is applied to indutry because these methods are always only efficient in theory. I encourage more explorations in this topic. Regarding this work, I believe many experimental settings are not reasonable or not as expected.

strengths:
1. The authors tackle an important issue in improving the efficiency of ViTs by proposing a method that dynamically selects the most relevant tokens for computation. This approach aligns with similar techniques in the field, such as A-ViT, which also prioritize token selection for efficiency gains.

2. Additionally, the experimental results show promising performance.

weaknesses:
The main idea behind the proposed method makes sense to me, but I have several concerns, particularly regarding the experiments and their settings.

1) First, the authors mention in Line 234 that they continue training from a previous checkpoint for an additional 100 epochs. I’m curious whether this approach is justified. Why not simply start training from scratch? 

2) I also find the transfer learning setup a bit confusing, especially since the authors do not use fixed pretrained weights. Why is transfer learning considered in a work that focuses on efficiency? How does it contribute to the overall goals of the work? Clarity on this point would be help.

3) Regarding the baseline, I expected to see a vanilla ViT for comparison, but instead,  the authors reduced the original ViT layers to match FLOPs directly. This choice makes no sense to me if the goal is to demonstrate that A-MOD improves effiency. Additionally, I noticed that A-MOD significantly lowers the original performance of ViT. For instance, the original ViT-Tiny achieves ~74% accuracy, whereas A-MOD only reaches 69.76% and 71.8%. 
The abstract is is also very misleading; I initially was thinking the method reduces the number of tokens while improving accuracy on ImageNet by 2% (a huge improvement on ImageNet).

4) Is the same selection ratio applied across all layers, from shallow to deep? we know that shallow layers typically do not generate meaningful attention maps.

5) At a high level, A-MOD aims to reduce token number for computation, which is similar to the approach taken in the TO-ME paper (“Token Merging: Your ViT but Faster,” ICLR 2023). However, I noticed that the performance of A-MOD is much worse, and there are no speed tests provided to show efficiency gains.

6) In Fig. 4,  the last image of the first row, it appears that MOD selects almost all unimportant boundary patches and parts of the bird's head. Does this make sense? If it does, I would like to know why.

7) Generally, attention maps in shallow layers do not reliably indicate importance. Given this, how should we interpret the results shown in Fig. 5? Were the specific attention heads or examples chosen selectively, or could you provide additional much more examples (besides the bird and car used throughtout the paper) to support the findings?

8) Finally, in Line 420, the authors state, “attention maps do not always learn semantically meaningful scores.” This leads me to two questions: How did you conclude that this issue is only limited to larger models? If this only happens in larger models, why does A-MOD work for ViT-Tiny?

questions:
1. Are there any other reproducible results for MOD on ImageNet?

2. How are different heads within a single layer handled? Will different heads select different tokens for computation in a layer?

3. As indicated in the abstract, it’s unclear how this method improves training and inference efficiency on an ImageNet-scale dataset. Are these efficiency improvements achieved in practice (like in training and inference speed) or just in theory?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730628671903

thread_type:
official_review

<thread_object>
creator: reviewer_68jn

<content>
summary:
This paper proposes a parameter-free routing method to improve the routing mechanism in the mixture-of-depths (MoD) method. The basic idea is to directly use the attention scores in Transformer.

strengths:
- Using the transformer attention in the MoD routing makes sense.
- Extensive results do prove that the proposed A-MoD outperforms MoD.

weaknesses:
- The entire study is too narrow and limited. This paper is specifically targeted at MoD. However, MoD is just an arxiv paper. Does MoD represent the SoTA in terms of the Pareto frontier? All the experiments are mainly compared with MoD? What about other SoTA methods? BTW, MoD is not impressive in Table 1, where it is even inferior to a simple baseline, isoFLOP.

- Why is higher average attention in (4) corresponding to higher importance? What is the semantic meaning? If so, shouldn't the background tokens in Fig. 4 not be skipped since they have high attentions with many other background tokens? BTW, what are the different columns in Fig. 4 and Fig. 5?

questions:
Please address the issues pointed out in the weakness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730792141908

thread_type:
official_review

<thread_object>
creator: reviewer_l4xz

<content>
summary:
This paper introduces a parameter-free routing method for mixture-of-depths (MoD) models. 
It uses attention maps to discover the importance of each token. 
The authors also tested the results on a series of DeiT/ViT models.

strengths:
This paper explores the token importance evaluation through the existing attention maps and thus reduce the overhead of extra layers.

weaknesses:
1. There is no comparison with the current SOTA models since DeiT/ViT are relatively old. 
2. The idea is quite similar to the following papers. It would be great to include the comparison (such as accuracy and latency) with these papers. 
DynamicViT: Efficient Vision Transformers with Dynamic Token Sparsification and 
Token Merging: Your ViT But Faster

questions:
1. what's benefits of the proposed approach compared to traditional ConvNet? It would be great to include some comparison with ResNext etc
2. What's the advantage of the proposed approach against the transformer / convnet hybrid architecture such as LeViT: a Vision Transformer in ConvNet's Clothing for Faster Inference or FastViT: A Fast Hybrid Vision Transformer using Structural Reparameterization
</content>
</thread_object>

</thread>
<thread>
creation date:
1732285826097

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response to reviewers

comment:
General Response 

We thank all the reviewers for taking the time and providing us with valuable feedback to improve our work.  

We would like to highlight the main contributions of this paper as appreciated by the reviewers: 

1. Reviewer 68Jn: A-MoD makes sense and extensive results show it is better.  
2. Reviewer N1Y6: Experimental results show promising performance. 
3. Reviewer je9B: The method is clearly explained and easy to implement. 

We would like to provide a summary of our response to address issues raised by the reviewers (L4xZ, N1Y6, je9B) as well as the updates made to the draft. Individual points are addressed in direct responses.  When possible, we have also addressed these points in the paper.  

1. **Comparison to more baselines** (Reviewer L4xZ, 68Jn, je9B): We compare our method, A-MoD, to state-of-the-art token merging and token pruning methods including Dynamic ViT [3], A-ViT[4] and Token Merging (ToMe)[2]. A-MoD outperforms Dynamic ViT and A-ViT. Moreover, despite standard training A-MoD is comparable to ToMe which additionally benefits from further training with distillation (see Table 6 in the Appendix). We also provide a Pareto-front curve for FLOPs vs accuracy in Figure 16. Please see Appendix A.8 in our updated draft. We have also stated a summary of the results below: 

| **Model**  | **Method**              | **Top-1 Acc (%)** | **FLOPs (G)** |
|------------|-------------------------|-------------------|---------------|
| **DeiT-T** | A-MoD                  | 71.8              | 0.9           |
|            | A-ViT                  | 71.0              | 0.8           |
|            | Dynamic ViT            | 70.9              | 0.9           |
|            | ToMe (with distillation) | 71.69*           | 0.93          |
| **DeiT-S** | A-MoD                  | 78.66             | 3.42          |
|            | A-ViT                  | 78.6              | 3.6           |
|            | Dynamic ViT            | 78.3              | 3.4           |
|            | ToMe (with distillation) | 79.68*           | 3.43          |
*with distillation 

2. **Efficiency Results: How many tokens can be dropped without sacrificing performance?** (Reviewer N1Y6): To highlight the efficiency of A-MoD, we compare it with the baseline DeiT-S and report the top-1 accuracy on ImageNet. A-MoD can reduce the number of FLOPs by up to 18% without dropping performance, with standard training and no additional tricks. See Appendix A.11, Table 7. The Pareto-front curve in Figure 16 also confirms the effectiveness of A-MoD. 
- DeiT-S baseline (4.6G FLOPS): 79.6% 
- A-MoD (3.8G FLOPS): 79.63% 

3. **Inference time optimizations** (Reviewer je9B, N1Y6): We acknowledge the reviewer comments that our analysis is limited to FLOPs and not inference time or throughput. We report inference time results for A-MoD below and in Appendix A.10 (Fig. 17) in the updated draft. GPU hardware and low code are not yet optimized for dynamic compute. This is ongoing efforts [5] and we compare the models simply in PyTorch  as is (e.g., no TensorRT, quantization, …).  

4. **Training from scratch** (Reviewer N1Y6): We also provide results for training from scratch on ImageNet and observe that A-MoD outperforms standard routing. See Appendix A.12, Table 8. 
- DeiT-S after 300 epochs: 
  - A-MoD 76.63% 
  - MoD 75.9% 
- ViT-Base after 160 epochs 
  - A-MoD 73.66% 
  - MoD 72.47% 
 
5. **A-MoD with Flash Attention** (Reviewer je9B): In response to reviewer je9B, we show that A-MoD can be integrated with Flash Attention without the need to realize the NxN attention matrix explicitly. The detailed derivation is included in Appendix A.8 (Algorithm 1) of our updated draft. 

6. **A-MoD beyond classification tasks** (Reviewer 68Jn , je9B): We also provide results with A-MoD on object detection tasks. Preliminary experiments suggest MoD can be applied to object detection with only a marginal drop in performance, with only 50% FLOPs (which has not been done before). The performance of A-MoD is on par with MoD. Longer training runs are in progress and will be updated for the camera ready version. Please see Appendix A.13, Table 9 in the updated draft. 

| **Metric**                               | **DETR MOD-50** | **DETR A-MOD-50** | **DETR-Baseline** |
|------------------------------------------|-----------------|-------------------|-------------------|
| **mAP**                                  | 39.6%           | 38.6%             | 39.9%             |
| **GFLOPS (Total)**                       | 83.2            | 83.2              | 86.56             |
| **GFLOPS (Transformer Encoder/Decoder)** | 7.747           | 7.745             | 10.745            |

We hope that our response has addressed the concerns of all the reviewers. We are happy to continue this discussion and further improve our work. Based on our response, we also please request the reviewers to reconsider their scores. 

Thank you 

Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper introduces a parameter-free routing method for mixture-of-depths models. The proposed approach aims to optimize the trade-off between computation and performance, which is especially relevant for large-scale models. The authors demonstrate the effectiveness of their method through experiments on DeiT/ViT models.

However, after the review process, all reviewers voted to reject this submission, with concerns such as limited experiments, insufficient explanation of the experimental results, and a lack of fair comparisons with existing works.