REVIEW 4 major objections 5 minor 24 references
Attention Is All You Need For Mixture-of-Depths Routing
T0 review · 4 major / 5 minor · reviewed 2026-08-10 · deepseek-v4-flash
Pith's one-line read A-MoD, a parameter-free Mixture-of-Depths router, derives token importance from the previous layer's attention maps and reports gains up to 2% accuracy over standard learned routing on ImageNet.
desk verdict A neat, well-tested idea for attention-based MoD routing, but the main comparison is confounded by output scaling, so the central claim needs a cleaner experiment. read the letter →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
What carries the argument
The central object is the attention map of the previous layer, $A^{l-1}_h \in \mathbb{R}^{N\times N}$, whose column sums are averaged across heads to form a token-importance score. This score is then used in a threshold router: keep tokens with $r_i \ge P_\beta(R^l)$, skip the rest. The mechanism replaces the learned linear projection $W^l_r$ of standard routing and, unlike standard routing, does not multiply the layer output by the routing score, which the paper credits for faster adaptation from pretrained checkpoints.
What would settle it
A direct test would be to shuffle the attention-derived scores across tokens within each layer, keep the same capacity, and retrain or re-evaluate: if accuracy is unchanged, the routing signal itself is not doing the work. A second test would measure A-MoD against standard routing on a large model whose attention maps are known to concentrate on a background or register token; the paper's own Fig. 15 suggests such a regime exists, and if A-MoD still wins there, the explanation cannot be that attention tracks semantic importance.
Extended reading notes
Core claim
A-MoD computes the importance score of each token as the average attention it receives from all queries across all heads in the preceding layer: $r_i = \frac{1}{HN}\sum_{h=1}^{H}\sum_{j=1}^{N} a^{(l-1)}_{h,ji}$ (Eq. 4). Tokens whose score falls below the $\beta$-th percentile are skipped in the current layer (Eq. 5). The paper argues that this parameter-free signal is a better indicator of token importance than a learned linear router, and supports this with leave-one-out token-importance correlations showing A-MoD scores consistently correlate with importance while standard router scores sometimes correlate negatively. Because the attention maps come from the pretrained model, A-MoD starts from higher zero-training accuracy and converges faster when adapting pretrained checkpoints.
Load-bearing premise
The argument rests on the assumption that the average attention a token receives in the previous layer is a reliable measure of how important that token is for the current layer, an assumption the paper states but does not prove and which the authors note can fail in larger models where attention concentrates on single patches.
Editorial extensions
If this is right
- MoD adaptation from pretrained checkpoints becomes a router-free procedure: no router parameters to initialize, train, or stabilize.
- Finetuning and transfer learning need fewer epochs to reach a given accuracy, with reported speedups up to about 2x on small datasets.
- For small vision transformers (DeiT-Tiny, DeiT-Small), A-MoD can exceed both standard-routing MoD and isoFLOP ViT baselines on ImageNet at 50% and 12.5% capacity.
- Token selections become interpretable: routed patches align with object-relevant regions and routing scores correlate with leave-one-out importance, unlike standard routing.
- A-MoD remains implementable under Flash Attention, so the parameter-free router can be used with memory-efficient attention kernels.
Reading between the lines
- If averaged attention is a general importance signal, the same trick could be applied to MoD in language models, where attention maps are already computed; the paper only demonstrates vision transformers.
- The authors' own caveat about attention maps lacking semantic meaning in larger models suggests A-MoD's advantage may shrink at scale; a router that mixes attention scores with a small learned correction could combine both signals.
- A-MoD effectively recovers something like token pruning without a separate scoring head, so it could be compared against token-pruning methods under identical training budgets; the paper's comparison relies on published baselines with different training recipes.
- Because the router adds no parameters, the method makes the FLOP-accuracy trade-off of MoD essentially a property of the pretrained attention itself, which may change how MoD capacity is chosen during adaptation.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The manuscript introduces A-MoD, a routing mechanism for Mixture-of-Depths Vision Transformers that computes token importance scores by averaging the previous layer's attention maps over heads and token positions (Eq. 4) and uses these scores to decide which tokens are processed by the current MoD layer (Eq. 5). The authors compare A-MoD with the standard learned linear router of Raposo et al. (Eqs. 1-2) in finetuning from ImageNet-pretrained DeiT-Tiny, DeiT-Small, ViT-Base, and ViT-Large at 50% and 12.5% capacity, and in transfer learning to Stanford Cars, Oxford Pets, and Flowers102. They report that A-MoD mostly outperforms standard routing and isoFLOP baselines, converges faster, requires no router parameters, and can be adapted from pretrained checkpoints with little or no training. They also report a leave-one-out correlation analysis, a Flash-Attention variant, DETR results, and comparisons to token-pruning methods. The central claim is that information already present in attention maps is sufficient and beneficial for MoD routing.
Significance. The paper addresses a real practical problem: MoD routers add parameters, are hard to train, and must be learned when adapting dense pretrained transformers. A parameter-free router built from existing attention maps would be a simple and useful contribution, and the paper tests it across a broad range of model sizes, capacities, and datasets. The Flash-Attention adaptation is also a thoughtful engineering contribution. If the empirical claims were clean, the paper would likely be of interest to the efficient-vision and conditional-computation communities. However, the headline comparison is currently confounded: A-MoD changes both the routing source and the output-scaling rule relative to the standard router, and the reported comparisons are single runs. The correlation analysis has a train/eval mismatch. These issues prevent the paper, in its present form, from establishing that attention-based routing itself is the cause of the observed gains. The contribution remains promising and the required fixes are experimental, not conceptual.
major comments (4)
- [§3.2–3.3, Eqs. (2) and (5), Fig. 11] The main comparison in Tables 1, 4, and 5 changes two variables at once. Standard routing in Eq. (2) multiplies the processed block output by the router score: x_i^l = r_i f_l(X^{l-1})_i + x_i^{l-1}; A-MoD in Eq. (5) drops this multiplication: x_i^l = f_l(X^{l-1})_i + x_i^{l-1}. The paper's own ablation in Fig. 11 shows that adding the r_i multiplication to A-MoD lowers accuracy and slows convergence, so the absent scaling is independently beneficial. No experiment holds the output rule fixed and varies only the routing source (e.g., a standard router without r_i multiplication trained with a straight-through estimator, or an A-MoD with r_i multiplication compared directly to standard MoD under the same multiplication rule). Until such a condition is added, the conclusion that attention maps are better routers than learned routers is not established; the gains in Table 1 may be entirely or partly due to the different output rule.
- [§4.1–4.2, Tables 1, 4, and 5] All reported accuracies are single runs with no error bars, seeds, or significance tests. The headline 'mostly outperforms' includes margins as small as 0.07 percentage points (ViT-Base, 12.5% capacity, Table 1) and many 0.1–0.5 point differences in the transfer tables, which are within typical run-to-run variation for ImageNet finetuning. The authors should report mean and standard deviation over at least three seeds for the central ImageNet comparisons, or provide paired bootstrap confidence intervals for the A-MoD versus MoD differences. Without this, the quantitative claims 'up to 2% higher' and 'mostly outperforms' are not reliably supported.
- [§4.4] The leave-one-out correlation analysis does not evaluate the deployed model. The text says token importance is measured by omitting the token 'in the vanilla transformer,' whereas the routing scores are taken from the trained MoD models. Because the MoD model has skipped tokens and different activations, importance measured in the dense vanilla model need not reflect the effect of skipping a token in the MoD model. The correlation should be recomputed by removing tokens at MoD layers inside the trained MoD model itself, and compared with a random-router baseline; otherwise Fig. 7 and the claim that 'A-MoD consistently selects important tokens' are not supported.
- [§3.3, Fig. 15, related work] The paper's core assumption—that averaged attention maps estimate token importance—is acknowledged in Section 3.3 to be an assumption, and the paper itself cites Darcet et al. (2024) and shows in Fig. 15 that attention maps of larger models can concentrate on a single patch without semantic meaning. Given that A-MoD nevertheless performs well on ViT-Base and ViT-Large, the authors need to address whether the empirical gains in those cases are actually due to semantically meaningful routing or to the output-scaling difference identified above. A quantitative check (e.g., routing quality versus attention-map entropy or semantic-segmentation agreement) would help, but at minimum the discussion should explain how the method's premise survives these observations.
minor comments (5)
- [§1 and §4.3] There is a missing space in 'domain.Our' in the introduction, and 'learnign' appears in Section 4.3; the paper would benefit from another proofreading pass.
- [Table 3] The zero-training comparison in Table 3 is not a fair head-to-head: standard MoD starts with a randomly initialized router, while A-MoD's scores come from the pretrained attention maps. This should be stated explicitly, or a random-attention baseline should be added to show the effect of the pretrained attention itself.
- [§A.9, Table 6] The comparison with ToMe is favorable to ToMe because it trains with distillation, as the authors note; the table would be more informative if the training settings were matched (or if ToMe's non-distillation result were included).
- [§A.8, Algorithm 1] In Algorithm 1, the accumulation step 'R ← R + Pij/li' should specify that Pij is normalized per query row and that Atemp is reinitialized for each query block; as written, the pseudocode is ambiguous about the loop scope.
- [§4.5] The 'MoD only in later layers' ablation is reported only for Stanford Cars and two architectures; stating whether the same behavior holds on the other transfer datasets would make the proposed remedy more compelling.
Circularity Check
A-MoD is an externally evaluated, parameter-free routing method; the main results are not fitted and no load-bearing argument reduces to its own inputs.
full rationale
The paper's central derivation is self-contained: A-MoD defines routing scores directly from the previous layer's attention maps (Eq. 4) and uses them in the MoD output rule (Eq. 5); no parameter is fitted to the target benchmark and no prediction is constructed from the data it claims to predict. The comparisons against standard MoD, isoFLOP baselines, and token-pruning methods on ImageNet, transfer tasks, and from-scratch training are external evaluations. The leave-one-out token-importance correlation is an independent proxy (loss change upon token removal) and is not defined in terms of the attention scores, so that validation is not circular. The only notable design concern is a treatment confound: standard MoD (Eq. 2) multiplies the block output by the learned score r_i while A-MoD (Eq. 5) does not, so the reported gains could partly stem from the absent scaling rather than from attention-based selection; the paper's own ablation (Fig. 11) shows that adding the multiplication to A-MoD reduces accuracy. This is an experimental-identification issue, not a circularity by construction, and it does not make the central result equivalent to its inputs. There are no load-bearing self-citations: references to Raposo et al. define the baseline MoD framework, and Darcet et al. is cited to acknowledge a limitation of the attention-importance assumption. The paper also explicitly states its assumption that attention maps estimate token importance rather than hiding it as a derived result. Overall, the derivation chain is not circular.
Assumptions & free parameters
free parameters (3)
- MoD capacity C =
0.125, 0.5, 0.7
- MoD layer placement =
alternate layers; later-layer variant
- Finetuning learning rate =
1e-5 for ImageNet, 0.01 for transfer
assumptions (5)
- domain assumption Softmax attention scores from the previous layer measure pairwise token relevance.
- domain assumption Average column attention aggregated over heads is a good proxy for token importance.
- domain assumption MoD's premise that tokens need different amounts of compute and that skipping tokens is beneficial.
- domain assumption Leave-one-out loss difference is a valid estimate of token importance.
- standard math Transformer attention as softmax(QK^T/sqrt(d)) works as described.
Cite this review
Pith. "Pith review of Attention Is All You Need For Mixture-of-Depths Routing." pith.science (2026). https://pith.science/paper/EFWPRL4G
@misc{pith2026241220875,
author = {Pith},
title = {Pith review of: Attention Is All You Need For Mixture-of-Depths Routing},
year = {2026},
howpublished = {\url{https://pith.science/paper/EFWPRL4G}},
note = {Machine review of arXiv:2412.20875}
}
read the original 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 2x faster transfer learning.
Figures
Figures from the paper (16 more)
Reference graph
Works this paper leans on
-
[1]
Neural machine translation by jointly learning to align and translate
Dzmitry Bahdanau. Neural machine translation by jointly learning to align and translate. arXiv preprint arXiv:1409.0473,
-
[9]
(2022)) to column first in order to aggregate row wise scores efficiently
It must be noted that we swap the order of tiling from row first (as in Dao et al. (2022)) to column first in order to aggregate row wise scores efficiently. Thus in each iteration only the Q can be cached while the K and V need to be loaded for each tile. Our algorithm introduces a small memory overhead of O(N ) due to additional temporary variables. 16 ...
work page 2022
-
[10]
M.I. Jordan and R.A. Jacobs. Hierarchical mixtures of experts and the EM algorithm. InProceedings of 1993 International Conference on Neural Networks (IJCNN-93-Nagoya, Japan), volume 2, pp. 1339–1344 vol.2,
work page 1993
-
[12]
Decoupled weight decay regularization
I Loshchilov and F Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101,
-
[14]
Mixture-of-depths: Dynamically allocating compute in transformer-based lan- guage models
David Raposo, Sam Ritter, Blake Richards, Timothy Lillicrap, Peter Conway Humphreys, and Adam Santoro. Mixture-of-depths: Dynamically allocating compute in transformer-based lan- guage models. arXiv preprint arXiv:2404.02258,
-
[15]
doi: 10.1007/s11263-015-0816-y. Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In International Conference on Learning Representations ,
-
[16]
The computational limits of deep learning
Neil C Thompson, Kristjan Greenewald, Keeheon Lee, and Gabriel F Manso. The computational limits of deep learning. arXiv preprint arXiv:2007.05558,
arXiv 2007
-
[17]
Yolov10: Real-time end-to-end object detection
Ao Wang, Hui Chen, Lihao Liu, Kai Chen, Zijia Lin, Jungong Han, and Guiguang Ding. Yolov10: Real-time end-to-end object detection. arXiv preprint arXiv:2405.14458,
Show all 24 references
-
[18]
Emergent abilities of large language models
Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yo- gatama, Maarten Bosma, Denny Zhou, Donald Metzler, et al. Emergent abilities of large language models. arXiv preprint arXiv:2206.07682,
-
[19]
Algorithm 1 A-MoD with Flash Attention, modified from Algorithm 1 in Dao et al
18 Figure 16: Visualizing routing scores and token selection of each MoD layer of a DeiT-Tiny MoD with 50% capacity for standard routing (top) and attention routing (bottom) finetuned on Stanford Cars. Algorithm 1 A-MoD with Flash Attention, modified from Algorithm 1 in Dao et...
2022
-
[20]
Fine- tuning with A-MoD: Results comparing A-MoD with standard routing and isoFLOP baselines for 12.5% capacity on ImageNet-1k
13 (a) DeiT-T (b) DeiT-S (c) ViT-B (d) ViT-L Figure 9: A-MoD achieves better performance and faster convergence on ImageNet-1k. Fine- tuning with A-MoD: Results comparing A-MoD with standard routing and isoFLOP baselines for 12.5% capacity on ImageNet-1k. A A PPENDIX A.1 M ODE...
2015
-
[21]
A.8 A-MoD WITH FLASH ATTENTION Flash Attention has been proposed by Dao et al. (2022) as a method to implement attention without the need to compute theN ×N attention map explicitly, reducing hardward communication overhead and thus speeding up computation considerably. Flash ...
2022
-
[23]
We compare with the baseline results provided in Table 11 in Bolya et al
to validate the performance of A-MoD. We compare with the baseline results provided in Table 11 in Bolya et al. and Table 3 in Yin et al. (2022). However, we note that ToMe (Bolya et al.) trains their models with distillation while the other methods do not, which aids ToMe. Re...
2022
-
[1991]
Mixture of nested experts: Adaptive processing of visual tokens
Gagan Jain, Nidhi Hegde, Aditya Kusupati, Arsha Nagrani, Shyamal Buch, Prateek Jain, Anurag Arnab, and Sujoy Paul. Mixture of nested experts: Adaptive processing of visual tokens. arXiv preprint arXiv:2407.19985,
-
[1993]
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei
doi: 10.1109/IJCNN.1993.716791. Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361,
1993
-
[2013]
Token merging: Your vit but faster
Daniel Bolya, Cheng-Yang Fu, Xiaoliang Dai, Peizhao Zhang, Christoph Feichtenhofer, and Judy Hoffman. Token merging: Your vit but faster. In The Eleventh International Conference on Learning Representations. Weilin Cai, Juyong Jiang, Fan Wang, Jing Tang, Sunghun Kim, and Jiayi...
-
[2014]
Conditional computation in neural networks for faster models
Emmanuel Bengio, Pierre-Luc Bacon, Joelle Pineau, and Doina Precup. Conditional computation in neural networks for faster models. arXiv preprint arXiv:1511.06297,
-
[2015]
Estimating or propagating gradients through stochastic neurons for conditional computation
Yoshua Bengio, Nicholas L ´eonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv:1308.3432,
-
[2016]
Train- ing compute-optimal large language models
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Train- ing compute-optimal large language models. arXiv preprint arXiv:2203.15556,
-
[2017]
Scaling laws for fine-grained mixture of experts
Jan Ludziejewski, Jakub Krajewski, Kamil Adamczewski, Maciej Pi ´oro, Michał Krutul, Szymon Antoniak, Kamil Ciebiera, Krystian Kr ´ol, Tomasz Odrzyg ´o´zd´z, Piotr Sankowski, et al. Scaling laws for fine-grained mixture of experts. InICLR 2024 Workshop on Mathematical and Empi...
2024
-
[2018]
Mixtral of experts
Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bam- ford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. Mixtral of experts. arXiv preprint arXiv:2401.04088,
-
[2021]
William Fedus, Jeff Dean, and Barret Zoph
URL https: //openreview.net/forum?id=YicbFdNTTy. William Fedus, Jeff Dean, and Barret Zoph. A review of sparse expert models in deep learning. arXiv preprint arXiv:2209.01667, 2022a. William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion paramet...
-
[2022]
ST-MoE: Designing stable and transferable sparse expert models
Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du, Yanping Huang, Jeff Dean, Noam Shazeer, and William Fedus. ST-MoE: Designing stable and transferable sparse expert models. arXiv preprint arXiv:2202.08906,
-
[2024]
On the benefits of learning to route in mixture-of-experts models
Nishanth Dikkala, Nikhil Ghosh, Raghu Meka, Rina Panigrahy, Nikhil Vyas, and Xin Wang. On the benefits of learning to route in mixture-of-experts models. InProceedings of the 2023 Conference on Empirical Methods in Natural Language Processing , pp. 9376–9396,
2023
Reviewed August 10, 2026 · model on record in the stance chip above.
Discussion (0). Continue with ORCID to comment.