Gating Network (Router)

The router computes a logit score for every expert for a given input token. These raw scores are then typically normalized using a softmax function to produce a probability distribution over the available experts. The selection is then made based on these probabilities, most commonly by choosing the top-k experts with the highest scores. This mechanism allows the model to learn complex, non-linear assignment policies.

The gating network is the architectural component that directly enables sparse activation. By selecting only a small subset of experts (e.g., 2 out of 64) for each token, the router ensures that the total computational cost is proportional to the number of active parameters, not the total. This is the fundamental principle that allows MoE models to scale to trillions of parameters without a proportional increase in inference FLOPs.

An untrained router tends to collapse to a state where it always selects the same few experts, starving the rest. To prevent this, training incorporates auxiliary loss functions (load-balancing losses) that penalize uneven assignment. A common approach is to minimize the coefficient of variation of the token counts across experts. This ensures all parameters are utilized during training and prevents a small subset of experts from becoming computational bottlenecks.

The router is trained jointly with the experts via standard backpropagation. The gradients flow through the discrete routing decision using techniques like the straight-through estimator. Over time, a phenomenon called token-expert affinity emerges, where the router learns to consistently send tokens with specific syntactic or semantic properties to the same experts. For example, one expert may specialize in prepositions while another handles named entities.

Gating networks can operate in two modes. Hard routing makes a discrete, binary decision to send a token to a specific expert (e.g., via top-k or argmax), which is computationally efficient. Soft routing computes a weighted average of all expert outputs based on the router's probability scores, which is dense and computationally expensive. Modern MoE models like Mixtral and the Switch Transformer exclusively use hard routing for its sparsity benefits.

The gating network itself introduces a non-trivial latency cost. The router must compute scores for all experts, perform a top-k selection, and then orchestrate the all-to-all communication that sends tokens to the correct expert devices. In distributed systems, this communication overhead can dominate the total inference time if not carefully optimized with fused kernels and efficient network primitives.

The most common routing strategy where the gating network selects the k experts with the highest scores for each token. For a given input token x, the router computes logits, applies a softmax, and keeps only the top-k values, setting the rest to zero. This enforces strict sparsity—only k experts are activated per token, regardless of the total number of experts. Typical values are k=1 (Switch Transformer) or k=2 (Mixtral). The choice of k represents a direct trade-off between model capacity and computational cost.

Without intervention, gating networks often collapse into a degenerate state where they route most tokens to a single expert, defeating the purpose of the MoE architecture. Load balancing is enforced through an auxiliary loss function added to the primary training objective. This loss penalizes uneven token distribution across experts, typically using the coefficient of variation of expert assignment counts. The auxiliary loss coefficient is a critical hyperparameter—too high and it overrides useful routing specialization, too low and experts become underutilized.

A refinement of standard top-k gating that adds tunable Gaussian noise to the expert logits before the top-k selection. The noise term is typically parameterized as Noise = Normal(0, 1) * Softplus(W_noise * x), where W_noise is a learnable weight matrix. This serves two purposes: it acts as an exploration mechanism during training, preventing premature routing convergence, and it provides a natural form of load balancing by occasionally perturbing the router away from overused experts. The noise is typically disabled during inference.

Expert capacity defines a hard limit on how many tokens can be routed to a single expert in one forward pass. It is calculated as: Capacity = (tokens_per_batch / num_experts) * capacity_factor. When an expert reaches capacity, excess tokens are dropped—they bypass the expert entirely and pass through a residual connection. The capacity factor (typically 1.0–1.5) provides a buffer. Dropped tokens represent a direct performance-accuracy trade-off: higher capacity factors reduce dropped tokens but increase peak memory usage and computational imbalance.

In distributed MoE inference, different experts reside on different GPUs. After the router assigns tokens to experts, an all-to-all collective communication pattern is executed: tokens are scattered from their current device to the device hosting their assigned expert, processed, and the results are gathered back. This communication is often the primary bottleneck in MoE inference latency. Optimized implementations use fused MoE kernels that combine routing, permutation, and sparse matrix multiplication into a single operation to minimize kernel launch overhead and memory movement.

Introduced by the Switch Transformer architecture, this simplified routing strategy sends each token to exactly one expert (k=1). The gating network computes logits and selects the expert with the maximum score via argmax. This maximizes sparsity and computational efficiency, enabling models to scale to trillions of parameters with constant per-token FLOPs. The trade-off is that each token only benefits from a single expert's specialization, which can limit representational capacity compared to top-k>1 routing. Switch Transformers rely heavily on auxiliary loss for load balancing.