What is a Mixture-of-Experts model, really?

If you have ever gone window shopping for today‚ open-weight language models, you have probably seen model descriptions like this:

30B total parameters, but only 3B active per token.

At first glance, this sounds really strange. In a typical dense model, more parameters mean more weight. If the model has 30 billion parameters, these parameters are part of the computations. Hence, what does it mean that only 3 billion of them are active? What about the other 27 billion parameters then? Do they sit idle?

It is the idea of Mixture-of-Experts, or MoE, models.

Hi I'm Darren. During my internship with the AI Practice team of GovTech, I attempted to investigate the inner workings of such a model. Coming into the internship, I found myself being a little like a fish out of water. More precisely, a cybersecurity fish surrounded by artificial intelligence fish. So I did what any malware analyst would do: try my best to reverse engineer what goes on inside.

Before discussing MoE models, let us go over the basics of transformers.

A language model processes textual prompts in three major steps: breaking the text into tokens and embedding tokens to vector space. Afterward, the vectors undergo a series of operations that involve going through the stack of transformer layers. Transformers contain two key operations:

1. Attention
2. A feed forward network or FFN

As its name implies, the attention operation allows tokens to interact with each other.

For instance, in the sentence:

The Eiffel Tower is in the city of

the last token requires information from the previous tokens. 'City of' alone is insufficient to make sense. The connection to 'Eiffel Tower' must be established. And this process of mixing information happens thanks to attention.

FFN, in contrast, rewrites information about individual tokens. The expectation is that FFN can incorporate knowledge of our world into the model in a way that will allow us to use the token information for predictions. Put another way, attention accumulates intra-related context about the original text, while FFN accumulate external contexts into the original text.

In contemporary SOTA transformer models, most parameters belong to the feed forward networks within transformer blocks. It makes sense; the FFN is responsible for acquiring knowledge in our model. The number of parameters tends to vary between 55-65%. That means that if the whole model required 32 GB of VRAM, roughly 20GB was dedicated to feed forward network parameters!

Hence, if we are attempting to decrease the weight of the model, a good place to start looking would be FFN parameters.

Changes made by Mixture of Experts

MoE is among these attempts. Whereas a usual transformer contains only one FFN block, MoE model uses multiple blocks named experts. Then, for every token, a learned small neural network called a router chooses which experts should process it.

Thus, when a model claims that:

30B total parameters, but only 3B active per token.

it means that the model owns all 30 billion of parameters, but every token uses a small portion of them.

In the model that I experimented with, each MoE block contained 128 experts. But only 8 experts processed a particular token. Other experts remained inactive at the moment, although they are available to deal with either other tokens, the same token in another layer, etc. Thus, the advantage of MoE models is that they can have a huge capacity without using all of it in computation. But they also pose a question:

What does our router learn to choose between different experts?

A naive hypothesis is that each expert specialises in a particular topic. For example, there can be geography experts, code experts, grammar experts, and so on. If the prompt involves the Eiffel Tower, the router would select French or geography experts.

An alternative hypothesis is that the experts do not focus on a topic, but specialise in computation. Experts in early layers build token representations. Experts in middle layers retrieve facts from memory. Lastly, experts in late layers turn internal information into concrete token choices.

But how can we test the hypotheses? Mechanistic interpretability framework allows running some interesting experiments.

The Experiment

My prompt was the following:

The Eiffel Tower is in the city of

With a correctly working router, a model gives a sensible output:

The Eiffel Tower is in the city of Paris, France. The Eiffel Tower is a famous landmark and is one of the most recognisable structures in the world.

However, to run my test, I made a single change. Instead of selecting the top 8 experts, the router had to select the bottom 8 experts. In other words, I was forcing the model to use experts that it did not want to use. However, all else remained unchanged: attention was computed normally, residual stream flowed as before, selected experts received hidden states as before.

But this intervention could be applied to only certain layers. So what happens when we change the expert selection in different layers?

All layers wrong-routed

First, I wrong-routed all layers in MoE blocks.

As a result, the output was the following:

.getAction.getAction.getAction.getAction...

This is what you would expect from an incorrect experiment. If the selection of all experts is wrong, the model generates meaningless gibberish. Thus, this case proves our hypothesis that the router is essential, but is not informative by itself. We need to look closer at partial routing failures!

Only early layers wrong-routed

In this experiment, I wrong-routed only the first layers, i.e., layers 0 to 6. Overall, it is only 7 out of 48 layers.

The output turned out as follows:

的000000000000000000000000

Not a slightly worse completion of the prompt. Not even the English sentence. It consists of Chinese characters followed by many zeros.To me, this result is rather unexpected since almost all layers (41 of them!) in MoE blocks worked properly. Why was the output so bad? I assume that MoE blocks of the early layers build up some kind of representation of the input. Once this representation is corrupted, further processing becomes meaningless.The result is similar to trying to execute a random sequence of instructions in a computer program. Its behavior will be unpredictable since there is likely no sensible code to execute at all.

Middle layers wrong-routed

Now, let me wrong-route layers 24 to 40.

The output became the following:

Paris. The Eiffel Tower is in the city of Paris. The Eiffel Tower is in the city of Paris

This output is significantly different from previous cases.

It is clear that the model remembers the correct answer - 'Paris'. Nevertheless, once it retrieves the answer, it starts repeating it endlessly.

This result tells us that by layer 24, the token‚ 'Paris'‚ has been built into the internal representation of the model. Even if the next few layers operate improperly, they still manage to preserve this answer. In other words, there is some fact-retrieval operation performed by the middle layers.

But this case reveals something else. Although the model manages to retrieve the fact successfully, it cannot continue generating the prompt. It circles the same sentence repeatedly.

Thus, we can conclude that middle layers also help to organise thoughts and choose the order of words.

Late layers wrong-routed

Lastly, I wrong-routed the last 7 layers, i.e., layers 40 to 47.

Here is the result:

____. ____ is the city of ____.

Unlike previous results, it does not consist of random gibberish. In fact, it retains the structure of the answer. It is clear that the model recognises that there should be words in the places that I highlighted. But the generation stops here since it does not generate tokens.

This suggests that the late layers also serve a special purpose. They take an internal representation and convert it into actual words.

Conclusion

Early layers build the initial token representation. Middle layers accumulate knowledge from it and retrieve facts. At the same time, they help in organising thoughts and choosing next tokens to emit. Finally, late layers finalise the result by converting an internal representation to tokens.In conclusion, MoE layers seem to perform multiple functions.

What I have learned from this internship

It was such a blast to join the exciting area of research, interpretability, and try to contribute to it with some novel results.

Entering the field anew in this internship was challenging since the rate at which new concepts and techniques are introduced into interpretability is very fast. Many beginners could easily find themselves struggling with catching up with recent developments in this sphere. Hopefully, this blog post brought that aha! moment to you as it did for me.

There's of course much left to discuss. In particular, in the interest of teaching, I glossed over some subtleties with writing code with today's interpretability frameworks. How can write cross-model intervention code? How can we scale these interventions across models? To discuss these problems, feel free to check out the code at this github repository (Note: permission is required to access this direct link).

And finally, in this internship at GovTech, I would like to thank my mentors, Shangru and Ryan, who were always ready to help me with technical difficulties and check in on me at every step of this journey. Big thank you guys!

Note: This article was written by Darren Lim, who has since completed his internship, and published on his behalf by Ryan Lin.