A Tour of MoE: 9. The Debate on Gate Normalization

MoE环游记：9、门控归一化之争 苏剑林 · 2026-06-17 · Machine Learning probability loss function gradient MoE

Tracing back through the history of MoE, we find that in the early years, the Router in MoE—when acting as a Gate multiplied onto the Expert—almost always used Softmax activation, and to this day it remains one of the standard forms of MoE. However, to accommodate Loss-Free load balancing, DeepSeek changed the activation function to Sigmoid, and demonstrated that this too is a rather competitive scheme, which sparked deeper thinking and experimentation around the form of the Router.

Even when restricting the discussion to Softmax, there are two slightly different approaches: should we do Softmax first and then pick the Top-$k$, or pick the Top-$k$ first and then Softmax? The latter can also be understood as performing an additional normalization after selecting the Top-$k$, i.e., Re-Norm. So, should the Gate's activation function be normalized, and if so, should normalization happen before or after Top-$k$ selection (Re-Norm)? This is the topic we'll discuss in this post.

Problem Statement

We know that the general form of MoE is

\begin{equation}\boldsymbol{y} = \sum_{i\in \mathop{\text{argtop}}_k \boldsymbol{\rho}} \rho_i \boldsymbol{e}_i \label{eq:moe-1}\end{equation}more

Here $\boldsymbol{\rho}$ actually plays two roles: when used to select the Top-$k$ Experts, it acts as the Router; when multiplied onto the Expert, it acts as the Gate. From the design standpoint of MoE, the core role of $\boldsymbol{\rho}$ is clearly the Router, while the Gate's function is to supply it with gradients during training.

The question we want to discuss can also be understood as: how do we more scientifically construct $\boldsymbol{\rho}=(\rho_1, \rho_2, \cdots, \rho_n)$ so that the Router obtains better gradients? For a long time, the standard answer has been Softmax, i.e.,

\begin{equation}\rho_i = \frac{e^{s_i}}{\sum_{j=1}^n e^{s_j}} \end{equation}

where $\boldsymbol{s}=(s_1, s_2, \cdots, s_n)$ is the logits projected directly by a linear layer. However, although this "standard" answer is widely used, the author has not found a satisfying explanation for it—it seems everyone simply accepted it and carried on using it, which at one point left the author quite puzzled about the training mechanics of MoE.

Other Choices

As mentioned at the outset, DeepSeek tried Sigmoid activation in its Loss-Free load balancing scheme, and later used it in DeepSeek-V3 as well; its success shows that non-Softmax activations can also work quite well. This has inspired people to try more general approaches—for example, ReMoE uses ReLU activation, while from the geometric perspective in 《MoE环游记：1、从几何意义出发》, any non-negative activation function is permitted.

Besides this, there's also the Re-Norm option in the MoE formulation, which changes $\eqref{eq:moe-1}$ to

\begin{equation}\boldsymbol{y} = \frac{\sum\limits_{i\in \mathop{\text{argtop}}_k \boldsymbol{\rho}} \rho_i \boldsymbol{e}_i}{\sum\limits_{i\in \mathop{\text{argtop}}_k \boldsymbol{\rho}} \rho_i} \label{eq:moe-2}\end{equation}

That is, performing an additional normalization on the selected Top-$k$ $\rho_i$. For Softmax, this is equivalent to using $\boldsymbol{s}$ to select the Top-$k$, setting the unselected ones to $-\infty$, and then applying Softmax. The benefit of Re-Norm is more numerically stable forward computation, but note that when using Re-Norm, $k$ must be greater than 1; otherwise $\boldsymbol{\rho}$ would receive no gradient at all, making it untrainable.

Looking across current practices from various parties, these MoE variants all perform roughly the same, with no single one clearly dominating. Since practice can't distinguish a winner, let's examine theoretically which form is actually more principled.

Design Principle

Our goal is to find a first-principles foundation closer to the essence of the problem, and use it to derive the gating mechanism of current MoE models.

So the first question is naturally: what is this "principle"? For simplicity, let's first consider $k=1$. We know that the most important feature of MoE is sparsity: a Router first determines which Experts to activate, and then only those Experts are computed, thereby increasing the parameter count while controlling computational cost. If this were the only consideration, the naive model would be

