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

We thank the reviewer for the detailed feedback.

As suggested by the reviewer, our method A-MoD is a simple routing mechanism that can be easily integrated in dynamic computation models, specifically into MoD which is the focus of our work. Our main contribution is to highlight that there is no need for a dedicated router in MoD models, which needs additional tuning and can lead to training instability. The information required to perform routing is already present in the network which we leverage from attention maps, via A-MoD.

1. Integration to Flash Attention: A-MoD is in fact directly compatible with Flash Attention without having to compute the NxN matrix explicitly. Token importance score can be aggregated iteratively together with Flash Attention tiling strategy. This requires additional O(N) memory overhead for intermediate variables. We provide an algorithm for the same in Appendix A.8 of our updated draft.

2. Non-classification tasks: We also provide additional experiments on object detection with A-MoD as suggested by the reviewer (please see our general response and Table 9 in the updated draft).

3. Maximum achievable speedup: In our experiments for MoD with standard routing, we find that placing MoDs in each layer does not improve the performance and often leads to unstable training. Intuitively, in earlier layers it is often not yet clear, which tokens are relevant, and which can be skipped [6]. Thus, achieving the theoretical larger speedups is empirically difficult. Note that a denser placement of A-MoD layers, especially in later layers, could be achieved by re-using prior attention matrices (e.g., from 2 layers before) and relying on the fact that embeddings change slowly [5].
   Nonetheless, as shown in our comparisons with SOTA baselines, A-MoD is competitive and even better than other SOTA token merging and token pruning methods like Dynamic ViT[3] and A-ViT [4] and ToMe[2] (see our general response for results).

4. Closer IsoFLOPs: We design our isoFLOP models by reducing the depth of the vanilla ViT architectures such that our isoFLOP models can also benefit from the pretrained weights of the baseline models.

• We would like to clarify that, in majority of the cases (5/8), the isoFLOP models have a greater number of FLOPs than the A-MoD models, and yet, A-MoD is able to outperform them as shown in Table 1. We agree, in (2/8) cases, the number of FLOPs for the isoFLOP model <90%. Please refer to Figure 16 in the Appendix for a Pareto-curve of accuracy vs FLOPs.
• Additionally, we show an exact comparison between A-MoD and the isoFLOP baseline on DeiT-Small for 2.6GFlops by reducing the embedding dimension of the isoFLOP model to get more finegrained control over the FLOPs.
  • A-MoD (12.5% Capacity, 2.6GFlops): 76.98%
  • MoD (12.5% Capacity, 2.6GFlops): 76.07%
  • IsoFLOP (2.6GFlops): 75.73%

5. Pareto front visualization and throughput: Please see general response and the updated draft for accuracy vs FLOPs (Appendix A.9, Figure 16), accuracy vs throughput (A.10, Figure 17) and run-time analysis of our method.

Q1. We provide an algorithm compatible with Flash Attention in Appendix A.8.

Q2. Please see our general response for results on object detection (See Table 9 in Appendix).

Q3. We provide a Pareto curve in Figure 16 in the Appendix and make an exact isoFLOP comparison for DeiT-S.

• A-MoD (12.5% Capacity, 2.6GFlops): 76.98%
• MoD (12.5% Capacity, 2.6GFlops): 76.07%
• IsoFLOP (2.6GFlops): 75.73%

Q4. Throughput results are provided in Figure 17 in the Appendix .

Q5. A-MoD does not affect the scalability of Tranformer models w.r.t. the input sequence length and is compatible with efficient attention modules like Flash Attention (see Algorithm 1 in Appendix A.8). Hence, we believe that A-MoD still performs well in the regime with >4k tokens. However, as our paper focuses on the vision domain, where large sequence lengths are not typical, we leave this exploration to future work.

[1] Jain, Gagan, et al. "Mixture of nested experts: Adaptive processing of visual tokens." arXiv preprint arXiv:2407.19985 (2024).

[2] Bolya, Daniel, et al. "Token Merging: Your ViT But Faster." The Eleventh International Conference on Learning Representations.

[3] Rao, Yongming, et al. "Dynamicvit: Efficient vision transformers with dynamic token sparsification." Advances in neural information processing systems 34 (2021): 13937-13949.

[4] Yin, Hongxu, et al. "A-vit: Adaptive tokens for efficient vision transformer." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.

[5] Liu, Zichang, et al. "Deja vu: Contextual sparsity for efficient llms at inference time." International Conference on Machine Learning. PMLR, 2023.

