Sticky Routing: Training MoE Models for Memory-Efficient Inference

# Sticky Routing: Training MoE Models for Memory-Efficient Inference
Source: [https://arxiv.org/html/2607.08780](https://arxiv.org/html/2607.08780)
###### 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 system\-level \(caching heuristics\) or post\-hoc \(router fine\-tuning\), leaving the root cause unchanged during pretraining\. We proposeStickyMoE, 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\.StickyMoErequires no architectural changes, adds a single hyperparameterλ\\lambda, 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 thatStickyMoEreduces the expert switch rate by up to59%while simultaneously*improving*perplexity on the medium model, and reduces cache misses by up to3\.92×\\mathbf\{3\.92\\times\}, 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)\.

## 1Introduction

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\(Shazeeret al\.,[2017](https://arxiv.org/html/2607.08780#bib.bib1); Feduset al\.,[2022](https://arxiv.org/html/2607.08780#bib.bib2); Lepikhinet al\.,[2021](https://arxiv.org/html/2607.08780#bib.bib3)\)\. By replacing dense feed\-forward network \(FFN\) sublayers with a collection ofNNexpert 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 tokens in a sequence routinely activate entirely different experts\(Zhouet al\.,[2025](https://arxiv.org/html/2607.08780#bib.bib7)\)\. 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 and others,[2024](https://arxiv.org/html/2607.08780#bib.bib10); Zhuet al\.,[2025](https://arxiv.org/html/2607.08780#bib.bib13)\)\. 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 and others,[2024](https://arxiv.org/html/2607.08780#bib.bib10); Kamahoriet al\.,[2024](https://arxiv.org/html/2607.08780#bib.bib11); Zhanget al\.,[2025](https://arxiv.org/html/2607.08780#bib.bib12)\)\. 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\(Zhuet al\.,[2026](https://arxiv.org/html/2607.08780#bib.bib8)\), which is more targeted but still treats the symptom rather than the cause: the expert representations have already been shaped by a training regime that never incentivised locality, and a brief fine\-tuning pass can only partially re\-align the routing surface\. A third line of work, most closely represented by Oracle\-MoE\(Zhouet al\.,[2025](https://arxiv.org/html/2607.08780#bib.bib7)\), redesigns the routing architecture so that the router operates in an attention\-derived “oracle space” that is inherently more semantically stable across tokens\. Oracle\-MoE trains from scratch, which is a step in the right direction, but the locality property is structural—embedded in the choice of routing input—rather than explicitly optimised\. We propose a different approach:StickyMoE, which instils routing locality directly as a training objective, requiring no architectural changes and no post\-hoc correction\.

#### Our proposal\.

We argue that the cleanest and most general solution is to treat routing temporal locality as an explicit training objective\. Specifically, we introduce a*routing consistency loss*— a differentiableℓ2\\ell\_\{2\}penalty on consecutive gate distributions — to directly optimise for temporal locality in MoE routing:

ℒcons=1T−1​∑t=2T‖𝐠t−𝐠t−1‖22,\\mathcal\{L\}\_\{\\text\{cons\}\}\\;=\\;\\frac\{1\}\{T\-1\}\\sum\_\{t=2\}^\{T\}\\bigl\\\|\\mathbf\{g\}\_\{t\}\-\\mathbf\{g\}\_\{t\-1\}\\bigr\\\|\_\{2\}^\{2\},\(1\)where𝐠t∈ℝN\\mathbf\{g\}\_\{t\}\\in\\mathbb\{R\}^\{N\}is the softmax gate probability vector produced by the router for tokenttandTTis the sequence length\. The total training loss is

ℒ=ℒCE\+λ​ℒcons\+μ​ℒbal,\\mathcal\{L\}=\\mathcal\{L\}\_\{\\text\{CE\}\}\+\\lambda\\,\\mathcal\{L\}\_\{\\text\{cons\}\}\+\\mu\\,\\mathcal\{L\}\_\{\\text\{bal\}\},\(2\)whereℒCE\\mathcal\{L\}\_\{\\text\{CE\}\}is the standard cross\-entropy language modelling loss,ℒbal\\mathcal\{L\}\_\{\\text\{bal\}\}is the load\-balancing auxiliary loss fromFeduset al\.\([2022](https://arxiv.org/html/2607.08780#bib.bib2)\), andλ,μ≥0\\lambda,\\mu\\geq 0are scalar hyperparameters\. This formulation:

- •requires*no architectural modification*—it applies to any standard top\-kkMoE without changing the router structure;
- •is*differentiable everywhere*and adds negligible computational overhead;
- •is*architecture\-agnostic*, attaching equally to small experimental models and large\-scale MoEs;
- •*co\-trains*with the language modelling objective, so that expert representations and routing decisions co\-evolve toward local consistency from the very first update\.

#### Contributions\.

We make the following specific contributions:

1. 1\.We propose the routing consistency loss, a simple differentiable training objective that directly penalises expert switching between adjacent tokens, requiring no architectural changes and adding a single hyperparameterλ\\lambda\([section4](https://arxiv.org/html/2607.08780#S4)\)\.
2. 2\.We show empirically thatℒcons\\mathcal\{L\}\_\{\\text\{cons\}\}andℒbal\\mathcal\{L\}\_\{\\text\{bal\}\}are complementary: expert utilisation entropy remains above1\.921\.92bits out oflog2⁡4=2\.0\\log\_\{2\}4=2\.0bits across all settings, confirming that the consistency loss does not induce expert collapse \([section5](https://arxiv.org/html/2607.08780#S5)\)\.
3. 3\.We propose a*soft\-hard*variant that addresses long\-range routing drift by combining the per\-step soft penalty with a segment\-level anchor constraint, achieving stronger locality guarantees without additional quality cost \([section4](https://arxiv.org/html/2607.08780#S4)\)\.
4. 4\.We present controlled experiments on small and medium MoE language models trained on WikiText\-2, comparing against a vanilla MoE baseline, a Hard\-Window ablation, a simulated ReMoE post\-hoc baseline, and a simplified Oracle\-MoE reimplementation, measuring perplexity, switch rate, cache hit rate, and utilisation entropy \([section5](https://arxiv.org/html/2607.08780#S5)\)\.

#### Paper organisation\.

[Section2](https://arxiv.org/html/2607.08780#S2)reviews MoE architectures and the expert\-swapping bottleneck\.[Section3](https://arxiv.org/html/2607.08780#S3)surveys related work and positionsStickyMoEprecisely\.[Section4](https://arxiv.org/html/2607.08780#S4)presents the method in full detail\.[Section5](https://arxiv.org/html/2607.08780#S5)describes the experimental setup and presents results\.[Section6](https://arxiv.org/html/2607.08780#S6)discusses findings, limitations, and broader implications\.[Section7](https://arxiv.org/html/2607.08780#S7)concludes and outlines future directions\.

## 2Background

### 2\.1Mixture\-of\-Experts Language Models

A standard MoE transformer replaces each FFN sublayer with a collection ofNNexpert networks\{ℰi\}i=1N\\\{\\mathcal\{E\}\_\{i\}\\\}\_\{i=1\}^\{N\}and a routerℛ\\mathcal\{R\}\. Given the hidden representation𝐡t∈ℝd\\mathbf\{h\}\_\{t\}\\in\\mathbb\{R\}^\{d\}of tokentt, the router computes gate logitsℓt=𝐖r​𝐡t∈ℝN\\boldsymbol\{\\ell\}\_\{t\}=\\mathbf\{W\}\_\{r\}\\mathbf\{h\}\_\{t\}\\in\\mathbb\{R\}^\{N\}and selects the top\-kkexperts by probability:

𝐠t\\displaystyle\\mathbf\{g\}\_\{t\}=softmax​\(ℓt\),\\displaystyle=\\mathrm\{softmax\}\(\\boldsymbol\{\\ell\}\_\{t\}\),\(3\)𝒮t\\displaystyle\\mathcal\{S\}\_\{t\}=top​\-​k​\(𝐠t\),\\displaystyle=\\mathrm\{top\}\\text\{\-\}k\(\\mathbf\{g\}\_\{t\}\),\(4\)𝐲t\\displaystyle\\mathbf\{y\}\_\{t\}=∑i∈𝒮tgt,i​ℰi​\(𝐡t\)\.\\displaystyle=\\sum\_\{i\\in\\mathcal\{S\}\_\{t\}\}g\_\{t,i\}\\,\\mathcal\{E\}\_\{i\}\(\\mathbf\{h\}\_\{t\}\)\.\(5\)The router weight matrix𝐖r∈ℝN×d\\mathbf\{W\}\_\{r\}\\in\\mathbb\{R\}^\{N\\times d\}is the only additional parameter relative to a dense FFN layer\. The expert networksℰi\\mathcal\{E\}\_\{i\}are typically two\-layer FFNs with the same structure as the dense FFN they replace\.

### 2\.2Load\-Balancing Auxiliary Loss

Without explicit regularisation, MoE routers tend to collapse: a small subset of experts receives disproportionately many tokens, while the remainder are rarely activated and fail to develop useful representations\(Feduset al\.,[2022](https://arxiv.org/html/2607.08780#bib.bib2)\)\. To prevent this, we include the standard load\-balancing auxiliary loss:

ℒbal=N​∑i=1Nfi⋅pi,\\mathcal\{L\}\_\{\\text\{bal\}\}=N\\sum\_\{i=1\}^\{N\}f\_\{i\}\\cdot p\_\{i\},\(6\)where the sum is over allNNexperts,fif\_\{i\}is the fraction of top\-kkassignments to expertii\(counting allkkslots\), andpi=1T​∑t=1Tgt,ip\_\{i\}=\\frac\{1\}\{T\}\\sum\_\{t=1\}^\{T\}g\_\{t,i\}is the mean gate probability assigned to expertii, both computed over the current batch\. The productfi⋅pif\_\{i\}\\cdot p\_\{i\}is minimised when routing is uniform across experts; scaling byNNensures the loss magnitude is independent of the number of experts\.

### 2\.3Expert Swapping on Memory\-Constrained Hardware

LetMfastM\_\{\\text\{fast\}\}denote the fast\-memory capacity \(VRAM or on\-device SRAM\) andwew\_\{e\}the size of a single expert’s weight tensor\. A device can hold at most⌊Mfast/we⌋\\lfloor M\_\{\\text\{fast\}\}/w\_\{e\}\\rfloorexpert weight matrices simultaneously alongside the attention and embedding parameters\. When the router at decoding stepttselects an expert not currently in fast memory, the system must evict a cached expert and load the new one from slow memory—a*cache miss*\. The per\-step latency is:

τt=τcompute\+𝟙​\[misst\]⋅τload,\\tau\_\{t\}=\\tau\_\{\\text\{compute\}\}\+\\mathbb\{1\}\[\\text\{miss\}\_\{t\}\]\\cdot\\tau\_\{\\text\{load\}\},\(7\)whereτload≫τcompute\\tau\_\{\\text\{load\}\}\\gg\\tau\_\{\\text\{compute\}\}on bandwidth\-limited devices\. The expected fraction of steps that incur a cache miss is determined by the*expert switch rate*:

SR=1T−1​∑t=2T𝟙​\[𝒮t≠𝒮t−1\]\.\\text\{SR\}=\\frac\{1\}\{T\-1\}\\sum\_\{t=2\}^\{T\}\\mathbb\{1\}\[\\mathcal\{S\}\_\{t\}\\neq\\mathcal\{S\}\_\{t\-1\}\]\.\(8\)This single\-sequence definition illustrates the concept; the full empirical estimator averaging over layers is defined formally in[section5](https://arxiv.org/html/2607.08780#S5)\. Reducing SR is therefore directly equivalent to reducing expected inference latency under memory constraints\.

## 3Related Work

### 3\.1System\-Level Expert Offloading

A large body of work treats expert weights as a two\-level memory hierarchy and focuses on minimising the latency cost of loading experts from slow to fast memory at inference time\. MoE\-Infinity\(Xue and others,[2024](https://arxiv.org/html/2607.08780#bib.bib10)\)introduces a sparsity\-aware expert cache on personal machines, keeping attention weights in GPU memory and streaming expert weights from host RAM\. Fiddler\(Kamahoriet al\.,[2024](https://arxiv.org/html/2607.08780#bib.bib11)\)uses the CPU itself for expert computation to avoid PCIe transfers\. DuoServe\(Zhanget al\.,[2025](https://arxiv.org/html/2607.08780#bib.bib12)\)separately optimises prefill and decode phases using a learned prediction model for expert activation paths\. EdgeMoE\(Zhuet al\.,[2025](https://arxiv.org/html/2607.08780#bib.bib13)\)proposes importance\-driven expert scheduling specifically for edge settings\. A related thread in the literature studies locality in deployed models:Lianget al\.\([2025](https://arxiv.org/html/2607.08780#bib.bib9)\)conduct an empirical study showing that local routing consistency varies dramatically across model families, and that some models are fundamentally ill\-suited to offloading\-based deployment\.

All of the above methods are*inference\-time*interventions\. They exploit whatever locality the trained model’s router provides but cannot alter the router’s inductive bias\.StickyMoEis orthogonal and complementary: it can be used in conjunction with any of these systems, supplying a router that already tends toward consistency\.

### 3\.2Post\-Hoc Router Adaptation

ReMoE\(Zhuet al\.,[2026](https://arxiv.org/html/2607.08780#bib.bib8)\)is the most directly related post\-hoc method\. It fine\-tunes only the router parameters of a pretrained MoE checkpoint to increase short\-horizon expert reuse, using a combination of the language modelling loss and a locality\-aware regulariser on the gate distribution\. ReMoE reports meaningful reductions in cache miss rate and throughput improvements in vLLM\-based serving\. Its key advantage is practicality: it can be applied to any pretrained model without retraining from scratch\. Its key limitation is that the expert representations themselves are fixed—shaped by a training process that never incentivised locality—so the router is being asked to change its decisions in a representation space that was not designed with locality in mind\.StickyMoEavoids this mismatch by instilling the locality bias from the first training step, allowing the expert representations and routing decisions to co\-adapt\. In our controlled experiments, we were unable to reproduce meaningful SR reductions from post\-hoc router fine\-tuning; we attribute this to the representation mismatch discussed in[section6](https://arxiv.org/html/2607.08780#S6)\.

### 3\.3Training\-Time Architectural Redesign

Oracle\-MoE\(Zhouet al\.,[2025](https://arxiv.org/html/2607.08780#bib.bib7)\)is the closest prior work to ours in motivation and in training regime\. It identifies the same root cause—temporal routing inconsistency—and proposes a structural solution: routing tokens not in the standard hidden\-state space but in an “oracle space” derived from attention scores, which is empirically more semantically stable across consecutive tokens\. Oracle\-MoE is trained from scratch and achieves state\-of\-the\-art inference speed on edge devices without sacrificing task performance\. However, it requires a non\-trivial architectural modification \(replacing the router input with attention\-derived features\), making it harder to apply to existing architectures and training pipelines\.StickyMoEachieves comparable locality through a lightweight auxiliary loss that requires no architectural change, making it more accessible and easier to combine with other MoE innovations\.

### 3\.4Routing Regularisation in MoE

The Switch Transformer\(Feduset al\.,[2022](https://arxiv.org/html/2607.08780#bib.bib2)\)and ST\-MoE\(Zophet al\.,[2022](https://arxiv.org/html/2607.08780#bib.bib4)\)introduce load\-balancing losses to prevent expert collapse\. Expert Choice routing\(Zhouet al\.,[2022](https://arxiv.org/html/2607.08780#bib.bib5)\)inverts the routing direction to guarantee balanced utilisation\. These works regularise the*spatial*distribution of routing decisions across tokens in a batch, but do not address the*temporal*coherence of routing across adjacent tokens in a sequence\. Our consistency loss is orthogonal to and composable with all of these\.

### 3\.5Temporal Coherence in Related Architectures

Mixture of Depths\(Raposoet al\.,[2024](https://arxiv.org/html/2607.08780#bib.bib14)\)routes tokens to*skip layers entirely*, creating a different form of computation\-aware sparse routing\. Pre\-gated MoE\(Hwang and others,[2024](https://arxiv.org/html/2607.08780#bib.bib15)\)predicts which experts will be needed before the full computation, enabling prefetch to hide loading latency\. Neither method explicitly trains for inter\-token routing consistency, but both share the broader goal of reducing unnecessary computation or memory movement\. In the time\-series domain,Erbaset al\.\([2025](https://arxiv.org/html/2607.08780#bib.bib16)\)assigns experts to predetermined temporal segments based on known signal physics—a hard form of the stickiness we achieve through soft regularisation\.

### 3\.6Summary of Differences

[Table1](https://arxiv.org/html/2607.08780#S3.T1)provides a structured comparison ofStickyMoEwith the most relevant prior works along the dimensions most pertinent to edge deployment\.

Table 1:Comparison of methods that address the expert\-swapping bottleneck in memory\-constrained MoE inference\. “Training\-time” indicates whether the method modifies the pretraining procedure\. “Arch\. change” indicates whether a modification to the model architecture is required\. “Any checkpoint” indicates whether the method can be applied to an arbitrary pretrained MoE without retraining\.

## 4Method

### 4\.1Routing Consistency Loss

LetLLdenote the number of MoE layers in the model\. For layerℓ\\elland a sequence ofTTtokens, let𝐠t\(ℓ\)∈ℝN\\mathbf\{g\}\_\{t\}^\{\(\\ell\)\}\\in\\mathbb\{R\}^\{N\}be the softmax gate probability vector produced by the router of layerℓ\\ellfor tokentt\. We define the per\-layer consistency loss as:

ℒcons\(ℓ\)=1T−1​∑t=2T‖𝐠t\(ℓ\)−𝐠t−1\(ℓ\)‖22\.\\mathcal\{L\}\_\{\\text\{cons\}\}^\{\(\\ell\)\}\\;=\\;\\frac\{1\}\{T\-1\}\\sum\_\{t=2\}^\{T\}\\bigl\\\|\\mathbf\{g\}\_\{t\}^\{\(\\ell\)\}\-\\mathbf\{g\}\_\{t\-1\}^\{\(\\ell\)\}\\bigr\\\|\_\{2\}^\{2\}\.\(9\)The full consistency loss aggregates over all MoE layers:

ℒcons=1L​∑ℓ=1Lℒcons\(ℓ\)\.\\mathcal\{L\}\_\{\\text\{cons\}\}\\;=\\;\\frac\{1\}\{L\}\\sum\_\{\\ell=1\}^\{L\}\\mathcal\{L\}\_\{\\text\{cons\}\}^\{\(\\ell\)\}\.\(10\)
#### Interpretation\.

Theℓ2\\ell\_\{2\}distance between consecutive gate distributions𝐠t\\mathbf\{g\}\_\{t\}and𝐠t−1\\mathbf\{g\}\_\{t\-1\}is zero if and only if the router assigns identical probability mass to all experts for both tokens\. Minimisingℒcons\\mathcal\{L\}\_\{\\text\{cons\}\}therefore directly penalises abrupt changes in the routing distribution\. Because𝐠t\\mathbf\{g\}\_\{t\}is a probability simplex vector, the maximum possible value of‖𝐠t−𝐠t−1‖22\\\|\\mathbf\{g\}\_\{t\}\-\\mathbf\{g\}\_\{t\-1\}\\\|\_\{2\}^\{2\}is22\(when one token is routed entirely to expertiiand the next entirely to expertj≠ij\\neq i\)\. The loss thus has a natural\[0,2\]\[0,2\]range, makingλ\\lambdainterpretable across different model sizes and expert counts\.

#### Gradient flow\.

The gate probabilities𝐠t=softmax​\(𝐖r​𝐡t\)\\mathbf\{g\}\_\{t\}=\\mathrm\{softmax\}\(\\mathbf\{W\}\_\{r\}\\mathbf\{h\}\_\{t\}\)are differentiable in both𝐖r\\mathbf\{W\}\_\{r\}\(the router weight\) and𝐡t\\mathbf\{h\}\_\{t\}\(the input hidden state\)\. Minimisingℒcons\\mathcal\{L\}\_\{\\text\{cons\}\}therefore:

1. 1\.directly updates𝐖r\\mathbf\{W\}\_\{r\}to produce more similar gate distributions for semantically adjacent tokens, and
2. 2\.backpropagates through𝐡t\\mathbf\{h\}\_\{t\}into the preceding attention and embedding layers, encouraging the model to produce hidden representations whose local geometry is better aligned with the router’s decision boundaries\.

This second effect is the key advantage over post\-hoc router fine\-tuning:StickyMoEtrains the*representation*as well as the*routing function*to be locally consistent\.

### 4\.2Total Training Objective

The full training loss is:

ℒ=ℒCE\+λ​ℒcons\+μ​ℒbal,\\mathcal\{L\}\\;=\\;\\mathcal\{L\}\_\{\\text\{CE\}\}\+\\lambda\\,\\mathcal\{L\}\_\{\\text\{cons\}\}\+\\mu\\,\\mathcal\{L\}\_\{\\text\{bal\}\},\(11\)whereℒCE\\mathcal\{L\}\_\{\\text\{CE\}\}is the standard next\-token cross\-entropy loss,ℒbal\\mathcal\{L\}\_\{\\text\{bal\}\}is the load\-balancing loss from[eq\.6](https://arxiv.org/html/2607.08780#S2.E6), andλ,μ≥0\\lambda,\\mu\\geq 0are hyperparameters\. We recommendμ=0\.01\\mu=0\.01followingFeduset al\.\([2022](https://arxiv.org/html/2607.08780#bib.bib2)\)and treatλ\\lambdaas the primary experimental variable\.

### 4\.3Interaction with Load Balancing

A natural concern is whetherℒcons\\mathcal\{L\}\_\{\\text\{cons\}\}andℒbal\\mathcal\{L\}\_\{\\text\{bal\}\}conflict: a trivial minimiser ofℒcons\\mathcal\{L\}\_\{\\text\{cons\}\}is to route all tokens to the same expert, which would eliminate switching but collapse expert diversity\. In practice,ℒbal\\mathcal\{L\}\_\{\\text\{bal\}\}effectively prevents this\. For allλ≤0\.5\\lambda\\leq 0\.5explored in our experiments, utilisation entropy remains above1\.921\.92bits out of a maximum oflog2⁡4=2\.0\\log\_\{2\}4=2\.0bits, confirming that expert usage stays near\-uniform throughout training \(see[tables6](https://arxiv.org/html/2607.08780#S5.T6)and[7](https://arxiv.org/html/2607.08780#S5.T7)\)\. The two objectives are therefore complementary in the regime we consider:ℒcons\\mathcal\{L\}\_\{\\text\{cons\}\}shapes*when*experts are used across the sequence, whileℒbal\\mathcal\{L\}\_\{\\text\{bal\}\}ensures*which*experts are used remains balanced\.

### 4\.4Soft\-Hard Variant

The soft consistency loss enforces a*chain constraint*: each token is penalised for differing from its immediate predecessor\. While effective at suppressing abrupt local switches, this formulation is vulnerable to gradual drift — the router may take small steps at each transition yet migrate far from its starting expert over a long span, incurring cache misses that the per\-step penalty never directly penalises\.

To address this, we introduce a complementary*anchor constraint*, inspired by commitment losses in discrete representation learning\(Van Den Oordet al\.,[2017](https://arxiv.org/html/2607.08780#bib.bib28)\)\. The sequence is partitioned into non\-overlapping windows ofWWtokens\. The first token of each window routes freely; its gate distribution𝐠s​\(t\)\(ℓ\)\\mathbf\{g\}\_\{s\(t\)\}^\{\(\\ell\)\}then becomes a fixed anchor for the remainder of the window\. Every subsequent token is penalised for deviating from this anchor, with the penalty weight increasing linearly with distance from the window start:

ℒhard\(t,ℓ\)\\displaystyle\\mathcal\{L\}\_\{\\text\{hard\}\}^\{\(t,\\ell\)\}=t−s​\(t\)W⋅‖𝐠t\(ℓ\)−𝐠s​\(t\)\(ℓ\)‖22,\\displaystyle\\;=\\;\\frac\{t\-s\(t\)\}\{W\}\\cdot\\bigl\\\|\\mathbf\{g\}\_\{t\}^\{\(\\ell\)\}\-\\mathbf\{g\}\_\{s\(t\)\}^\{\(\\ell\)\}\\bigr\\\|\_\{2\}^\{2\},\(12\)ℒhard\\displaystyle\\mathcal\{L\}\_\{\\text\{hard\}\}=1L​T​∑ℓ=1L∑t=1Tℒhard\(t,ℓ\),\\displaystyle\\;=\\;\\frac\{1\}\{LT\}\\sum\_\{\\ell=1\}^\{L\}\\sum\_\{t=1\}^\{T\}\\mathcal\{L\}\_\{\\text\{hard\}\}^\{\(t,\\ell\)\},wheres​\(t\)=W​⌊t/W⌋s\(t\)=W\\lfloor t/W\\rflooris the start of the window containing tokentt\. Anchor tokens \(wheret=s​\(t\)t=s\(t\)\) contribute zero through thet−s​\(t\)W\\frac\{t\-s\(t\)\}\{W\}factor and require no special treatment\. The sum over layers mirrors the aggregation inℒcons\\mathcal\{L\}\_\{\\text\{cons\}\}\([eq\.10](https://arxiv.org/html/2607.08780#S4.E10)\), ensuring the hard loss operates uniformly across all MoE layers\. The linear ramp is a deliberate design choice: early in the window the correct expert for the segment is uncertain, so flexibility is preserved; by the window’s end the router should be fully committed\. At window boundaries the anchor resets and the router is free to reassign without penalty\.

The soft and hard losses are geometrically complementary\. The soft loss controls the*step size*of routing changes — each transition is small\. The hard loss controls the*radius*of routing changes — every token stays close to the window anchor\. A small\-step path can still drift far from its origin; a radius constraint alone does not prevent jagged local switching within the window\. Together they enforce both fine\-grained smoothness and coarse\-grained commitment, which is precisely what a memory\-constrained cache requires: not just that adjacent tokens agree, but that an entire span can be served by the same resident expert\.

The total training loss for the soft\-hard variant is:

ℒ=ℒCE\+λ​ℒcons\+μ​ℒbal\+α​ℒhard,\\mathcal\{L\}\\;=\\;\\mathcal\{L\}\_\{\\text\{CE\}\}\+\\lambda\\,\\mathcal\{L\}\_\{\\text\{cons\}\}\+\\mu\\,\\mathcal\{L\}\_\{\\text\{bal\}\}\+\\alpha\\,\\mathcal\{L\}\_\{\\text\{hard\}\},\(13\)whereλ\\lambda,μ\\mu, andα\\alphaare scalar hyperparameters controlling the relative weight of the consistency, load\-balancing, and anchor terms respectively\. Settingα=0\\alpha=0recovers the soft\-only objective in[eq\.11](https://arxiv.org/html/2607.08780#S4.E11)\.

### 4\.5Algorithm

Algorithm[1](https://arxiv.org/html/2607.08780#algorithm1)summarizes theStickyMoEtraining procedure\.

Input:Training corpus

𝒟\\mathcal\{D\}, MoE model with

LLlayers and

NNexperts, hyperparameters

λ,μ,α,W\\lambda,\\mu,\\alpha,W
Output:Trained MoE with locality\-aware router

1exInitialise model parameters

θ\\thetaas in standard MoE;

foreach*mini\-batch ofBBsequences of lengthTTsampled from𝒟\\mathcal\{D\}*do

Forward pass: compute

\{𝐠t\(ℓ\)\}\\\{\\mathbf\{g\}\_\{t\}^\{\(\\ell\)\}\\\}for all

t,ℓt,\\ell;

Compute

ℒCE\\mathcal\{L\}\_\{\\text\{CE\}\}\(cross\-entropy on next\-token prediction\);

Compute

ℒcons\\mathcal\{L\}\_\{\\text\{cons\}\}via[eq\.10](https://arxiv.org/html/2607.08780#S4.E10);

Compute

ℒbal\\mathcal\{L\}\_\{\\text\{bal\}\}via[eq\.6](https://arxiv.org/html/2607.08780#S2.E6);

if*soft\-hard variant*then

Compute

ℒhard\\mathcal\{L\}\_\{\\text\{hard\}\}via[eq\.12](https://arxiv.org/html/2607.08780#S4.E12);

ℒ←ℒCE\+λ​ℒcons\+μ​ℒbal\+α​ℒhard\\mathcal\{L\}\\leftarrow\\mathcal\{L\}\_\{\\text\{CE\}\}\+\\lambda\\,\\mathcal\{L\}\_\{\\text\{cons\}\}\+\\mu\\,\\mathcal\{L\}\_\{\\text\{bal\}\}\+\\alpha\\,\\mathcal\{L\}\_\{\\text\{hard\}\};

end if

else

ℒ←ℒCE\+λ​ℒcons\+μ​ℒbal\\mathcal\{L\}\\leftarrow\\mathcal\{L\}\_\{\\text\{CE\}\}\+\\lambda\\,\\mathcal\{L\}\_\{\\text\{cons\}\}\+\\mu\\,\\mathcal\{L\}\_\{\\text\{bal\}\};

end if

Backpropagate

∇θℒ\\nabla\_\{\\theta\}\\mathcal\{L\}and update

θ\\theta;

end foreach

Algorithm 1StickyMoEtraining procedure\.

## 5Experiments

### 5\.1Dataset

All experiments use the WikiText\-2 raw character dataset\(Merityet al\.,[2017](https://arxiv.org/html/2607.08780#bib.bib17)\)\(wikitext\-2\-raw\-v1\), a standard language modelling benchmark derived from verified Good and Featured Wikipedia articles\. We use the raw \(un\-tokenised\) variant to avoid pre\-processing artefacts\. Text is tokenised with the GPT\-2 byte\-pair encoding \(BPE\) vocabulary\(Radfordet al\.,[2019](https://arxiv.org/html/2607.08780#bib.bib23)\)using theGPT2Tokenizerfrom the HuggingFace Transformers library\(Wolfet al\.,[2020](https://arxiv.org/html/2607.08780#bib.bib24)\), yielding a vocabulary of\|𝒱\|=50,257\|\\mathcal\{V\}\|=50\{,\}257tokens\. Each split is tokenised as a single flat sequence; document boundaries are not marked with special tokens\.

Table 2:WikiText\-2 token counts after GPT\-2 BPE tokenisation\.#### Batching\.

The flat token sequence is partitioned into non\-overlapping fixed\-length chunks ofT=256T=256tokens\. Each training batch consists ofB=16B=16such chunks drawn uniformly at random without replacement per epoch\. Targets are the same chunk shifted by one position, with the last token of each chunk predicting the first token of the immediately following chunk\. Tokenised arrays are cached as 32\-bit integer NumPy files after the first run to avoid redundant tokenisation\.

### 5\.2Model Architecture

All variants share a GPT\-style\(Radfordet al\.,[2018](https://arxiv.org/html/2607.08780#bib.bib25)\)causal transformer in which every FFN sublayer is replaced by a Mixture\-of\-Experts layer\. We evaluate two model sizes:

Table 3:Model configurations\.dffd\_\{\\text\{ff\}\}is the hidden dimension of each expert FFN\. Parameters include the embedding matrix, which is shared with the output projection via weight tying\(Press and Wolf,[2017](https://arxiv.org/html/2607.08780#bib.bib26)\)\.Each MoE layer uses a linear gate𝐖g∈ℝdmodel×N\\mathbf\{W\}\_\{g\}\\in\\mathbb\{R\}^\{d\_\{\\text\{model\}\}\\times N\}to produce a routing probability distribution𝐠t\\mathbf\{g\}\_\{t\}over theNNexperts via softmax\. The top\-k=2k=2experts by gate probability are activated; their outputs are weighted by the renormalised gate values and summed\. Learned absolute positional embeddings are used with a maximum sequence length of 256\. All weights are initialised with𝒩​\(0,0\.02\)\\mathcal\{N\}\(0,0\.02\); residual projection weights are additionally scaled by1/2​L1/\\sqrt\{2L\}followingRadfordet al\.\([2019](https://arxiv.org/html/2607.08780#bib.bib23)\)\. Both model sizes are tractable on a single consumer GPU \(NVIDIA GTX 1080 Ti, 11 GB VRAM\), making the full experiment pipeline reproducible without large\-scale compute resources\.

Note that the absolute perplexity values reflect the limited model capacity and short training schedule \(10,000 steps\); our focus is on the*relative*differences between variants rather than state\-of\-the\-art language model performance\.

### 5\.3Training Variants and Baselines

All variants share the base objective

ℒ=ℒCE\+μ​ℒbal,\\mathcal\{L\}=\\mathcal\{L\}\_\{\\text\{CE\}\}\+\\mu\\,\\mathcal\{L\}\_\{\\text\{bal\}\},\(14\)with load\-balancing coefficientμ=0\.01\\mu=0\.01throughout\. The variants we evaluate are as follows\.

#### Vanilla MoE \(baseline\)\.

Standard top\-kkMoE trained with[eq\.14](https://arxiv.org/html/2607.08780#S5.E14)only\. It serves as the control condition against which all other variants are measured\.

#### StickyMoE— Soft Consistency\.

The primary proposed method augments[eq\.14](https://arxiv.org/html/2607.08780#S5.E14)with the routing consistency loss from[eq\.10](https://arxiv.org/html/2607.08780#S4.E10):

ℒ=ℒCE\+μ​ℒbal\+λ​ℒcons\.\\mathcal\{L\}=\\mathcal\{L\}\_\{\\text\{CE\}\}\+\\mu\\,\\mathcal\{L\}\_\{\\text\{bal\}\}\+\\lambda\\,\\mathcal\{L\}\_\{\\text\{cons\}\}\.\(15\)We sweepλ∈\{0\.01,0\.05,0\.1,0\.2,0\.5\}\\lambda\\in\\\{0\.01,0\.05,0\.1,0\.2,0\.5\\\}on both model sizes\.

#### StickyMoE— Soft\-Hard Variant\.

The combined soft\-hard method described in[section4\.4](https://arxiv.org/html/2607.08780#S4.SS4)\. We fixλ=0\.1\\lambda=0\.1and sweepα∈\{0\.05,0\.1,0\.5,1\.0\}\\alpha\\in\\\{0\.05,0\.1,0\.5,1\.0\\\}with window sizeW=4W=4\.

#### Hard\-Window \(ablation baseline\)\.

Rather than an auxiliary loss, this variant constrains routing at the logit level by biasing the router toward experts it has recently used\. For each tokenttand layerℓ\\ell, a binary bonus is added to the router logits before softmax:

ℓ~t,e\(ℓ\)=ℓt,e\(ℓ\)\+β⋅𝟏​\[e∈ℋ​\(t,W\)\],\\tilde\{\\ell\}\_\{t,e\}^\{\(\\ell\)\}=\\ell\_\{t,e\}^\{\(\\ell\)\}\+\\beta\\cdot\\mathbf\{1\}\\bigl\[e\\in\\mathcal\{H\}\(t,W\)\\bigr\],\(16\)whereℋ​\(t,W\)\\mathcal\{H\}\(t,W\)is the set of experts that appeared in any top\-kkslot over the precedingWWpositions\[t−W,…,t−1\]\[t\{\-\}W,\\ldots,t\{\-\}1\], andβ=10\.0\\beta=10\.0is the bias strength\. The window mask is computed from a detached copy of the current\-step routing assignments so that no gradient flows through the constraint\. We sweepW∈\{2,4,8\}W\\in\\\{2,4,8\\\}\. This baseline is*not*a proposed method — its purpose is to ablate the contribution of the differentiable soft loss by replacing it with a hard, non\-differentiable logit\-level constraint, isolating whether the gradient signal itself is necessary or whether any routing pressure suffices\.

#### ReMoE \(simulated post\-hoc baseline\)\.

This variant simulates the ReMoE approach\(Zhuet al\.,[2026](https://arxiv.org/html/2607.08780#bib.bib8)\)of post\-hoc router fine\-tuning via two phases\. InPhase 1, the full model is trained forSSsteps with[eq\.14](https://arxiv.org/html/2607.08780#S5.E14)only \(identical to the Vanilla MoE baseline\)\. InPhase 2, all parameters except the gate linear layers\{𝐖g\(ℓ\)\}ℓ=1L\\\{\\mathbf\{W\}\_\{g\}^\{\(\\ell\)\}\\\}\_\{\\ell=1\}^\{L\}are frozen, and only the router weights are fine\-tuned for0\.1×S0\.1\\times Sadditional steps with the consistency loss atλ=0\.1\\lambda=0\.1and a reduced learning rate of10−410^\{\-4\}\. Because expert FFN weights are fixed in Phase 2, the model cannot adapt representations to the new routing distribution, providing a direct ablation of the representation co\-adaptation advantage claimed forStickyMoE\.

#### Oracle\-MoE \(simplified architectural baseline\)\.

A simplified reimplementation of Oracle\-MoE\(Zhouet al\.,[2025](https://arxiv.org/html/2607.08780#bib.bib7)\)\. In the standard router, logits are computed from the token hidden state:

Standard:𝐫t=𝐖g​𝐡t\.\\text\{Standard:\}\\quad\\mathbf\{r\}\_\{t\}=\\mathbf\{W\}\_\{g\}\\,\\mathbf\{h\}\_\{t\}\.In the oracle variant, the router input is replaced with the pre\-projection attended values𝐚t\\mathbf\{a\}\_\{t\}from the preceding multi\-head attention sublayer:

Oracle:𝐫t=𝐖g​𝐚t,\\text\{Oracle:\}\\quad\\mathbf\{r\}\_\{t\}=\\mathbf\{W\}\_\{g\}\\,\\mathbf\{a\}\_\{t\},where𝐚t\\mathbf\{a\}\_\{t\}is the weighted sum of value vectors before the output projection\. Because𝐚t\\mathbf\{a\}\_\{t\}aggregates information over the full context via the attention mechanism, it is more temporally stable than𝐡t\\mathbf\{h\}\_\{t\}, and the resulting routing decisions exhibit greater locality as an emergent property\. No consistency loss is added\. Note that this is a simplified approximation of the full Oracle\-MoE architecture\(Zhouet al\.,[2025](https://arxiv.org/html/2607.08780#bib.bib7)\)\.

[Table4](https://arxiv.org/html/2607.08780#S5.T4)summarises the loss coefficients for each variant\.

Table 4:Loss coefficients and routing hyperparameters per variant\.

### 5\.4Training Protocol

All variants are trained with AdamW\(Loshchilov and Hutter,[2019](https://arxiv.org/html/2607.08780#bib.bib19)\)for 10,000 steps\. Full hyperparameters are listed in[table5](https://arxiv.org/html/2607.08780#S5.T5)\.

Table 5:Optimiser hyperparameters shared across all variants\.The learning rate follows a cosine decay schedule with linear warmup, decaying to 10% of the peak value by the end of training\. All runs use a fixed random seed \(42\)\.

### 5\.5Metrics

We report four metrics evaluated on the full WikiText\-2 validation split\.

#### Perplexity \(PPL\)\.

Standard token\-level perplexity:

PPL=exp⁡\(−1n​∑t=1nlog⁡P​\(xt∣x<t\)\)\.\\text\{PPL\}=\\exp\\\!\\left\(\-\\frac\{1\}\{n\}\\sum\_\{t=1\}^\{n\}\\log P\(x\_\{t\}\\mid x\_\{<t\}\)\\right\)\.\(17\)We also reportΔ​PPL%\\Delta\\text\{PPL\}\\%, the percentage change relative to the Baseline, to isolate the quality cost of routing constraints\.

#### Expert Switch Rate \(SR\)\.

The fraction of consecutive token pairs\(t,t\+1\)\(t,t\{\+\}1\)within a sequence where the top\-1 expert assignment changes, averaged over all MoE layers:

SR=1L​T′​∑ℓ=1L∑t=1T′𝟏​\[arg⁡max⁡𝐠t\(ℓ\)≠arg⁡max⁡𝐠t\+1\(ℓ\)\],\\text\{SR\}=\\frac\{1\}\{LT^\{\\prime\}\}\\sum\_\{\\ell=1\}^\{L\}\\sum\_\{t=1\}^\{T^\{\\prime\}\}\\mathbf\{1\}\\\!\\left\[\\arg\\max\\mathbf\{g\}\_\{t\}^\{\(\\ell\)\}\\neq\\arg\\max\\mathbf\{g\}\_\{t\+1\}^\{\(\\ell\)\}\\right\],\(18\)whereLLis the number of MoE layers andT′=T−1T^\{\\prime\}=T\-1\. SR=0=0means the same expert is used for every token; SR=1=1means a different expert is used at every step\.

#### LRU Cache Hit Rate \(CHR\)\.

Expert accesses are simulated against a per\-sequence LRU cache of capacityC=2C=2expert slots — reflecting a realistic edge device budget where two expert weight sets fit in fast memory simultaneously\. A hit occurs when the top\-1 expert for tokenttis already resident in the cache:

CHR=cache hitstotal expert accesses\.\\text\{CHR\}=\\frac\{\\text\{cache hits\}\}\{\\text\{total expert accesses\}\}\.\(19\)CHR directly proxies the fraction of expert weight loads that can be served from fast on\-chip memory rather than DRAM, and is the most hardware\-relevant metric we report\.

#### Utilisation Entropy \(UE\)\.

Utilisation entropy is computed per layer and then averaged\. For each MoE layerℓ\\ell, all expert assignments across sequence positions and both top\-kkslots are pooled into an empirical count distribution over theNNexperts\. Letp^i\(ℓ\)\\hat\{p\}\_\{i\}^\{\(\\ell\)\}denote the fraction of top\-kkassignments to expertiiat layerℓ\\ell\. The per\-layer entropy and its mean are:

UE\(ℓ\)=−∑i=1Np^i\(ℓ\)​log2⁡p^i\(ℓ\),UE=1L​∑ℓ=1LUE\(ℓ\)\.\\text\{UE\}^\{\(\\ell\)\}=\-\\sum\_\{i=1\}^\{N\}\\hat\{p\}\_\{i\}^\{\(\\ell\)\}\\log\_\{2\}\\hat\{p\}\_\{i\}^\{\(\\ell\)\},\\qquad\\text\{UE\}=\\frac\{1\}\{L\}\\sum\_\{\\ell=1\}^\{L\}\\text\{UE\}^\{\(\\ell\)\}\.\(20\)UE lies in\[0,log2⁡N\]\[0,\\,\\log\_\{2\}N\]bits, wherelog2⁡N\\log\_\{2\}Nindicates perfectly uniform usage and0indicates complete collapse to a single expert\. Averaging per\-layer entropies rather than pooling across layers treats each layer independently — a collapsed layer is not masked by a well\-balanced one, and vice versa\.

### 5\.6Main Results

Tables[6](https://arxiv.org/html/2607.08780#S5.T6)and[7](https://arxiv.org/html/2607.08780#S5.T7)report results on WikiText\-2 for the small and medium models respectively\. We highlight six key observations\.

Table 6:Results on WikiText\-2\(Small model,≈\\approx8\.8M parameters, 4 experts, top\-2 routing\)\. PPL = perplexity \(lower is better\)\.Δ\\DeltaPPL% = percentage change relative to Baseline \(negative = improvement\)\. SR = expert switch rate \(lower is better\)\. CHR = LRU cache hit rate withC=2C\{=\}2expert slots \(higher is better\)\. Ent = utilisation entropy in bits \(higher indicates more balanced expert usage\)\.Highlightedrow is the recommended operating point\.Table 7:Results on WikiText\-2\(Medium model,≈\\approx22M parameters, 4 experts, top\-2 routing\)\. Columns as in[table6](https://arxiv.org/html/2607.08780#S5.T6)\.Highlightedrow is the recommended operating point\.#### Soft consistency improves or preserves perplexity at lowλ\\lambda\.

On the small model, Softλ=0\.01\\lambda\{=\}0\.01reduces perplexity relative to the Baseline \(245\.3→243\.9245\.3\\to 243\.9,−0\.5%\-0\.5\\%\), and Softλ=0\.05\\lambda\{=\}0\.05matches the Baseline within\+0\.1%\+0\.1\\%\. The effect is stronger on the medium model, where Softλ=0\.05\\lambda\{=\}0\.05yields the best perplexity of any method \(263\.5263\.5,−4\.1%\-4\.1\\%relative to Baseline\)\. This is a notable finding: at moderateλ\\lambda, the consistency loss acts as a beneficial regulariser rather than a quality penalty, shaping representations that are both more locally coherent and more predictive\. Only at aggressiveλ≥0\.2\\lambda\\geq 0\.2does a clear quality–locality trade\-off emerge\.

#### Switch rate is substantially reduced across both scales\.

Softλ=0\.5\\lambda\{=\}0\.5reduces SR from0\.710\.71to0\.300\.30on the small model and from0\.710\.71to0\.290\.29on the medium model — a consistent∼\\sim59% reduction at both scales\. The smooth monotone relationship betweenλ\\lambdaand SR \([fig\.1](https://arxiv.org/html/2607.08780#S5.F1)\) means practitioners can select an operating point to match their hardware memory budget without retraining from scratch\.

#### ReMoE fails to reduce switch rate at either scale\.

Post\-hoc router fine\-tuning achieves SR of0\.7050\.705on the small model and0\.7090\.709on the medium model, compared to baselines of0\.7090\.709and0\.7130\.713respectively — a negligible change of less than0\.5%0\.5\\%at both scales\. The per\-layer analysis \([tables9](https://arxiv.org/html/2607.08780#A3.T9)and[8](https://arxiv.org/html/2607.08780#S5.T8)\) makes this even starker: ReMoE’s per\-layer switch rates are virtually identical to the Baseline at every layer\. This confirms the central claim of this paper — that routing locality cannot be meaningfully retrofitted post\-hoc, because the expert representations themselves were never trained to support it\.

#### Hard baseline achieves locality at excessive quality cost\.

Hard \(W=2W\{=\}2\) reduces SR to0\.4990\.499on the small model but at a\+5\.6%\+5\.6\\%perplexity cost — more than five times the cost of Softλ=0\.1\\lambda\{=\}0\.1at a comparable switch rate\. More strikingly, larger windows do*not*improve locality: Hard \(W=8W\{=\}8\) produces higher switch rates than Hard \(W=2W\{=\}2\) at every layer on both models, suggesting that the logit bias becomes ineffective when the window is too wide to maintain consistent pressure\. The Hard baseline is therefore both expensive in quality and unreliable in locality, motivating the soft differentiable approach\.

#### Soft\-Hard variant provides the best locality at moderate cost\.

On the medium model, Soft\-Hard \(λ=0\.1\\lambda\{=\}0\.1,α=1\.0\\alpha\{=\}1\.0,W=4W\{=\}4\) achieves the lowest switch rate of any method \(0\.2740\.274\) with only\+0\.4%\+0\.4\\%perplexity cost, and the highest cache hit rate \(0\.8890\.889\)\. The per\-layer profile shows L1 dropping to0\.1080\.108— an85%85\\%reduction from the Baseline — while perplexity remains within1%1\\%of the Baseline\. This confirms that the anchor constraint and the soft penalty are complementary: the soft loss prevents local switching while the anchor prevents long\-range drift within a window\.

#### Expert utilisation is preserved throughout\.

Utilisation entropy remains above1\.921\.92bits across allStickyMoEvariants on both models, out of a maximum oflog2⁡4=2\.0\\log\_\{2\}4=2\.0bits\. Even atλ=0\.5\\lambda\{=\}0\.5, where SR is reduced by nearly60%60\\%, all four experts remain actively used\. The consistency loss does not cause the model to collapse routing to a single expert — it redistributes*when*experts are used, not*which*experts exist\.

### 5\.7Pareto Frontier and Inference\-Time Impact

StickyMoErequires no modification to the forward pass at inference time — the trained router naturally produces more temporally consistent gate distributions, and any standard caching strategy \(LRU, LFU, or learned\) benefits from the improved locality without additional engineering effort\.

![Refer to caption](https://arxiv.org/html/2607.08780v1/figures/pareto_small_chr.png)\(a\)Small model \(∼\{\\sim\}8\.8M params\)
![Refer to caption](https://arxiv.org/html/2607.08780v1/figures/pareto_medium_chr.png)\(b\)Medium model \(∼\{\\sim\}22M params\)

Figure 1:Quality–locality Pareto frontier \(cache hit rate vs\. perplexity\) on WikiText\-2 for small and medium models\. Each point represents one method configuration; higher cache hit rate and lower perplexity are both desirable \(upper\-left is better\)\.StickyMoESoft and Soft\-Hard variants trace a smooth Pareto frontier, consistently achieving higher cache hit rates at comparable or better perplexity than the Hard baseline, ReMoE, and Oracle\-MoE\. ReMoE fails to meaningfully improve cache hit rate at either scale, clustering near the Baseline point\.Figure[1](https://arxiv.org/html/2607.08780#S5.F1)illustrates the quality–locality Pareto frontier\.StickyMoESoft and Soft\-Hard variants strictly dominate ReMoE across all operating points: because ReMoE fails to move SR meaningfully, it offers no cache efficiency benefit whatsoever despite the additional fine\-tuning cost\. Oracle\-MoE achieves competitive locality but requires architectural changes to the router input;StickyMoEmatches or exceeds its cache hit rate without any structural modification\. The Hard baseline traces a separate, dominated frontier — high locality is achievable but only at perplexity costs that make the model uncompetitive\.

The cache impact is directly quantified by our LRU simulations \(C=2C\{=\}2expert slots\)\. On the small model,StickyMoEatλ=0\.1\\lambda\{=\}0\.1reduces SR from0\.710\.71to0\.470\.47, raising CHR from0\.540\.54to0\.750\.75— a1\.82×1\.82\\timesreduction in cache misses\. The effect is more pronounced on the medium model: atλ=0\.5\\lambda\{=\}0\.5, SR drops from0\.710\.71to0\.290\.29, raising CHR from0\.540\.54to0\.880\.88— a3\.92×3\.92\\timesreduction in cache misses, translating directly into lower per\-token latency on bandwidth\-limited edge hardware\.

### 5\.8Ablation Studies

Table 8:Per\-layer expert switch rate on WikiText\-2\(Medium model,≈\\approx22M parameters\)\. SR per layer \(lower is better\)\. Layer 0 is closest to the input\.#### Per\-layer consistency profiles\.

Tables[9](https://arxiv.org/html/2607.08780#A3.T9)and[8](https://arxiv.org/html/2607.08780#S5.T8)show per\-layer switch rates across all methods\. Two consistent patterns emerge across both model sizes\. First, L0 is the most resistant layer: even atλ=0\.5\\lambda\{=\}0\.5, L0 SR remains above0\.490\.49on both models, compared to reductions to0\.220\.22and below at L1–L3\. This is expected — L0 processes token embeddings before any contextual mixing via attention, making adjacent tokens intrinsically more dissimilar at that layer\. Second, the consistency loss distributes its effect non\-uniformly: L1 and L2 see the largest absolute reductions, suggesting these layers are most amenable to locality regularisation once some contextual integration has occurred\. The Soft\-Hard variant atα=1\.0\\alpha\{=\}1\.0achieves the single lowest per\-layer SR in either table: L1=0\.108=0\.108on the medium model, an85%85\\%reduction from the Baseline value of0\.7390\.739\.

#### λ\\lambdasensitivity\.

The medium model is substantially more robust to largeλ\\lambdathan the small model\. On the small model,λ=0\.5\\lambda\{=\}0\.5incurs\+3\.9%\+3\.9\\%PPL; on the medium model the sameλ\\lambdaincurs only−0\.9%\-0\.9\\%\(an improvement\)\. This scaling behaviour is encouraging: as model capacity grows, the consistency constraint becomes easier to satisfy without sacrificing representational quality, suggestingStickyMoEwill become more effective, not less, at larger scales\.

## 6Discussion

#### When doesStickyMoEhelp most?

The consistency loss is most beneficial when two conditions hold simultaneously: the target hardware has severely constrained fast memory relative to expert size, making each cache miss expensive, and the target domain has natural local coherence — long\-form text, code, conversations, documents — that the consistency loss can reinforce\. Our results suggest a third condition that was not anticipated:StickyMoEis more effective at larger model scales\. On the medium model, moderateλ\\lambdavalues actually*improve*perplexity relative to the baseline, suggesting that larger expert networks have sufficient capacity to satisfy the routing constraint without sacrificing representational quality\. This is an encouraging sign for deployment at scale\.

#### Why post\-hoc fine\-tuning fails\.

The most striking empirical finding is the complete ineffectiveness of ReMoE\-style post\-hoc router fine\-tuning\. Across both model sizes and all four layers, ReMoE produces switch rates within0\.5%0\.5\\%of the Baseline — a negligible change despite additional compute\. This is not a failure of the fine\-tuning procedure itself but a fundamental consequence of the representation mismatch problem: the expert weight matrices were trained without any locality incentive, and the hidden representations they produce do not cluster in a way that supports temporally stable routing\. Fine\-tuning only the router cannot overcome this because the router’s input space is fixed\.StickyMoEavoids this by shaping both representations and routing decisions jointly from the first training step\.

#### The L0 barrier\.

Across all methods and both model sizes, layer L0 is consistently the most resistant to locality improvement\. Even atλ=0\.5\\lambda\{=\}0\.5, L0 switch rate remains above0\.490\.49, while L1–L3 drop to0\.170\.17–0\.240\.24\. We attribute this to the nature of L0 inputs: token embeddings carry token\-specific identity information and have not yet been mixed with contextual information via attention\. Adjacent tokens are therefore genuinely dissimilar at L0 in a way that the consistency loss cannot easily overcome without large perplexity cost\. This suggests that future work targeting L0 specifically — perhaps through embedding\-level smoothing or subword\-aware routing — could yield further locality gains\.

#### Relationship to language structure\.

The non\-uniform distribution of locality gains across layers is consistent with the view that higher layers represent more abstract, semantically stable information while lower layers process surface\-level token features\.StickyMoEeffectively asks lower layers to behave more like higher layers from a routing perspective\. The fact that this is achievable at moderateλ\\lambdawithout quality loss suggests that the routing surface in these layers is more flexible than the representations themselves — the router can be steered toward consistency without the underlying features needing to change dramatically\. Consistent with this, L1 SR drops from0\.7390\.739to0\.1770\.177atλ=0\.5\\lambda\{=\}0\.5on the medium model while PPL simultaneously*improves*by0\.9%0\.9\\%— evidence that the routing surface at L1 is genuinely more flexible than the representations it operates on\.

#### Towards cross\-layer routing consistency\.

The consistency loss proposed in this paper operates within each layer across consecutive tokens, which is most directly beneficial during prefill and batched decode\. For single\-token autoregressive decoding — the dominant inference mode on memory\-constrained edge devices — the forward pass traverses all layers sequentially for each token, making cross\-layer memory access the primary bottleneck\. A complementary cross\-layer consistency loss that encourages the same expert to be selected at consecutive layers for the same token would directly target this bottleneck\. However, such a loss only reduces cache misses if expert weights are shared across layers, since routing consistently to expertiiat layersℓ\\ellandℓ\+1\\ell\+1only avoids a memory load whenℰi\(ℓ\)=ℰi\(ℓ\+1\)\\mathcal\{E\}\_\{i\}^\{\(\\ell\)\}=\\mathcal\{E\}\_\{i\}^\{\(\\ell\+1\)\}\. In standard MoE models with independent per\-layer expert matrices, cross\-layer routing consistency provides no cache benefit\. We provide a full analysis of this access pattern, the parameter sharing requirement, and a proposed combined training objective in[appendixD](https://arxiv.org/html/2607.08780#A4)\.

#### Limitations\.

The consistency loss penalises switches uniformly regardless of whether a switch occurs at a natural semantic boundary — a sentence break, topic shift, or dialogue turn — or in the middle of a coherent span\. At a sentence boundary, switching experts is semantically appropriate and penalising it unnecessarily constrains the model\. A boundary\-aware variant that gates the consistency loss using punctuation or BOS signals would relax the penalty precisely where switching is warranted and strengthen it where it is not\. We expect this would improve the quality–locality trade\-off particularly at higherλ\\lambdavalues where the current formulation begins to hurt perplexity\. We leave this to future work\.

## 7Conclusion

We introducedStickyMoE, a training\-time routing consistency loss for Mixture\-of\-Experts language models that directly optimises for temporal locality in expert activation patterns\. By penalising abrupt expert switches between adjacent tokens during pretraining,StickyMoEproduces routers that are intrinsically locality\-aware without requiring any architectural modifications or post\-hoc adaptation\. Experiments on small and medium MoE models demonstrate thatStickyMoEreduces the expert switch rate by up to59%while simultaneously*improving*perplexity on the medium model, and reduces cache misses by up to3\.92×\\mathbf\{3\.92\\times\}— a result that stands in stark contrast to post\-hoc router fine\-tuning, which fails to meaningfully reduce switch rate at either scale, confirming that routing locality cannot be retrofitted and must be instilled during training\.

These results instantiate a broader principle: inference\-time memory access patterns are a legitimate training objective\. Just as quantisation\-aware training\(Jacobet al\.,[2018](https://arxiv.org/html/2607.08780#bib.bib20)\)has become standard practice by demonstrating that training with deployment constraints in the loop consistently outperforms post\-hoc compression, we argue that locality\-aware training should be the default starting point for MoE models intended for memory\-constrained deployment\. The post\-hoc stack — caching heuristics, router fine\-tuning, 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\.

#### Future work\.

Promising directions include: scalingStickyMoEto larger MoE models \(1B\+ parameters\) where the regularisation benefit observed at medium scale may be even more pronounced; boundary\-aware consistency losses that relax the penalty at sentence and paragraph boundaries; combiningStickyMoEwith structured expert grouping so that switches within a group are cheaper than switches across groups; and investigating whether experts trained withStickyMoEdevelop more interpretable and topically coherent specialisations, given that they process longer contiguous spans of semantically related tokens — a testable prediction that consistent routing implies more coherent and persistent expert assignments across a document\.

## References

- I\. Erbas, F\. Demirkiran, K\. Swaminathan, N\. Wang, N\. I\. Nizam, S\. T\. Radev, K\. E\. Maghraoui, X\. Intes, and V\. Pandey \(2025\)EvidenceMoE: a physics\-guided mixture\-of\-experts with evidential critics for advancing fluorescence light detection and ranging in scattering media\.arXiv preprint arXiv:2505\.21532\.Cited by:[§3\.5](https://arxiv.org/html/2607.08780#S3.SS5.p1.1)\.
- W\. Fedus, B\. Zoph, and N\. Shazeer \(2022\)Switch transformers: scaling to trillion parameter models with simple and efficient sparsity\.InJournal of Machine Learning Research,Vol\.23,pp\. 1–39\.Cited by:[§1](https://arxiv.org/html/2607.08780#S1.SS0.SSS0.Px3.p1.7),[§1](https://arxiv.org/html/2607.08780#S1.p1.1),[§2\.2](https://arxiv.org/html/2607.08780#S2.SS2.p1.10),[§3\.4](https://arxiv.org/html/2607.08780#S3.SS4.p1.1),[§4\.2](https://arxiv.org/html/2607.08780#S4.SS2.p1.5)\.
- R\. Hwanget al\.\(2024\)Pre\-gated MoE: an algorithm\-system co\-design for fast and scalable mixture\-of\-expert inference\.InProceedings of the International Symposium on Computer Architecture \(ISCA\),Cited by:[§3\.5](https://arxiv.org/html/2607.08780#S3.SS5.p1.1)\.
- B\. Jacob, S\. Kligys, B\. Chen, M\. Zhu, M\. Tang, A\. Howard, H\. Adam, and D\. Kalenichenko \(2018\)Quantization and training of neural networks for efficient integer\-arithmetic\-only inference\.InProceedings of the IEEE Conference on Computer Vision and Pattern Recognition \(CVPR\),Cited by:[§7](https://arxiv.org/html/2607.08780#S7.p2.1)\.
- K\. Kamahori, Y\. Gu, K\. Zhu, and B\. Kasikci \(2024\)Fiddler: CPU\-GPU orchestration for fast inference of mixture\-of\-experts models\.InarXiv preprint arXiv:2402\.07033,Cited by:[§1](https://arxiv.org/html/2607.08780#S1.SS0.SSS0.Px2.p1.1),[§3\.1](https://arxiv.org/html/2607.08780#S3.SS1.p1.1),[Table 1](https://arxiv.org/html/2607.08780#S3.T1.4.5.4.1)\.
- D\. Lepikhin, H\. Lee, Y\. Xu, D\. Chen, O\. Firat, Y\. Huang, M\. Krikun, N\. Shazeer, and Z\. Chen \(2021\)GShard: scaling giant models with conditional computation and automatic sharding\.InInternational Conference on Learning Representations \(ICLR\),Cited by:[§1](https://arxiv.org/html/2607.08780#S1.p1.1)\.
- J\. Liang, S\. Wang, M\. Tian, Y\. Li, D\. Tang, and Z\. Wei \(2025\)Not all models suit expert offloading: on local routing consistency of mixture\-of\-expert models\.arXiv preprint arXiv:2505\.16056\.Cited by:[§3\.1](https://arxiv.org/html/2607.08780#S3.SS1.p1.1)\.
- I\. Loshchilov and F\. Hutter \(2019\)Decoupled weight decay regularization\.InInternational Conference on Learning Representations \(ICLR\),Cited by:[§5\.4](https://arxiv.org/html/2607.08780#S5.SS4.p1.1)\.
- S\. Merity, C\. Xiong, J\. Bradbury, and R\. Socher \(2017\)Pointer sentinel mixture models\.InInternational Conference on Learning Representations \(ICLR\),Cited by:[§5\.1](https://arxiv.org/html/2607.08780#S5.SS1.p1.1)\.
- O\. Press and L\. Wolf \(2017\)Using the output embedding to improve language models\.InProceedings of the 15th Conference of the European Chapter of the Association for Computational Linguistics: Volume 2, Short Papers,pp\. 157–163\.Cited by:[Table 3](https://arxiv.org/html/2607.08780#S5.T3),[Table 3](https://arxiv.org/html/2607.08780#S5.T3.2.1)\.
- A\. Radford, K\. Narasimhan, T\. Salimans, I\. Sutskever,et al\.\(2018\)Improving language understanding by generative pre\-training\.Cited by:[§5\.2](https://arxiv.org/html/2607.08780#S5.SS2.p1.1)\.
- A\. Radford, J\. Wu, R\. Child, D\. Luan, D\. Amodei, I\. Sutskever,et al\.\(2019\)Language models are unsupervised multitask learners\.OpenAI blog1\(8\),pp\. 9\.Cited by:[§5\.1](https://arxiv.org/html/2607.08780#S5.SS1.p1.1),[§5\.2](https://arxiv.org/html/2607.08780#S5.SS2.p2.6)\.
- D\. Raposo, S\. Ritter, B\. Richards, T\. Lillicrap, P\. C\. Humphreys, and A\. Santoro \(2024\)Mixture of depths: dynamically allocating compute in transformer language models\.InarXiv preprint arXiv:2404\.02258,Cited by:[§3\.5](https://arxiv.org/html/2607.08780#S3.SS5.p1.1)\.
- S\. Roller, S\. Sukhbaatar, J\. Weston,et al\.\(2021\)Hash layers for large sparse models\.advances in neural information processing systems34,pp\. 17555–17566\.Cited by:[§D\.4](https://arxiv.org/html/2607.08780#A4.SS4.p2.2),[§D\.4](https://arxiv.org/html/2607.08780#A4.SS4.p4.1)\.
- N\. Shazeer, A\. Mirhoseini, K\. Maziarz, A\. Davis, Q\. Le, G\. Hinton, and J\. Dean \(2017\)Outrageously large neural networks: the sparsely\-gated mixture\-of\-experts layer\.InInternational Conference on Learning Representations \(ICLR\),Cited by:[§1](https://arxiv.org/html/2607.08780#S1.p1.1)\.
- A\. Van Den Oord, O\. Vinyals,et al\.\(2017\)Neural discrete representation learning\.Advances in neural information processing systems30\.Cited by:[§4\.4](https://arxiv.org/html/2607.08780#S4.SS4.p2.2)\.
- T\. Wolf, L\. Debut, V\. Sanh, J\. Chaumond, C\. Delangue, A\. Moi, P\. Cistac, T\. Rault, R\. Louf, M\. Funtowicz,et al\.\(2020\)Transformers: state\-of\-the\-art natural language processing\.InProceedings of the 2020 conference on empirical methods in natural language processing: system demonstrations,pp\. 38–45\.Cited by:[§5\.1](https://arxiv.org/html/2607.08780#S5.SS1.p1.1)\.
- J\. Xueet al\.\(2024\)MoE\-Infinity: efficient MoE inference on personal machines with sparsity\-aware expert cache\.InarXiv preprint arXiv:2401\.14361,Cited by:[§1](https://arxiv.org/html/2607.08780#S1.SS0.SSS0.Px1.p1.1),[§1](https://arxiv.org/html/2607.08780#S1.SS0.SSS0.Px2.p1.1),[§3\.1](https://arxiv.org/html/2607.08780#S3.SS1.p1.1),[Table 1](https://arxiv.org/html/2607.08780#S3.T1.4.4.3.1)\.
- Y\. Zhang, G\. Pinkert, N\. Yang, Y\. Li, and D\. Yuan \(2025\)DuoServe\-moe: dual\-phase expert prefetch and cache scheduling for efficient moe llm inference\.arXiv preprint arXiv:2509\.07379\.Cited by:[§1](https://arxiv.org/html/2607.08780#S1.SS0.SSS0.Px2.p1.1),[§3\.1](https://arxiv.org/html/2607.08780#S3.SS1.p1.1),[Table 1](https://arxiv.org/html/2607.08780#S3.T1.4.6.5.1)\.
- J\. Zhou, F\. Dong, R\. Huang, H\. Cao, M\. Chen, Y\. Yang, A\. Chen, M\. Dong, Y\. Wang, D\. Li, D\. A\. Clifton, Q\. Lv, R\. Zhu, C\. Zhang, F\. Yang, T\. Lu, N\. Gu, and L\. Shang \(2025\)Oracle\-MoE: locality\-preserving routing in the oracle space for memory\-constrained large language model inference\.InProceedings of the 42nd International Conference on Machine Learning \(ICML\),Cited by:[§1](https://arxiv.org/html/2607.08780#S1.SS0.SSS0.Px1.p1.1),[§1](https://arxiv.org/html/2607.08780#S1.SS0.SSS0.Px2.p1.1),[§3\.3](https://arxiv.org/html/2607.08780#S3.SS3.p1.1),[Table 1](https://arxiv.org/html/2607.08780#S3.T1.4.9.8.1),[§5\.3](https://arxiv.org/html/2607.08780#S5.SS3.SSS0.Px6.p1.4),[§5\.3](https://arxiv.org/html/2607.08780#S5.SS3.SSS0.Px6.p1.5)\.
- Y\. Zhou, T\. Lei, H\. Liu, N\. Du, Y\. Huang, V\. Zhao, A\. M\. Dai, Z\. Chen, Q\. V\. Le, and J\. Laudon \(2022\)Mixture\-of\-experts with expert choice routing\.InAdvances in Neural Information Processing Systems \(NeurIPS\),Cited by:[§3\.4](https://arxiv.org/html/2607.08780#S3.SS4.p1.1)\.
- G\. Zhu, M\. Li, H\. Dai, X\. Liu, W\. Wang, K\. Li, L\. Chen, W\. Wang,et al\.\(2025\)Enabling moe on the edge via importance\-driven expert scheduling\.arXiv preprint arXiv:2508\.18983\.Cited by:[§1](https://arxiv.org/html/2607.08780#S1.SS0.SSS0.Px1.p1.1),[§3\.1](https://arxiv.org/html/2607.08780#S3.SS1.p1.1),[Table 1](https://arxiv.org/html/2607.08780#S3.T1.4.7.6.1)\.
- X\. Zhu, X\. Liao, T\. Jiang, Y\. Zhang, L\. Wang, and L\. Xiao \(2026\)ReMoE: boosting expert reuse through router fine\-tuning in memory\-constrained moe llm inference\.arXiv preprint arXiv:2605\.27081\.Cited by:[§1](https://arxiv.org/html/2607.08780#S1.SS0.SSS0.Px2.p1.1),[§3\.2](https://arxiv.org/html/2607.08780#S3.SS2.p1.1),[Table 1](https://arxiv.org/html/2607.08780#S3.T1.4.8.7.1),[§5\.3](https://arxiv.org/html/2607.08780#S5.SS3.SSS0.Px5.p1.5)\.
- B\. Zoph, I\. Bello, S\. Kumar, N\. Du, Y\. Huang, J\. Dean, N\. Shazeer, and W\. Fedus \(2022\)ST\-MoE: designing stable and transferable sparse expert models\.InarXiv preprint arXiv:2202\.08906,Cited by:[§3\.4](https://arxiv.org/html/2607.08780#S3.SS4.p1.1)\.

## Appendix AReproducibility

All code, configuration files, and dependency versions are committed to the project repository\. Exact package versions are specified inuv\.lock\. The WikiText\-2 dataset is downloaded from the HuggingFace Hub \(Salesforce/wikitext, splitwikitext\-2\-raw\-v1\) and cached locally as NumPy arrays\. The full experiment pipeline is reproducible by running:

```
uv sync
    uv run python download_data.py
    bash run_experiments.sh
```

## Appendix BSwitch Rate Pareto Frontier

Figure[2](https://arxiv.org/html/2607.08780#A2.F2)shows the quality–locality Pareto frontier using switch rate as the locality metric, complementing the cache hit rate frontier in the main text\.

![Refer to caption](https://arxiv.org/html/2607.08780v1/figures/pareto_small_sr.png)\(a\)Small model \(∼\{\\sim\}8\.8M params\)
![Refer to caption](https://arxiv.org/html/2607.08780v1/figures/pareto_medium_sr.png)\(b\)Medium model \(∼\{\\sim\}22M params\)

Figure 2:Quality–locality Pareto frontier \(switch rate vs\. perplexity\) on WikiText\-2 for small and medium models\. Each point represents one method configuration; lower switch rate and lower perplexity are both desirable \(lower\-left is better\)\.StickyMoESoft and Soft\-Hard variants trace a smooth Pareto frontier, consistently achieving lower switch rates at comparable or better perplexity than the Hard baseline, ReMoE, and Oracle\-MoE\. ReMoE fails to meaningfully reduce switch rate at either scale, clustering near the Baseline point\.
## Appendix CPer\-Layer Switch Rates: Small Model

Table 9:Per\-layer expert switch rate on WikiText\-2\(Small model,≈\\approx8\.8M parameters\)\. SR per layer \(lower is better\)\. Layer 0 is closest to the input\.
## Appendix DCross\-Layer Routing Consistency: A Forward\-Looking Analysis

### D\.1The Decode\-Time Memory Access Pattern

The consistency loss proposed in the main paper operates*within*a layer across consecutive tokens:

ℒcons\(ℓ\)=1T−1​∑t=2T‖𝐠t\(ℓ\)−𝐠t−1\(ℓ\)‖22\.\\mathcal\{L\}\_\{\\text\{cons\}\}^\{\(\\ell\)\}=\\frac\{1\}\{T\-1\}\\sum\_\{t=2\}^\{T\}\\bigl\\\|\\mathbf\{g\}\_\{t\}^\{\(\\ell\)\}\-\\mathbf\{g\}\_\{t\-1\}^\{\(\\ell\)\}\\bigr\\\|\_\{2\}^\{2\}\.\(21\)This formulation is well\-motivated for*prefill*\(where many tokens are processed in parallel, one layer at a time\) and for batched decode\. For single\-token autoregressive decoding — the dominant inference mode on memory\-constrained edge devices — the memory access pattern is different\. The forward pass processes one token at a time and traverses all layers sequentially:

tokent:ℰ\(0\)→ℰ\(1\)→ℰ\(2\)→⋯→ℰ\(L−1\),\\text\{token \}t:\\quad\\mathcal\{E\}^\{\(0\)\}\\;\\to\\;\\mathcal\{E\}^\{\(1\)\}\\;\\to\\;\\mathcal\{E\}^\{\(2\)\}\\;\\to\\;\\cdots\\;\\to\\;\\mathcal\{E\}^\{\(L\-1\)\},\(22\)so consecutive memory accesses are*across layers for the same token*, not across tokens for the same layer\. The access sequence for two consecutive tokens under single\-token decode is:

whereσt\(ℓ\)\\sigma\_\{t\}^\{\(\\ell\)\}denotes the expert index selected for tokenttat layerℓ\\ell\. Each step requires loading a potentially different expert weight matrix\. Under this access pattern, the token\-level consistency loss reduces misses at step 5 relative to step 1 \(same layer, adjacent tokens\) but does nothing to reduce misses at step 2 relative to step 1 \(adjacent layers, same token\)\.

### D\.2Cross\-Layer Consistency and the Parameter Sharing Requirement

A natural extension is a*cross\-layer consistency loss*that encourages the same expert to be selected at consecutive layers for the same token:

ℒcross=1T​∑t=1T∑ℓ=1L−1‖𝐠t\(ℓ\)−𝐠t\(ℓ\+1\)‖22\.\\mathcal\{L\}\_\{\\text\{cross\}\}=\\frac\{1\}\{T\}\\sum\_\{t=1\}^\{T\}\\sum\_\{\\ell=1\}^\{L\-1\}\\bigl\\\|\\mathbf\{g\}\_\{t\}^\{\(\\ell\)\}\-\\mathbf\{g\}\_\{t\}^\{\(\\ell\+1\)\}\\bigr\\\|\_\{2\}^\{2\}\.\(23\)However, there is a critical prerequisite: cross\-layer routing consistency only reduces cache misses if the*same expert index at different layers refers to the same weight tensor in memory*\. In a standard MoE model where each layer maintains its own independent set of expert matrices\{ℰi\(ℓ\)\}i=1N\\\{\\mathcal\{E\}\_\{i\}^\{\(\\ell\)\}\\\}\_\{i=1\}^\{N\}, expert 2 at layerℓ\\elland expert 2 at layerℓ\+1\\ell\+1are entirely different weight tensors\. Routing consistently to expert 2 across both layers still requires two separate memory loads — no cache benefit is obtained\.

Cross\-layer consistency therefore becomes practically useful only under one of the following architectural conditions:

1. 1\.Full cross\-layer weight sharing\.Expertiiis the same weight matrix at every layer:ℰi\(ℓ\)=ℰi\\mathcal\{E\}\_\{i\}^\{\(\\ell\)\}=\\mathcal\{E\}\_\{i\}for allℓ\\ell\. Consistent routing to expertiiacross all layers requires only a single load, amortised overLLlayers of computation\.
2. 2\.Grouped weight sharing\.Layers are partitioned into blocks of sizeKK\(e\.g\.K=2K=2orK=4K=4\), and all layers within a block share the same expert pool\. Consistent routing within a block reduces loads; routing across block boundaries still requires new loads\.
3. 3\.Universal expert pool\.All layers draw from one shared set of expert matrices — a design explored in recent parameter\-efficient architectures\. Any cross\-layer routing consistency directly translates to cache reuse\.

### D\.3The Combined Objective

Under a grouped weight sharing architecture with block sizeKK, the full memory\-aware training objective combining token\-level and cross\-layer consistency would be:

ℒ=ℒCE\+λ​ℒcons\+μ​ℒbal\+γ​ℒcross,\\mathcal\{L\}=\\mathcal\{L\}\_\{\\text\{CE\}\}\+\\lambda\\,\\mathcal\{L\}\_\{\\text\{cons\}\}\+\\mu\\,\\mathcal\{L\}\_\{\\text\{bal\}\}\+\\gamma\\,\\mathcal\{L\}\_\{\\text\{cross\}\},\(24\)whereγ≥0\\gamma\\geq 0controls the cross\-layer consistency weight\. The two consistency terms are complementary and target different bottlenecks:ℒcons\\mathcal\{L\}\_\{\\text\{cons\}\}reduces misses across tokens at the same layer \(prefill and batched decode\), whileℒcross\\mathcal\{L\}\_\{\\text\{cross\}\}reduces misses across layers for the same token \(single\-token decode\)\.[Figure3](https://arxiv.org/html/2607.08780#A4.F3)illustrates the two access patterns and the corresponding consistency objectives\.

ℒcons\\mathcal\{L\}\_\{\\text\{cons\}\}ttt\+1t\{\+\}1t\+2t\{\+\}2t\+3t\{\+\}3L3L2L1L0Token axis\(a\)Token\-level consistency \(prefill / batched decode\)\. Arrows show penalised pairs underℒcons\\mathcal\{L\}\_\{\\text\{cons\}\}at layer L0\.ℒcross\\mathcal\{L\}\_\{\\text\{cross\}\}ttt\+1t\{\+\}1t\+2t\{\+\}2t\+3t\{\+\}3L3L2L1L0Token axis\(b\)Cross\-layer consistency \(single\-token decode\)\. Arrows show penalised pairs underℒcross\\mathcal\{L\}\_\{\\text\{cross\}\}for tokentt\.
Figure 3:The two memory access patterns and corresponding consistency objectives\. Blue arrows \(left\) show token\-level consistency targeting prefill and batched decode\. Red arrows \(right\) show cross\-layer consistency targeting single\-token autoregressive decode\. Both objectives are complementary and can be combined in[eq\.24](https://arxiv.org/html/2607.08780#A4.E24)\.
### D\.4Limiting Case: Pre\-defined Expert Trajectories

The consistency objectives proposed in this paper nudge the router toward stable routing decisions but still allow it to adapt freely\. A natural limiting case is to remove the router entirely and fix expert assignments before training begins — defining a deterministic*trajectory*σt=\(σt\(0\),σt\(1\),…,σt\(L−1\)\)\\sigma\_\{t\}=\(\\sigma\_\{t\}^\{\(0\)\},\\sigma\_\{t\}^\{\(1\)\},\\ldots,\\sigma\_\{t\}^\{\(L\-1\)\}\)through the network for every tokentt, whereσt\(ℓ\)∈\{1,…,N\}\\sigma\_\{t\}^\{\(\\ell\)\}\\in\\\{1,\\ldots,N\\\}is the expert index at layerℓ\\ell\.

The idea of fixed routing is not new:Rolleret al\.\[[2021](https://arxiv.org/html/2607.08780#bib.bib27)\]show that assigning tokens to experts via a deterministic hash function, with no learned router, produces surprisingly competitive models\. Our trajectory framing differs in two key respects\. First, assignments are coordinated*across layers*to define a complete network path per token, rather than independent per\-layer hashes\. Second, trajectories are chosen to maximise cache locality — consecutive tokens are assigned the same trajectory, so the entire sequence of expert weight loads for tokenttis identical to that of tokent−1t\{\-\}1, yielding zero cache misses within a coherent span\.

Under this design the cache miss rate during single\-token autoregressive decode is determined entirely by how often the trajectory changes between consecutive tokens — a quantity that is fully under the designer’s control rather than an emergent property of learned routing\. In the extreme case where all tokens in a document share the same trajectory, cache misses are eliminated entirely at the cost of routing flexibility\. This represents the upper bound on locality achievable by any routing strategy, and serves as a useful reference point for evaluating how close learned consistency objectives likeStickyMoEcome to the theoretical optimum\.

Whether a model trained with fixed locality\-maximising trajectories can match the quality of a model with learned routing is an open empirical question\. The hash routing results ofRolleret al\.\[[2021](https://arxiv.org/html/2607.08780#bib.bib27)\]suggest it is plausible, at least at small to moderate scale\. We leave a systematic investigation of this trade\-off to future work\.

### D\.5Implications and Future Work

The analysis above identifies a gap in the currentStickyMoEformulation: it is most directly beneficial for prefill and batched decode, while single\-token autoregressive decode — the dominant mode on memory\-constrained edge devices — is better served by cross\-layer consistency under a parameter\-sharing architecture\. We leave the empirical validation ofℒcross\\mathcal\{L\}\_\{\\text\{cross\}\}and its interaction with grouped weight sharing to future work, noting that the training\-time intervention principle extends naturally to this setting: just as token\-level locality must be instilled during training rather than retrofitted, cross\-layer locality requires architectural choices \(parameter sharing\) and training objectives \(ℒcross\\mathcal\{L\}\_\{\\text\{cross\}\}\) that are made from the outset\.