MoE Odyssey: 1. Starting from Geometric Intuition

MoE环游记：1、从几何意义出发 苏剑林 · 2025-02-08 · Machine Learning model geometry sparse MoE

A couple of years ago, moved by a moment of inspiration, I started a series called "The Path to Better Transformers," sharing various improvements to mainstream Transformer architectures along with some of my own thoughts, and it seemed to resonate with a portion of my readers. Starting with this post, we'll take up the same style and turn our attention to MoE (Mixture of Experts), another currently mainstream architecture.

The popularity of MoE hardly needs elaborating: the recently trending DeepSeek-V3 uses an MoE architecture, GPT-4 is rumored to be MoE as well, and quite a few recent domestic models have also adopted MoE. However, although research on MoE goes back a long way, its application remained lukewarm for a long time. It was roughly starting from Mixtral of Experts at the beginning of last year that MoE gradually began to attract attention, its notable advantage being a large parameter count while keeping both training and inference costs significantly lower.

That said, MoE also comes with some tricky problems — training instability, load imbalance, and underwhelming performance, among others — and these are the main reasons it didn't catch on in its early years. However, with the increased attention it has received over the past couple of years, these problems have largely been resolved, and we'll go through them one by one in what follows.

Problem Definition

First, let me note that the discussion here reflects my own way of understanding MoE. I will cite relevant references where necessary, but I won't attempt a systematic history tracing the origins of the MoE architecture — I hope readers will forgive me for that.

We know that the Transformer model is built from Attention layers and MLP layers, and MoE replaces the MLP layers in the model. MLP layers come in two flavors, FFN (FeedForward Network) and GLU (Gated Linear Unit); GLU is the mainstream choice these days, but for simplicity we'll use FFN as our running example.

\begin{equation}\boldsymbol{y} = f(\boldsymbol{x}\boldsymbol{W}^{(A)})\boldsymbol{W}^{(B)}\end{equation}

Here $\boldsymbol{x}\in\mathbb{R}^{d}$ is the input vector (a row vector), $\boldsymbol{W}^{(A)}\in\mathbb{R}^{d\times D},\boldsymbol{W}^{(B)}\in\mathbb{R}^{D\times d}$ are the two parameter matrices, and $f$ is an element-wise activation function. Let $n$ be an integer that divides $D$ evenly; then the above can be equivalently written using block matrices as

\begin{equation}\boldsymbol{y} = f\big(\boldsymbol{x}\begin{bmatrix}\boldsymbol{W}^{(A)}_1 & \boldsymbol{W}^{(A)}_2 & \cdots & \boldsymbol{W}^{(A)}_n\end{bmatrix}\big)\begin{bmatrix}\boldsymbol{W}^{(B)}_1 \\ \boldsymbol{W}^{(B)}_2 \\ \vdots \\ \boldsymbol{W}^{(B)}_n\end{bmatrix} = \sum_{i=1}^n \underbrace{f(\boldsymbol{x}\boldsymbol{W}^{(A)}_i)\boldsymbol{W}^{(B)}_i}_{\boldsymbol{v}_i}\end{equation}

Here $\boldsymbol{W}^{(A)}_i = \boldsymbol{W}^{(A)}_{[:,(i-1)c:ic]}, \boldsymbol{W}^{(B)}_i = \boldsymbol{W}^{(B)}_{[(i-1)c:ic,:]},c= D/n$, where the slicing follows Python conventions. This shows that the FFN can be equivalently expressed as the sum of $n$ vectors $\boldsymbol{v}_1,\boldsymbol{v}_2,\cdots,\boldsymbol{v}_n$, each vector representing the output of a small model $f(\boldsymbol{x}\boldsymbol{W}^{(A)}_i)\boldsymbol{W}^{(B)}_i$, with each small model requiring the same amount of computation — these small models are precisely the "Experts" in MoE.

# The Question MoE Raises

Let's continue our discussion from last time on model architecture, moving from Dense to MoE.

Actually, this "sequel" wasn't originally planned. But after finishing the previous post 《Transformer升级之路：Dense前传》, I realized that focusing purely on Dense models can only take the discussion so far — because the biggest and most influential Chinese open-source models of the past two years, whether DeepSeek or Qwen, are basically all MoE models. So this piece serves both as a supplement to the previous post's discussion of "$d/n$" and, more importantly, as an exploration of a question that has puzzled me for a long time: what exactly is the problem that MoE is trying to solve?

