MoSE: Mixture of Slimmable Experts for Efficient and Adaptive Language Models

AI summary

Overview

Research area: Machine learning / large language model architecture, specifically Mixture-of-Experts (MoE) conditional computation and efficient inference.

Technical level: Advanced. The paper assumes familiarity with transformer feed-forward layers, top-k routing, sparsely-gated MoE training, and test-time adaptation.

Scope: The paper proposes an MoE variant, MoSE, in which each expert has a nested, slimmable width, so that a single pretrained model can be executed at a continuous range of accuracy-compute trade-offs at inference time.

What This Paper Is About

Standard Mixture-of-Experts models save computation by activating only a few experts per token, but any expert that is selected is executed at full capacity. This makes the accuracy-versus-computation curve discontinuous: to change the cost, you must change how many experts you activate, not how much of each expert you use. The paper's goal is to add a second, finer-grained axis of conditional computation by allowing each selected expert to run at a variable width, and to show that this yields a smoother Pareto frontier without retraining or changing expert parameters.

Key Contributions

1. The MoSE architecture. Each expert is a transformer FFN with an expansion ratio of 4 and a nested (ordered) slimmable structure. A width w activates only the first m(w) = ⌈w · 4d⌉ intermediate hidden units by slicing W_up and W_down, so all widths share parameters. MoSE decouples expert selection (the router) from expert capacity (the execution width).

2. A training recipe for slimmable experts under sparse routing. For each mini-batch the model is run twice, at full width w_max and at one random width sampled w ~ Uniform(w_min, w_max), and gradients flow through the average of the two language-modeling losses. Standard MoE auxiliary objectives (a load balancing loss and a router z-loss) are retained, so no changes to the standard MoE pipeline are required.

3. Three inference-time width allocation modes. (i) Uniform width across all activated experts; (ii) router-conditioned widths derived from router probabilities normalized over the selected experts, raised to a sharpness parameter γ and rescaled to a per-token budget Γ; (iii) test-time training (TTT), which learns γ (a single shared scalar or layer-specific values) on a short calibration stream while every model and expert weight stays frozen.

4. Empirical validation across scales and adaptation regimes. Experiments span GPT2-Small (55M), GPT2-Standard (322M), GPT2-Medium (1B) and GPT2-Large (2.5B), finer-grained routing (E64A8, E128A8), zero-shot reasoning benchmarks, and continual-pretraining adaptation of DeepSeek-V2-Lite (16B).

Main Findings

• Pre-training is stable. Under a 15B-token budget on OpenWebText with GPT2-Small, MoSE closely tracks the convergence trajectory of MoE, indicating that width-adaptive execution does not destabilize or hinder pre-training.

• Full-width MoSE matches or beats MoE. At w = 1.0, GPT2-Small with 3B tokens reaches OpenWebText perplexity 39.02 versus 41.82 for MoE, and 154.85 versus 160.13 on WikiText-103. At 15B tokens the two are close (31.26 versus 31.03 on OpenWebText; 121.50 versus 120.48 on WikiText-103), and the authors report that MoSE with uniform execution already matches or improves on the MoE baseline across all settings.

• TTT gives the best results in Table 1 (reported at lower FLOPs). GPT2-Small at 3B tokens: OpenWebText 38.48, WikiText-103 153.24, LAMBADA accuracy 0.164, LAMBADA perplexity 110.129, WSC 0.5201. GPT2-Standard at 15B tokens: OpenWebText 20.38, WikiText-103 76.99, LAMBADA accuracy 0.338, LAMBADA perplexity 16.430, WSC 0.5531. The paper's convention is that the star symbol marks lower FLOPs at comparable or better performance.

• TTT shifts the Pareto frontier. Across GPT2-Small, GPT2-Standard and GPT2-Medium at E8A2 (n = 8, k = 2), MoSE with test-time training achieves lower perplexity than uniform-width execution at comparable MFLOPs per token, and the advantage remains stable as the base model grows.

• Pareto gains scale with pre-training data. Comparing 3B versus 15B tokens for GPT2-Small and GPT2-Standard at E8A2, more data shifts the frontier downward at every compute budget while the TTT advantage over other inference modes stays stable, which the authors read as evidence that width identification captures structure in the routing-compute trade-off rather than overfitting a data regime.

