The Promise of Sparsity
Training a sparse Mixture of Experts (MoE) model is incredibly difficult. We attempted to scale our internal models up to an 8x7B architecture. The goal was to increase total parameter count (56B) while keeping active inference parameters low (only routing to 2 experts per token, or 14B active parameters).
Routing Collapse
During training, the gating network suffered from "routing collapse." Because neural networks are lazy, the router found that sending 95% of the tokens to just Expert 1 and Expert 2 minimized the loss the fastest. The other six experts were completely starved of gradients and became dead weight.
The Balancing Loss
To fix this, we had to introduce a strict load-balancing loss penalty to the training objective. This mathematical penalty forces the router to distribute tokens evenly across all experts across the batch.
# Simplified load balancing loss
def load_balancing_loss(router_probs, num_experts):
# router_probs shape: [batch_size, sequence_length, num_experts]
mean_probs = router_probs.mean(dim=(0, 1)) # Average probability per expert
# We want mean_probs to be uniform (1 / num_experts)
# The loss penalizes deviation from a uniform distribution
loss = num_experts * torch.sum(mean_probs * mean_probs)
return lossWe also added a small amount of Gumbel noise to the routing logits to encourage exploration early in training. MoE models are fantastic for scaling inference compute, but their training dynamics are incredibly fragile compared to dense models.