A Tour of MoE: 7. A Minimalist Explanation of Dynamic Activation

MoE环游记：7、动态激活极简解 苏剑林 · 2026-02-23 · Machine Learning linear duality optimization MoE dynamics

In the previous post, MoE Wandering Notes: 6. Optimal Assignment for Balance, we achieved load balancing by solving the following optimal assignment problem:

\begin{equation}\max_{x_{i,j}\in\{0,1\}} \sum_{i,j} x_{i,j}s_{i,j} \qquad\text{s.t.}\qquad \sum_j x_{i,j} = k,\quad \sum_i x_{i,j} = \frac{mk}{n}\end{equation}

where $\sum_j x_{i,j} = k$ says that each Token activates exactly $k$ Experts, and $\sum_i x_{i,j} = mk/n$ says that each Expert gets activated exactly $mk/n$ times. On closer thought, however, neither training nor inference actually requires the former; what we really need is the latter, which already implies that "on average each Token activates $k$ Experts" together with per-Expert load balance—and that alone is enough to achieve MoE's goal. So in this post we consider the simplified problem

\begin{equation}\max_{x_{i,j}\in\{0,1\}} \sum_{i,j} x_{i,j}s_{i,j} \qquad\text{s.t.}\qquad \sum_i x_{i,j} = \frac{mk}{n}\label{eq:target-dyn}\end{equation}more

Dynamic Activation

We'll assume readers are already familiar with the previous post, so the derivation here will be relatively brief; if any step is unclear, feel free to go back and review. As before, we start from the relaxed version of the objective $\eqref{eq:target-dyn}$:

\begin{equation}\max_{x_{i,j}\in[0,1]} \sum_{i,j} x_{i,j}s_{i,j} \qquad\text{s.t.}\qquad \sum_i x_{i,j} = \frac{mk}{n}\end{equation}

and then consider its equivalent $\max\text{-}\min$ form:

\begin{equation}\max_{x_{i,j}\in[0,1]}\min_{\beta_j} \sum_{i,j} x_{i,j}s_{i,j} - \sum_j \beta_j\left(\sum_i x_{i,j} - \frac{mk}{n}\right)\end{equation}

Similarly, the order of $\max$ and $\min$ here can be swapped; swapping and rearranging a bit gives

\begin{equation}\min_{\beta_j}\max_{x_{i,j}\in[0,1]} \sum_{i,j} x_{i,j}(s_{i,j} - \beta_j) + \frac{mk}{n} \sum_j \beta_j\label{eq:relax-min-max}\end{equation}

The $\max$ step can be carried out first, giving