Trade-off

For those familiar with MoE (Mixture of Experts), the standard answer to "what problem does MoE solve" is usually something like: "it decouples parameter count from computational cost, achieving a larger parameter count without increasing computational cost (inference cost)." This is indeed correct, but it can feel more like a description of what MoE does rather than a deeper explanation of why it does it. In other words, this "standard answer" is essentially restating MoE's basic setup in different words, but it hasn't explained why we need to decouple parameter count from computational cost in the first place, nor why the specific "sparse" solution that MoE offers is a good one.

Since we're already in an era where MoE has proven its worth through large-scale practice, we no longer need to prove from scratch that "decoupling parameter count from computational cost, and doing so via sparsification" is a viable path — the facts speak for themselves. What we want to figure out here is why it works, i.e., can we find a set of relatively self-consistent logical starting points from which MoE emerges as a natural, or at least reasonable, choice?

Guided by this question, let's revisit the Trade-off section from the previous post. There we established the following relationship between test loss $L$, parameter count $N$, and data volume (measured in Tokens) $D$:
$$ L(N,D) = L_{\infty} + \frac{A}{N^{\alpha}} + \frac{B}{D^{\beta}} $$
Under the constraint of a fixed compute budget $C$, we obtained the compute-optimal solution:
$$ N_C = N_{\text{opt}}\left(\frac{C}{C_0}\right)^{a},\qquad D_C = D_{\text{opt}}\left(\frac{C}{C_0}\right)^{b} $$
Here $a = \frac{\beta}{\alpha+\beta}, b = \frac{\alpha}{\alpha+\beta}$, and according to Chinchilla's original paper, $\alpha,\beta$ are close to each other, hence $a,b$ are both close to $1/2$. This means that under the same compute budget, the optimal parameter count $N$ and optimal data volume $D$ should grow at roughly the same rate as compute increases.

Then, from Scaling Law's basic setup — where $C=6ND$ (that is, compute is (approximately) proportional to the product of parameter count and data volume) — we further derived
$$ D_C/N_C \propto C^{b - a} = C^{\frac{\alpha - \beta}{\alpha+\beta}} $$

Since $\alpha < \beta$ typically holds (as verified across many subsequent papers), this means: as the compute budget increases, the optimal ratio of data volume to parameter count grows correspondingly, i.e., the amount of data needed per parameter increases with the compute budget.

Bottleneck

We know that once a model's architecture and parameter count are fixed, its capacity is essentially capped. When the data volume is not particularly large, more data mainly helps the model learn existing knowledge and patterns more thoroughly, so the loss decreases quickly. But once data volume exceeds a certain threshold, the loss reduction gained from additional data becomes limited by the parameter count, and the marginal benefit of more data starts to shrink noticeably — a phenomenon we might call a "bottleneck effect."

The conclusion from the previous section — "the optimal ratio of data volume to parameter count grows with the compute budget" — tells us that, at least within the range covered by current experiments, this bottleneck effect has not yet become severe enough that we need to worry about it. In other words, we can basically use "how much data to prepare in proportion to the parameter count" as a criterion, without worrying too much about the risk of the model's capacity being insufficient to hold the data — a model with more parameters can basically always accommodate proportionally more data.

However, this doesn't mean the bottleneck effect won't eventually become severe. In fact, we can find hints of this even within the same period as Chinchilla. For example, the paper 《Data and Parameter Scaling Laws for Neural Machine Translation》 noted that in translation tasks, once the parameter count is fixed, increasing the data volume can even cause the test loss to rise instead of fall — a phenomenon it calls "Data Saturation":

The performance of NMT models has often been shown to scale directly with training data size, at least for reasonably-sized datasets that manageably fit into GPU memory. However, this pattern has been shown to break down at the multiple GPU scale for very large datasets, exhibiting a "data-saturation" effect. In this regime, the largest models' performance stops improving with additional training data, and even starts to slightly deteriorate, which suggests that the largest of the current architectures are already too powerful for the amount of naturally-available NMT data.

