Temporally Extended Mixture-of-Experts Models

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.

• The paper introduces a novel RL-based controller that persistently allocates experts across tokens, reducing switch rates from about 55–58% to below 5%.
• It leverages the options framework in hierarchical reinforcement learning to model expert selection as temporally extended actions, balancing accuracy with hardware memory constraints.
• The method achieves up to 90% accuracy retention on large-scale LLMs, enabling memory-efficient inference and scalable expert pool expansion.

Temporally Extended Mixture-of-Experts: Reinforcement Learning for Persistent Expert Allocation

Introduction and Motivation

Mixture-of-Experts (MoE) architectures have become the principal approach for scalable capacity in large-scale LLMs, leveraging expert sparsity to decouple total parameter count from inference-time compute. However, current MoE implementations activate expert sets at nearly every token, inducing excessive expert-churn—a design misaligned with the realities of hardware limitations, specifically memory constraints. Once the expert pool outgrows GPU memory, token-wise switching fundamentally impedes offloading optimizations and memory-efficient inference pipelines.

The authors make a formal and algorithmic connection between MoE expert selection and the options framework in hierarchical reinforcement learning (HRL), framing expert set allocation as a temporally extended action. Integration of an option-critic controller allows the model to persist with a chosen set of experts across multiple tokens, incurring an explicit deliberation cost only when switching is deemed beneficial. This approach opens MoE serving and training to systematic memory-resource trade-offs, efficient chunked training, and scalable continual expansion of the expert pool.

Figure 1: Temporally extended control dramatically reduces switch rates in MoEs from over 50% to below 5% without significant accuracy loss, enabling practical memory savings and efficient deployment.

Formalization: MoE Routing as an Options s-MDP

The MoE expert routing problem is modeled as a semi-Markov Decision Process (s-MDP) in the options framework. Each possible binary expert mask (top-$k$ selection per layer) constitutes an “option,” with a policy over hidden states selecting the mask and a learned termination function deciding when to switch.

• Expert Mask as Option: The option set $\Omega^{(\ell)}$ for each layer $\ell$ is the set of $k$-hot expert masks over the $N$ available experts.
• Temporal Extension: The controller, parameterized via option-critic with RL, decides, per token, whether to persist with the current mask or incur a deliberation cost $\eta$ to sample and switch to a new mask.
• Rewards, Critic, Advantage: Controller training employs the per-token reverse KL from student (controlled MoE) to teacher (static MoE) as reward, sharing gradient pathways with both the controller and intra-option LLM parameters for joint optimization.

This reframing allows temporal chunking and direct trade-off control between performance (accuracy) and operational cost (switches/reloads).

Method: Option-Critic MoE Routing Controller

The controller receives token-wise layer hidden states (and the current option) and outputs:

• Termination Head: A stochastic Bernoulli parameterization, conditioned on current state and mask embedding, modulates the switch probability.
• Mask Selection Head: When switching, the next option is selected by sampling (via a Plackett-Luce distribution, implemented with Gumbel-top-$k$) from scores computed over all possible experts. This produces a new sparse binary expert mask.
• Critic Infrastructure: Option-value ($Q_\Omega$) and state-value ($V_\Omega$) heads are used for RL training, enabling advantage computation for both intra-option policy and controller updates.

Mask switch rates are directly regularized by the deliberation cost $\eta$; increasing $\Omega^{(\ell)}$0 increases persistence and reduces memory load at a known cost to token-level adaptivity.

Empirical Results

Evaluation is performed on gpt-oss-20b (32 experts/layer, top-4 active), comparing MoE variants on standard benchmarks (MATH, MMLU, MMMLU) and against expert pruning baselines (frequency, reconstruction, Wanda, random).

Key Observations:

• Switch Rate Reduction: Controller reduces average switch rates from baseline values ($\Omega^{(\ell)}$155–58%) to as low as $\Omega^{(\ell)}$2, even under strong accuracy constraints.
• Retention of Accuracy: For representative expert budget $\Omega^{(\ell)}$3 and low deliberation cost $\Omega^{(\ell)}$4, controlled models retain up to $\Omega^{(\ell)}$5 of the unpruned base model’s accuracy.
• Superior to Pruning: All static pruning baselines (incl. reconstruction and frequency) exhibit catastrophic accuracy drops when reducing active expert set size, demonstrating the unique efficacy of the temporal approach for aggressive memory reduction.
• Configurability: The deliberation cost can be tuned post-hoc to calibrate the trade-off between switch rate (i.e., memory / bandwidth cost) and retained task accuracy.

Figure 3: Controller training curves for $\Omega^{(\ell)}$6, showing rapid convergence in reward and monotonic switch rate reduction as the policy adapts to the deliberation cost.

Figure 5: Loss curves for $\Omega^{(\ell)}$7, indicating stable critic learning and effective value prediction over the RL trajectory.

Figure 7: Switch rate curves for $\Omega^{(\ell)}$8, confirming that after early rapid drop, switch rates remain stably suppressed for remainder of training.

Figure 2: Repetition and perplexity rates for $\Omega^{(\ell)}$9, ruling out error modes where switch suppression induces degenerate outputs.

Qualitative rollouts on reasoning tasks confirm that only the temporally extended controller maintains coherent, task-consistent sequence generation; baselines degrade into gibberish or repetitive outputs when forced to use severely restricted static expert sets.

Systemic and Theoretical Implications

The framework proposed has strong implications for both theoretical and practical systems:

• Principled Trade-off: Explicit deliberation cost as a regularizer integrates hardware reality into the RL objective, allowing direct optimization for latency/memory-constrained deployments.
• Extensibility: Temporal extension naturally enables continual expansion—new experts can be added on disk without increasing compute per token, provided they are only mapped into the mask upon switch.
• Chunk-wise Training: The persistent mask over several tokens allows training and inference systems to offload inactive experts, dramatically reducing memory pressure and enabling larger models on commodity hardware.
• Foundation for Memory-Efficient Pretraining: While designed as a post-training intervention, the philosophy suggests that temporally aware routing could be embedded into pretraining (possibly via RL or other differentiable routing policies) to produce models that inherently avoid token-wise expert churn.

Directions for Future Work

• Cross-layer Synchronization: The current per-layer policy yields staggered switch points; a joint multi-layer controller could synchronize chunk boundaries for maximal chunked computation and memory efficiency.
• Hardware-Aware Costing: Integrating real measured expert loading times or system latency into the cost function would make the trade-off more actionable for deployment.
• Expanded Evaluation: Analysis on larger prompt sets, code tasks, and more challenging reasoning benchmarks would further establish the approach’s robustness.

Conclusion

This work establishes a rigorous framework for temporally extended MoEs, aligning expert routing decisions with principles of temporal abstraction in RL. By coupling an option-critic controller with explicit deliberation cost, expert set persistence becomes tractable, yielding dramatic switch rate and memory reduction without commensurate accuracy loss. This structurally bridges the gap between current model design and the operational requirements of scalable, memory-efficient LLM deployment. The methodology fundamentally reframes MoE control as sequential decision-making under resource constraints, opening new avenues for large-scale, dynamic, and expandable model architectures.