Attention-Aware Routing: Coupling Routing and Attention in MoEs

Despoina Kosmopoulou; Anastasios Tsetsilas; Efthymios Georgiou; Giannis Karamanolakis; Swastik Roy; Alexandros Potamianos

cs.AISep 17, 2026

Attention-Aware Routing: Coupling Routing and Attention in MoEs

Authors: Despoina Kosmopoulou, Anastasios Tsetsilas, Efthymios Georgiou, Giannis Karamanolakis, Swastik Roy, Alexandros Potamianos

Organizations: National Technical University of Athens · University of Bern · Amazon AGI · Archimedes RU, Athena RC

Abstract

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.

Figures & tables

Figure 1: Attention-Aware Routing. For the current token $i$ , attention weights to the previous $M$ tokens are averaged across heads and assembled into a sliding window $\mathbf{w}$ , then transformed into the augmented window $\tilde{\mathbf{w}}=\phi(\mathbf{w})$ . The attention router $A_{l}$ maps $\tilde{\mathbf{w}}$ to expert logits $\mathbf{g}^{\text{attn}}$ , which are merged and gated with the standard router’s logits $\mathbf{g}^{\text{std}}=W_{r}\mathbf{h}$ to produce the final routing decision over experts $E_{1},\ldots,E_{N}$ .

Configuration	GSM8K	BBH	MMLU	MMLU-Math	HumanEval	IFEval
Configuration	8-shot CoT acc.	0-shot acc.	5-shot acc.	5-shot acc.	pass@10	loose acc.
OLMoE-0125-SFT	51.33	36.80	53.96	31.01	41.46	56.19
+ routing-only SFT	51.99 ${\scriptstyle\pm 0.76}$	36.17 ${\scriptstyle\pm 0.45}$	53.91 ${\scriptstyle\pm 0.01}$	30.45 ${\scriptstyle\pm 0.56}$	41.46 ${\scriptstyle\pm 0.80}$	56.56 ${\scriptstyle\pm 1.09}$
+ AAR S , $M=40$	55.36 ${\scriptstyle\pm 0.74}$	36.71 ${\scriptstyle\pm 0.46}$	53.96 ${\scriptstyle\pm 0.20}$	32.91 ${\scriptstyle\pm 0.54}$	40.00 ${\scriptstyle\pm 1.26}$	55.64 ${\scriptstyle\pm 0.41}$
+ AAR S , $M=20$	54.89 ${\scriptstyle\pm 1.02}$	36.30 ${\scriptstyle\pm 0.42}$	53.93 ${\scriptstyle\pm 0.10}$	33.48 ${\scriptstyle\pm 0.49}$	40.98 ${\scriptstyle\pm 1.80}$	56.67 ${\scriptstyle\pm 1.36}$
+ AAR ${}^{\text{T}}_{S}$ , $M=20$	54.57 ${\scriptstyle\pm 1.10}$	36.33 ${\scriptstyle\pm 0.61}$	53.78 ${\scriptstyle\pm 0.14}$	32.29 ${\scriptstyle\pm 0.64}$	40.85 ${\scriptstyle\pm 1.14}$	56.34 ${\scriptstyle\pm 1.88}$

Table 1: AAR compared against trained baselines. GSM8K consistently improves, with full AAR achieving the best results.

S=\{9,\ldots,15\}

for OLMoE.

S_{Q}=\{14,\ldots,17\}

for Qwen1.5. Best in bold.

Method	MATH-500
OLMoE
Routing-only SFT	14.00 $\pm$ 1.83
AAR ( $L9$ – $15$ )	16.64 $\pm$ 0.91
Qwen3.6-MoE
Routing-only SFT	78.47 $\pm$ 1.30
AAR ( $L23,27,31$ )	79.60 $\pm$ 1.51

Table 2: MATH-500 accuracy for AAR

{}^{\text{TF}}

and routing-only SFT. OLMoE results are averaged over 5 seeds and Qwen3.6-MoE results over 3 seeds.

Figure 2: Normalized autocorrelation function (ACF) of the attention window signal $w^{l}_{i}(t)$ at layers 4, 9, and 12, for example window sizes $\in\{20,64,128\}$ . The ACF is averaged over query tokens and prompts from GSM8K (5-shot). The dashed vertical line marks lag $=20$ . Across all layers and window sizes, the ACF decays below $0.1$ by lag $\sim\!20$ , revealing an intrinsic correlation length of $20$ tokens in the attention signal. A window of $M=20$ is thus sufficient to capture most structured information.

Configuration	$M$	GSM8K
routing-only SFT	–	51.99 ${\scriptstyle\pm 0.76}$
AAR T	10	51.99 ${\scriptstyle\pm 1.13}$
	20	54.10 ${\scriptstyle\pm 0.52}$
	128	53.74 ${\scriptstyle\pm 0.91}$
	512	52.96 ${\scriptstyle\pm 0.76}$

Table 3: Window size ablation for AAR T , applied in all layers. Best in bold.

