Interviewer asks: We need a language model that works well in code, mathematics, and creative writing.
How to achieve multi-domain performance?
You: I'll increase the number of attention heads.
interview ends…
➡️ What you missed?
Attention heads capture patterns, not domain expertise.
More heads = richer representations in a single pass.
More experts = dedicated subnets for different types of knowledge.
The correct answer: Mixture of Experts (MoE).
Let's break down how MoE differs from standard transformers:
Transformers and MoE differ in the decoder block:
- Transformers use a fully connected feed-forward network.
- MoE uses experts — also fully connected feed-forward networks, but smaller in size compared to transformers.
During inference, a subset of experts is selected. This speeds up inference in MoE.
Since the network contains multiple decoder layers:
- the text passes through different experts at different layers
- the selected experts also differ for different tokens
But how does the model decide which experts are better suited?
This is done by the router. Let's break it down further.
Task 1) Note this pattern at the beginning of training:
- the model selects "Expert 2"
- the expert slightly improves
- it might be selected again
- the expert trains further
- it's selected again
- it continues to train
- and so on
Many experts remain undertrained.
We solve this in two steps:
- Add noise to the output of the fully connected layer of the router so that other experts can receive higher logits.
- Set all logits except the top-K to -infinity. After softmax, their values become zero.
This way, other experts also get the opportunity to train.
Task 2) Some experts might receive more tokens than others — this leads to some experts remaining undertrained.
This is prevented by limiting the number of tokens that a single expert can process.
If an expert reaches the limit, the input token is redirected to the next most suitable expert.
MoE contains more parameters to load. However, only a portion of them are activated, as only a limited number of experts are selected.
This leads to faster inference. Mixtral 8x7B from MistralAI is a well-known language model built on MoE.
Visuals on first comment that compares transformers and MoE again
••••••••••••••••••••••••••••••••••••••
🤖 Data & ML | @DataXplore