Attention-Aware Routing: a MoE router improves when it can see attention

Changing only the router parameters of a Mixture-of-Experts model, without touching a single weight in the rest of the transformer, is enough to gain 3.37 percentage points on GSM8K. That is the central result of Attention-Aware Routing, a paper that appeared today on arXiv and rests on a simple observation: a MoE router decides with very little information about the context.

In a conventional MoE, each token goes through a router that chooses which experts to send it to and with what weight. That decision is made from the token's hidden state, a vector summarising what the model knows about that token at that layer, but one that barely encodes explicitly how attention is being distributed over the recent window. The authors start there: if attention is the mechanism by which the model spreads its focus over the context, the router should be able to consult it.

What AAR proposes

Attention-Aware Routing, AAR, adds to the router a set of features extracted from a sliding window of attention weights. They come in two kinds:

Temporal: how attention evolves over the last positions.
Spectral: the frequency structure of those weights, which summarises repetitive or scattered patterns.

That summary of the model's contextual state is kept separate from the hidden state, so the router receives two distinct signals instead of one. The experimental design is deliberately narrow: the base transformer stays fully frozen and only the routing parameters are trained. Any improvement is therefore attributed solely to the router, with no noise from general fine-tuning.

With that setup, AAR improves GSM8K by 3.37 percentage points over a baseline that also does SFT only on the router, on OLMoE, the open MoE from the Allen Institute for AI. The comparison is fair because both train exactly the same parameters. The only thing that changes is what information the router sees.

Routing and attention form a circuit

The most interesting part of the work is not the number but the mechanistic analysis. The authors show that a routing change at layer l propagates through the residual stream and amplifies the attention sinks at layer l+1. Attention sinks are those positions, typically the first token, that absorb a disproportionate share of attention and act as a release valve for the mechanism.

In other words, modifying only the router reshapes the next layer's attention without anyone touching the attention mechanism. Routing and attention are not independent modules that add up, but a coupled circuit in which a local decision about experts alters how the model looks at its context. For anyone working on interpretability, this complicates the idea of studying each component in isolation.

Shorter wrong answers

There is a useful side effect. AAR reduces what the paper calls long diverging generation: incorrect answers get shorter, while correct ones keep their length. In mathematical reasoning, a model that goes wrong tends to ramble through chains that lead nowhere. Having the router help cut them off earlier reduces generated tokens and, in a real deployment, cost.

Not valid at every layer

The last finding is a warning. AAR is strongly depth-sensitive: applying it indiscriminately across all layers can degrade factual retrieval, while mathematical reasoning responds differently, according to the abstract itself. The per-layer detail is in the body of the paper, but the message is clear: choosing which layers to apply it to is part of the method, not an optional tweak.

Who it is relevant for

This work changes nothing for those consuming models via API. It does matter to:

Teams training or fine-tuning open MoEs such as OLMoE, Mixtral or the Qwen3 MoE variants, because it proposes a parameter-cheap improvement that can be tried without retraining the full model.
Interpretability researchers, for the evidence that routing and attention sinks are connected through the residual stream.
* Anyone optimising inference, for the reduction in length of wrong outputs, which translates into less wasted compute.

Open questions remain: the paper is evaluated on OLMoE, a relatively small model, and we do not know whether the effect scales to MoEs with dozens of experts and hundreds of billions of parameters. The cost of computing spectral features at every inference step is also unclear.

At ElephantPink we find the coupled-circuit finding more valuable than the three GSM8K points, because it explains something that until now was treated as two separate problems. If it holds up in larger models, it will be one more reason to stop thinking of the router as a simple selector.