Mixture-of-experts models gain transparent routing memory state
A Persistent State for Auditable Mixture-of-Experts Routing
Artificial Intelligence
Summary
Mixture-of-Experts models decide which parts of a neural network to use for each input, but they usually don't keep a clear record of their routing choices across layers. The authors introduce a method called Scratchpad-Augmented Mixture-of-Experts (SA-MoE), which gives routers a small memory state to keep track of routing decisions over time. This memory helps routers make better choices by remembering past influences and can be inspected directly during the model’s operation. Their experiments show this change is useful and doesn't add much extra computation.
What this means in practice
- •For machine learning engineers: Build interpretable MoE models that record routing history for debugging and understanding token routing decisions across layers.
- •For natural language processing developers: Improve token routing consistency and long-term dependency handling in sparse language models by adding persistent routing memory.
Authors
Abdurrahman Javat, Allan Kazakov
Abstract
Mixture-of-Experts (MoE) models repeatedly route tokens to sparse subsets of experts, but conventional routers expose no routing-specific record of how cross-layer influences accumulate. We introduce Scratchpad-Augmented Mixture-of-Experts (SA-MoE), which gives each router access to a low-dimensional persistent state that is not provided to the experts. Learned layerwise writes update this state, and their realized post-update changes exactly decompose the state-mediated contribution to any later routing margin, forming a routing ledger. Across sparsely upcycled SmolLM2- and Gemma-based models and three independent training seeds per architecture, this pathway adds less than 1% analytical forward compute and is strongly used by trained routers: local removal of its router contribution changes the selected Top-2 expert set in 87.6% and 69.9% of decisions, respectively. Relative to a matched latest-write-only control, persistent accumulation increases long-horizon future-routing accessibility by 19.4 and 12.2 percentage points, with positive effects in every seed. More than 90% of absolute ledger contribution comes from non-recent writes in both families, and full-forward suppression of ledger-selected writes changes later routing and output distributions. The ledger is an exact provenance object for the persistent-state pathway, not a complete causal explanation of routing. Sensitivity-aware scores better predict full-forward intervention effects, and post-hoc methods recover related cross-layer attribution without architectural modification. SA-MoE instead makes one routing-specific computational history explicit and directly inspectable within the model's natural forward computation.