Mixture of Experts (MoE) replaces the single FFN with multiple "expert" FFNs and a router that selects which experts process each token. Only the top-K experts activate per token — typically K=2 out of 8–64 experts.
Key challenges: (1) Load balancing — prevent all tokens routing to the same expert, fixed with auxiliary loss. (2) Expert collapse — some experts never activate. (3) Communication — experts on different GPUs need token routing across devices.
Python — simplified MoE layer#
Mixtral 8×7B has 46.7B parameters but uses 12.9B per token. Solving the two equations gives about 5.63B per expert and 1.63B shared (attention, embeddings). With 8 experts and top-2 routing, each expert should see 25% of tokens; keeping it near that is what the load-balancing loss is for.
- Active parameters set compute, but all parameters must sit in memory; MoE saves FLOPs, not GPU memory.
- “8×7B” is not 56B: the experts share attention and embeddings.
- Routing can collapse onto a few experts, and experts do not specialise into neat human topics.
total, active, experts, top_k = 46.7e9, 12.9e9, 8, 2 # Mixtral 8x7B, published counts per_expert = (total - active) / (experts - top_k) # total = shared + 8E, active = shared + 2E shared = active - top_k * per_expert share_per_expert = top_k / experts # tokens each expert sees under perfect balance
Assumes import numpy as np. Configurations named after real models use their published shapes; other numbers are made up or simulated, as the comments say.
The same idea with a library
import torch, torch.nn as nn, torch.nn.functional as F
class MoELayer(nn.Module):
def __init__(self, d_model, n_experts=8, top_k=2):
super().__init__()
self.top_k = top_k
self.router = nn.Linear(d_model, n_experts, bias=False)
self.experts = nn.ModuleList([
SwiGLU_FFN(d_model) for _ in range(n_experts)
])
def forward(self, x):
# x: (batch, seq, d_model)
logits = self.router(x) # (B, T, E)
weights, indices = logits.topk(self.top_k) # top-K experts
weights = F.softmax(weights, dim=-1) # normalize
out = torch.zeros_like(x)
for k in range(self.top_k):
for e in range(len(self.experts)):
mask = indices[..., k] == e
if mask.any():
out[mask] += weights[..., k:k+1][mask] * \
self.experts[e](x[mask])
return out- N. Shazeer et al., “Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer”, ICLR, 2017
- W. Fedus, B. Zoph & N. Shazeer, “Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity”, Journal of Machine Learning Research 23(120), 2022
- A. Q. Jiang et al., “Mixtral of Experts”, arXiv:2401.04088, 2024
Last reviewed October 2026 · How this site checks its content · Report a mistake