Attention aware routing improves reasoning by linking routing to model context
Attention-Aware Routing: Coupling Routing and Attention in MoEs
Artificial Intelligence
Summary
Large language models use expert groups to process information, but usually decide which expert to use based only on limited context. The authors suggest a method called Attention-Aware Routing that looks at a bigger summary of what the model is focusing on, without changing the main model itself. This helps the model solve math problems better and changes how it pays attention layer by layer inside. It also stops wrong answers from getting too long and shows the importance of picking where in the model to use this technique.
What this means in practice
- •For machine learning engineers: Improve accuracy in specialized transformer models by integrating attention-based routing techniques without retraining base model weights.
- •For natural language processing developers: Reduce errors and improve reasoning length control in language model outputs by selectively applying advanced routing in deeper network layers.
Authors
Despoina Kosmopoulou, Anastasios Tsetsilas, Efthymios Georgiou, Giannis Karamanolakis, Swastik Roy, Alexandros Potamianos
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.