Seyed Masoud Hosseini · Overview · Study log · Weekly summaries · Ideas · Search · Transcript · RSS feed

Language Modeling from Scratch · Lecture 4 of 17 · 1:22:04

Lecture 4: Mixture of Experts

Stanford CS336 Language Modeling from Scratch | Spring 2025 | Lecture 4: Mixture of experts on YouTube

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 K experts, typically K=2 or 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

  1. Why does sparse activation let a model gain more parameters without increasing training or inference flops?
  2. What problem does an auxiliary load-balancing loss solve, and what happens to training if you omit it entirely?
  3. How does DeepSeek V3's auxiliary-loss-free balancing differ mechanically from the switch-transformer-style balancing loss?
  4. Why can token dropping make inference nondeterministic even at temperature zero?
  5. What tradeoff motivates slicing experts into smaller, fine-grained pieces instead of using a handful of full-size experts?

Vocabulary

sparse (adjective)
Using only a small active portion of a larger system at a time.
Sparse activation lets a model grow without more compute per token.
router (noun)
A component that decides which expert or path should handle each input.
The router sends each token to its best-matching experts.
expert (noun)
One of several smaller sub-networks in a mixture-of-experts model.
Each token is routed to only a few experts, not all of them.
affinity (noun)
How well two things match or suit each other.
The router ranks experts by their affinity for each token.
gating weight (noun)
A number that controls how strongly an expert's output contributes to the final result.
The top-K experts' outputs are combined using gating weights.
load balancing (noun)
Distributing work evenly across resources so none are overloaded or idle.
Load balancing keeps every expert receiving a fair share of tokens.
collapse (verb)
To fall into a much smaller, less useful state than intended.
Without balancing, training can collapse onto just one expert.
auxiliary loss (noun)
An extra training objective added alongside the main one to encourage a desired behavior.
An auxiliary loss encourages even routing across experts.
bias term (noun)
An extra learned value added to shift a result up or down.
DeepSeek V3 learns a per-expert bias to nudge routing.
parallelism (noun)
Running different parts of a computation at the same time on different hardware.
Expert parallelism spreads experts across separate devices.
non-differentiable (adjective)
Not smooth enough to compute an exact gradient for, making standard training methods harder to apply.
Routing decisions are non-differentiable, so heuristics are used instead.
upcycle (verb)
To reuse an already trained model as the starting point for a new one.
Some teams upcycle a dense model into a sparse MoE model.
device-limited routing (noun)
A rule restricting which devices a token's chosen experts can be spread across.
Device-limited routing reduces communication cost across machines.
dominate (verb)
To be the most common or important choice among alternatives.
Mixture-of-experts architectures now dominate frontier models.
fine-grained (adjective)
Divided into many small, detailed pieces.
DeepSeek used fine-grained experts sliced smaller than usual.
shared expert (noun)
An expert that processes every token, in addition to the ones chosen by the router.
A shared expert absorbs common processing every token needs.
token dropping (noun)
Discarding a token's assignment when its chosen expert is overloaded.
Token dropping can make outputs vary depending on batch composition.
stability (noun)
The quality of staying steady and not breaking down during training.
Computing the router softmax in float32 improves stability.
ablation (noun)
An experiment that removes a component to test its effect on results.
Ablations showed fine-grained experts helped consistently.
innovation (noun)
A new and original idea or method.
DeepSeek's key innovation was slicing experts smaller.

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 →