Since the corresponding parameters and data volume in that paper are relatively small (parameter count on the order of $10^8$, data volume around $10^9$ tokens), it doesn't immediately follow that this experience directly applies to today's LLMs. But it at least illustrates a general phenomenon: for a fixed model architecture, once the parameter count is set, capacity is capped, and blindly increasing data volume beyond that cap won't help — in the mild case it merely wastes compute, and in the severe case it may even hurt performance.

Increase

So a question naturally arises: if we've already collected far more data than the current model's parameter count can "digest," what should we do?

The straightforward answer is: increase the parameter count $N$ accordingly. But there's a problem here: as we discussed in 《Transformer升级之路：Dense前传》, for the currently mainstream Dense (non-MoE) models, $N \propto d^2$ where $d$ is the model width, while compute $C \propto N D \propto d^2 D$. This means that if $D$ is fixed and we want to increase $N$ by a certain multiple, then correspondingly $C$ needs to increase by the same multiple.

However, in real-world scenarios, both compute and data are typically hard constraints. That is, the total compute budget is usually roughly fixed by the training duration and the number of available GPUs, while the available data volume for training is also basically capped by the total scale of internet corpora after cleaning (excluding synthetic data for now). In other words, when we want to increase the parameter count $N$ to prevent it from becoming the bottleneck, we hope to keep both $C$ and $D$ fixed — but for standard Dense models, this is impossible, because increasing $N$ necessarily increases $C=6ND$ as well.

Now we can start to see where MoE might come in: it needs to find a way to increase the parameter count $N$ without changing the compute $C$ and data volume $D$ (much) — this is exactly the sense in which MoE "decouples compute from parameter count." Once we see it this way, MoE stops looking like a technique invented purely for the sake of "reducing compute," but rather as a natural requirement that emerges once compute and data are both capped while we still want to increase the parameter count.

Sparsification

Once we've identified the goal — increasing $N$ while holding $C$ and $D$ roughly fixed — the next question is: how do we actually achieve this?

The compute formula $C\propto N D$ already tells us that if $D$ is fixed and we want $C$ to remain (roughly) unchanged while $N$ increases, then effectively only a fraction of the $N$ parameters can be participating in the computation for any given token — meaning we must "activate" only part of the total parameters at a time. This is precisely the core idea of "sparsification": instead of letting every parameter participate in every forward pass (as in Dense models), we let a subset of parameters handle each input, so that the total parameter count $N$ can grow largely independently of the per-token compute.

More concretely, this is realized via the "expert" mechanism: we split (part of) the parameters into multiple modules called "experts," and for a given input, only a small number of the most relevant experts are activated for computation, while the rest are skipped. This way, the total parameter count $N$ can be made very large (equal to the sum of parameters across all experts), while the activated parameter count per token — and hence the compute — stays roughly at the previous, smaller scale.

This is exactly what MoE (Mixture of Experts) does, and now we can see clearly the logical chain from "compute-data trade-off" to "MoE": once both compute $C$ and data $D$ are fixed while we still want the total parameter count $N$ to increase, we are forced towards sparsification, i.e., activating only part of the parameters at a time — and MoE is currently the most mature and effective way of implementing this sparsification.

Of course, one might reasonably ask: why must sparsification take the specific form of "experts"? Are there other forms of sparsification that might work just as well or better? This is indeed a meaningful question, but for now let's set it aside and return to it after we've digested the more basic logic of "why sparsification (in some form) is needed."

Summary

In this post, starting from the Trade-off between compute and data volume discussed in the previous post, and combined with the practical constraints that both compute and data are essentially capped, I tried to establish a chain of reasoning explaining why we need MoE: since increasing the parameter count under a Dense architecture would require compute to grow correspondingly, but compute and data are both essentially fixed, we're compelled to look for a way to decouple parameter count from compute — and sparsification (in the form of MoE) is exactly such a solution.

Can we approximate the sum of $n$ vectors by summing only $k$ of them? That would bring the computational cost down to $k/n$.

Sorting by norm

We actually already explored this problem in The Path of Low-Rank Approximation (III): CR, which, written as a mathematical formula, is

