Experiments / Does a sparse MoE beat its dense counterpart at the same active parameters?

Does a sparse MoE beat its dense counterpart at the same active parameters?

Three hypotheses, one survived.

PyTorch MoE Routing LLM
View notebook ↗

The question

I replaced the dense FFN in a 12-layer LLM with a from-scratch MoE layer: a router, top-k expert selection, and a load-balancing auxiliary loss. I went in with three hypotheses:

  1. MoE converges faster than the dense model at the same active parameters.
  2. The load-balancing loss spreads tokens evenly across experts; without it, a few experts dominate.
  3. More, finer-grained experts (4, 8, 16, 32 at proportionally smaller hidden dims, same active parameters) converge faster.

Setup

12 layers, embedding size 768, sequence length 256, trained on 768k tokens of Project Gutenberg text. Dense and MoE variants trained under the same budget.

Results

Only hypothesis 2 survived.

  • Without the auxiliary loss, expert 3 swallowed most tokens, climbing to around 70% of routing. With it, tokens spread across all four experts.
  • MoE did not converge faster than dense, despite carrying more total parameters.
  • All fine-grained configurations, from 4 to 32 experts, produced nearly identical loss curves.
Token routing across four experts during training without the load-balancing loss: one expert rises to about 70 percent.
Without the load-balancing loss.
Token routing across four experts during training with the load-balancing loss: the load stays spread.
With it.
Validation loss across expert configurations, which land on nearly identical curves.
Validation loss across expert configurations.

What I learned

Two of my three hypotheses failed, and the likely reason is scale. The results I was testing against come from runs on billions of tokens; at 768k tokens there is little room for experts to specialize. The experiment cleanly separated what held (the routing mechanism behaves as theory predicts) from what did not transfer down to small scale (the convergence and fine-graining benefits).