Seyed Masoud Hosseini · Overview · Study log · Ideas · Transcript · RSS feed

Language Modeling from Scratch · Lecture 3 of 17 · 1:27:03

Lecture 3: Architectures and Hyperparameters

Stanford CS336 Lang. Modeling from Scratch | Spring 2025 | Lec. 3: Architectures, Hyperparameters on YouTube

Study guide

What this lecture covers

Rather than deriving a transformer from first principles, this lecture takes an empirical approach: having surveyed roughly nineteen dense model releases from the past year plus the major earlier models, it asks which architectural choices have converged to consensus and which are still genuinely varied. It covers normalization placement and type, activation functions, serial versus parallel transformer blocks, and position embeddings, then moves to hyperparameters (feedforward ratio, head dimensions, aspect ratio, vocabulary size, regularization), and closes with training-stability tricks and attention variants aimed at inference efficiency.

After watching, you should be able to explain why nearly all modern LLMs use pre-norm, RMSNorm, no bias terms, SwiGLU-style gated activations, and RoPE position embeddings, and be able to reason about hyperparameter choices like the feedforward-to-model-dimension ratio using the same evidence the lecture cites, rather than picking numbers arbitrarily.

Key ideas

  • Pre-norm dominance: placing layer normalization before each attention/MLP sublayer (rather than after, as in the original transformer) removes the need for careful learning-rate warmup and produces far more stable training, which is why virtually every modern model uses it.
  • RMSNorm over LayerNorm: dropping the mean-centering and bias-shift steps of LayerNorm gives a comparable or better model while doing less memory movement, which matters because normalization operations take a large share of runtime despite being a tiny share of total flops.
  • Gated activations (SwiGLU, GeGLU): adding an elementwise gate to the MLP consistently improves loss over plain ReLU/GELU MLPs, at the cost of an extra weight matrix, which is why gated models shrink the hidden dimension by roughly 2/3 to keep parameter counts matched.
  • RoPE (rotary position embeddings): encodes relative position by rotating query/key vectors by an angle proportional to position, so inner products depend only on relative distance; nearly all surveyed models have converged on it.
  • Consensus hyperparameter ratios: feedforward dimension is about 4x the model dimension for standard MLPs (about 2.6x-2.7x for gated ones); head dimension times number of heads roughly equals the model dimension; the aspect ratio (model dimension per layer) clusters around 128; vocabulary sizes have grown from 30-50k to 100-250k as models become multilingual.
  • Weight decay's real role: in single-epoch pre-training, weight decay isn't preventing overfitting (train/validation gap is unaffected) — it interacts with the learning-rate schedule to produce lower training loss late in training.
  • Stability interventions: z-loss (penalizing the softmax normalizer for drifting from 1) and QK-norm (normalizing queries and keys before the attention softmax) have become common ways to prevent gradient-norm spikes in large training runs.
  • KV-cache-aware attention variants: multi-query attention (MQA) and grouped-query attention (GQA) share key/value heads across multiple query heads to reduce the memory-bound cost of autoregressive decoding.

Walkthrough

Pre-norm versus post-norm (6:09)

The lecture reviews the original transformer's post-norm design (layer norm applied after each residual addition) versus the now-standard pre-norm design (layer norm applied before each sublayer, inside the residual branch). Early papers showed post-norm required careful warmup to avoid instability, while pre-norm trains stably without it because the residual stream keeps a clean identity path for gradients. A newer "double norm" variant, used by models like Grok and Gemma 2, adds a further layer norm after the sublayer as well, which some evidence suggests is even more stable at large scale.

RMSNorm and dropping bias terms (11:12)

Most models (LLaMA, PaLM, Chinchilla, T5) have moved from LayerNorm to RMSNorm, which skips mean subtraction and the learned bias/shift term. The lecture cites profiling showing that matrix multiplications account for about 99.8% of a transformer's flops but normalization operations still consume roughly 25% of runtime, because they involve heavy memory movement rather than compute — explaining why a change that saves negligible flops still measurably speeds up training. The same logic motivates dropping bias terms from linear layers almost everywhere, which also appears to improve training stability for reasons not fully understood.

Gated activations: SwiGLU and friends (19:17)

