Source-linked AI summary
Temporally Extended Mixture-of-Experts Models
Zeyu Shen, Peter Henderson
TL;DR
MoE expert switching can undermine memory-efficient serving once models exceed GPU memory. The paper introduces temporally extended MoE layers with an option-critic controller and reports sharply reduced switching while retaining benchmark capability. It presents this framework as a path toward managing expert-loading trade-offs in large MoEs.
Problem
MoE models switch experts nearly every token, creating loading costs that challenge serving when expert weights exceed available GPU memory.
Method
A lightweight option-critic controller with deliberation costs learns when to switch expert sets and which new set to load.
Results
Switching falls from over 50% to under 5%, or even 1%, on gpt-oss-20b with configurable performance trade-offs.
Takeaways & Limitations
Temporally extended MoEs can balance generation quality against expert-set transfer latency and may support memory-efficient serving and continual learning.
Takeaways & Limitations
The deliberation cost is a tunable hyperparameter rather than measured hardware latency, so deployment-specific cost-quality trade-offs remain ungrounded.
Abstract
from arXiv · showhide
Mixture-of-Experts models, now popular for scaling capacity at fixed inference speed, switch experts at nearly every token. Once a model outgrows available GPU memory, this churn can render optimizations like offloading and pre-fetching ineffective. We make the case that the options framework in reinforcement learning is a perfect match to tackle this problem, and argue for temporally extended mixture-of-experts layers. Building on the option-critic framework with deliberation costs, we add a controller to each layer that learns when to switch expert sets and which to load. By applying this to gpt-oss-20b with low-rank adapters and a self-distillation reward, our method reduces switch rates from over 50% to below 5% while retaining up to 90% of base-model accuracy on MATH, MMLU, and MMMLU. This shows that even existing pre-trained models can be converted to temporally extended MoEs with lightweight training, with the deliberation cost allowing model trainers to trade off switching rates against capability. We hope this opens a principled path, grounded in the options framework, for memory-efficient serving and continual learning in ever-growing MoE models.
1 Introduction
MoE routing keeps per-token computation sparse but can switch experts almost every token, making memory-constrained serving costly. The paper frames switching as a temporally extended decision and trains a lightweight controller to trade capability against switching cost.
- Sparse MoE routing keeps inference-time computation flat as total model parameters grow, but requires access to potentially many experts.
- Once experts exceed GPU memory, on-demand loading from host memory or disk adds latency, while current architectures largely ignore switching costs.
- The options framework models expert-set persistence and switching as temporally extended decisions with an explicit deliberation cost.
- The proposed controller learns when to switch expert sets and which set to load using option-critic optimization with deliberation costs.
- Over 50% to under 5% switching, or even 1%, is achieved on gpt-oss-20b with lightweight training and configurable capability trade-offs.
- Explicit deliberation costs are presented as a principled path for managing the trade-off between expert-loading overhead and model capability.
2 Preliminaries
The paper formulates temporally extended expert routing using semi-Markov decisions and options. Its option-critic controller optimizes both expert-selection policies and termination decisions, extending options when the current expert set is more valuable than switching.
- MDP, s-MDP, and Options: A semi-Markov decision process allows high-level actions to persist for variable durations while accumulating rewards before the next decision.
- MDP, s-MDP, and Options: An option combines an initiation set, an intra-option policy, and a termination function, while a policy over options selects which option starts.
- MoE Routing and Expert Masks: In the MoE formulation, a binary expert mask specifies allowed experts, and top-k routing is restricted to that mask.
- MoE Routing and Expert Masks: The binary expert mask is treated as the option, with k̂ allowed experts and k̃ activated experts.
- Option-Critic Architecture: The option-critic framework jointly optimizes intra-option policies and termination functions to maximize expected discounted return.
- Option-Critic Architecture: When the current option has higher value than switching, the termination update decreases its termination probability and extends its duration.
3 Motivation
Current frontier MoEs show little temporal continuity: expert sets change at nearly every token. This creates opportunities for temporally extended routing to reduce memory demands during serving and training while preserving fixed active computation.
- Current Mixture-of-Experts LLMs are not Temporally Extended: Three frontier open-source MoEs were evaluated on 100 prompts from each of ten categories, generating 256 tokens while recording expert activations.
- Current Mixture-of-Experts LLMs are not Temporally Extended: Average switch rates are close to 1 across all evaluated models, indicating expert selection changes at nearly every token.
- Missed Opportunities: Temporal continuity could reduce serving memory requirements by keeping only a fixed active expert subset instead of accommodating arbitrary per-token selections.
- Missed Opportunities: For gpt-oss-20b, keeping 16 experts reduces VRAM by approximately 4.7 GiB (37%), while keeping 8 reduces it by approximately 7.1 GiB (55%).
- Missed Opportunities: During training, fixed expert masks can partition responses into contiguous chunks so only the current mask’s experts participate in forward and backward computation.
- Missed Opportunities: Continual learning can add new experts without increasing per-token compute or active memory, while the controller learns when to route to them.
4 Method
The method gives each MoE layer a persistent expert mask and a lightweight controller that learns when to switch masks and which experts to load. Option-critic training with deliberation costs coordinates switching decisions with self-distillation rewards while fine-tuning the model.
- Temporally extended routing: Each layer maintains an expert mask that restricts routing to an allowed subset while a controller learns when and how to switch it.The controller simultaneously fine-tunes MoE parameters through intra-option policy updates.
- Controller architecture: Each option is an expert mask, represented with a DeepSets encoder that combines learned expert embeddings through an MLP.Separate encoders are used for each layer, and the resulting representation feeds controller heads.
- Options execution: Active options persist across tokens until termination, after which a newly selected option replaces the current one.The router masks disallowed experts before top-˜k routing.
- Controller architecture: The controller factorizes the joint control problem into independent per-layer controllers for tractability, with each conditioning on its layer hidden state and current mask.A joint mask over all layers is possible in principle, but the method instead uses L independent controllers.
- Option selection: The selection head produces candidate logits, samples ˆk experts with the Plackett–Luce distribution, and uses Gumbel-top-ˆk sampling for vectorized implementation.The sampled ordered tuple induces the expert-mask set used as the new option.
- Training objective: Deliberation costs make termination preferable only when an alternative sufficiently outweighs the current option, encouraging switches whose expected quality gain justifies their cost.The method uses per-token dense rewards and a reverse-KL reward from a frozen teacher to train the controller and student.
5 Experiments
The experiments evaluate temporally extended MoE controllers on gpt-oss-20b using option-critic training with deliberation costs, comparing them against pruning baselines across accuracy, switch rates, and temporal continuity. Results show improving rewards and perplexity, lower switch rates at higher deliberation costs, competitive benchmark accuracy, and visibly more continuous expert selection, with single-run stability limiting the claims.
- Experimental setup: The setup uses gpt-oss-20b with 24 layers, 32 experts per layer, top-4 routing, and training on four 140GB H200 GPUs.The model is dequantized to bf16 for training, while retaining native MXFP4 quantization.
- Training dynamics: Higher deliberation costs yield lower converged switch rates, while rewards increase and perplexity decreases during training, with clearer improvements for k̂ = 8.The switch rate initially decreases as the value networks learn and then stabilizes at a level determined by η.
- Benchmark evaluation: With k̂ = 16 and η = 0.02, the controller achieves accuracy close to the unpruned base model and substantially outperforms all pruning baselines across benchmarks.The reported trade-off between performance, deliberation cost, and mask size may improve with a full post-training run.
- Caveat: One training run is reported per configuration, so curve variance reflects within-run stochasticity rather than across-run training stability.The authors claim only that at least one run produced a well-performing temporally extended MoE.
- Temporal continuity: The trained controller produces significantly stronger temporal continuity in expert selection across layers and deliberation costs.The analysis records expert masks for 256 generated tokens and reports that different layers can exhibit different continuity levels.
6 Discussion and Conclusion
The paper introduces temporally extended MoEs as an options-based approach to dynamic expert loading and reports a quality–latency trade-off. It presents the framework as a first step toward MoE architectures that minimize expert switching by design.
- Contribution: The framework addresses dynamic expert loading by applying the options framework to temporally extended MoE models.The method balances generation quality against the latency cost of transferring expert sets.
- Future direction: The authors propose making temporal extension a core objective during post-training and potentially during pre-training.They suggest that inherently temporally extended architectures could minimize expert switches by design.
A1 Related Works
Related work addresses MoE efficiency through pruning, caching, prefetching, and offloading, while this paper frames expert-mask selection as temporally extended control. The options framework supplies the formal connection between switching costs and semi-Markov decisions.
- MoE architectures: Modern MoE models increase parameter capacity while activating only a sparse subset of experts per token.Examples include models with 128 or 512 experts per layer but only 4 or 10 active experts per token.
- MoE efficiency: Existing MoE efficiency methods mainly prune experts or use caching, prefetching, and offloading-aware serving strategies.These approaches reduce the expert set or predict which experts to keep on GPU and fetch from host memory.
- Options and s-MDPs: The paper casts expert-mask selection as a temporally extended control problem in which switching can be regularized by an explicit latency cost.This connects expert-mask termination to options and induces a semi-Markov decision process.
- Related language-model work: Prior work also applies options or hierarchical reinforcement learning to language-model training and reasoning.The cited examples include macro-actions for RLHF, hierarchical GRPO, and internal reinforcement learning.
A2 Additional Implementation Details
The implementation provides a full option-critic training procedure with per-layer critics, termination decisions, importance weighting, and advantage normalization. Additional details describe reconstruction-loss baseline selection and the computational compromise required for 32-expert layers.
- Training procedure: The full algorithm trains option selection, termination, critic, and intra-option policy components using token-level rewards and teacher-mixed sampling.The inputs include deliberation cost, discount and GAE parameters, learning rates, and a teacher model.
- Controller architecture: Each MoE layer has an independent controller with DeepSets encoding and separate termination, option-value, expert-selection, and state-value heads.The controller uses learned expert-set embeddings and hidden dimension 1024 for most heads.
- Gradient computation: The intra-option policy gradient uses importance weights and normalized advantages, with the shared advantage computed from Monte Carlo returns.RMS normalization preserves advantage signs while stabilizing scale, according to the implementation discussion.
- Baseline implementation: The reconstruction-loss baseline replaces exhaustive subset search with greedy forward selection for N = 32 experts.Exhaustive enumeration was feasible for N = 8 in prior work but is computationally infeasible in this setting.
A3 Additional Experimental Results
Additional analyses show that controller training learns value estimates, adapts switching behavior to deliberation cost, avoids repetition, and produces temporally continuous options across layers.
- Loss curves: Value loss steadily decreases, indicating that the value and option-value heads learn state and state-option values effectively.
- Additional switch rate curves: The 95th-percentile switch rate and switch-rate variability increase during training, indicating that the termination head distinguishes when to switch.
- Training stability: Repetition remains stable within a healthy range for k̂ = 8 and k̂ = 16, showing that constrained routing avoids catastrophic repetition during training.
- Temporal continuity with controller: Options demonstrate significant temporal continuity across layers and deliberation costs η = 0.03 and η = 0.04.
A4 Limitations and Future Directions
The paper presents temporally extended MoEs as feasible, but identifies deployment, pre-training, cross-layer coordination, evaluation coverage, and attribution as open limitations.
- From philosophy to deployment: Current experiments validate graceful degradation in generation quality rather than end-to-end memory savings, which remain substantial systems work.
- From philosophy to deployment: The deliberation cost is a tunable hyperparameter rather than measured hardware latency, limiting direct actionability for a specific deployment target.
- Realizing temporal continuity in pre-training: The framework performs post-training on a pre-trained MoE, leaving temporal structure in the pre-training objective as future work.
- Per-layer vs. cross-layer options: Per-layer independent options make learning tractable, but simultaneous cross-layer switching would better support single-event loading and chunk-wise training.
- Evaluation scope: Evaluation covers MATH, MMLU, and MMMLU using 200 randomly selected questions per benchmark, excluding code, long-form instruction following, and open-ended conversation.
- Disentangling temporal extension from self-distillation: The study does not isolate temporal extension from on-policy self-distillation, leaving their relative contributions unresolved.
A5 Examples
The examples section examines a MATH number-theory rollout and compares methods on computing an LCM, while illustrating the factorial-based expression used in the solution.
- Example setup: The rollout studies a MATH number-theory problem with k̂ = 16 experts, temperature 0.5, and a 2048-token limit.The method uses deliberation cost 0.02 and the checkpoint at training step 120.
- LCM objective: The example compares methods by computing the least common multiple of the eight factorial-sum terms.
- Factorial reduction: Each term satisfies i! + (i + 1)! = (i + 2) · i! for i = 1,…,8.
- LCM computation: The solution proceeds by factoring the terms and taking the maximum exponent of each prime across them.