Sparse Mixture-of-Experts with
Gradient-Aware Expert Lifecycle

1. Overview

MoE-Stream augments the Stream SSM backbone with sparse Mixture-of-Experts feed-forward layers and gradient-aware expert lifecycle management. Unlike Stream's minimalist design (which has no MLP at all), MoE-Stream adds specialized capacity through expert networks while maintaining Stream's clean single-loss training and O(n) SSM recurrence.

2. Architecture

2.1 MoE-Stream Block

Each block stacks an SSMBlock (identical to Stream's core recurrence) followed by a sparse MoE FFN, with layer-normalized residual connection:

2.2 MoE Feed-Forward Network

The MoEFFN replaces Stream's absent MLP with a top-k sparse mixture of expert sub-networks. Each expert is a 2-layer MLP with GELU activation:

Routing is performed by a learned linear gate followed by top-2 selection with normalized weights.

2.3 Config

3. Gradient-Aware Expert Lifecycle Management

The core innovation in MoE-Stream is the expert lifecycle system, which distinguishes rare but valuable experts from truly useless ones using gradient-based impact measurements. This section details each component.

3.1 Expert Value Scoring

Each expert $i$ maintains two exponential moving averages updated every forward pass:

The frequency EMA tracks how often the expert is selected. The impact EMA uses backward hooks on each expert's output layer to compute per-token loss reduction: $-(\partial \mathcal{L} / \partial y) \cdot y$, which is a first-order approximation of how much the expert's output contributes to reducing the loss. The combined value score is:

3.2 Impact Hooks

Gradient-based impact is captured via PyTorch backward hooks registered on the last Linear layer of each expert:

3.3 Birth-Age Survival Boost

Recently replaced experts receive a temporary value score bonus to prevent immediate re-death before they have a chance to learn:

This provides a score boost of up to 0.35 at birth, decaying linearly to 0 once the expert has seen 2× the minimum token threshold.

3.4 Dead Expert Detection

The `find_dead_experts` method identifies candidates for replacement using an adaptive per-layer threshold:

Safety constraints:

• Experts must have seen at least expert_min_tokens total tokens (probation period)
• The single best expert in each layer is never replaceable
• At most $⌊n_{experts} / 2⌋$ experts can be replaced in a single check

3.5 Expert Replacement

When an expert is identified for replacement, the system clones the best-performing expert's weights with small Gaussian noise rather than random initialization:

3.6 Exploration Bias

After replacement, an exploration bias is injected into the router logits to encourage the new expert to receive tokens during early training. This bias decays exponentially (factor 0.99 per training step), allowing the router to naturally integrate the new expert over time.

4. Training

4.1 Loss Function

MoE-Stream maintains Stream's clean single-loss approach with one additional term for load balancing:

where $f_i$ is the batch-local fraction of tokens routed to expert $i$, and $P_i$ is the average router probability. The load balancing loss uses batch-local counts (not cumulative history) to ensure responsiveness to recent routing patterns.

4.2 Post-Backward Lifecycle Update

After each backward pass, `flush_grad_impacts()` is called via a hook registered on the loss tensor, ensuring that gradient-based impact EMAs are updated immediately after gradients are computed but before the next forward pass.

4.3 Training Config (stream_light equivalent)

5. Test Report

The following sections document smoke tests, utilization logging, and simulation results for the MoE-Stream expert lifecycle system.

5.1 Expert Utilization Log Format

During training, expert utilization statistics are logged at evaluation intervals. The `log_expert_utilization()` method produces per-layer, per-expert reports:

5.2 Value Score Distribution

The value score distribution across experts in a 4-expert, 4-layer configuration typically shows a bimodal pattern: 2-3 high-value core experts (score > 0.5) and 1-2 low-value candidates for replacement (score < threshold). The gradient-based metric naturally captures the distinction between frequency and impact.

5.3 Replacement Safety Validation

The replacement system includes multiple safeguards verified through simulation:

5.4 Load Balancing Validation

The load balancing loss uses batch-local token fractions:

This design avoids the hysteresis of cumulative count-based balancing, ensuring the router can rapidly adapt to changing expert quality after replacements. The coefficient $\alpha = 0.01$ provides gentle balancing pressure without overwhelming the prediction loss.

6. Architecture Comparison

7. Implementation Details

7.1 Code Structure

MoE-Stream is implemented as a standalone module in moe_stream.py with no dependencies on Stream's model.py or VECTOR's model_vector.py. The SSMBlock is duplicated for independence, ensuring each variant can evolve independently.

7.2 Key Components

7.3 Configuration

The MoEStreamConfig dataclass provides all hyperparameters with sensible defaults (8 experts, top-2, min 5000 tokens before replacement). Training is orchestrated through train.py with model_type='moe_stream' and CLI overrides via configurator.py.

8. Pending Work & Known Issues