Mixture-of-Experts (MoE) architectures have become a standard tool for scaling language models without proportionally increasing per-token compute. The appeal for edge deployment is clear: if a model activates only 2 of 64 experts per token, the effective working set in fast memory should be a small fraction of total parameters. In practice, this promise breaks down. Standard routers make independent decisions for every token, and consecutive tokens routinely activate entirely different experts. On memory-constrained hardware, each such switch triggers a weight eviction and reload from CPU RAM or flash storage, where bandwidth is one to two orders of magnitude lower than GPU memory bandwidth. The theoretical sparsity advantage evaporates.
This paper, arXiv:2607.08780, proposes StickyMoE, a training-time fix to this problem. The core argument is that existing remedies, whether system-level caching heuristics or post-hoc router fine-tuning, treat the symptom rather than the cause. If the router was never trained to be cache-friendly, retrofitting locality after the fact is structurally limited. The authors make a clean case that routing temporal locality must be instilled during pretraining, when expert representations and routing decisions can co-adapt.
Key Contributions
The central contribution is a differentiable routing consistency loss that penalises abrupt expert switches between adjacent tokens. This loss, weighted by a single hyperparameter lambda, is added to the standard language modelling objective during pretraining. No architectural changes are required. The router is still a linear projection over the hidden state; the only change is that the training signal now includes a penalty for switching experts between token positions t and t+1.
The paper evaluates three loss variants:
- Soft: a continuous KL-based penalty on the routing distribution between adjacent tokens, differentiable end-to-end.
- Hard: a discrete consistency constraint over a sliding window of W tokens, applied as a regulariser.
- Soft-Hard: a combination of both, with separate weighting hyperparameters.
The authors also benchmark against ReMoE, a post-hoc router fine-tuning baseline, and an Oracle-MoE that represents an upper bound on locality achievable by forced expert assignment. This framing is useful: it lets the reader see where StickyMoE sits on the quality-locality frontier relative to both practical alternatives and a theoretical ceiling.
Methodology and Experimental Setup
Experiments are conducted on two GPT-style causal transformer models, a small (~8.8M parameter) and a medium (~22M parameter) variant, both trained on WikiText-2. Every FFN sublayer is replaced with a top-2 MoE layer over 4 experts. Training runs for 10,000 steps on a single NVIDIA GTX 1080 Ti, which is a deliberate choice: the authors want the pipeline to be reproducible without large-scale compute. Absolute perplexity values are correspondingly modest, and the paper is explicit that relative differences between variants are what matter here.
The primary evaluation metrics are perplexity (language modelling quality), expert switch rate (fraction of consecutive token pairs that activate different experts), and LRU cache hit rate with a fixed cache of 2 expert slots. The cache hit rate is the most practically relevant metric for edge deployment, directly measuring how often a required expert is already resident in fast memory.
Results
The headline numbers are striking. On the medium model, Soft StickyMoE with lambda=0.05 achieves a 4.1% perplexity improvement over the baseline while reducing the switch rate from 0.713 to 0.538. At lambda=0.20, the switch rate drops to 0.368 with a 2.1% perplexity improvement, and cache hit rate rises from 0.538 to 0.832, a 3.92x reduction in cache misses. The model is simultaneously better at language modelling and dramatically more cache-friendly. This Pareto dominance over post-hoc fine-tuning is the paper's strongest empirical result.
Several observations are worth highlighting:
- ReMoE (post-hoc fine-tuning) achieves almost no reduction in switch rate at either scale, confirming the authors' central thesis that locality cannot be retrofitted.
- The Hard variants hurt perplexity significantly on the medium model (up to +6.8% at W=2), suggesting that discrete consistency constraints are too blunt an instrument at this scale.
- The Soft-Hard combination offers modest additional locality over Soft alone but at a perplexity cost, making it a less attractive operating point.
- On the small model, the perplexity-locality trade-off is less favourable: Soft at lambda=0.01 is the only configuration that improves perplexity, and larger lambda values degrade quality noticeably. The medium model appears to have enough capacity to absorb the consistency constraint more gracefully.
The implication is that the benefit of locality-aware training may scale with model capacity, which motivates the authors' call for experiments at 1B+ parameters.
Limitations and Open Questions
The paper is candid about its limitations. The consistency loss penalises all expert switches uniformly, including those that occur at natural semantic boundaries: sentence breaks, topic shifts, paragraph transitions. At these points, switching experts is semantically appropriate, and penalising it unnecessarily constrains the model. A boundary-aware variant that gates the penalty using punctuation or BOS signals is an obvious next step, and the authors acknowledge it would likely improve the quality-locality trade-off at higher lambda values where the current formulation begins to hurt perplexity.
A second limitation, discussed in the appendix, concerns the deployment mode that matters most on edge devices: single-token autoregressive decode. The token-level consistency loss is most directly beneficial for prefill and batched decode. For single-token decode, cross-layer expert consistency under parameter-sharing architectures would be more relevant. The authors sketch a cross-layer consistency loss but leave empirical validation to future work.
The experimental scale is also a real constraint. WikiText-2 with 10,000 training steps on a 22M parameter model is far from the regime where MoE models are actually deployed. The results are internally consistent and the methodology is sound, but extrapolating to production-scale MoE models (Mixtral, DeepSeek-MoE, etc.) requires caution. The authors are aware of this and frame it as a motivation for future scaling experiments rather than a gap they can fill here.
The broader point the paper makes is worth taking seriously. The analogy to quantisation-aware training is apt: once it became clear that training with deployment constraints in the loop consistently outperforms post-hoc compression, QAT became standard practice. If the scaling results hold, there is a reasonable case that locality-aware training should be the default starting point for MoE models intended for memory-constrained deployment. The post-hoc stack of caching heuristics and offloading schedulers remains useful as a complement, but not as a substitute for a model that was never asked to be cache-friendly in the first place.
Code is available at github.com/alikayyam/sticky_moe. The full paper is on arXiv.