[6] Puigcerver, Joan, et al. "From Sparse to Soft Mixtures of Experts." The Twelfth International Conference on Learning Representations.

I'd like to thank the authors for their comprehensive and concise responses. The following points have been addressed by their rebuttal:

• Integration with Flash Attention
• Fairer comparison with IsoFLOP baseline
• Pareto curve visualizations
• Throughput analysis

My concerns remain regarding the limited scope of the study (as shared by other reviewers), though I appreciate the inclusion of initial object detection results. I'd like the method to be evaluated on more dense prediction tasks such as segmentation or image generation. My concern regarding the reduced attractiveness of A-MoD compared to traditional MoD or MoE paradigms remains. I will update my score accordingly

We thank the reviewer for acknowledging our response and their positive comments.

Regarding the additional concerns,

1. We are the first to provide vision experiments on ImageNet classification for the MoD architecture, based on which we develop A-MoD to improve token selection based on attention. Our preliminary results suggest that we are able to extend the benefits of A-MoD to other vision tasks including detection. A comprehensive investigation of segmentation and detection lies outside the scope of this work, and we would dedicate time for it as part of future work.

2. We would also like to highlight that both MoD and A-MoD suffer from trainability issues in the case where each layer is accelerated with token skipping. However, if we keep the first few layers dense (no tokens skipped), we can leverage attention maps from those layers to implement A-MoDs in each of the successive layers. Similarly, standard routers would learn better in later layers allowing better speed up for MoDs when initial layers are kept dense. Please also note that MoEs are a different paradigm that are orthogonal to MoDs and can provide additional speed-up in combination with MoDs [1].

We are happy to clarify any further questions the reviewer might have and continue this discussion.

[1] Raposo, David, et al. "Mixture-of-Depths: Dynamically allocating compute in transformer-based language models." arXiv preprint arXiv:2404.02258 (2024).

We thank the reviewer for taking the time and providing detailed feedback for our work.

As pointed out by the reviewer our method A-MoD shows promising performance across provided experiments.

Experimental Setup: Training from a pretrained checkpoint

Our experiments are designed to showcase that A-MoD is a superior form of routing in terms of accuracy, when compared with the standard dedicated router as well as isoFLOP baselines. We build on the work of Raposo et al. [1], who have shown that MoD architectures are efficient, i.e., Pareto-optimal. We further show that the drawbacks of MoD architectures like training instability and router tuning can be easily eliminated with A-MoD.

1. We choose to start our experiments from a pretrained model to highlight that attention routing requires minimal training when adapting to a MoD layer from a standard dense architecture (see also Table 3 in the appendix). Our experimental setup starts from a model with pretrained weights and verifies the performance of A-MoD on a variety of datasets including finetuning and transfer learning.

   • Our results show that A-MoD outperforms isoFLOP baselines on ImageNet-1k and also converges faster on transfer learning tasks.
   • We also provide results for training from scratch:
     • DeiT-S after 300 epochs
       • A-MoD 76.63%
       • MoD 75.9%
     • ViT-Base after 160 epochs
       • A-MoD 73.66%
       • MoD 72.47%

2. We choose the transfer learning setting to test the adaptability of MoD models and find that A-MoD models converge faster than standard routing.

Efficiency and comparison to vanilla ViT

3. We have updated the abstract to make our experimental setup and results clear. Our main contribution is that A-MoD is able to outperform standard routing and its isoFLOP counterparts. Additional comparisons with other SOTA baselines also show that A-MoD is comparable and even better than these methods. While the focus of our work is not on efficiency, we provide additional results on efficiency as well where A-MoD can achieve the same performance as the baseline with 18% fewer FLOPs (please also see general response and Table 7 in the Appendix):

• DeiT-S baseline (4.6G FLOPS): 79.6%
• A-MoD (3.8G FLOPS): 79.63%

We also provide a Pareto-curve for accuracy vs FLOPs in Figure 16 and accuracy vs inference time in Fig 17 that highlights the pareto-optimality of A-MoD.

Layerwise capacity

4. We apply the same capacity to each MoD layer. Further, we have ablated the placement of the MoD layers by keeping the first few layers dense to allow improved feature learning and observe an improvement in comparison with the isoFLOP baseline as shown in Figure 7.
   • We believe that the capacity of each layer can be different in each layer to further improve A-MoD. However, finding the optimal per-layer capacity values is far from trivial, and we consider it as part of future work.

Comparison with SOTA baselines

