Top-k Gating

Top-k Gating enforces sparse activation, where for each input token, only a fixed, small number (k) of experts are selected from a large pool. This creates a conditional computation graph where the total FLOPs are proportional to k * (parameters per expert) rather than the full parameter count. For example, in a 1-trillion parameter model with 128 experts and k=2, only about 1.6% of parameters are active per token.

It performs hard routing by selecting the experts with the top k scores from the gating network. The selected experts process the token independently, and their outputs are combined, typically via a weighted sum based on the router scores. This is distinct from soft routing or mixture models where outputs are always blended. The discrete assignment enables efficient expert parallelism on distributed hardware.

A naive Top-k Gating network can collapse, always selecting the same few experts. To prevent this, load balancing techniques are critical:

• Auxiliary Loss (Load Loss): An added training loss term that penalizes the variance in routing probability across experts.
• Noise Top-k Gating: Adds tunable Gaussian noise to logits before the top-k operation to encourage exploration during training.
• Capacity Factor: A buffer hyperparameter that sets a soft limit on tokens per expert to prevent overflows and dropped tokens.

While sparse, Top-k Gating introduces unique overheads:

• Router Latency: The forward pass of the small gating network and the top-k selection operation.
• All-to-All Communication: In expert-parallel setups, selected tokens must be scattered to the devices hosting their experts and gathered back, incurring significant network cost.
• Kernel Fusion: Performance depends on fused MoE kernels that combine routing, permutation, and the sparse matrix multiplications of multiple experts to minimize memory movement.

Successful training leads to token-expert affinity, where experts develop specialized skills. For example, in language models, different experts may activate for:

• Specific syntactic structures (e.g., questions, negations).
• Semantic domains (e.g., scientific, legal, or conversational vocabulary).
• Linguistic features (e.g., function words, named entities). This emergent specialization is key to the model's efficiency and performance gains.

Deploying Top-k Gating requires tuning key parameters that trade off quality, speed, and memory:

• k (Number of Experts Selected): Typically 1 or 2. Higher k improves quality but increases computation and communication costs linearly.
• Capacity Factor: A multiplier (e.g., 1.25) on the theoretical minimum expert capacity. Higher values reduce dropped tokens but increase memory and compute for padding.
• Number of Experts: Ranges from 8 (e.g., Mixtral 8x7B) to 128+ in massive models. More experts increase model capacity and potential specialization but complicate parallelism.

Specialized, low-level GPU kernels are critical for performant top-k gating. Fused MoE Kernels combine multiple operations into one.

• Operation Fusion: Merges the routing logic, token permutation, and the sparse matrix multiplications across experts into a single kernel launch.
• Benefits: Dramatically reduces memory movement and kernel launch overhead, which are major bottlenecks.
• Libraries: Available in frameworks like FlashAttention and xFormers, and are a core component of optimized inference servers like TensorRT-LLM and Triton.