\begin{equation}\mathop{\text{argmin}}_{\lambda_1,\lambda_2,\cdots,\lambda_n\in\{0,1\}}\left\Vert\sum_{i=1}^n \lambda_i \boldsymbol{v}_i - \sum_{i=1}^n\boldsymbol{v}_i\right\Vert^2\quad\text{s.t.}\quad \sum_{i=1}^n \lambda_i = k\end{equation}

Let $\gamma_i = 1 - \lambda_i$, then it can also be written as

\begin{equation}\mathop{\text{argmin}}_{\gamma_1,\gamma_2,\cdots,\gamma_n\in\{0,1\}}\left\Vert\sum_{i=1}^n \gamma_i \boldsymbol{v}_i\right\Vert^2\quad\text{s.t.}\quad \sum_{i=1}^n \gamma_i = n - k\end{equation}

An exact solution to this problem is fairly difficult to obtain, but there is a simple approximate solution: when $\boldsymbol{v}_i$ are pairwise orthogonal, we have

\begin{equation}\left\Vert\sum_{i=1}^n \gamma_i \boldsymbol{v}_i\right\Vert^2 = \sum_{i=1}^n \gamma_i^2 \Vert\boldsymbol{v}_i\Vert^2 = \sum_{i=1}^n \gamma_i \Vert\boldsymbol{v}_i\Vert^2\end{equation}

The obvious optimal solution to the above is to set the $n-k$ $\gamma_i$ with the smallest norm $\Vert\boldsymbol{v}_i\Vert$ to 1, which is equivalent to picking out the $k$ vectors with the largest norms to approximate the sum of $n$ vectors. When $\boldsymbol{v}_i$ does not satisfy pairwise orthogonality, we still use it as an approximate solution. The geometric intuition is also quite clear: the larger a vector's norm, the less likely it is to be canceled out during summation, and hence the more prominent its contribution.

In addition, in Path to Low-Rank Approximation (III): CR we also discussed a probability-based sampling approximation, where the optimal sampling probability — derived under the assumption of minimum variance — also turns out to be proportional to the vector norm. So, overall, sorting by vector norm is a simple yet effective strategy.

The first hints of MoE

Now that we have a strategy in place — "pick the $k$ vectors with the largest norm" — a closer look reveals that it's not actually practical: to pick out the $k$ vectors with the largest norm, we'd have to compute the norms of all the vectors first, which in turn means computing all of the $\boldsymbol{v}_i$ in advance. But our original goal was precisely to reduce the amount of computation for $\boldsymbol{v}_i$!

Fixing the Contradiction

To resolve this contradiction, we need to redesign each Expert model so that its norm can be computed cheaply. What does this mean exactly? First, we normalize $\boldsymbol{v}_i$ to get $\boldsymbol{e}_i = \boldsymbol{v}_i/\Vert\boldsymbol{v}_i\Vert$, so that every $\boldsymbol{e}_i$ now has the same norm. Next, we define

\begin{equation}\underbrace{[\rho_1,\rho_2,\cdots,\rho_n]}_{\boldsymbol{\rho}} = h(\boldsymbol{x}\boldsymbol{W}^{(R)})\quad\in\mathbb{R}_{\geq 0}^n\end{equation}

Here $\boldsymbol{W}^{(R)}\in\mathbb{R}^{d\times n}$ is a parameter matrix and $h(\cdot)$ is a $\mathbb{R}\to\mathbb{R}_{\geq 0}$ activation function — in plain terms, this is just a linear transformation from $d$ dimensions to $n$ dimensions followed by an activation function, so the computational cost is quite small. This part of the model is called the "Router" in MoE.

What does $\boldsymbol{\rho}$ do? It predicts the norm of each Expert! In other words, we treat $\rho_i$ as the norm of the $i$-th Expert, and $\rho_i \boldsymbol{e}_i$ is the complete Expert, decomposed into two parts: the norm $\rho_i$, which is cheap to compute, and the direction $\boldsymbol{e}_i$, which is expensive to compute. To reduce the computational cost, we first compute $\boldsymbol{\rho}$, pick out the top $k$ largest ones, and only then compute the corresponding $\boldsymbol{e}_i$, finally multiplying by $\rho_i$ and summing:

It looks like your message only contains a placeholder (\begin{equation}\boldsymbol{y} = \sum_{i\in \mathop{\text{argtop}}_k \boldsymbol{\rho}} \rho_i \boldsymbol{e}_i\end{equation}) with no surrounding text to translate. Could you please share the full text you'd like translated, including all the placeholders in context? That way I can preserve the markdown structure and placeholder positions accurately.

