Temporally Extended Mixture-of-Experts Models
| Authors | Zeyu Shen & Peter Henderson |
| Year | 2026 |
| HF Upvotes | 3 |
| arXiv | 2604.20156 |
| Download | |
| HF Page | View on Hugging Face |
Abstract
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.
Engineering Breakdown
Plain English
This paper tackles a critical problem in Mixture-of-Experts (MoE) models: when they grow larger than available GPU memory, traditional expert switching (happening at nearly every token) becomes inefficient because memory prefetching and offloading optimizations break down. The authors propose using reinforcement learning's options framework to make expert switching temporally extended—meaning experts stay active for multiple tokens rather than switching constantly. Applied to a 20B parameter model with low-rank adapters, their method drops expert switch rates from over 50% to below 5% while maintaining 90% of the base model's accuracy on reasoning benchmarks (MATH, MMLU, MMMLU), making the model practical to run on constrained hardware.
Core Technical Contribution
The core novelty is applying the option-critic framework from reinforcement learning to MoE layers, where each layer learns a policy for when to switch expert sets rather than using fixed routing logic. The authors add a learnable controller to each MoE layer that treats expert selection as a hierarchical decision problem: which expert set to commit to, and for how long to use it before reconsidering. They incorporate deliberation costs into the RL objective to penalize unnecessary switches, and combine this with self-distillation rewards to maintain model accuracy. This is fundamentally different from prior MoE work, which either switches experts at every token or uses static expert assignments—here, the switch decision itself becomes a learned, continuous optimization problem.
How It Works
The mechanism works by extending MoE layers with an option-critic controller that maintains a commitment to an expert set across multiple tokens. At each layer, the controller observes the hidden state and decides either to continue with the current expert set or switch to a new one, with a learned penalty for switching. When a switch occurs, the system must load a new expert set into GPU memory; the controller learns to minimize these loads while preserving model performance. The RL training uses two losses: a deliberation cost that penalizes frequent switches (weighted by memory access cost), and a self-distillation loss where the model's outputs are compared to a teacher model running on unrestricted compute. Low-rank adapters (LoRA) are applied to make expert sets more compact and memory-efficient. During inference, the controller runs forward passes through its policy network to decide commitment duration, effectively batching multiple token computations per expert set load.
Production Impact
For engineers operating large MoE models on constrained hardware, this approach directly solves the memory thrashing problem that makes offloading ineffective—you can now run 20B+ parameter models on GPUs with insufficient memory to hold all experts simultaneously. The practical benefit is 10x reduction in expert switches (50% to 5%), which means GPU-to-CPU data movement and cache misses drop dramatically, improving end-to-end latency and throughput. Integration into a production pipeline requires: (1) retraining the option-critic controllers on your specific model and hardware, (2) benchmarking accuracy/latency trade-offs on your actual workload (the paper shows 10% accuracy drop, which may be unacceptable for some applications), and (3) careful tuning of the deliberation cost hyperparameter, which controls the switch/accuracy trade-off. The approach is most valuable when inference hardware is the bottleneck; if compute is abundant, traditional dense models or full-capacity MoE may be better.
Limitations and When Not to Use This
The method assumes that clustering tokens into expert-set groups doesn't degrade model quality too severely—the 10% accuracy loss on benchmarks is non-trivial and may be unacceptable for safety-critical applications or high-quality text generation. It's unclear how well this generalizes to other model architectures beyond the tested gpt-oss-20b, or whether the option-critic framework scales to much larger models (100B+) where the control overhead itself becomes significant. The paper doesn't address online learning or continual adaptation of the controllers, so they're frozen at inference time; dynamic adjustment based on actual hardware state or latency measurements isn't discussed. Finally, there's no analysis of worst-case scenarios: the paper doesn't show whether adversarial inputs or out-of-distribution prompts could trigger high switch rates and degrade the method's benefits.
Research Context
This work bridges two previously separate research areas: scaling MoE models (popularized by Google's Switch Transformers and Meta's work on sparse models) and hierarchical RL with options (Sutton et al., 2017). Prior MoE research focused on improving routing quality or load balancing; this paper shifts the problem to temporal coherence in expert selection, which is novel. The option-critic framework is well-established in RL but rarely applied to neural architecture optimization, so this represents a successful transfer of RL ideas to a classical deep learning scaling problem. The self-distillation + RL training approach opens doors for future work on jointly optimizing model capacity and inference efficiency, potentially inspiring similar hierarchical decision-making in other memory-constrained scenarios like multi-token speculation or dynamic batching.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::