The lecture surveys activation functions (ReLU, GELU) and then gated linear units (GLU, GeGLU, SwiGLU), which multiply the MLP's hidden activations elementwise by a second linear projection of the input before the output projection. Citing Noam Shazeer's original GLU variants paper and a follow-up ablation study, it shows GLU variants consistently outperform their ungated counterparts on held-out loss, which is why the course's assignment specifies SwiGLU. Because the extra gating matrix adds parameters, gated models typically shrink the feedforward hidden size by a factor of about 2/3 to keep total parameter count comparable to an ungated MLP.

Serial versus parallel blocks, and RoPE (28:22)

Most transformers compute attention and the MLP serially within a block, but a minority (GPT-J, PaLM, Cohere's Command models) compute them in parallel from the same input and sum the results, trading some expressiveness for systems efficiency. The lecture then covers position embeddings in more depth: sinusoidal and learned absolute embeddings leak or ignore relative-position information, while RoPE rotates query and key vectors — in blocks of two dimensions, each rotating at a different rate — by an angle proportional to token position, so that inner products between any two positions depend only on their relative distance. This property, combined with strong empirical results even at small scale, has made RoPE close to universal in recent models.

Hyperparameter rules of thumb (40:31)

Drawing on a spreadsheet of real model configurations and Jared Kaplan's scaling-law paper, the lecture identifies several conventions: feedforward dimension is about 4x the model dimension for plain MLPs and about 2.6x-2.7x for gated ones (T5-11B's outlier 64x ratio was walked back in its successor); head dimension times number of heads is usually close to 1x the model dimension; the aspect ratio (model dimension divided by number of layers) clusters near 128 across several orders of magnitude of model scale; and vocabulary size has grown from 30-50k tokens for early monolingual models to 100-250k for modern multilingual, production-oriented ones. On regularization, dropout has largely disappeared from pre-training, and weight decay persists not to control overfitting (train/validation gaps are unaffected by it) but because it interacts with learning-rate decay schedules to produce measurably lower training loss late in training.

Stability tricks: z-loss and QK-norm (1:05:51)

The lecture identifies softmax operations — the output softmax and the attention softmax — as the main source of training instability (gradient-norm spikes) in large models. Two interventions have become common: z-loss, an auxiliary loss term (pioneered for LLMs by PaLM) that penalizes the log of the softmax normalizer for drifting from zero, keeping the softmax numerically well-behaved; and QK-norm, which applies layer normalization to queries and keys before computing attention scores, an idea borrowed from vision transformer training and now used by models like Gemma 2, DCLM and OLMo 2. A related but less common technique, logit soft-capping (used in Gemma 2), clips attention logits with a tanh-based function but was found by one comparison to hurt perplexity slightly, unlike QK-norm.

Attention variants for inference: MQA, GQA, and long context (1:15:00)

The lecture explains why standard multi-head attention becomes memory-bound during autoregressive decoding: the KV cache must repeatedly move previously computed keys and values in and out of memory, one token at a time, making arithmetic intensity poor. Multi-query attention (MQA) shares a single key/value head across all query heads to cut this memory traffic sharply; grouped-query attention (GQA) is a middle ground, sharing keys/values across groups of query heads rather than all of them. The lecture closes with a newer trick used in models like Llama 4: alternating blocks where most layers use sliding-window attention with RoPE (bounded local context) and one layer in every few uses full attention with no position embedding at all, which helps both systems efficiency and extrapolation to very long context lengths.

Before you watch

  • Watch Lectures 1 and 2 first — this lecture assumes familiarity with the course's assignment-one transformer implementation and the flop/memory accounting from Lecture 2.
  • A prior transformer architecture course (the lecture references Stanford CS224N) helps with the review of the original transformer, multi-head attention, and layer norm, since this lecture moves quickly past those basics to focus on variations.

Check your understanding

  1. Why does pre-norm training remove the need for learning-rate warmup that post-norm training required, and what does the residual stream have to do with it?
  2. RMSNorm and dropping bias terms save well under 1% of a transformer's flops — so why do they still measurably speed up training?
  3. Walk through why RoPE's rotation-based approach guarantees that attention scores depend only on relative position, not absolute position.
  4. Why does weight decay improve training loss in single-epoch pre-training even though it isn't preventing overfitting, according to the lecture's learning-rate-schedule explanation?
  5. Why does the KV cache make autoregressive decoding memory-bound rather than compute-bound, and how do MQA and GQA address that specific bottleneck?

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 2: PyTorch and Resource Accounting · Lecture 4: Mixture of Experts →