This is the basic formula for the MoE model. Since only the Top-$k$ part is kept in the computation, it is essentially a form of sparse model, whereas the original FFN, or the model when $k=n$, is usually referred to as the corresponding dense model.

Summary of the approach

Whether or not the reader is already familiar with MoE, the derivation above may feel a bit unfamiliar, since this is a route to understanding MoE that I worked out on my own, so to speak, behind closed doors. But because its geometric meaning is clearer, it should in essence be easier to understand.

Let's lay out the whole train of thought once more.

1. A standard dense FFN can be equivalently rewritten as the sum of $n$ expert vectors $\boldsymbol{v}_1,\boldsymbol{v}_2,\cdots,\boldsymbol{v}_n$;
2. To save computation, we try to pick out $k$ of these vectors and sum them to approximate the original sum of $n$ vectors;
3. After converting this into a math problem and solving it, we find that the selection rule is to pick the $k$ vectors with the largest norms;
4. Directly computing the norms of all $n$ experts and then selecting $k$ of them doesn't actually save any computation, so we need to redesign the experts;
5. We normalize $\boldsymbol{v}_i$ to get $\boldsymbol{e}_i$, and then use a separate small model (the router) to predict the norm $\rho_i$, so that the final expert becomes $\rho_i \boldsymbol{e}_i$;
6. At this point, we can first compute all the $\rho_i$, pick out $k$ of them, and only then compute $\boldsymbol{e}_i$, thereby achieving the goal of saving computation.

Why do it this way

Some readers might be wondering: why go through this seemingly complicated process? Wasn't the original MoE already easy enough to understand? The general form of MoE is

\begin{equation}\boldsymbol{y} = \sum_{i\in \mathop{\text{argtop}}_k \boldsymbol{\rho}} \rho_i \boldsymbol{v}_i\end{equation}

This means the normalization by $\boldsymbol{v}_i$ before the sum is missing, and $\rho_i$ no longer has any meaning as a norm — it is purely a scoring model used to rank the Experts (i.e., the Router). But why would multiplying $\rho_i$ onto the Experts allow the Router to learn to rank the Experts correctly? The author has found that only Sparse Backpropagation for MoE Training offers an explanation for this, though it's still not entirely intuitive.

Under the geometric perspective of this article, many of these questions suddenly become "obvious." Once we reparameterize the Experts as $\rho_i \boldsymbol{e}_i$, the Dense model corresponds to summing over all $\rho_i \boldsymbol{e}_i$, while MoE corresponds to selecting the Top-$k$ out of $\rho_i$ and summing over those, which constitutes a theoretically justified approximation of the Dense model. We haven't even needed to think about how the Router chooses Experts — we simply try, at every step, to approximate the Dense model as closely as possible. This can be seen as the best possible way to have your cake and eat it too: both a large parameter count and a small computational cost.

Now the geometric meaning of $\rho_i$ is a norm rather than a probability, so the activation function $h(\cdot)$ no longer needs to satisfy any normalization requirement. Besides softmax, we could consider using Sigmoid or ReLU, or the smooth Top-$k$ approximation introduced in Softmax's Sequel: Finding a Smooth Approximation to Top-K. Using a non-normalized activation function for the Router helps avoid vicious competition among Experts during $k > 1$, and sometimes yields better results.

One final addendum: earlier we defined $\boldsymbol{e}_i = \boldsymbol{v}_i/ \Vert\boldsymbol{v}_i\Vert$ with the goal of making all $\boldsymbol{e}_i$ have the same norm. In practice, this doesn't have to be implemented as an L2 normalization specifically — any equivalent operation will do, such as RMSNorm with the gamma parameter fixed to 1, which is more in line with our usual coding conventions.

Summary

In this post we set out from the problem of finding the best approximation to a Dense model, and used it to derive and understand MoE, arriving at a particular form of MoE. It adds one extra Normalize step compared to existing MoE designs, but this makes the geometric meaning of MoE much more apparent. Of course, whether or not we include the Normalize step, the road for MoE has only just begun — plenty more difficulties still lie ahead.