• Routing configuration matters, and MoSE exploits it. Increasing the number of activated experts and the width budget (E8A2, E8A4, E16A4 on GPT2-Small) consistently shifts the frontier downward; the paper states that width identification complements routing capacity.

• Routing transfers at inference. A checkpoint trained under E16A4 degrades gracefully when evaluated at E16A2, E16A3 and E16A4. Notably, E16A3 achieves essentially the same performance as E16A4 at reduced compute, suggesting diminishing returns from activating more experts beyond that point in this experiment.

• Learned width allocation beats probability-based allocation. In the inference-mode ablation, the normalized-probability mode (γ = 1.0) improves over uniform widths, confirming that router confidence carries useful allocation signal, but its Pareto frontier is uniformly dominated by the learned TTT variants. The gap is largest at low compute budgets; a shared γ captures much of the gain and layer-wise γ adds expressiveness.

• Learned γ transfers across datasets. Sharpness parameters calibrated on OpenWebText and reused unchanged on LAMBADA (GPT2-Standard, 15B tokens) give lower perplexity than the non-adapted variant across the full compute range, with layer-wise γ forming the dominant frontier. Accuracy curves show mild non-monotonicity, which the paper attributes to accuracy being a discrete metric sensitive to sampling noise and prediction discretization.

• Large-model and fine-grained gains. On GPT2-Large (2.5B parameters, 50B tokens), MoSE achieves 32.8% lower FLOPs at matched perplexity. In finer-grained routing, it yields 32.0% and 40.8% FLOPs reductions for E64A8 and E128A8 respectively.

• Continual pre-training on DeepSeek-V2-Lite (16B). Starting from a pretrained non-slimmable MoE checkpoint and continuing on OpenWebText, MoSE improves the inference-time Pareto frontier by 38.1% at comparable perplexity, corresponding to roughly 920 MFLOPs/token absolute savings.

• Gains carry over to downstream reasoning (GPT2-Medium). HellaSwag accuracy rises from 0.335 for MoE to 0.385 for MoSE (TTT) while MFLOPs/token fall from 402.7 to 272.5; PIQA rises from 0.675 to 0.710 while compute falls from 402.7 to 319.1; SIQA rises from 0.290 to 0.305 with compute falling from 402.7 to 319.1.

• Stated limitation. The authors report no full from-scratch training at frontier MoE scale; empirical validation covers GPT-style models trained from scratch, the 2.5B/50B setting, finer-grained routing, zero-shot benchmarks, and DeepSeek-V2-Lite (16B) continual pre-training.

Methodology in Plain English

The researchers start from a standard decoder-only transformer where feed-forward blocks are replaced by MoE layers containing n experts and a router. A router selects a sparse top-k set of experts per token, and the token's output is the router-weighted sum of the selected experts' outputs — exactly as in ordinary MoE.

The change is inside each expert. Because an FFN's cost is dominated by its 4x intermediate hidden dimension, they define a discrete set of width multipliers between w_1 > 0 and w_r = 1.0 and treat the expert's up-projection and down-projection as nested slices of the full matrices. Running an expert at width w simply uses the first ⌈w · 4d⌉ intermediate units, so all widths share the same parameters and a width of 1.0 recovers the standard MoE exactly.

To make one model good at every width, they train with two forward passes per step: one at full width and one at a randomly sampled width, backpropagating through the average of the two losses, plus the usual load-balancing and router z-loss terms. This keeps training overhead low relative to strategies that sample more than two widths.

At inference, three strategies assign widths to activated experts. The simplest uses one global width. The second turns router probabilities into widths: probabilities are normalized over selected experts, raised to a sharpness exponent γ, renormalized, scaled by a per-token budget Γ measured in full-expert widths, and clipped into the valid range. The third, test-time training, treats γ as the only learnable quantity — either one shared scalar or one per layer — and fits it on a short calibration stream by minimizing language-modeling loss under the budget, leaving all other weights frozen. The fitted γ is then reused for subsequent inputs.

Methodologically, the interesting move is that the router's own probabilities are used as a free signal of how much each expert matters, and that a single scalar controls how sharply compute concentrates on the most probable experts. The paper's comparisons are expressed as Pareto frontiers of validation perplexity against MFLOPs per token, produced by sweeping the width budget with model parameters fixed.

Why This Matters

