Technical Summary: Sticky Routing for Memory-Efficient MoE Inference

1. Problem Statement: The Expert-Swapping Bottleneck

Mixture-of-Experts (MoE) models are designed to scale Large Language Models (LLMs) by activating only a sparse subset of experts per token. Theoretically, this allows for massive parameter counts with low per-token computation, making them attractive for edge deployment where only a few experts need to reside in fast memory (VRAM/SRAM) at any given time.

However, the paper identifies a critical failure mode in practice: temporal routing inconsistency. Standard MoE routers make independent decisions for each token based solely on that token's hidden representation. Consequently, consecutive tokens frequently activate entirely different experts. On memory-constrained hardware, each switch necessitates evicting cached expert weights and loading new ones from slow storage (CPU RAM/Flash). Since PCIe/NVMe bandwidth is orders of magnitude lower than GPU memory bandwidth, these frequent "expert swaps" dominate inference latency, often rendering MoE models slower than comparably-sized dense models.

Existing solutions are deemed insufficient:

• System-level approaches (e.g., caching heuristics, prefetchers) are reactive; they exploit whatever locality exists but cannot fix intrinsically inconsistent routers.
• Post-hoc approaches (e.g., fine-tuning only the router) fail because the expert representations were shaped by a training regime that never incentivized locality. The router cannot easily re-align decisions in a representation space not designed for it.
• Architectural redesigns (e.g., Oracle-MoE) often require significant structural changes to the model or training pipeline.

2. Methodology: StickyMoE

The authors propose StickyMoE, a training-time approach that treats routing temporal locality as an explicit objective. The core innovation is a differentiable routing consistency loss ($L_{cons}$) that penalizes abrupt changes in expert assignment between adjacent tokens.

The Routing Consistency Loss

The loss function minimizes the $\ell_2$ distance between the softmax gate probability vectors ($g_t$) of consecutive tokens:
$$L_{cons} = \frac{1}{T-1} \sum_{t=2}^{T} \| g_t - g_{t-1} \|_2^2$$
Where $g_t \in \mathbb{R}^N$ is the router's output for token $t$.

This loss is integrated into the total training objective alongside standard cross-entropy ($L_{CE}$) and load-balancing ($L_{bal}$) losses:
$$L = L_{CE} + \lambda L_{cons} + \mu L_{bal}$$

Key Mechanisms

• Co-adaptation: Unlike post-hoc methods, StickyMoE trains the router and the expert representations simultaneously. Minimizing $L_{cons}$ updates the router weights ($W_r$) and backpropagates through the hidden states ($h_t$), encouraging the model to learn representations that are naturally consistent across semantically coherent spans.
• Soft-Hard Variant: To address "long-range drift" (where small step-by-step changes accumulate into a distant switch), the authors introduce an anchor constraint. The sequence is divided into windows; the first token of a window sets an anchor, and subsequent tokens are penalized for deviating from this anchor, with the penalty weight increasing linearly with distance from the window start.
• Compatibility: The method requires no architectural changes and adds negligible computational overhead. It is compatible with any standard top-$k$ MoE and existing load-balancing techniques.

3. Key Contributions

1. Routing Consistency Loss: A simple, differentiable training objective that directly penalizes expert switching between adjacent tokens, requiring only a single hyperparameter ($\lambda$).
2. Complementarity with Load Balancing: Empirical evidence shows that consistency loss and load-balancing loss are complementary. The consistency loss shapes when experts are used without causing "expert collapse" (utilization entropy remains high, near the theoretical maximum).
3. Soft-Hard Variant: A combination of per-step soft penalties and segment-level anchor constraints that guarantees strong locality without additional quality costs.
4. Comprehensive Evaluation: Controlled experiments comparing StickyMoE against vanilla MoE, post-hoc router fine-tuning (ReMoE), and architectural baselines (Oracle-MoE) on small and medium models.

4. Experimental Results

Experiments were conducted on WikiText-2 using small (8.8M params) and medium (22M params) MoE models.

• Expert Switch Rate (SR): StickyMoE reduces the expert switch rate by up to 59% (e.g., from 0.71 to 0.29 on the medium model).
• Cache Hit Rate (CHR): The reduction in switching translates to a 3.92× reduction in cache misses (CHR increased from ~0.54 to ~0.88) on the medium model with a 2-expert slot cache.
• Perplexity (Quality):
  • At moderate $\lambda$ values (e.g., 0.05), StickyMoE improves perplexity compared to the baseline (e.g., -4.1% on the medium model), suggesting the consistency loss acts as a beneficial regularizer.
  • Aggressive $\lambda$ values introduce a trade-off, but the Pareto frontier of StickyMoE strictly dominates other methods.
• Failure of Post-Hoc Fine-Tuning: The ReMoE baseline (fine-tuning only the router on a pre-trained model) failed to reduce the switch rate meaningfully (<0.5% change), confirming that locality cannot be retrofitted once representations are fixed.
• Layer-wise Analysis: The consistency loss is most effective in deeper layers (L1–L3). The input layer (L0) remains resistant to locality improvements, likely because token embeddings lack contextual mixing.

5. Significance and Claims

The paper argues that routing temporal locality is a legitimate training objective that should be instilled during pretraining rather than treated as a system-level afterthought.

• Paradigm Shift: Just as quantization-aware training has become standard, the authors posit that "locality-aware training" should be the default for MoE models intended for memory-constrained deployment.
• Root Cause Resolution: StickyMoE addresses the root cause (inconsistent routing decisions) rather than the symptoms (cache misses). By co-training representations and routing decisions, it avoids the representation mismatch that plagues post-hoc methods.
• Practicality: The method is architecture-agnostic and requires no changes to the inference pipeline. The trained router naturally produces consistent gate distributions, allowing standard caching strategies (LRU, LFU) to function significantly more efficiently without additional engineering.

The authors conclude that while system-level optimizations and post-hoc adjustments remain useful complements, they cannot substitute for a model that was never trained to be cache-friendly in the first place. StickyMoE provides a path to realizing the theoretical memory efficiency of MoE architectures on edge devices.