Attention-Aware Routing: Coupling Routing and Attention in MoEs
On the map: Mixture-of-experts → Expert routing strategies
Abstract (arXiv)
In Mixture-of-Experts language models, the router typically selects and weights experts based on the token's hidden state, utilizing limited contextual information. We propose Attention-Aware Routing (AAR), which augments the router with temporal and spectral features extracted from a sliding window of attention weights that represent a summary of the model's contextual state, disentangled from the hidden state. Keeping the base transformer entirely frozen, we train only the routing parameters, isolating routing as the sole variable. AAR improves GSM8K by +3.37 pp over a routing-only SFT baseline on OLMoE. Beyond performance, we show that routing and attention form a coupled circuit: routing changes at layer l propagate through the residual stream to amplify attention sinks at layer l+1, reshaping attention without any direct update to the attention mechanism itself. Further, AAR reduces long diverging generation, with incorrect answers getting shorter, while correct answers remain unchanged in length. Finally, AAR is strongly depth-sensitive: applying it indiscriminately across layers can degrade factual retrieval, whereas mathematical reasoning gains persist when it is introduced deeper in the network. This sensitivity exposes a retrieval--reasoning tension across depth and makes layer-selective AAR a controlled probe of the routing-relevant information carried by attention at different layers.
Key facts. No citing papers in the map yet (data through 14 Oct 2026). Placed in Expert routing strategies, one of 12 directions in Mixture-of-experts. Builds on 18 of its 24 references that are in the map. First posted on arXiv 17 Sep 2026.
Builds on
18 of 24 references.
Closest ancestors · 4
Same direction, newest first
- 2024Layerwise Recurrent Router for Mixture-of-Experts13 Aug
- 2023HyperRouter: Towards Efficient Training and Inference of Sparse Mixture of Experts12 Dec
- 2022Mixture-of-Experts with Expert Choice Routing18 Feb
- 2021Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity11 Jan
Other references · 14
- 2021Evaluating Large Language Models Trained on Codecited by 8,320
- 2020Measuring Massive Multitask Language Understandingcited by 7,963
- 2023Let's Verify Step by Stepcited by 3,789
- 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,673
- 2017Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layercited by 3,509
- 2023Efficient Streaming Language Models with Attention Sinkscited by 2,186
- 2022Challenging BIG-Bench Tasks and Whether Chain-of-Thought Can Solve Themcited by 1,807
- 2020GShard: Scaling Giant Models with Conditional Computation and Automatic Shardingcited by 1,631
- 2025Why do LLMs attend to the first token?cited by 127
- 2025Between Underthinking and Overthinking: An Empirical Study of Reasoning Length and correctness in LLMscited by 117
- 2025Demystifying the Roles of LLM Layers in Retrieval, Knowledge, and Reasoningcited by 15
- 2025Decoupling Knowledge and Reasoning in LLMs: An Exploration Using Cognitive Dual-System Theorycited by 5
Built on by
Not cited yet.
Similar, not cited
Close in wording, no citation either way
- 2026TriRoute: Unified Learned Routing for Joint Adaptive Attention, Experts, and KV-Cache Allocation
- 2026Multi-Head Attention as a Source of Catastrophic Forgetting in MoE Transformers
- 2026Attention Sink Forges Native MoE in Attention Layers: Sink-Aware Training to Address Head Collapse
- 2022Mixture of Attention Heads: Selecting Attention Heads Per Token
- 2025Input Domain Aware MoE: Decoupling Routing Decisions from Task Optimization in Mixture of Experts
- 2026Mixture of Layers with Hybrid Attention
- 2025Probing Semantic Routing in Large Mixture-of-Expert Models
- 2025Mixture of Sparse Attention: Content-Based Learnable Sparse Attention via Expert-Choice Routing