\begin{equation}\newcommand{argmax}{\mathop{\text{argmax}}}\boldsymbol{f}\left(\boldsymbol{e}_{\argmax\boldsymbol{\rho}}\right)\end{equation}

That is, from the Router $\boldsymbol{\rho}$, pick out the one with the highest score and activate the corresponding Expert. This form works perfectly fine for inference, but during training, the Router would receive no gradient at all and thus could not be updated. So we need to design a way to give the Router gradients. How should we design gradients for the Router? To answer this, we first need to clarify: what kind of Router do we actually want?

Since only 1 Expert can be activated, we naturally want that Expert to be the best-performing one. If we denote the loss function by $\ell$, our expectation can be written as

\begin{equation}\newcommand{argmin}{\mathop{\text{argmin}}}\argmax \boldsymbol{\rho} = \argmin\, [\ell(\boldsymbol{e}_1),\ell(\boldsymbol{e}_2),\cdots,\ell(\boldsymbol{e}_n)]\label{eq:target}\end{equation}

This is the design principle we're looking for.

Transforming the Objective

However, the objective $\eqref{eq:target}$ is not yet a loss function we can directly train with—it still needs further transformation. To this end, we construct two distributions. The first is a target distribution $\boldsymbol{q}=(q_1,q_2,\cdots,q_n)$ built from the loss function, defined as

\begin{equation}q_i = \frac{e^{-\ell(\boldsymbol{e}_i)/\tau}}{\sum_{j=1}^n e^{-\ell(\boldsymbol{e}_j)/\tau}}\end{equation}

This distribution has nothing to do with the Router; from the perspective of learning the Router, it is a "target distribution." The second distribution is a predicted distribution $\boldsymbol{p}$ constructed from $\boldsymbol{\rho}$. There are many possibilities here: for instance $\boldsymbol{\rho}$ itself could already be the distribution $\boldsymbol{p}$ (if $\boldsymbol{\rho}$ is already normalized), or $\boldsymbol{p}$ could be the Softmax of $\boldsymbol{\rho}$ (in which case $\boldsymbol{\rho}$ are logits), or some other normalization scheme besides Softmax. In any case, $\boldsymbol{p}$ is some probabilistic representation of the Router, with parameters $\boldsymbol{\theta}$.

We transform the objective $\eqref{eq:target}$ into bringing $\boldsymbol{p}$ and $\boldsymbol{q}$ closer together, thereby providing a gradient for $\boldsymbol{\theta}$. To do this, we consider minimizing the KL divergence

\begin{equation}KL(\boldsymbol{p}\Vert \boldsymbol{q}) = \sum_{i=1}^n p_i \log \frac{p_i}{q_i}\end{equation}

Rearranging slightly gives

\begin{equation}KL(\boldsymbol{p}\Vert \boldsymbol{q}) = - \mathcal{H}(\boldsymbol{p}) + \frac{1}{\tau}\sum_{i=1}^n p_i \ell(\boldsymbol{e}_i) - \log \sum_{i=1}^n e^{-\ell(\boldsymbol{e}_i)/\tau}\end{equation}

We can see this objective has three terms. The first is the negative entropy $-\mathcal{H}(\boldsymbol{p})$; minimizing it means maximizing entropy, which effectively encourages the model to explore thoroughly—one could argue load balancing already plays a similar role, so let's set this aside for now. The third term has nothing to do with $\boldsymbol{p}$, i.e., nothing to do with $\boldsymbol{\theta}$, so the equivalent loss function is $\mathcal{L} = \sum_{i=1}^n p_i \ell(\boldsymbol{e}_i)$

Straight-Through Estimation

Taking the gradient of the equivalent loss gives

\begin{equation}\nabla_{\boldsymbol{\theta}}\mathcal{L} = \sum_{i=1}^n \nabla_{\boldsymbol{\theta}} p_i \cdot \ell(\boldsymbol{e}_i) = \sum_{i=1}^n p_i \nabla_{\boldsymbol{\theta}} \log p_i \cdot \ell(\boldsymbol{e}_i) = \mathbb{E}_{i\sim \boldsymbol{p}} [\nabla_{\boldsymbol{\theta}}\log p_i \cdot \ell(\boldsymbol{e}_i)]\end{equation}

