Mixture of Experts

In this post, I will do a deep dive into Mixture of Experts, a Sparse neural network which replaces a transformers Feed Forward block with multiple parallel "Expert" sub networks, which helps in saving compute when we process tokens.

Why does it matter?

Large Scale AI models often contain Billions to Trillions of parameters. Parameters are basically the Weights & Biases that the model learns during the training phase. When we have to pass a single token through the network, we dont want to pass it through entire network, which can be computationally expensive and time consuming, which is basically what a Dense network does.

Mixture of Experts on the other hand actives Experts, usually sub-networks of the whole network, depending on the input token, through what is called Gate, which I will discuss in the later part of the article.

Dense vs. Sparse Networks

Parameters are the weights and bias that the network learned during training.

First of all, lets see how Dense vs. Sparse Network looks like.

Dense Network

In a dense network, every parameter of the network is used by the token. When the token reaches the feed forward block, the token is multiplied through the full weight matrix of the network, which can be quite expensive. We need to run each token through entire network every time. For example, if there is a 70 Billion parameter dense model, then it has to do 70 Billion parameters worth of matrix multiplation+addition for each token.

Every time we pass a token over each parameter, we need two operations - multiply, & add, we are computing two Floating point operations, so for each token we need 2 FLOPs and assume a 70 Billion parameter model which is Dense, we need 140 Billion FLOPs which is too expensive computationally.

Sparse Network

In a Sparse Network, when we pass the token through the network, few sub groups of the network called experts get activated. So we need not do the arithmetic operation on the entire network parameters. We do the arithmetic operation using the parameters of the activated experts. In the above Spark Network(Mixture of Experts) visual, you can see that Expert 1 and Expert 2 got activated by the Router/Gate based on the given input token X.

MoE replaces FFN in the Transformer block

Each transformer layer has two sub-layers - a multi head attention, feed forward network, a residual connection, and a normalization step. Mixture of Experts replaces the feed forward network of the transfomer block.

Feed Forward Network holds the factual knowledge, for example:

- The capital of India is -> New Delhi
- The closest star to Earth is -> Sun

These facts are learned by the model from traning data and converted to weights at each neuron. So when we pass a token through the FFN, it acts as a lookup table.

We replace this FFN block with Mixture of Experts because we need this factual knowledge and we only activate parts of the network, called experts based on the content of the input token. When we only activate certain(K) number of experts, we are making the lookup cheaper than the FFN. So Mixture of Experts is like a trick that scales the factual knowledge store

SwiGLU:

It is like a smooth version of ReLU

Each expert contains a router + experts. and inside each expert we use SwiGLU activation function. ReLU converts all values below 0 to 0 whereas SwiGLU creates two projections of the input, passes one projection through the Swish activation (Swish(x) = x * Sigmoid(x)), and multiplies the two projections using element wise before creating a final down projection(W_down). Compared to ReLU, which converts all negative values to 0, SwiGLU is much smoother and lets the little negative singal pass through the network.

Gate/Router weights:

Gate weights are computed for every new token x, that decide how much each selected expert contributes. And based on these weights, after the gate/router scores the experts and the Top-K number of experts are chosen, we use softmax on these selected scores to produce weights(that sum to 0, because sum of probabilities is 1). We can easily spot the confidence of the router, usually experts that get higher probability means that the router is much more confident on these experts. Example, a router is confident on the expert which gets 0.9 probability and less confident on the expert that got 0.1 probability.

How to choose the "K" no.of experts

There is no automated way to choose the K. It is a hyper parameter that needs to be experimented with.

Example of MoE applied over Next Word Prediction

As each token passes through the layer, the router sends this token to the selected experts. For example, as you can see in the above image, the token "capital" go through two experts, E1 & E6, and token India goes through two experts, E4 & E6. After the final layer, the last token's representation is converted into probability distribution, and the word with the highest probability which is New Delhi. with 0.71 gets selected.

Comparison of few Large Scale AI Models, if they were Dense vs. Sparse

We already discussed that the FFN will multiply+add the token through the entire FFN weights, whereas the MoE model routes the token through the weights of the activated experts only.

Active/token - Are the neurons which get active when we pass a token through the network.

MoE share is calculated based on the active/token neurons. (share = (active per token / total params) * 100%)

What about the Compute Saved?

One FLOP(floating point operation) is one multiply and addition,

and 1 GLOP = 1 billion FLOPs. (and even though out of scope 1 TFLOP = 1 Trillion FLOPs)

I already discussed earlier in the article about what a FLOP is, and I will again with example here.

y1 = w11 * x1 + w12 * x2 + w13 * x3

In the above equation, we are computing 7 FLOPs. - 4 multiplications, and 3 additions. and it is 0. To compute GLOPs, it is - 000000007 FLOPs * 1 Billion GFLOPs.

In the above image we can see that if Kimi K3 were a dense network, it consumes 5600 GFLOPs per token whereas if it were Mixture of Experts, it would consume 208 GFLOPs per token, which is 27x saved in computation.

Memory is still the same between Dense vs. Sparse, all the neurons still hold all the parameters. No change in that.

Load Balancing the Experts

During early phase of training, few experts may get slightly more traffic, and improve faster because they see more tokens, and they end up getting more traffic, like a Rich get Rich situation.

and some experts become dead because they get very less traffic.

Do overcome this, we add auxilliary load-balancing loss which penalizes the imbalance by multiplying per expert, the fraction of tokens sent to it, by averaging the routing probability and sum of all.

Here,

• \(f_i\) = fraction of tokens sent to the expert
• \(P_i\) = Mean Router probability
• \(N\) = Number of experts in the layer
• \(i\) = Expert index, from 1 to N
• \(\alpha\) = a small weight, which is a hyperparameter, that sets how strongly balancing is pushed.

Fine Grained Experts and Shared Experts

Large LLM Models uses a combination of fine gradined and shared experts.

A shared expert is an always on that every token passes through, without any routing. We need shared experts because some knowledge is common like it can be basic syntax, or frequent patterns, or some general structure in the data. Without a shared expert, each routed expert has to handle independently and independently store this common knowledge which might cause duplication and wasted capacity and also memory.

A fine grained expert is an expert choosen by the router to learn specific knowledge. These often go together in an Mixture of Experts.

How do we scale MoE across multiple GPUs?

We cannot fit all the experts in one GPU and we have to spread these experts across multiple GPUs. Each GPU keeps just a few of the MoE layer's experts. The router looks at every token and decides which expert each token has to go to, and then the correspending GPU will send the token across the network to whichever GPU holds that choosen expert and each GPU now has a pile ot tokens which computes the arthmetic and sends the computation back to the original GPU the token originated from.