{"id":"945742ec-11bc-4072-9466-9dc0ea419163","arxiv_id":"2412.20875","paper_version":1,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":6.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":3,"one_line_summary":"A-MoD routes tokens in Mixture-of-Depths vision transformers using attention maps from the prior layer, improving accuracy and convergence without extra trainable parameters.","lead":"This paper proposes A-MoD, a routing method for Mixture-of-Depths vision transformers that scores tokens by averaging the previous layer's attention maps instead of learning a separate router. This removes router parameters and speeds up adaptation of pretrained models, with reported accuracy gains of up to about 2% on ImageNet.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"A-MoD vs. standard MoD changes both the routing source and the output scaling rule (Eq. 2 vs. Eq. 5); the reported gains may come from dropping the r_i multiplication, not from attention-based routing.","rationale":"The reader's weakest assumption is that attention maps track token importance. I do not think that is the decisive weakness: the paper gives leave-one-out correlations (Fig. 7) and routing visualizations suggesting the attention signal is meaningful, and the zero-shot adaptation numbers in Table 3 are consistent with that. The decisive weakness is that the headline comparison is confounded. The authors clearly state in Section 3.3 that A-MoD does not multiply the output by routing scores, while standard routing does in Eq. (2); they even note that the multiplication is 'required' for standard routing. Because the ablation in Fig. 11 demonstrates that multiplication hurts A-MoD, the performance gap in Table 1 could be entirely explained by this output-scaling difference. The reader's rationale lists the scaling change as a limitation but keeps it secondary to the attention-proxy assumption; I would promote it to the primary concern. A matched-scaling control is cheap and decisive, so the conditional verdict stands, with the condition made explicit.","tokens_in":14876,"tokens_out":8934,"duration_ms":86606,"concrete_test":"Run Table 1's ImageNet finetuning with A-MoD modified to use Eq. (2) instead of Eq. (5): multiply the MoD block output by the attention-derived r_i for selected tokens, exactly as standard routing does. At minimum cover DeiT-Tiny and ViT-Base at C=50% with the same optimizer and schedule. If this multiplied A-MoD still matches or beats standard MoD, the attention-routing claim survives; if it drops toward or below standard MoD (as Fig. 11 suggests), the reported advantage is attributable to removing output scaling, not to attention-based routing.","verdict_should_be":"CONDITIONAL","load_bearing_attack":"The weakest point is the treatment contrast, not the attention-importance proxy. In Eq. (2), standard routing scales the processed token's block output by the learned router scalar: x_i^l = r_i f_l(X^{l-1})_i + x_i^{l-1}. In Eq. (5), A-MoD does not scale the output by its attention-derived score: x_i^l = f_l(X^{l-1})_i + x_i^{l-1}. Thus Table 1 compares two systems that differ in routing source and output scaling simultaneously. The paper's own ablation (Fig. 11) shows that adding the r_i multiplication to A-MoD slows convergence and reduces accuracy, so the absence of scaling is independently beneficial. No experiment holds output scaling fixed and varies only the routing source: a standard router without multiplication receives no gradient, and the scaled A-MoD in Fig. 11 is never compared against standard MoD under the same scaled-output rule. The from-scratch results in Table 8 share the confound. Therefore the central attribution 'attention maps are better routers than learned routers' is not yet established; the advantage could be an artifact of a different output rule.","agreement_with_reader":"partial"},"referee_report":{"model":"deepseek-v4-flash","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.","tokens_in":15109,"tokens_out":6132,"duration_ms":58406,"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":[{"comment":"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.","section":"§3.2–3.3, Eqs. (2) and (5), Fig. 11"},{"comment":"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.","section":"§4.1–4.2, Tables 1, 4, and 5"},{"comment":"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.","section":"§4.4"},{"comment":"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.","section":"§3.3, Fig. 15, related work"}],"minor_comments":[{"comment":"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.","section":"§1 and §4.3"},{"comment":"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.","section":"Table 3"},{"comment":"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).","section":"§A.9, Table 6"},{"comment":"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.","section":"§A.8, Algorithm 1"},{"comment":"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.","section":"§4.5"}],"recommendation":"major_revision","confidential_remarks":"The core idea is attractive and the experimental scope is broad, but the missing controlled comparison between routing source and output scaling is the main blocker. If the authors add a condition that varies only the routing source, report multiple seeds, and fix the correlation mismatch, the paper could become a solid contribution. No code is provided; for a parameter-free routing method, releasing code would substantially aid reproducibility and would help reviewers verify the Flash-Attention variant."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"Thanks for the report. Here's my read.\n\nThe idea is simple and well-presented: use the averaged column attention from the previous layer as the MoD routing score, instead of a learned linear router. It's parameter-free by construction, and the experiments span four architectures, two capacity levels, finetuning and transfer, plus a Flash Attention variant. The paper is honest about the main assumption (attention maps are a proxy for token importance) and even cites Darcet et al. showing cases where that fails. The leave-one-out correlation analysis is a good attempt to show the routing decisions are sensible. I'd give the authors credit for a clear, thorough empirical study.\n\nBut the central comparison is confounded. Standard MoD (Eq. 2) multiplies the block output by the router scalar r_i, which is needed to give the learned router a gradient. A-MoD (Eq. 5) drops that multiplication entirely. So Table 1 compares two systems that differ in both the routing source and the output rule. The paper's own ablation (Fig. 11) shows that adding the multiplication back to A-MoD slows convergence and reduces accuracy. That means the observed gains could come from removing the scaling, not from attention-based routing. No experiment isolates the routing source by holding the output rule fixed. As the stress-test note says, the claim 'attention maps are better routers' is not yet established.\n\nThere's also no code or data and no repeated-seed statistics. For a paper reporting 1–2% accuracy differences, single runs are not enough to be confident.\n\nThe token-pruning comparison is relegated to the appendix, and A-ViT/DynamicViT/ToMe are cited only there. That's a smaller issue, but the core idea of attention-based token importance has prior art.\n\nThis is a paper worth engaging with, not a desk reject. A careful referee could ask the authors to decouple the two variables: e.g., compare standard routing without output scaling (using a stop-gradient or a straight-through estimator) or compare scaled A-MoD against scaled standard routing. They also need error bars.\n\nMy recommendation: accept for peer review, but with the expectation of heavy revision. The idea has legs, but the headline result needs a fair apples-to-apples comparison.","headline":"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.","tokens_in":15663,"tokens_out":2767,"would_cite":false,"duration_ms":27339,"reading_group":"maybe","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"deepseek-v4-flash","headline":"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.","keywords":["Mixture-of-Depths","attention routing","token importance","vision transformers","conditional computation","parameter-free router","token pruning","ImageNet"],"falsifier":"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.","tokens_in":14695,"feed_emoji":"🎯","tokens_out":4981,"duration_ms":44659,"temperature":0.7,"pith_summary":"The paper claims that the trainable router in Mixture-of-Depths (MoD) vision transformers is unnecessary: token importance for an MoD layer can be read directly from the previous layer's attention maps, averaged over attention heads and source tokens. This attention-based router, called A-MoD, adds zero trainable parameters and no extra router layers. Across DeiT-Tiny, DeiT-Small, ViT-Base, and ViT-Large at 50% and 12.5% capacity, the authors report that A-MoD mostly outperforms standard learned routing on ImageNet, with gains up to 2% accuracy, and reaches peak accuracy faster during finetuning and transfer learning. If true, this means an MoD model can be adapted from an off-the-shelf pretrained transformer checkpoint without training a router, lowering the overhead of dynamic computation.","feed_headline":"Attention maps replace learned routers in Mixture-of-Depths","feed_subtitle":"A-MoD skips the trainable router, beating standard MoD by up to 2% on ImageNet with faster convergence.","key_machinery":"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.","core_discovery":"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.","pith_inferences":["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."],"forward_implications":["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."],"supporting_citations":[{"why":"Supplies the attention-map definition (Eq. 3) that A-MoD aggregates into routing scores.","marker":"Vaswani et al. (2017)"},{"why":"Introduces Mixture-of-Depths and the standard learned router (Eq. 1-2) that A-MoD is compared against and replaces.","marker":"Raposo et al. (2024)"},{"why":"Provides the Vision Transformer architectures (ViT-Base, ViT-Large) and pretrained checkpoints used for adaptation.","marker":"Dosovitskiy et al. (2021)"},{"why":"Provides the DeiT-Tiny and DeiT-Small models used as small-scale testbeds.","marker":"Touvron et al. (2021)"},{"why":"Cited by the paper as evidence that attention maps do not always carry semantically meaningful scores, marking the known failure mode of the core assumption.","marker":"Darcet et al. (2024)"},{"why":"Supplies the leave-one-out procedure used to define token importance and validate that A-MoD scores track importance.","marker":"Hastie et al. (2009)"},{"why":"Flash Attention integration: the paper's Algorithm 1 aggregates routing weights without materializing the NxN attention map.","marker":"Dao et al. (2022)"},{"why":"DETR experiments showing A-MoD and standard MoD perform comparably on object detection.","marker":"Carion et al. (2020b)"}],"fun_headline_variants":["Attention is all you need for MoD routing","No extra parameters: attention maps decide token skips","A-MoD: Use attention to route MoD without learned routers","Attention routing improves on learned routers in MoD","Parameter-free MoD routing from attention scores"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"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.","fun_headline_variants_meta":{"raw":{"variants":["Attention is all you need for MoD routing","No extra parameters: attention maps decide token skips","A-MoD: Use attention to route MoD without learned routers","Attention routing improves on learned routers in MoD","Parameter-free MoD routing from attention scores"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000673,"raw_usage":{"total_tokens":3070,"prompt_tokens":955,"completion_tokens":2115,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":571,"completion_tokens_details":{"reasoning_tokens":2039}},"tokens_in":571,"tokens_out":2115,"duration_ms":17413,"temperature":1.0,"reasoning_tokens":2039,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-10T23:08:09.477363+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"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.","supporting_citations":[{"cited_title":"(2022)) to column first in order to aggregate row wise scores efficiently","cited_arxiv_id":null,"evidence_quote":"Flash Attention integration: the paper's Algorithm 1 aggregates routing weights without materializing the NxN attention map."}],"review_version":1}