Seyed Masoud Hosseini · Overview · Study log · Ideas · Transcript · RSS feed
Language Modeling from Scratch · Lecture 4 of 17 · 1:22:04
Lecture 4: Mixture of Experts
Study guide
What this lecture covers
The lecture answers a practical question: why have mixture-of-experts (MoE) architectures become the default for state-of-the-art open models like DeepSeek, Qwen, Grok, and Llama 4? It builds up the idea from scratch, starting with the basic swap of a dense feed-forward block for a router plus many smaller expert blocks, then works through routing functions, training objectives, and systems tradeoffs.
This is a systems-and-architecture lecture in a course on building language models from scratch. By the end, you should be able to explain why sparse activation gives more parameters without more flops, how token-choice top-K routing works, why load balancing is necessary, and how DeepSeek's architecture evolved from its first MoE model to V3.
Key ideas
- Sparse activation: an MoE model replaces one large feed-forward block with many smaller "expert" blocks and a router that activates only a few per token, so parameter count grows without growing flops per token.
- Token-choice top-K routing: nearly all production models rank experts by affinity for each token and route it to the top
Kexperts, typicallyK=2or more; this has won out over expert-choice routing and global assignment schemes. - Fine-grained and shared experts: DeepSeek's innovation was slicing experts into smaller pieces (so you can afford many more of them) and optionally keeping one or a few "shared" experts that process every token for common structure.
- Load balancing is essential, not optional: without an auxiliary balancing loss, training collapses onto one or two experts and the rest go unused, wasting capacity even before any systems concern is considered.
- Auxiliary-loss-free balancing (DeepSeek V3): instead of an auxiliary loss term, DeepSeek V3 learns a per-expert bias that nudges routing scores up or down based on recent load, plus a smaller sequence-level auxiliary loss for inference-time robustness.
- Expert parallelism: because experts are sparsely activated, they can be placed on separate devices and tokens routed to the right device, giving a new axis of parallelism beyond data and model parallelism.
- Training instability: MoE routing is a non-differentiable, RL-like problem in practice solved with heuristic balancing losses; the router softmax is also a common source of instability, handled with float32 computation and an auxiliary z-loss.
- DeepSeek's architecture barely changed across versions: the core routing and expert design from DeepSeek's first MoE model carried through to V3, with additions like device-limited routing and auxiliary-loss-free balancing layered on top.
Walkthrough
Why mixture of experts now dominates (0:05)
The lecture opens by noting that most frontier models in 2025 use MoE architectures, and reviews evidence (from papers such as Fedus et al. 2022 and AI2's OLMoE) that at matched training flops, MoE models reach lower loss than dense models. The core mechanism is explained: the dense feed-forward block is replaced with a router and multiple smaller expert blocks, so a sparsely activated model gets more parameters without more flops per forward pass.
Design choices: routing, sizing, and training (14:15)
Three design axes are introduced: how tokens are routed to experts, how many experts there are and how big each one is, and how the non-differentiable routing decision gets trained. The lecture contrasts token-choice, expert-choice, and global-assignment routing, and explains that token-choice top-K has become the standard because it lets each token pick the expert best suited to it.
Top-K routing mechanics (19:18)
The router computes an inner product between the token's hidden state and learned per-expert vectors, applies a softmax (or sigmoid, in later models), and keeps only the top K scores as gating weights. Outputs from the selected experts are combined by a weighted sum and added back to the residual stream. Alternatives like hashing-based routing, RL-learned routing, and optimal-transport-based assignment are mentioned as historically explored but largely abandoned due to cost or instability.
Fine-grained and shared experts (31:26)
DeepSeek's key architectural idea is explained: slicing each expert into a smaller fraction of the standard feed-forward size lets a model afford many more experts at the same flop budget, and adding one or a few shared experts that always fire can absorb common processing. Ablations from DeepSeek and OLMoE show fine-grained experts helping consistently, while the benefit of shared experts is more mixed across replications.
Training the router: balancing losses (41:36)
Because routing decisions aren't differentiable, training relies on heuristic auxiliary losses rather than reinforcement learning or stochastic exploration, both of which were tried early on and abandoned. The classic switch-transformer balancing loss (comparing each expert's actual token fraction to its intended routing probability) is derived, along with device-level balancing losses for systems efficiency. DeepSeek V3's auxiliary-loss-free approach, which learns a per-expert bias term instead, is presented as a refinement, though a sequence-level auxiliary loss is still kept for inference-time robustness.
Systems concerns and stability tricks (58:48)
The lecture covers expert parallelism (placing experts on separate devices and routing tokens via all-to-all communication), token dropping when an expert is overloaded, and why this can make even "deterministic" inference produce different outputs depending on batch composition. Stability tricks are covered: computing router softmax in float32, adding a z-loss to keep softmax normalizers near one, and using upcycling (initializing an MoE from a trained dense model) as a cost-effective way to bootstrap a sparse model.
The DeepSeek V1 to V3 architecture walkthrough (1:10:01)
The lecture traces DeepSeek's architecture from its first MoE model (16B parameters, 2.8B active) through V2 (236B/21B active, adding device-limited routing and communication balancing losses) to V3 (671B/37B active, adding auxiliary-loss-free balancing and a normalized sigmoid gate). It closes with a brief tour of DeepSeek V3's non-MoE innovations: multi-head latent attention (compressing K/V into a smaller cached vector and merging projection matrices to avoid extra flops) and multi-token prediction as an auxiliary training objective.
Before you watch
- Be comfortable with the standard transformer feed-forward block and residual stream, since the lecture assumes you know what it's replacing.
- Review grouped-query and multi-head attention (GQA/MHA) from the previous lecture, since multi-head latent attention is explained as an alternative to them.
- Familiarity with softmax and basic gradient-based training will help with the routing and balancing-loss derivations.
Check your understanding
- Why does sparse activation let a model gain more parameters without increasing training or inference flops?
- What problem does an auxiliary load-balancing loss solve, and what happens to training if you omit it entirely?
- How does DeepSeek V3's auxiliary-loss-free balancing differ mechanically from the switch-transformer-style balancing loss?
- Why can token dropping make inference nondeterministic even at temperature zero?
- What tradeoff motivates slicing experts into smaller, fine-grained pieces instead of using a handful of full-size experts?
From the YouTube description
For more information about Stanford's online Artificial Intelligence programs visit: https://stanford.io/ai
To learn more about enrolling in this course visit: https://online.stanford.edu/courses/cs336-language-modeling-scratch
To follow along with the course schedule and syllabus visit: https://stanford-cs336.github.io/spring2025/
Percy Liang
Associate Professor of Computer Science
Director of Center for Research on Foundation Models (CRFM)
Tatsunori Hashimoto
Assistant Professor of Computer Science
View the entire course playlist: https://www.youtube.com/playlist?list=PLoROMvodv4rOY23Y0BoGoBGgQ1zmU_MT_
← Lecture 3: Architectures and Hyperparameters · Lecture 5: GPUs →
