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

arXiv: https://arxiv.org/abs/2412.20875 · HTML: https://mlsanity.com/paper/2412.20875/

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Builds on 15 of its 47 references that are in the map.
- First posted on arXiv 30 Dec 2024.

## Builds on (most cited first, 15 of 15)

- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [Emergent Abilities of Large Language Models](https://mlsanity.com/paper/2206.07682/) (2022)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Decoupled Weight Decay Regularization](https://mlsanity.com/paper/1711.05101/) (2017)
- [Mixtral of Experts](https://mlsanity.com/paper/2401.04088/) (2024)
- [YOLOv10: Real-Time End-to-End Object Detection](https://mlsanity.com/paper/2405.14458/) (2024)
- [ST-MoE: Designing Stable and Transferable Sparse Expert Models](https://mlsanity.com/paper/2202.08906/) (2022)
- [Mixture-of-Depths: Dynamically allocating compute in transformer-based language models](https://mlsanity.com/paper/2404.02258/) (2024)
- [A Survey on Mixture of Experts in Large Language Models](https://mlsanity.com/paper/2407.06204/) (2024)
- [The Computational Limits of Deep Learning](https://mlsanity.com/paper/2007.05558/) (2020)
- [Conditional Computation in Neural Networks for faster models](https://mlsanity.com/paper/1511.06297/) (2015)
- [A Review of Sparse Expert Models in Deep Learning](https://mlsanity.com/paper/2209.01667/) (2022)
- [Mixture of Nested Experts: Adaptive Processing of Visual Tokens](https://mlsanity.com/paper/2407.19985/) (2024)

## Built on by (most cited first, 7 of 7)

- [Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation](https://mlsanity.com/paper/2507.10524/) (2025)
- [ThinkingViT: Matryoshka Thinking Vision Transformer for Elastic Inference](https://mlsanity.com/paper/2507.10800/) (2025)
- [Self-Routing: Parameter-Free Expert Routing from Hidden States](https://mlsanity.com/paper/2604.00421/) (2026)
- [DND: Boosting Large Language Models with Dynamic Nested Depth](https://mlsanity.com/paper/2510.11001/) (2025)
- [DirMoE: Dirichlet-routed Mixture of Experts](https://mlsanity.com/paper/2602.09001/) (2026)
- [ProgResViT: Progressive Resolution and Width for Adaptive Vision Transformers](https://mlsanity.com/paper/2609.03216/) (2026)
- [Bounded Autonomy and Verifiable Safety for Agentic AI Enabled Automation](https://mlsanity.com/paper/2610.08815/) (2026)

## Abstract (arXiv)

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 2x faster transfer learning.