Expert Routing

Unlike dense models that activate all parameters, expert routing enforces sparse activation. For each input token, the gating network selects only a small, fixed subset (e.g., top-2) of the total experts for computation. This creates a conditional computational graph where the FLOPs are proportional to the active experts, not the total parameter count, enabling massive models with feasible inference costs.

A naive router can collapse, always selecting the same few experts, leaving others unused (load imbalance). To prevent this, training employs auxiliary load balancing losses (e.g., based on expert assignment variance) and techniques like Noise Top-k Gating. This ensures uniform utilization of all experts, maximizing parameter efficiency and hardware throughput by distributing the computational load evenly across devices in an expert-parallel setup.

For hardware efficiency, experts process tokens in batched matrix operations. Expert capacity is a fixed buffer size defining how many tokens an expert can accept per batch. If more tokens are routed to an expert than its capacity, the excess tokens are dropped. Dropped tokens are typically passed through the layer unprocessed, creating a critical trade-off: higher capacity reduces drops but increases memory and compute waste; lower capacity risks losing information.

Through training, experts develop specialized skills. The router learns token-expert affinity, mapping specific linguistic, semantic, or functional features to particular experts. For example, one expert may specialize in mathematical symbols, another in proper nouns. This emergent specialization is key to the model's quality, as it allows different parts of the parameter space to become domain experts, which the router calls upon contextually.

In distributed inference (Expert Parallelism), experts are sharded across GPUs. Routing triggers all-to-all communication: tokens are scattered from all devices to the devices hosting their assigned experts, and results are gathered back. This communication overhead is a major bottleneck. Optimized systems use fused MoE kernels that combine routing, permutation, and computation to minimize latency, making network bandwidth a critical scaling factor.

The gating network itself adds inference latency. It is a small neural network (often a linear layer) that must process every token to produce routing scores. While computationally cheap compared to expert FFNs, its execution and the subsequent top-k selection & permutation logic introduce fixed overhead. In high-throughput serving, this cost is amortized across many tokens, but it remains a key difference from the fixed execution path of dense models.

A neural network architecture where the model is composed of multiple sub-networks called experts. A gating network (router) dynamically selects a sparse subset of these experts to process each input, enabling massive model sizes (e.g., trillions of parameters) with a manageable, conditional computational cost per token.

• Key Innovation: Decouples model capacity from computational cost.
• Example: The Switch Transformer and Mixtral 8x7B are landmark MoE models.

The small, trainable neural network component within an MoE model responsible for expert routing. It receives an input token's hidden state and outputs a set of scores (logits) used to select which expert(s) should process that input.

• Function: Implements the routing policy (e.g., Top-k).
• Training Challenge: Must be co-trained with the experts and often requires auxiliary load balancing losses to prevent expert collapse.

The two primary hard routing strategies that determine how many experts process each token.

• Top-k Routing: Selects the k experts with the highest router scores. The token's output is a weighted sum of these k expert outputs. Common in models like Mixtral (k=2).
• Switch Routing: A special case of Top-k where k=1. Each token is routed to exactly one expert, simplifying communication and implementation, as used in the Switch Transformer.

Critical techniques to prevent router collapse, where the gating network learns to favor only a small subset of experts, leaving others untrained and underutilized.

• Load Balancing Loss: An auxiliary term added to the training objective. A common method is the Expert Balancing Loss, which minimizes the squared coefficient of variation of the routing probabilities across experts.
• Noise Top-k: Adds tunable Gaussian noise to router logits before gating, encouraging exploration during training.

A model parallelism strategy designed for MoE models, where different experts are placed on different physical devices (e.g., GPUs).

• Process: After the local router decides token assignments, an All-to-All communication operation scatters tokens across the device mesh to their assigned expert's device. After processing, another All-to-All gathers the outputs.
• Challenge: The variable, sparse communication pattern can become a bottleneck, requiring optimized networking and overlapping of compute and communication.

The practical outcome of expert routing: only a fraction of the model's total parameters are activated per token. This sparsity enables efficiency but introduces unique inference challenges.

• Core Kernel: Sparse Matrix Multiplication is performed only on the weight blocks of activated experts.
• System Optimizations: Include Fused MoE Kernels (combining routing and computation) and Continuous Batching for MoE to handle dynamic, uneven workloads across experts in a production serving system like vLLM.