Figure 3: Top: Absolute weight to the attention sink, per layer (averaged across prompt’s tokens). Bottom: the absolute difference, isolated. The different routing decisions of layer 9, because of AAR, lead to different attention weights produced in the attention submodule, at layer 10.

Condition	Metric	Base	AAR-T	Rel. $\Delta$
Both Correct	Mean	68.2	67.6	$-0.8\%$
	Median	54.0	54.0	$0.0\%$
	P99 (chars)	251.0	248.0	$-1.2\%$
Both Wrong	Mean	163.7	144.4	$-11.8\%$
	Median	78.0	73.0	$-6.4\%$
	P99 (chars)	8513.0	7242.0	$-14.9\%$

Table 4: Length analysis (tokens) on GSM8K averaged across seeds and context sizes

M

. Ramble Rate denotes responses exceeding 1,000 characters. AAR

{}^{\text{T}}

leaves correct responses nearly unchanged while reducing the length of incorrect reasoning, especially in the tail (P99). The answer length reduction also persists for answers that get corrected by AAR

{}^{\text{T}}

Figure 4: AAR ${}^{\text{T}}$ ( $M=20$ ) accuracy across MMLU subjects as a function of starting layer. Most subjects’ accuracy curves follow a bell shape.

MMLU Subject	$\Delta$ $M$ =10	$\Delta$ $M$ =20	$\Delta$ $M$ =64
Math	$+0.7$	$+0.6$	$+0.5$
Comp. Sci.	$+0.5$	$+0.4$	$-0.1$
Biology	$-4.5$	$-3.4$	$-3.8$
History	$-2.5$	$-2.4$	$-2.2$
Health	$-1.4$	$-1.8$	$-1.9$
Psychology	$-1.4$	$-1.7$	$-1.2$

Table 5: MMLU subcategory

\Delta

(absolute % accuracy difference from baseline, 5 seeds) when AAR

{}^{\text{T}}

is applied to all layers of OLMoE. Reasoning-oriented subjects (top) are improved; degradation is concentrated in retrieval-heavy subjects (bottom).

LLM Routing

Large language model routing selects models, tools, experts, or inference strategies to balance answer quality, cost, latency, and deployment constraints.

See all 342 papers

When Are Experts Misrouted? Counterfactual Routing Analysis in Mixture-of-Experts Language Models
May 8, 2026Youngsik Yoon, Siwei Wang, Wei Chen +1Mixture-of-Experts Language Models LLM Routing
Department of Computer Science and Engineering, POSTECH, South Korea · Microsoft Research Asia, Beijing, China · Graduate School of Artificial Intelligence, POSTECH, South Korea
Self-Routing: Parameter-Free Expert Routing from Hidden States
Apr 1, 2026Jama Hussein Mohamud, Drew Wagner, Mirco RavanelliParameter-Free Mixture-of-Experts Routing LLM Routing
Mila – Quebec AI Institute · Universite de Montreal · Concordia University
Path-Constrained Mixture-of-Experts
Mar 18, 2026Zijin Gu, Tatiana Likhomanenko, Vimal Thilak +2Mixture-of-Experts Language Models Sparse Mixture-of-Experts
Apple · †Google

May 8, 2026J/K move · Enter open · S save

When Are Experts Misrouted? Counterfactual Routing Analysis in Mixture-of-Experts Language Models

Youngsik Yoon, Siwei Wang, Wei Chen, Jungseul Ok

Department of Computer Science and Engineering, POSTECH, South Korea · Microsoft Research Asia, Beijing, China · Graduate School of Artificial Intelligence, POSTECH, South Korea

Mixture-of-Experts Language Models LLM Routing Mixture-of-Experts Models

Open paper AskPDF arXiv

Mixture-of-Experts (MoE) language models route each token to a small subset of experts, but whether the routes selected by a trained top-

k

router are good ones is rarely evaluated directly. Holding the model fixed, we compare each standard route against sampled equal-compute alternatives for the same token and score each by the next-token probability it assigns to the realized token in a verified reasoning trajectory. The result is sharply token-conditional: the standard router is well-aligned with route utility on confident tokens but uninformative on the fragile tokens that drive hard reasoning, where lower-loss equal-compute routes consistently exist inside the frozen model but are not selected. The same pattern holds across Qwen3-30B-A3B, GPT-OSS-20B, DeepSeek-V2-Lite, and OLMoE-1B-7B, and follows structurally from how standard top-

k

training evaluates routing decisions: the language modeling loss scores only the executed route, and load balancing depends only on aggregate routing statistics. A minimal router-only update to the final-layer router, leaving every expert and every other router frozen, is sufficient to shift pass@K on AIME 2024+2025 and HMMT 2025 for both Qwen3-30B-A3B and GPT-OSS-20B, suggesting that at least part of the failure reflects router-reachable misallocation rather than expert capacity alone.

Abstract

LLM Routing

Explore similar work

When Are Experts Misrouted? Counterfactual Routing Analysis in Mixture-of-Experts Language Models

Self-Routing: Parameter-Free Expert Routing from Hidden States

Path-Constrained Mixture-of-Experts