MoE v2: Expert Choice Routing + Shared Experts + Aux-Loss-Free Load Balancing.

Implements three key improvements to the Mixture of Experts architecture: expert choice routing (Zhou et al., 2022), shared expert slots (DeepSeekMoE), and aux-loss-free load balancing via trainable bias (DeepSeek-V3).

Key Innovations

1. Expert Choice Routing

Standard MoE: each token picks its top-K experts. Expert Choice: each expert picks its top-C tokens.

Standard:  token -> selects experts   (load imbalance risk)
Expert:    expert -> selects tokens   (perfect load balance)

This eliminates the need for load balancing auxiliary loss since each expert processes exactly C tokens by construction.

2. Shared Expert Slots

Some experts are "shared" (always active for every token), while others are "routed" (selected by expert choice). This ensures a base level of computation for every token while allowing specialization:

Output = SharedExperts(x) + RoutedExperts(x)

3. Aux-Loss-Free Load Balancing (DeepSeek-V3)

Traditional MoE uses auxiliary loss to encourage uniform expert utilization. However, this creates a trade-off: higher aux loss weight improves balance but hurts model quality.

DeepSeek-V3 introduces a bias-based approach that achieves load balance without auxiliary loss:

1. Add a trainable bias term b[i] to each expert's routing score
2. Router computes: scores = gate_logits + bias, then selects top-K
3. After each forward pass, update bias based on expert utilization: bias[i] -= lr * (utilization[i] - target_utilization) where target_utilization = 1/num_experts

This pushes the model toward uniform utilization without interfering with the main training objective. When an expert is overused, its bias decreases, making it less likely to be selected. When underused, bias increases.

Architecture

Input [batch, seq_len, input_size]
      |
      v
+----------------------------------+
|     Router (transposed scores)   |
|  score = softmax(W_r * x + b)^T |
|  Each expert picks top-C tokens  |
+----------------------------------+
      |
+-----+-----+
|             |
v             v
Shared      Routed
Experts     Experts
(always     (expert
active)     choice)
|             |
+-----+-----+
      |
      v
Output = shared + routed

Usage

# Standard with auxiliary loss (default)
model = MoEv2.build(
  input_size: 256,
  hidden_size: 512,
  num_shared_experts: 1,
  num_routed_experts: 4,
  tokens_per_expert: 4,
  load_balance: :aux_loss
)

# DeepSeek-V3 style: aux-loss-free with bias
model = MoEv2.build(
  input_size: 256,
  load_balance: :bias
)

# After training step, update bias for :bias mode
utilization = MoEv2.compute_utilization(router_logits, tokens_per_expert)
params = MoEv2.update_load_balance_bias(params, utilization, lr: 0.001)

References

• Zhou et al., "Mixture-of-Experts with Expert Choice Routing" (NeurIPS 2022)
• DeepSeek-AI, "DeepSeekMoE: Towards Ultimate Expert Specialization" (2024)
• DeepSeek-AI, "DeepSeek-V3 Technical Report" (2024) — aux-loss-free load balancing