AI News

Mixture-of-Experts language models usually route each token using its hidden state, which provides only limited contextual information. The paper introduces Attention-Aware Routing (AAR), which augments the router with temporal and spectral features computed from a sliding window of attention weights. The base Transformer is kept completely frozen, so only routing parameters are trained and routing can be studied in isolation. On OLMoE, AAR improves GSM8K performance by 3.37 percentage points over a routing-only supervised fine-tuning baseline. The authors find that routing and attention form a coupled circuit: a routing change at one layer can travel through the residual stream and amplify attention sinks in the next layer, changing attention without directly updating its mechanism. AAR also shortens incorrect, long-diverging generations while leaving correct-answer lengths unchanged. Its effects depend strongly on depth: applying it across all layers can harm factual retrieval, while mathematical-reasoning gains remain when AAR is introduced deeper in the network. The results expose a retrieval-reasoning tension across layers and position layer-selective AAR as a probe of the attention information relevant to routing.