Lost in Latent Space

The Gate That Sells the Model

For a brief stretch in 2021, the largest publicly disclosed neural network in the world had 1.6 trillion parameters and ran, per token, on roughly 1/64th of them. Google’s Switch Transformer was not really one model. It was sixty-four feed-forward networks sharing a hallway, and a small bouncer at the door deciding which one each token would visit. The bouncer — a single linear layer feeding a softmax over experts — was trained jointly with everyone else, and it had to learn, in a strict sense, the politics of the building.

This is the trade that mixture-of-experts architectures make. A dense transformer pays for every parameter on every token. An MoE pays for the routing decision, plus the parameters of whichever expert was picked. If the routing is sharp enough, the model can be enormous in parameter count while remaining cheap in floating-point operations. The economics work because activation memory and FLOPs scale with active parameters, while capacity scales with total parameters. Sparse conditional computation, in principle, gets you the second without paying for the first.

In practice, the routing is the entire game.

The standard formulation looks deceptively simple. Given a token’s hidden state, compute scores over E experts, take the top k (almost always k=1 or k=2), and dispatch the token to those experts. The experts’ outputs are weighted by the gate’s probabilities and summed. Each expert is just an ordinary feed-forward block — the same two-layer MLP that appears in a dense transformer — replicated E times with independent weights. There is nothing exotic about the machinery. The exotic part is what happens when sixty-four feed-forward networks share a hallway and a flawed bouncer.

The pathologies are predictable and severe. Early in training, one expert tends to be marginally better than the others; the gate discovers this and sends it more traffic; that expert improves further from the extra signal; the gate sends it everything. Within a few hundred steps a 64-expert model can collapse into a 1-expert model wearing a costume. Researchers call this routing collapse. The standard countermeasure is a load-balancing loss — an auxiliary term that penalizes the gate when the expected fraction of tokens routed to each expert deviates from uniform. The loss is added to the training objective with a small coefficient, typically around 0.01. It works well enough that nearly every production MoE uses some descendant of it.

What the load-balancing loss buys, at heart, is a negotiation between two objectives the model would otherwise resolve violently. The language modeling loss wants the gate to send each token to the expert that handles it best. The balance loss wants the gate to distribute tokens evenly so no expert is starved of gradient. The equilibrium is a routing function that is almost-but-not-quite content-aware — sharp enough to specialize experts to different linguistic regimes, soft enough that all experts stay alive.

The architectures that have actually shipped — Switch, GShard, Mixtral, DeepSeekMoE — are largely variations on how to make this negotiation more graceful. Expert dropout, capacity factors that cap how many tokens any one expert can accept per batch, top-k routing with k=2 so no single mistake is fatal, and recently shared-expert designs where a fixed common expert handles the “easy” tokens while routed experts specialize on the rest. Each is an engineering response to the same underlying fact: that a gate trained by gradient descent is a self-reinforcing decision-maker, and self-reinforcing decision-makers, left to themselves, become tyrants.

The trillion-parameter regime is, mechanically, the regime of well-managed gates. Everything else is a feed-forward layer.