Transformer and Mixture of Experts, explained visually! Mixture of Experts (MoE) is a popular architecture that uses different experts to improve Transformer models. Transformer and MoE differ in the decoder block: - Transformer uses a feed-forward network. - MoE uses experts, which are feed-forward networks but smaller compared to those Transformer. During inference, a subset of experts are selected. This makes inference faster in MoE. Also, since the network has multiple decoder layers: - The text passes through different experts across layers. - The chosen experts also differ between tokens. But how does the model decide which experts should be ideal? The router does that. It is a multi-class classifier that produces softmax scores over experts to select the top K experts. The router is trained with the network, and it learns to select the best experts. But it isn't straightforward. There are challenges! Challenge 1) Notice this pattern at the start of training: - Say, the model selects "Expert 2" - This expert gets a bit better - It may get selected again since it's the "best" - It learns more - It gets selected again in the next iteration - It learns more, and so on! This means many experts can go under-trained due to the overselection of a few experts! We solve this in two steps: - Add noise to the feed-forward output of the router so that other experts can get higher logits. - Set all but the top K logits to -infinity. After softmax, these scores become zero. This way, other experts also get the opportunity to train. Challenge 2) Some experts may get exposed to more tokens than others, leading to under-trained experts. We prevent this by limiting the number of tokens an expert can process. If an expert reaches the limit, the token is passed to the next best expert. Overall, MoEs have more parameters to load. But a fraction of them are activated during inference. This leads to faster inference. Mixtral 8x7B and Llama 4 are two popular MoE-based LLMs. Have you used MoEs in production yet?