🤖 machine learning

Temporally Extended Mixture-of-Experts Models

This paper proposes a temporally extended Mixture-of-Experts framework grounded in reinforcement learning's options framework, which utilizes a layer-wise controller to significantly reduce expert switching rates (from over 50% to below 5%) while preserving model accuracy, thereby enabling more memory-efficient serving and continual learning for large-scale models.

Original authors: Zeyu Shen, Peter Henderson

Published 2026-04-23

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Zeyu Shen, Peter Henderson

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Fidgety Chef"

Imagine a massive, high-end kitchen (the AI model) designed to cook any dish imaginable. To handle this, the kitchen doesn't have one giant chef; it has 32 different specialist chefs (called "experts").

The Old Way (Standard MoE): For every single word the AI says, the kitchen manager looks at the current sentence, picks the best 4 chefs for that specific word, fires the other 28, and then immediately hires 4 different chefs for the next word.
- The Result: The kitchen is in a state of constant chaos. Chefs are running in and out of the pantry every split second. If the kitchen is too small to hold all 32 chefs at once (which happens when models get huge), the manager has to keep running to the basement to fetch new chefs and send others back. This "fetching" takes time, slowing everything down and making the kitchen inefficient.

The New Idea: The "Stable Shift"

The Princeton researchers asked: Why change the team for every single word? Can't a team of chefs work together for a whole paragraph or a whole thought?

They introduced a new system called Temporally Extended Mixture-of-Experts.

Think of it like a TV Show Production:

The Old Way: The director changes the entire cast of actors for every single line of dialogue.
The New Way: The director assigns a specific "Scene" (an Option). For the duration of that scene, the same 4 actors stay on stage. They only swap out when the scene changes (e.g., moving from a kitchen scene to a courtroom scene).

How It Works: The "Smart Manager"

The paper introduces a lightweight Controller (a small AI brain) that sits on top of the main model. This controller acts like a Smart Stage Manager.

The Decision: Instead of picking a new team for every word, the manager looks at the flow of the conversation.
The Cost: The manager knows that swapping teams is expensive (it takes time and energy to load new chefs). This is called the "Deliberation Cost."
The Strategy: The manager learns a simple rule: "Only swap the team if the new topic is so different that the current team can't handle it, and the benefit of swapping is worth the cost of the delay."

If the AI is writing a story about a cat, the "Cat Team" stays on stage for the whole story. If the story suddenly switches to a rocket ship, the manager swaps in the "Space Team."

The Results: Less Chaos, Same Quality

The researchers tested this on a large AI model (gpt-oss-20b). Here is what happened:

Switching Rate: In the old system, the team changed 50% of the time (basically every other word). In the new system, the team changed less than 5% of the time.
Quality: Despite keeping the same team for much longer, the AI still got 90% of the original accuracy on hard math and logic tests.
The Benefit: Because the team stays put for longer, the kitchen doesn't need to run to the basement as often. This means:
- Faster Speed: Less time waiting for chefs to arrive.
- Cheaper Memory: You don't need to keep all 32 chefs in the main room; you can keep just the 4 active ones and store the rest in the basement, knowing you won't need them for a while.
- Future Proofing: As AI models get bigger and bigger (with thousands of experts), this method allows them to run on standard computers without crashing.

The "Deliberation Cost" Trick

The coolest part of the paper is a knob they can turn called the Deliberation Cost.

Low Cost: The manager is lazy and rarely swaps teams. The AI is very fast and memory-efficient, but might get slightly "stuck" if the topic changes too abruptly.
High Cost: The manager is very eager to swap. The AI is very flexible but slower.

This allows engineers to tune the AI based on their needs: "Do we want it super fast and cheap, or super flexible and precise?"

Summary

In short, this paper teaches AI models to stop micro-managing every single word. Instead, it teaches them to commit to a "team of experts" for a whole chunk of a conversation. This reduces the frantic "churn" of loading and unloading data, making giant AI models faster, cheaper to run, and easier to manage, all without losing their smarts.

1. Problem Statement

Current Large Language Models (LLMs) increasingly rely on Mixture-of-Experts (MoE) architectures to scale parameter counts while maintaining fixed inference compute costs. However, a critical bottleneck exists in memory-constrained serving and training:

High Switching Frequency: In standard MoEs, the router selects a sparse subset of experts for every single token. Empirical analysis shows that for frontier models (e.g., gpt-oss-20b, gpt-oss-120b, Qwen3-Next), the active expert set changes at nearly every token (switch rates >90%).
Inefficiency: When the total number of experts exceeds GPU memory, weights must be offloaded to host memory (RAM/SSD) and loaded on demand. The high frequency of switching renders optimizations like prefetching and caching ineffective, as the system cannot predict which experts will be needed next. This leads to significant latency and reduced throughput.
Missed Opportunities: The lack of temporal continuity prevents efficient "chunk-wise" training (where chunks of tokens share an expert set) and limits the ability to expand expert pools for continual learning without increasing per-token memory footprints.