\begin{equation}\left\{\begin{aligned}&\,x_{i,j}^* = 1, &\, s_{i,j} - \beta_j > 0 \\ &\,x_{i,j}^* = 0, &\, s_{i,j} - \beta_j < 0 \\ &\,x_{i,j}^* \in [0,1], &\, s_{i,j} - \beta_j = 0 \end{aligned}\right.\end{equation}

This tells us that an Expert gets activated as long as $s_{i,j} - \beta_j > 0$, and the number of activations is not fixed—which formally matches exactly what we designed by intuition in MoE Wandering Notes: 4. Put More Effort Where It's Hard, except this time it is derived from the more fundamental objective $\eqref{eq:target-dyn}$ rather than guessed.

Solving in One Step

Substituting $x_{i,j}^*$ back into equation $\eqref{eq:relax-min-max}$ gives $x_{i,j}^*(s_{i,j} - \beta_j) = \max(0, s_{i,j} - \beta_j)$, so the optimization objective for $\beta_j$ simplifies to

\begin{equation}\min_{\beta_j} \sum_{i,j} \max(0, s_{i,j} - \beta_j) + \frac{mk}{n} \sum_j \beta_j\label{eq:beta-obj}\end{equation}

Every term with $\beta_j$ stands alone, so this decomposes into $m$ independent sub-problems, which means we can drop the subscript $j$ for now:

\begin{equation}\min_{\beta} \frac{mk}{n}\beta + \sum_i \max(0, s_i - \beta)\end{equation}

Sort all the $s_i$ values in decreasing order as $s_{\sigma_1}\geq s_{\sigma_2} \geq \cdots \geq s_{\sigma_m}$, and suppose we already know $s_{\sigma_l}\geq\beta\geq s_{\sigma_{l+1}}$; then we can drop $\max$ to get

\begin{equation}\frac{mk}{n}\beta + \sum_{i=1}^l (s_{\sigma_i} - \beta) = \left\{\begin{aligned} &\,\sum_{i=1}^{mk/n} s_{\sigma_i} + \sum_{i=mk/n+1}^l \underbrace{(s_{\sigma_i} - \beta)}_{\geq 0},&\, l \geq mk/n \\ &\,\sum_{i=1}^{mk/n} s_{\sigma_i} - \sum_{\hphantom{ab}i=l+1\hphantom{ab}}^{mk/n} \underbrace{(s_{\sigma_i} - \beta)}_{\leq 0},&\, l \leq mk/n \\ \end{aligned}\right.\end{equation}

This shows that both $l > mk/n$ and $l < mk/n$ enlarge the objective, so the minimum is attained at $l=mk/n$, i.e., $\beta^*$ lies between the $mk/n$-th and $mk/n+1$-th largest elements of $s_i$. For convenience, we take the $mk/n+1$-th largest element. Restoring the subscript $j$, we get that for a given $j$, $\beta_j^*$ is the $mk/n+1$-th element when $s_{i,j}$ is sorted in decreasing order, i.e., the "$1-k/n$-quantile."

Note that now the only variable to solve for is $\boldsymbol{\beta}$, so there is no need for an alternating iteration step at all—a single Quantile computation directly gives the perfectly balanced optimal solution! We again call this "Quantile Balancing (QB)."

Practical Notes

Readers already familiar with MoE Wandering Notes: 6. Optimal Assignment for Balance will notice that this post is essentially a streamlined version of it. Yet simplicity here is not simplistic—it's actually more elegant and efficient than the Top-$k$ style of MoE: first, it activates an Expert based on whether $\boldsymbol{s} - \boldsymbol{\beta}$ is greater than 0, which is more efficient than picking Top-$k$; second, it obtains the load-balanced optimal solution with a single Quantile computation, which is also more efficient and elegant.

Of course, a one-shot optimal solution is very likely to overfit the current training batch. To avoid this, the improvement we consider here is to take an EMA (Exponential Moving Average) between the current batch's optimal solution and the historical solution (note that the Top-$k$ version from the previous post could also use an EMA). We again need to be careful about the information-leakage trap, so we must first use the old $\boldsymbol{\beta}$ to activate Experts, and only afterwards update $\boldsymbol{\beta}$.

The algorithm proceeds as follows (in the author's experiments, $\lambda=0.9$):

$$\begin{array}{|l|} \hline \text{Quantile Balancing (QB) dynamic activation version} \\[4pt] \hline \text{input: score matrix}\boldsymbol{s}\in\mathbb{R}^{m\times n}\text{, previous step}\boldsymbol{\beta}\in\mathbb{R}^n\text{, decay rate}\lambda \\ \text{output: assignment}\boldsymbol{x}\in\{0,1\}^{m\times n}\text{, new}\boldsymbol{\beta}\in\mathbb{R}^n \\[4pt] \hline \begin{array}{ll} 1: & x_{i,j}=1 \text{ if } s_{i,j} - \beta_j > 0 \text{ else } 0 \\ 2: & \boldsymbol{\beta} \leftarrow \lambda\boldsymbol{\beta} + (1-\lambda)\mathop{\text{desc_sort}}(\boldsymbol{s}, \text{axis=0})_{[mk/n:mk/n+1]} \\ 3: & \text{Output } \boldsymbol{x},\boldsymbol{\beta} \end{array} \\ \hline \end{array}$$

In practice, we again run into the problem that finding the $1-k/n$-quantile of $\boldsymbol{s}$ is expensive, since it requires finding the $mk/n+1$-th largest element among $m$ elements, where $m$ equals "global sample count × sequence length." Because of various parallelism strategies and gradient accumulation constraints, an exact implementation is usually unacceptable. So we again resort to a compromise: split the samples into the largest micro-batches we can afford, compute a separate $\boldsymbol{\beta}$ for each micro-batch, and average them to get the final result.

Expert Choice

This dynamic version of QB actually has a very simple interpretation, namely "Expert Choice":

Don't you want each Expert to be activated exactly $mk/n$ times? Then why not let the Experts choose the Tokens? That is, have each Expert pick its Top-$mk/n$ Tokens—isn't that just finding the $mk/n$ largest elements along $\text{axis=0}$?

However, naive Expert Choice has a fatal flaw: it requires comparing all Tokens globally, which leads to selection across samples and across sequences—violating causality (seeing future information during training) and breaking train-inference consistency (the inference result depends on batch size). This is precisely why the method has historically been hard to put into practice.

The cleverness of this post lies in reformulating it as a Bias term: taking each Expert's $mk/n+1$-th largest element (equivalently, the $1-k/n$-quantile) as the threshold $\boldsymbol{\beta}$, we re-express Expert Choice as Token Choice—turning "Expert picks Top-$mk/n$" into "activate whenever $s_{i,j} - \beta_j > 0$." More importantly, we delay the update of $\boldsymbol{\beta}$ until after the activation decision is made, which perfectly avoids the information-leakage problem.

Remarkably, this Bias formulation, which fixes the flaw in Expert Choice, emerges naturally from the dual of the original linear program—there is a certain aesthetic pleasure in that.

Initialization Strategy

To prevent information leakage, during training we need to first use the old $\boldsymbol{\beta}$ to decide which Experts each Token activates, and only then update $\boldsymbol{\beta}$. This makes the first few training steps quite sensitive to the initialization of $\boldsymbol{\beta}$. For example, if $\boldsymbol{s}$ are all positive but $\boldsymbol{\beta}$ is initialized to zero, then in the very first training step every Token would activate all Experts, and the computation would blow up.

In MoE Wandering Notes: 4. Put More Effort Where It's Hard we already provided a simulation script to estimate a suitable initialization, but here, since we already know that the optimal solution for $\boldsymbol{\beta}$ is the "$1-k/n$-quantile," we can make some reasonable assumptions about the initial Scores and derive an analytical formula for the initialization directly. Specifically, suppose the Router's initial Logits follow $\mathcal{N}(0,\sigma^2)$; then its $1-k/n$-quantile is

\begin{equation}\sigma \cdot \Phi^{-1}\left(1 - \frac{k}{n}\right)\end{equation}

This is exactly the optimal $\boldsymbol{\beta}$ for this distribution's Logits, and we use it as the initialization for $\boldsymbol{\beta}$, where $\Phi^{-1}$ is the quantile function of the standard normal distribution, also known as the inverse CDF.

In addition, the Router may apply an activation function. If the activation function is element-wise and monotonically increasing, such as Sigmoid, it doesn't change the ordering, so we simply apply the corresponding activation function to the formula above. Softmax is slightly more complex, since it exponentiates first and then normalizes; assuming the normalizing denominator is approximately constant, it behaves similarly to the Sigmoid case—we just exponentiate the formula above, then divide by a denominator estimated via quantile-based simulated sampling:

\begin{equation}\frac{\exp\left[\sigma \cdot \Phi^{-1}\left(1 - \frac{k}{n}\right)\right]}{\sum\limits_{i=1}^n \exp\left[\sigma \cdot \Phi^{-1}\left(1 - \frac{i}{n+1}\right)\right]}\end{equation}

As for estimating $\sigma$, this is also straightforward: suppose the Router's input dimension is $d$ and RMS Norm is applied, and suppose the Router's weight matrix is initialized with $\mathcal{N}(0,\tilde{\sigma}^2)$; then the Router Logits approximately follow a normal distribution with mean around 0 and variance around $\tilde{\sigma}^2 d$, i.e., $\sigma\approx \tilde{\sigma} \sqrt{d}$. These are all classic results from initialization theory reference].

Demo Code

Once again, here's a piece of demo code for readers to play with:

import numpy as np

def quantile_bias(s, k):
    """动态版QB，一步quantile即可求最优bias
    原理：https://kexue.fm/archives/11626
    """
    m, n = s.shape
    beta = np.quantile(s, 1 - k / n, axis=0)
    return beta

def max_min_avg_vio(s):
    """计算max_vio、min_vio和avg_vio
    其中 max_vio ≥ 0, avg_vio ≥ 0, -1 ≤ min_vio ≤ 0，三者都是越接近于0表示越均衡
    """
    m, n = s.shape
    f = (s > 0).mean(0)
    f = f / f.sum() * n - 1
    return f.max(), f.min(), np.abs(f).mean()

def avg_std_active(s):
    """计算每个expert被激活次数的平均值和标准差
    avg越接近k越好，std越接近0越好
    """
    m, n = s.shape
    f = (s > 0).mean(0) * n
    return f.mean(), f.std()

m, n, k = 100000, 256, 8
s = np.random.rand(m, n) + np.random.rand(n)  # 模拟一个不均匀的打分

b = quantile_bias(s, k)
max_min_avg_vio(s - b)  # 如无意外是全0
avg_std_active(s - b)  # 如无意外是 (k, 0)

from scipy.stats import norm

sigma = 1
s = np.random.randn(m, n) * sigma  # 模拟初始化
avg_std_active(s)  # 大致是(128, 0.4)，远大于预算k
b0 = np.zeros(n) + norm.ppf(1 - k / n) * sigma  # norm.ppf正是标准正态分布的分位数函数
avg_std_active(s - b0)  # 大致是(8, 0.1)，接近预算k

s2 = 1 / (1 + np.exp(-s))  # Sigmoid激活
b0_2 = 1 / (1 + np.exp(-b0))  # Sigmoid激活
avg_std_active(s2 - b0_2)  # 依然大致是(8, 0.1)，接近预算k

s3 = np.exp(s) / np.exp(s).sum(axis=1, keepdims=True)  # Softmax激活
q = (np.arange(n) + 1) / (n + 1)  # 均匀间隔点（分位数）
b0_3 = np.exp(b0) / np.exp(sigma * norm.ppf(q)).sum() # Exp激活并除以模拟的分母
avg_std_active(s3 - b0_3)  # 大致是(7.5, 0.1)，接近预算k

Gradient Descent

Finally, if you'd rather avoid Quantile altogether, gradient descent is still an option. First, denote the loss function for objective $\eqref{eq:beta-obj}$ as $\ell$:

\begin{equation}\min_{\beta_j} \underbrace{\sum_{i,j} \max(0, s_{i,j} - \beta_j) + \frac{mk}{n} \sum_j \beta_j}_{\text{denote}\ell}\end{equation}

This is differentiable, with gradient

\begin{equation}\frac{\partial\ell}{\partial\beta_j} = \frac{mk}{n} - \sum_{i=1}^m \chi(s_{i,j} - \beta_j > 0)\end{equation}

where $\chi$ is the indicator function, $\chi(\text{True})=1,\chi(\text{False})=0$. Computing this gradient is cheaper than a global Quantile, making it suitable for readers who are more cost-sensitive. Once we have the gradient, we can perform gradient descent; we again consider SignSGD, aligned with Loss-Free:

\begin{equation}\beta_j \leftarrow \beta_j - \gamma\mathop{\text{sign}}\left(\frac{\partial\ell}{\partial\beta_j}\right)\end{equation}

This can replace the Quantile-based update of $\boldsymbol{\beta}$, and since this is inherently a long-run iterative algorithm, the EMA can be dropped as well. In fact, this is exactly the formula $(9)$ that we proposed in MoE Wandering Notes: 4. Put More Effort Where It's Hard; here we've re-derived it from the perspective of "optimal assignment + dual objective + gradient descent."

Summary

This post continues the exploration of Quantile Balancing (QB) from the previous one. By dropping the constraint "each Token can activate at most $k$ Experts" from the optimal assignment problem, we obtain a significantly simplified solution—load balancing is achieved with just a single Quantile computation, and each Token only needs to check the sign of its bias score to decide which Experts to activate, eliminating the overhead of Top-$k$ sorting altogether.