Routing Networks & Gating

9 min read

Learn

Quick Reading

Estimated 9 mins

Prereq

Intermediate

Basic ML concepts helpful

Interactive

Static Playbook

Static guide & reference tables

The gating network as a small decision-maker

Every conditional-computation architecture — Mixture of Experts, Mixture of Depths, early-exit networks — relies on a small component that decides where a token's computation should go. This component, the router or gating network, is usually tiny compared to the thing it's routing to: often just a single linear layer followed by a softmax or top-k selection. Its output determines which expert processes a token, whether a layer is skipped, or which branch of a larger architecture gets used. Despite being small, the router's quality bounds the quality of the entire sparse system — a bad router wastes capacity no matter how good the experts behind it are.

Softmax gating and top-k selection

The canonical gating mechanism, introduced for Mixture of Experts, computes $g = \text{softmax}(W_g \cdot x)$ over all available experts, producing a probability-like distribution. Two variants are common: dense (soft) gating, where every expert's output is combined using its full softmax weight (expensive — defeats the purpose of sparsity), and sparse (top-k) gating, where only the k highest-scoring experts are kept and the rest are zeroed out before combining, often after renormalizing the kept weights so they sum to 1. Top-1 routing (a single expert per token) is the cheapest and simplest; top-2 is the most common choice in practice because it gives the router some redundancy and smoother gradients without doubling cost relative to top-1.

topk_gating.py

python

import torch
import torch.nn.functional as F

def topk_gate(x, gate_weight, k=2):
    """
    x: (batch, d_model)
    gate_weight: (d_model, num_experts)
    Returns: expert indices and renormalized weights, shape (batch, k)
    """
    logits = x @ gate_weight                       # (batch, num_experts)
    topk_vals, topk_idx = torch.topk(logits, k, dim=-1)
    topk_weights = F.softmax(topk_vals, dim=-1)     # renormalize over the chosen k only
    return topk_idx, topk_weights

Router collapse is the default failure mode

Left untrained with no corrective pressure, gating networks tend to collapse onto a small favored subset of experts early in training and never recover — a self-reinforcing loop where popular experts get more gradient signal, become better, and get chosen even more. This is why almost every production routing network pairs the gate with an auxiliary balancing mechanism, not because balance is intrinsically desirable, but because without it the sparse capacity the architecture paid for in parameters simply never gets used.

Load balancing techniques

Several mechanisms address router collapse. Auxiliary load-balancing loss adds a term, typically proportional to the product of expert selection frequency and average router probability, that's minimized when load is spread evenly across experts. Noisy gating (as in the original sparsely-gated MoE paper) adds learned noise to router logits during training, encouraging exploration of less-favored experts. Expert capacity limits cap how many tokens any single expert can accept per batch; tokens that overflow a full expert are either dropped (their layer output becomes zero / identity) or routed to the next-best expert. Capacity limits also matter for an entirely practical reason: they bound worst-case memory and compute per device in a distributed, expert-parallel deployment.

Gating design choices

Mechanism	Purpose	Trade-off
Top-1 routing	Minimal compute per token	Less redundancy, more collapse-prone
Top-2 routing	Smoother gradients, some redundancy	~2x expert compute vs top-1
Auxiliary balancing loss	Encourage even expert usage	Competes with the main task loss
Noisy gating	Exploration during training	Adds training-time stochasticity
Capacity limits	Bound per-expert load for serving	Dropped/overflow tokens lose some computation

Beyond MoE: gating as a general pattern

The same gating machinery generalizes well past MoE feed-forward layers. Switch Transformer simplified routing to top-1 for efficiency at scale. Mixture of Depths uses a gate to decide run-vs-skip per layer. Multi-modal architectures use gates to decide how much weight to give a vision branch versus a text branch per token. Even classical ensemble methods have a gating analogue — stacked models with a learned combiner are doing a soft version of the same thing. The unifying idea is always: a small, cheap function learns to allocate a scarce resource (expert capacity, depth, modality weight) based on the input, instead of allocating it uniformly by default.

What's Next?

Good gating is ultimately a search problem — picking the right structure to route through — and that same search framing shows up, at a different timescale, in how architectures themselves get discovered and tuned rather than hand-designed.

I build these systems professionally.

Whether it's a RAG pipeline, analytics migration, or AI workflow — let's talk.

Work With Me →