The key here is using $\nabla_{\boldsymbol{\theta}} p_i = p_i \nabla_{\boldsymbol{\theta}} \log p_i$ to isolate a single term $p_i$, which allows the summation to be converted into an expectation, and thus achieves MoE's sparse computation goal via sampling. Some readers may already recognize this—it's exactly REINFORCE from policy gradient methods! (See 《从采样看优化：可导优化与不可导优化的统一视角》 and 《殊途同归的策略梯度与零阶优化》.)

The problem with REINFORCE is its high variance. Intuitively, this is because it places $p_i$ outside the loss function $\ell$; if possible, we'd rather have a "reparameterized" form where $p_i$ sits inside $\ell$. To derive such a form, we exploit REINFORCE's invariance to baseline subtraction to get

\begin{equation}\begin{aligned} \mathbb{E}_{i\sim \boldsymbol{p}} [\nabla_{\boldsymbol{\theta}}\log p_i \cdot \ell(\boldsymbol{e}_i)] =&\, \mathbb{E}_{i\sim \boldsymbol{p}} [\nabla_{\boldsymbol{\theta}}\log p_i \cdot (\ell(\boldsymbol{e}_i) - \ell(\boldsymbol{0}))] \\[4pt] \approx&\, \mathbb{E}_{i\sim \boldsymbol{p}} [\nabla_{\boldsymbol{\theta}}\log p_i \cdot \langle\nabla_{\boldsymbol{e}_i} \ell(\boldsymbol{e}_i), \boldsymbol{e}_i - \boldsymbol{0}\rangle] \\[4pt] = &\, \mathbb{E}_{i\sim \boldsymbol{p}} [\nabla_{\boldsymbol{\theta}} \langle\nabla_{\boldsymbol{e}_i} \ell(\boldsymbol{e}_i), \log p_i \cdot \boldsymbol{e}_i\rangle] \\[4pt] = &\, \mathbb{E}_{i\sim \boldsymbol{p}} [\nabla_{\boldsymbol{\theta}} \ell((\log p_i + \color{skyblue}{[}1 - \log p_i \color{skyblue}{]_{\text{sg}}}) \cdot\boldsymbol{e}_i)] \\[4pt] = &\, \nabla_{\boldsymbol{\theta}} \mathbb{E}_{i\sim \boldsymbol{p}} [\ell((\log p_i + \color{skyblue}{[}1 - \log p_i \color{skyblue}{]_{\text{sg}}}) \cdot\boldsymbol{e}_i)] \\ \end{aligned}\end{equation}

where the approximation symbol $\approx$ comes from a first-order Taylor expansion at $\boldsymbol{e}_i$, and $\color{skyblue}{[}\cdot\color{skyblue}{]_{\text{sg}}}$ denotes Stop Gradient. In the end, we obtain a Straight-Through Estimator (STE) that uses $1$ in the forward pass and $\log p_i$ in the backward pass, providing gradients to the Router.

Final Form

Although STE gives us a workable training scheme, because of the mismatch between forward and backward propagation, it often only achieves suboptimal results. Here's a remarkably clever improvement: change each Expert to $p_i\boldsymbol{e}_i$! Repeating the derivation above, we get

\begin{equation}\begin{aligned} \mathbb{E}_{i\sim \boldsymbol{p}} [\nabla_{\boldsymbol{\theta}}\log p_i \cdot \ell(p_i\boldsymbol{e}_i)] =&\, \mathbb{E}_{i\sim \boldsymbol{p}} [\nabla_{\boldsymbol{\theta}}\log p_i \cdot (\ell(p_i\boldsymbol{e}_i) - \ell(\boldsymbol{0}))] \\[4pt] \approx&\, \mathbb{E}_{i\sim \boldsymbol{p}} [\nabla_{\boldsymbol{\theta}}\log p_i \cdot \langle\nabla_{p_i\boldsymbol{e}_i} \ell(p_i\boldsymbol{e}_i), p_i\boldsymbol{e}_i - \boldsymbol{0}\rangle] \\[4pt] = &\, \mathbb{E}_{i\sim \boldsymbol{p}} [\nabla_{\boldsymbol{\theta}} \langle\nabla_{p_i \boldsymbol{e}_i} \ell(p_i \boldsymbol{e}_i), p_i \boldsymbol{e}_i\rangle] \\[4pt] = &\, \mathbb{E}_{i\sim \boldsymbol{p}} [\nabla_{\boldsymbol{\theta}} \ell(p_i \boldsymbol{e}_i)] \\[4pt] = &\, \nabla_{\boldsymbol{\theta}} \mathbb{E}_{i\sim \boldsymbol{p}} [\ell(p_i \boldsymbol{e}_i)] \\[4pt] \end{aligned}\end{equation}