2. Methodology: Temporally Extended MoEs

The authors propose reframing expert routing as a temporally extended decision-making problem using the Options Framework from Reinforcement Learning (RL).

Core Concept

Instead of selecting experts token-by-token, the model learns to select a set of experts (an "Option") that persists for a variable number of tokens. Switching to a new set incurs a deliberation cost ( $\eta$ ), representing the latency/overhead of loading new weights.

Technical Components

Semi-Markov Decision Process (s-MDP) Formulation:
- State: The LLM hidden state $h_t^{(\ell)}$ at layer $\ell$ .
- Option ( $\omega$ ): A binary mask selecting a subset of $\hat{k}$ experts out of $N$ total experts.
- Termination: A learned policy decides when to terminate the current option (switch experts) based on the state and current option.
- Reward: The goal is to maintain performance close to the base model while minimizing switches. The reward is defined as the negative Reverse KL Divergence between the student model (with controller) and a frozen teacher model (original base model).
Controller Architecture:
- A lightweight controller is added to each MoE layer.
- Inputs: Current hidden state $h_t$ and the current expert mask $\omega_{t-1}$ .
- DeepSets Encoder: Encodes the set of active expert indices into a permutation-invariant embedding.
- Heads:
  - Termination Head ( $\beta$ ): Predicts the probability of switching ( $d_t=1$ ).
  - Selection Head: When switching, selects a new set of $\hat{k}$ experts using a Plackett-Luce distribution (approximated via Gumbel-top- $k$ trick).
  - Critic Heads: Estimate State-Value ( $V_\Omega$ ) and Option-Value ( $Q_\Omega$ ) to guide learning.
Training via Option-Critic:
- The authors adapt the Option-Critic architecture with deliberation costs.
- Objective: Maximize expected return while penalizing switches. The gradient for termination includes a term $(Q_\Omega - V_\Omega + \eta)$ . The controller only switches if the expected gain in value ( $Q_\Omega - V_\Omega$ ) exceeds the deliberation cost $\eta$ .
- Intra-option Learning: The underlying LLM parameters (experts and router) are fine-tuned simultaneously (via LoRA) to adapt to the new routing constraints, effectively performing on-policy self-distillation.

3. Key Contributions

Design Philosophy: Introduced "Temporally Extended MoEs," identifying that excessive switching is a primary bottleneck for memory-efficient serving and that the Options Framework is the natural mathematical fit for solving this.
Algorithmic Framework: Formalized expert loading as an s-MDP and adapted the Option-Critic algorithm to train a lightweight controller that learns when to switch and which experts to load, balancing performance against switching costs.
Efficiency & Scalability: Demonstrated that existing pre-trained models can be converted to this architecture with minimal training (using LoRA), enabling memory-efficient serving, chunk-wise training, and continual learning without retraining from scratch.

4. Experimental Results

The method was evaluated on gpt-oss-20b (32 experts per layer, top-4 routing) using the Nemotron Post-Training Dataset.

Switch Rate Reduction:
- Baseline: Standard MoE switch rates were ~58% (switching almost every token).
- Proposed Method: Switch rates were reduced to <5% (and as low as 1.2% with higher deliberation costs) while maintaining a fixed active expert budget ( $\hat{k}$ ).
Performance Retention:
- Despite drastically reducing the number of switches, the model retained ~90% of the base model's accuracy on benchmarks like MATH, MMLU, and MMMLU.
- Example: On MATH with $\hat{k}=16$ and $\eta=0.02$ , accuracy was 64.0% (vs. 71.5% base), significantly outperforming static pruning baselines (e.g., Frequency selection: 53.5%, Random: 15.0%).
Comparison to Baselines:
- Outperformed static pruning methods (Frequency, Reconstruction Loss, Wanda) and random selection.
- Avoided the "catastrophic repetition" failure mode often seen in constrained MoE routing, maintaining coherent generation.
Temporal Continuity: Visualizations confirmed that expert masks now persist for long sequences of tokens, validating the "temporally extended" behavior.

5. Significance and Future Impact

Memory-Efficient Serving: By reducing switch rates, the active expert set becomes predictable. This allows systems to keep only the necessary $\hat{k}$ experts on GPU, offloading the rest, and significantly reducing VRAM requirements (e.g., ~37-55% reduction for gpt-oss-20b).
Training Efficiency: Enables "temporal chunking" where training batches can be processed with a fixed expert set, allowing inactive experts to be offloaded during forward/backward passes.
Continual Learning: Provides a mechanism to add new experts to the model pool without increasing inference latency or memory footprint, as the controller learns to route to them only when beneficial.
Principled Trade-off: The deliberation cost $\eta$ offers a tunable hyperparameter for practitioners to explicitly trade off between model capability (accuracy) and system efficiency (switching latency).

In conclusion, this work bridges Reinforcement Learning theory and LLM architecture, offering a principled path to manage the growing complexity of MoE models in resource-constrained environments.