Impact on research. MoE scaling has largely treated expert capacity as a fixed hyperparameter, so precision in the compute-quality trade-off has come only from choosing k. MoSE shows that a second, orthogonal axis — per-expert width — can be trained cheaply inside the standard MoE pipeline, and that a very low-dimensional adaptation (a scalar per layer, or even one shared scalar) is enough to steer allocation at inference. That reframes the question from "which experts" to "how much of each expert," and it also connects MoE research to the slimmable-network and once-for-all supernet literature that previously lived mostly outside large-scale MoE. The continual-pretraining results on DeepSeek-V2-Lite suggest the method can be bolted onto already-trained checkpoints rather than requiring fresh pre-training.

Real-world applications:

• Serving LLM inference under fluctuating traffic: one checkpoint can be run wide when latency budget allows and narrow under load, instead of maintaining several separately trained models.
• Edge and on-device deployment, where the same weights can be executed at reduced widths to fit tight memory or compute envelopes.
• Speculative-style decoding, which the authors explicitly raise: a lightweight draft pass could be realized by executing the same model at reduced widths, avoiding separate draft-model training or alignment.
• Agentic pipelines, where an agent could pick a width based on perceived task difficulty or uncertainty — small widths for easy steps, more compute for high-uncertainty ones.

Industry relevance. The reported savings are concrete: 32.8% lower FLOPs at matched perplexity on GPT2-Large, 38.1% on DeepSeek-V2-Lite (about 920 MFLOPs/token), and 32.0% to 40.8% in fine-grained routing regimes — all at comparable perplexity, with accuracy gains on HellaSwag, PIQA and SIQA. For anyone paying per-token inference costs, that is directly monetizable. The fact that adaptation requires only fitting a scalar while weights stay frozen also means the method is compatible with frozen, licensed, or otherwise immutable checkpoints.

Future Directions

• Frontier-scale from-scratch training. The authors name this explicitly as the remaining limitation: they do not report full from-scratch training at frontier MoE scale, only the 2.5B/50B setting, fine-grained routing, and continual pre-training on DeepSeek-V2-Lite (16B).
• Self-speculative decoding as an implementation. The discussion proposes using reduced widths as an implicit draft model for speculative decoding; whether this delivers wall-clock speedups rather than just FLOP reductions is not established in the paper.
• Task- and uncertainty-conditioned width control. The agentic use case — deciding width from perceived task difficulty or uncertainty — is raised in the discussion but not evaluated, so the mechanism for choosing width from a confidence signal remains open.
• Safety and robustness across execution widths. The impact statement notes that slimmable behavior adds a conditional axis and argues future safety-alignment and robustness evaluations should account not just for base-model behavior but for behavior across different execution widths and compute regimes. The paper reports no such evaluations.

Target Audience

This paper is most useful to MoE architecture researchers, LLM inference and serving engineers, and efficiency-focused practitioners who already understand transformer FFNs, top-k routing and MoE training losses. Those working on slimmable or elastic networks, model compression, or "train once, deploy many" supernets will find the connection to large-scale MoE valuable. Readers seeking an accessible introduction to MoE will need background reading first, since the methodology section uses router probabilities, width multipliers and budget constraints without much scaffolding. The code is available at https://github.com/tnurbek/mose, and the paper is released under a CC BY-NC-SA 4.0 license.

Authors’ abstract

Mixture-of-Experts (MoE) models scale large language models efficiently by sparsely activating experts, but once an expert is selected, it is executed fully. Hence, the trade-off between accuracy and computation in an MoE model typically exhibits large discontinuities. We propose Mixture of Slimmable Experts (MoSE), an MoE architecture in which each expert has a nested, slimmable structure that can be executed at variable widths. This enables conditional computation not only over which experts are activated but also over how much of each expert is utilized. Consequently, a single pretrained MoSE model can support a more continuous spectrum of accuracy-compute trade-offs at inference time. We present a simple and stable training recipe for slimmable experts under sparse routing, combining multi-width training with standard MoE objectives. During inference, we explore strategies for runtime width determination, including a lightweight test-time training mechanism that learns how to map router confidence/probabilities to expert widths under a fixed budget. Experiments on GPT-style models, various routing regimes, zero-shot downstream reasoning benchmarks, and continual pre-training adaptation of DeepSeek model show that MoSE matches or improves standard MoE at full width and consistently shifts the compute-quality frontier toward lower inference FLOPs. The code can be found at: https://github.com/tnurbek/mose.