This transformation process is quite elegant and worth savoring carefully. By changing the Expert from $\boldsymbol{e}_i$ to $p_i\boldsymbol{e}_i$, we eliminate the Stop Gradient, achieve consistency between forward and backward passes, and in principle raise the ceiling on the model's performance.

Now we can answer the question posed at the beginning:

If we want a top-down probabilistic derivation, then when the Router acts as a Gate, it should be normalized—but it should not use Re-Norm.

To Sample or Not to Sample

One detail worth noting is: $\mathbb{E}_{i\sim \boldsymbol{p}}$ implies we should sample from $\boldsymbol{p}$, but in practice we usually directly select the Top-$k$. How should we understand this discrepancy?

This is really a trade-off between diversity and stability. Random sampling encourages the model to explore more thoroughly, but sampling increases gradient variance and introduces additional instability; directly selecting the Top-$k$ is more stable, but risks trapping the model in a suboptimal solution or even causing model collapse. Fortunately, various load balancing strategies today are already quite mature and to some extent already encourage the model to explore broadly, so selecting Top-$k$ remains the mainstream approach for now.

If we want to sample while retaining stability, we can slightly extend the Top-$k$ approach rather than fully opening up to sampling—for example, first selecting the Top-$k+c$, then randomly picking $k$ out of those $k+c$ Experts, or adding slight noise to the logits $\boldsymbol{p}$ before selecting the Top-$k$. This adds randomness while not straying too far from the original Top-$k$, balancing exploration with stability.

Related Work

To be clear, the derivation in this post isn't new—it's something the author distilled and adapted from Liyuan Liu's article 《Sparse Backpropagation for MoE Training》, which also has a preceding part 《Bridging Discrete and Backpropagation: Straight-Through and Beyond》 and a follow-up 《GRIN: GRadient-INformed MoE》.

Although these are already articles from 2023-24, if you want to deepen your understanding of MoE Routers, I still highly recommend reading this trilogy—it provides a unified probabilistic framework for designing gradients for various discretization operations. Of course, the probabilistic framework also has its limitations, namely that it is quite formal, which can feel a bit constraining in practice.

For example, when $k = 2$, if we extend the earlier result in a parallel fashion, we should get

\begin{equation}\mathbb{E}_{i,j\sim \boldsymbol{p}} [\nabla_{\boldsymbol{\theta}}\log p_i p_j \cdot \ell(p_i p_j (\boldsymbol{e}_i + \boldsymbol{e}_j))] \approx \nabla_{\boldsymbol{\theta}} \mathbb{E}_{i,j\sim \boldsymbol{p}} [\ell(p_i p_j (\boldsymbol{e}_i + \boldsymbol{e}_j))]\end{equation}

That is, taking the joint distribution $p_i p_j$ and the sum of pairs of experts $\boldsymbol{e}_i + \boldsymbol{e}_j$ as the basic unit, converting Top-2 into Top-1. However, the MoE we've been using all along is actually of the form $\ell(p_i \boldsymbol{e}_i + p_j\boldsymbol{e}_j)$, and it's not so easy to find a precise probabilistic derivation for this.

At this point, perhaps a more "relaxed" way of understanding things is to simply treat it as an analogue of MaxPooling, without insisting on a probabilistic interpretation—or one could alternatively choose to understand it via the geometric meaning described in 《MoE环游记：1、从几何意义出发》. Overall, the probabilistic framework only proves that a certain scheme is feasible, but in principle it does not rule out the feasibility of other schemes.

Summary

This post attempted to start from first principles to explore the design of the Router and Gate in MoE, providing a probabilistic explanation for the normalization of gating.