5. We provide a comparison with state-of-the-art token pruning/merging methods including Dynamic ViT, A-ViT and Token Merging and find that A-MoD is comparable (even better) than these methods (please see the general response for results and Table 6 in the updated draft).

Token importance from attention maps

6. Figure 4 aims to differentiate between the tokens picked by standard routing (top row) and A-MoD (bottom row). Each column denotes a MoD layer. Standard routing struggles to identify the bird and only picks the background tokens, while A-MoD learns to pick the appropriate tokens representing the bird’s head starting from the third layer as the attention features become more meaningful in the later layers. Additional images are provided in the Appendix (see Figure 15).

7. In Figure 5, the attention maps are represented for every head in the last layer. Figure 5 highlights that A-MoD attention maps can better identify the objects in the image. Please see the Figure in our updated draft.

8. We mention that attention maps do not always learn semantic meaning (in line with the findings of [2]). However, as highlighted in Figure 6, attention scores do have a very high correlation with leave-one-out token importance in every layer validating their use as routing weights.

We thank the reviewer for acknowledging our response,

Q1. We are the first to introduce MoDs for vision classification tasks and benchmark it on ImageNet. Prior to our work, MoDs had only been introduced on language tasks. On vision tasks, we observe that standard MoDs suffer from routing instability and hence we propose A-MoD for improved routing. As highlighted in our Pareto-curve, A-MoD is even better that Dynamic-ViT and A-ViT (and is competetive with ToMe) which are state-of-the-art token-merging and token-pruning methods.

Q2. We average across all the attention heads as highlighted in Equation 4, where we introduce A-MoD. We use an average across all attention heads to get an overall importance of tokens which encompasses information across all attention heads instead of individual head which can be used to skip tokens in the transformer block. Moreover, this estimated token importance has a very high correlation (almost 1) with the leave-one-out token importance as shown in Section 4.4.

Q3. We conduct a profiling of the model using the PyTorch profiler and for a further breakdown of the inference time of individual operations as reported in Figure 19 and find that isoFLOP baselines require more GPU time. We attribute this speed up to the fact that A-MoD needs to process fewer overall tokens, in spite of having more depth as MoD layers are placed in every other layer.

Updates to the draft: As outlined in our general response, we have provided updates to the draft by adding Appendix sections A.8-A.13 along with modifying the abstract as suggested by the reviewer. We also provide Figure 1 to highlight the Pareto-optimality of A-MoD in comparison to MoD and isoFLOP baselines on ImageNet.

1. Training from scratch: Due to the limited time during the discussion phase and the long duration of training from scratch (at least 300 epochs), we are unable to provide isoFLOP runs for training from scratch as of now. However, the experiments are already running and Wwe will update them in the camera ready draft.

2. MoD select boundary patches: The reviewer makes a valid point, that standard routing MoD struggles to choose the correct patches and chooses boundary patches instead which is precisely the motivation of our work. We introduce attention routing which shows improved choosing of tokens as shown in Figure 5.

3. IsoFLOP baselines are better on transfer learning tasks in Table 5: While isoFLOP baselines are slightly better than MoD models (both MoD and A-MoD), Table 5 shows that the MoD architectures are still able to compete with standard isoFLOP baselines. Moreover, as shown in our ablation in Figure 8, A-MoD training is improved by allowing the initial layers to be dense and these A-MoD models are now able to match their isoFLOP counterparts, suggesting that while we show that the standard MoD model with alternate layers skipped are competitive, an improved layer placement of MoDs can enhance their performance further.

4. IsoFLOP models: [1] use isoFLOP models by reducing dimension as well as reducing model depth. We choose to reduce model depth such that we can leverage pretrained weights for isoFLOP models as well. However, we also provide an isoFLOP run for DeiT-S with reduced embedding dimension and training this model from scratch:

   • A-MoD (12.5% Capacity, 2.6GFLOPs): 76.98%
   • MoD (12.5% Capacity, 2.6GFLOPs): 76.07%
   • IsoFLOP (2.6GFLOPs): 75.73%

[1] Raposo, David, et al. "Mixture-of-Depths: Dynamically allocating compute in transformer-based language models." arXiv preprint arXiv:2404.02258 (2024).

We are happy to clarify any further questions the reviewer might have and continue this discussion.

We thank the reviewer for the valuable feedback.

