Sticky Routing: Training MoE Models for Memory-Efficient Inference

Jul 16, 2026

Upload by Levi

Mixture-of-Experts (MoE) models activate only a sparse expert subset per token, but consecutive tokens often route to different experts, triggering frequent expert weight swaps between fast memory and slow storage on edge devices. StickyMoE introduces a differentiable routing consistency loss that penalizes abrupt expert switches across adjacent tokens, encouraging stable expert assignments over semantically coherent spans. The method requires no architectural changes, adds a single hyperparameter λ, and enables co-adaptation during pretraining, reducing switch rates up to 59% while improving perplexity and reducing cache misses up to 3.92×.

LanguageEnglish

CategoryResearch & Report

FormatPDF

Pages24

Reading Time1 h 0 min

Introduction

The expert-swapping problem

Why existing solutions are insufficient

Our proposal and training objective

Contributions

arXiv :2607 .08780v1 [ cs .LG] 12 Jun 2026 Sticky Routing: Training MoE Models for Memory-Efficient Inference Ali Kayyam BrainChip Inc., Laguna Hills, CA [akayyam@brainchip. com](akayyam@brainchip. com) Abstract Mixture-of-Experts (MoE) models activate only a sparse subset of experts per token, yet consecutive tokens frequently activate different experts—causing constant weight swapping between slow storage and fast memory on edge devices. Existing remedies are either systemlevel (caching heuristics) or post-hoc (router fine-tuning), leaving the root cause unchanged during pretraining. We propose StickyMoE, a differentiable routing consistency loss that penalises abrupt expert switches between adjacent tokens, encouraging the router to maintain the same expert assignment across semantically coherent spans. StickyMoE requires no architectural changes, adds a single hyperparameter λ, and unlike post-hoc methods, allows expert representations and routing decisions to co-adapt from the first training step. Experiments on small and medium MoE language models show that StickyMoE reduces the expert switch rate by up to 59% while simultaneously improving perplexity on the medium model, and reduces cache misses by up to 3.92 × , Pareto-dominating post-hoc fine-tuning on the quality–locality frontier. Routing temporal locality is most efficiently instilled at training time. Code is available at: [https://github.com/alikayyam/sticky_moe.git](https://github.com/alikayyam/sticky_moe.git). 1 Introduction The Mixture-of-Experts (MoE) architecture has emerged as one of the most practically important techniques for scaling large language models (LLMs) without a proportional increase in per-token computation [Shazeer et al. , 2017 , Fedus et al. , 2022 , Lepikhin et al. , 2021] . By replacing dense feed-forward network (FFN) sublayers with a collection of N expert FFNs and a learned sparse router, MoE models can possess billions of parameters while activating only a small subset—typically 1 or 2 experts out of 8–128—for each input token. This property makes MoE architectures theoretically attractive for edge deployment: if only 2 of 64 experts need to be resident in fast memory at any moment, the effective working-set size is a small fraction of the total parameter count. The expert-swapping problem. In practice, this promise is undermined by the temporal inconsistency of standard MoE routers. Because the router makes an independent routing decision for every token based solely on that token’s hidden representation, consecutive tokensin a sequence routinely activate entirely different experts [Zhou et al. , 2025] . On a system with limited GPU VRAM or on-device SRAM, each such switch requires evicting the currently cached expert weights and loading new ones from CPU RAM or flash storage. PCIe and NVMe bandwidth is one to two orders of magnitude lower than GPU memory bandwidth; a single expert weight transfer can therefore dominate the latency of the forward pass itself [Xue et al. , 2024 , Zhu et al. , 2025] . The result is that MoE models, despite their theoretical sparsity advantage, can be slower than comparably-sized dense models on memory-constrained hardware. Why existing solutions are insufficient. Prior work addresses this problem from two directions. System-level approaches maintain an expert cache in fast memory and use eviction policies (LRU, LFU) or learned prefetchers to reduce cache miss rates [Xue et al. , 2024 , Kamahoriet al. , 2024 , Zhang et al. , 2025] . While practically useful, these methods are fundamentally reactive: they try to exploit whatever locality the model’s router happens to produce, but cannot improve routers that are intrinsically inconsistent. Post-hoc approaches fine-tune only the router of an already-trained model to increase expert reuse [Zhu et al. , 2026], which is more targeted but still treats the symptom rather than the cause: the expert representations have already been shaped b

FAQs

What problem does StickyMoE address in MoE inference?

StickyMoE targets temporal inconsistency in standard MoE routers, where consecutive tokens frequently activate different experts, causing repeated weight swapping between slow and fast memory on memory-limited devices.

How does StickyMoE improve routing behavior during training?

StickyMoE adds a differentiable routing consistency loss that penalizes abrupt changes in the router’s gate distributions between adjacent tokens, promoting consistent expert assignments across coherent spans.

Does StickyMoE require architectural changes or only router training?

StickyMoE requires no architectural modifications and applies to standard top-k MoE routers by adding the proposed training objective; it co-trains with the language modeling loss from the first update.

Jul 16, 2026

Upload by Levi