Attention-Aware Routing: Coupling Routing and Attention in MoEs

Attention-Aware Routing proposes augmenting the Mixture-of-Experts (MoE) router by incorporating temporal and spectral features derived from a sliding window of attention weights, which summarize the model's contextual state, independent of the token's hidden state. The method keeps the base transformer frozen, training only the routing parameters to isolate the variable under study. Experimental results show that Attention-Aware Routing improves GSM8K performance by 3.37 percentage points over a routing-only SFT baseline on OLMoE. Furthermore, routing and attention are demonstrated to form a coupled circuit where routing decisions at one layer propagate through the residual stream to amplify attention sinks in the subsequent layer, effectively reshaping attention without directly modifying the attention mechanism itself. The approach also mitigates long diverging generation by shortening incorrect answers while preserving correct answer lengths. Depth-sensitivity is observed: indiscriminate application of AAR degrades factual retrieval, whereas deeper application preserves mathematical reasoning.

Attention-Aware Routing augments the MoE router with temporal and spectral features from a sliding window of attention weights to summarize contextual state. The base transformer remains frozen during training; only routing parameters are updated. Performance improvement on GSM8K was observed, with a 3.37 pp gain over a routing-only SFT baseline on OLMoE. Routing changes propagate through the residual stream to amplify attention sinks in the subsequent layer. Applying AAR indiscriminately across layers degrades factual retrieval but preserves mathematical reasoning when applied deeper.

The research suggests that optimizing routing in MoE architectures is not an isolated decision based solely on token hidden states, but rather requires integrating broader contextual awareness captured by attention patterns across time and space. The finding that routing and attention form a coupled circuit implies a deep structural relationship where expert selection influences the flow of contextual information through the network layers, necessitating joint optimization. The depth-sensitivity points toward a critical tension between local routing decisions and global context propagation; applying changes indiscriminately creates brittleness in factual recall but allows for nuanced reasoning when localized deeper modifications are made. This suggests that model efficiency and factual fidelity are governed by how temporal attention patterns modulate expert allocation throughout the transformer stack, offering a controlled probe into where routing-relevant information resides within attention mechanisms. The implication is that understanding routing requires mapping the dynamic interplay between selection (routing) and contextual representation (attention) across hierarchical levels to achieve robust generalization.