Module 8.1
A mixture-of-experts layer has several small feed-forward networks instead of one big one — but only the top few "light up" per token. Step through Nimbus's 4 tokens and watch a real router decide which experts fire.
The math
Without the balance term, the router can collapse onto a favorite expert early in training and never recover — the loss below penalizes exactly that.
The four tokens, together
| Expert | Used by | fᵢ (fraction of tokens) | Pᵢ (avg. gate mass) |
|---|
The code
def route(h, W_gate, k=2):
logits = h @ W_gate # one score per expert
probs = softmax(logits)
top_idx = np.argsort(probs)[-k:] # indices of the k largest
top_probs = probs[top_idx]
weights = top_probs / top_probs.sum() # renormalize to sum to 1
return top_idx, weights, probs
def load_balancing_loss(all_top_idx, all_probs, n_experts):
f = [np.mean([i in idx for idx in all_top_idx]) for i in range(n_experts)]
P = np.mean(all_probs, axis=0) # avg. gate mass per expert
return n_experts * np.sum(np.array(f) * P)This loss term gets added to the ordinary training loss with a small weight — enough to discourage collapse, not enough to override what the router is actually learning about which tokens suit which expert.
The Nimbus thread
Nobody assigns "billing questions go to expert 2" — the router learns it, because tokens like "quota" and "export" end up producing hidden states that the gate happens to score highly for the same expert, over and over across training. The table above is a toy version of exactly that pattern: experts 0 and 2 dominate for this batch of Nimbus tokens, expert 1 barely fires — the kind of imbalance the load-balancing loss exists to catch before it hardens into a permanently idle expert.
So what
More experts means more capacity, not more compute per token — only the top-k ever run.