Noisy Top-k Gating for Load Balancing

A standard top-k gating mechanism makes routing decisions deterministically for a given input and set of model weights. While this sounds desirable, it can lead to routing patterns that quickly become entrenched during training. Certain experts may consistently receive the highest scores from the gating network, attracting the majority of tokens, while others are starved of data. This imbalance hinders the model's ability to develop a diverse set of specialized experts and can even cause some experts to collapse, contributing nothing to the final output.

Noisy top-k gating introduces a simple and computationally inexpensive solution to this problem: add a small amount of random noise to the gating network's output logits before selecting the top-k experts. This technique acts as a form of regularization, encouraging the model to explore different routing paths and preventing it from relying too heavily on a small subset of experts.

The Mechanism of Noisy Gating

The core idea is to perturb the deterministic scores produced by the gating network. For each token, after the gating network calculates the logits (the raw, unnormalized scores) for each expert, we add a random value sampled from a noise distribution, typically a Gaussian distribution.

The mathematical formulation builds directly on the standard gating process. Let $h(x)$ represent the vector of logits produced by the gating network for an input token $x$, where $h(x) = x \cdot W_g$ and $W_g$ is the gate's weight matrix. In noisy gating, we compute a modified set of logits, $h_{noisy}(x)$:

$$h_{noisy}(x) = h(x) + \text{Noise}$$

The noise term is usually drawn from a normal distribution with a mean of zero. The standard deviation of this noise is a tunable hyperparameter. A common implementation, proposed in the original Sparsely-Gated MoE paper, scales the noise using a separate, learnable weight matrix, $W_{noise}$:

$$\text{Noise} = \text{StandardNormal}() \cdot \text{softplus}(x \cdot W_{noise})$$

Here, StandardNormal() generates random values from $\mathcal{N}(0, 1)$, and the softplus function ensures the scaling factor is always positive. The top-k selection is then performed on these noisy logits. This process is only active during training. During inference, the noise is disabled to ensure deterministic and stable outputs.

The diagram below illustrates how this addition of noise can alter a routing decision.

In standard gating, Experts 1 and 3 are chosen. With noise, the logit for Expert 3 is perturbed just enough to make Expert 4 the second-highest choice, thus altering the routing path.

Impact on Load Balancing and Training

The primary benefit of adding noise is the improvement in load distribution across experts. By "shaking up" the scores, tokens that are on the margin between being routed to a popular expert versus a less popular one are occasionally sent to the latter. This prevents any single expert from becoming a bottleneck and ensures all experts receive a sufficient variety of training examples to learn meaningful specializations.

This mechanism directly complements the auxiliary load balancing loss discussed in the previous chapter. While the loss function penalizes imbalance, noisy gating proactively discourages it from forming in the first place. The result is often a more stable training process with smoother loss curves and a lower likelihood of expert collapse.

Implementation and Practical Approach

In practice, implementing noisy gating is straightforward. It requires adding only a few lines of code within the gating module of your MoE layer.

Here is a simplified example in a PyTorch-like structure:

import torch
import torch.nn as nn
import torch.nn.functional as F

class NoisyTopkRouter(nn.Module):
    def __init__(self, d_model, num_experts, top_k):
        super().__init__()
        self.top_k = top_k
        self.gate_linear = nn.Linear(d_model, num_experts)
        self.noise_linear = nn.Linear(d_model, num_experts)

    def forward(self, x):
        # x shape: (batch_size * sequence_length, d_model)
        logits = self.gate_linear(x)

        # Add noise only during training
        if self.training:
            noise = self.noise_linear(x)
            noise_std = F.softplus(noise)
            noisy_logits = logits + (torch.randn_like(logits) * noise_std)
        else:
            noisy_logits = logits

        # Select top-k experts
        top_k_logits, indices = torch.topk(noisy_logits, self.top_k, dim=-1)

        # Create a sparse routing mask
        zeros = torch.full_like(noisy_logits, float('-inf'))
        sparse_logits = zeros.scatter(-1, indices, top_k_logits)
        router_output = F.softmax(sparse_logits, dim=-1)

        return router_output, indices

A significant consideration is the magnitude of the injected noise.

• Too little noise: The effect on load balancing will be minimal, and the model may still suffer from routing imbalance.
• Too much noise: The routing decisions can become overly random, disrupting the learning process. The gating network may struggle to learn meaningful routing patterns if its decisions are consistently drowned out by noise.

The use of a learnable W_{noise} matrix allows the model to adapt the noise level on a per-token basis, which is generally more effective than using a single, fixed noise hyperparameter for the entire training run.

Ultimately, noisy top-k gating is a simple yet effective technique for improving the stability and performance of MoE models. It encourages exploration in the routing space, leading to better load distribution and more specialized expert performance, without incurring a significant computational overhead.