1. Scope of our work and comparison with SOTA: As stated by the reviewer, the focus of this work is to introduce a new routing method leveraging attention for dynamic computation in Mixture-of-Depth architectures. The results in [1] show that MoD architecture can indeed achieve Pareto-optimal performance. Hence, we focus on improving the training and convergence of MoDs using attention-based routing as successfully shown by our experiments. Yet, we also provide a Pareto-curve for accuracy vs FLOPs in Figure 16 and accuracy vs inference time in Fig 17 that highlights the pareto-optimality of A-MoD.

   • Regarding Experimental Results:
     • A-MoD outperforms standard MoD and isoFLOP baselines in 7/8 settings in Table 1.
     • Please see general response for results on additional experiments, in particular, comparison to state-of-the-art methods on token pruning and merging (See also Figure 16 and Table 6 in the Appendix).

2. Regarding attention and token importance, we highlight:

   • The score in Equation 4 aggregates the interaction between a token and every other token. Intuitively, this means, that a higher score implies more tokens interact with a given token, and such a token is assigned higher token importance. With this, background tokens do not get higher scores in later layers, as also confirmed in Figure 5.
   • In Figure 4, each column is an MoD layer for standard routing on top and attention routing on the bottom. We have also updated the figure to highlight the same.
   • In Figure 5, each column represents an attention head of the last MoD layer in the model for standard routing on top and attention routing on the bottom. We have also updated the figure to highlight the same.

[1] Raposo, David, et al. "Mixture-of-Depths: Dynamically allocating compute in transformer-based language models." arXiv preprint arXiv:2404.02258 (2024).

We thank the reviewer for the valuable feedback.

1. No comparison with the current SOTA: We would like to clarify that the focus of our paper is not to claim a new ‘global’ SOTA method but to improve the field of dynamic computation models, which is still very young but holds great potential.

   • Here, our main contribution is that A-MoD offers a parameter-free routing mechanism for dynamically skipping tokens, preventing training instabilities and allowing for faster convergence. Since MoDs were designed for transformers and our method is built on leveraging the attention map for routing, we only focus on transformers in this work.
   • Specifically, we focus on understanding the effects of attention routing on training, convergence and Pareto-front comparisons for standard transformer architectures. Given the success of A-MoD on transformers it can be integrated into other hybrid architectures. Technically, we could also leverage internal representations within CNN to determine a routing. Both points are out of scope for this paper, but we aim to explore this as part of future work.
   • We also provide a Pareto-curve of accuracy vs FLOPs in Figure 16 in the Appendix to highlight the pareto-optimality of A-MoD.

2. Include the comparison (such as accuracy and latency) with these papers (Dynamic ViT[3], ToMe[2]): Please see our general response for a comparison with the mentioned existing SOTA methods for the token-pruning and token-merging and for additional experiments to verify the performance of A-MoD. We also provide a Pareto-curve for accuracy vs FLOPs in Figure 16 and accuracy vs inference time in Fig 17 that highlights the pareto-optimality of A-MoD.

[1] Raposo, David, et al. "Mixture-of-Depths: Dynamically allocating compute in transformer-based language models." arXiv preprint arXiv:2404.02258 (2024).

[2] Bolya, Daniel, et al. "Token Merging: Your ViT But Faster." The Eleventh International Conference on Learning Representations.

[3] Rao, Yongming, et al. "Dynamicvit: Efficient vision transformers with dynamic token sparsification." Advances in neural information processing systems 34 (2021): 13937-13949.

References

[1] Jain, Gagan, et al. "Mixture of Nested Experts: Adaptive Processing of Visual Tokens." The Thirty-eighth Annual Conference on Neural Information Processing Systems

[2] Bolya, Daniel, et al. "Token Merging: Your ViT But Faster." The Eleventh International Conference on Learning Representations.

[3] Rao, Yongming, et al. "Dynamicvit: Efficient vision transformers with dynamic token sparsification." Advances in neural information processing systems 34 (2021): 13937-13949.

[4] Yin, Hongxu, et al. "A-vit: Adaptive tokens for efficient vision transformer." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.

[5] Cao, Shiyi, et al. "MoE-Lightning: High-Throughput MoE Inference on Memory-constrained GPUs." arXiv preprint arXiv:2411.11217 (2024).

We provide additional updates to the draft, summarized as follows:

1. We added Figure 1 to the paper to highlight that A-MoD is Pareto-optimal on ImageNet compared to MoD and isoFLOP baselines (Please not that this changes the subsequent Figure numbers accordingly).
2. As per the suggestions of Reviewer N1Y6, we conduct a profiling of the model using the PyTorch profiler for inference time of individual operations as reported in Figure 19.