Stanford CS336: Language Modeling from Scratch | Spring 2025 | Architectures, Hyperparameters

Key Concepts

• Transformer Architecture: The fundamental structure for language models, including attention mechanisms, residual connections, and normalization layers.
• Pre-Norm vs. Post-Norm: Placement of LayerNorm before (pre) or after (post) the attention/feed-forward blocks in a transformer.
• RMS Norm: A simplified version of LayerNorm that omits mean adjustment and bias, improving speed and memory efficiency.
• Gated Linear Units (GLUs): Activation functions that use a gating mechanism to modulate the output of a linear transformation, such as SwiGLU and GeGLU.
• Serial vs. Parallel Layers: Arrangement of attention and MLP blocks in a transformer, either sequentially (serial) or simultaneously (parallel).
• Rotary Position Embeddings (RoPE): A relative position embedding technique that uses rotations to encode positional information in a way that is invariant to absolute position.
• Hyperparameter Tuning: Selecting optimal values for parameters that control the training process, such as hidden dimension, feed-forward size, and vocabulary size.
• Weight Decay: A regularization technique that penalizes large weights, often used to improve training stability and generalization.
• Stability Tricks: Techniques used to prevent gradient explosion and improve the stability of training large language models, such as z-loss and QK norm.
• Grouped-query Attention (GQA) and Multi-query Attention (MQA): Variants of multi-head attention that reduce the number of key and value heads to improve inference efficiency.
• Sparse Attention: Attention mechanisms that only attend to a subset of the input sequence, reducing computational cost and enabling longer context lengths.

1. Transformer Architecture Variations

1.1. Pre-Norm vs. Post-Norm

• The original transformer used post-norm, placing LayerNorm after each subcomponent (multi-head attention and feed-forward network).
• Modern LLMs predominantly use pre-norm, placing LayerNorm before each subcomponent.
• Reasoning: Pre-norm is more stable to train, avoids gradient attenuation, and allows for removing learning rate warm-up.
• Evidence: Papers by Salazar and Ngyuen, and Xiong (2020) compare pre-norm and post-norm, showing pre-norm allows removing warmup and is more stable.
• Double Norm: A recent innovation involves adding LayerNorm both before and after the feed-forward network (e.g., Grok, Gemma 2, Olmo 2).

1.2. LayerNorm vs. RMS Norm

• LayerNorm standardizes activations by subtracting the empirical mean and dividing by the standard deviation, scaling by gamma and shifting by beta.
• RMS Norm simplifies LayerNorm by dropping the mean adjustment and bias term.
• Reasoning: RMS Norm is faster and just as good as LayerNorm, reducing memory movement overhead.
• Evidence: Ivanov et al (2023) show that normalization operations account for 25% of runtime due to memory movement, even though they are only 0.17% of the flops.
• Example: Narang et al (2020) show that RMS Norm increases steps per second and achieves lower final loss compared to the vanilla transformer.

1.3. Bias Terms

• Most modern transformers omit bias terms in linear layers.
• Reasoning: Bias terms don't significantly improve performance and can negatively impact optimization stability.

1.4. Activations and Gated Linear Units (GLUs)

• Common activation functions include ReLU, GeLU, Swish, and their gated variants.
• ReLU: max(0, x)
• GeLU: x * CDF(Gaussian)
• Gated Linear Units (GLUs): Gate the output of a linear transformation with an entrywise linear term.
  • ReGLU: ReLU + GLU
  • GeGLU: GeLU + GLU (used by T5v1.1, Gemma 2, Gemma 3)
  • SwiGLU: Swish + GLU (very popular, used by LLaMA-family, PaLM, OlMo)
• Swish: x * sigmoid(x)
• Reasoning: GLUs consistently work well and improve performance.
• Evidence: Shazeer's original paper evaluates GLU variants and shows they perform better than ReLU. Narang et al (2020) also show that GLU variants achieve lower losses.
• Parameter Sizing: Gated models usually make the hidden size slightly smaller (by a factor of 2/3) to maintain the same total number of parameters.
• Exceptions: GPT-3 uses GeLU, Nemotron 340B uses squared ReLU, and Falcon 211b uses ReLU.

1.5. Serial vs. Parallel Layers

• Serial: Attention and MLP blocks are computed sequentially.
• Parallel: Attention and MLP blocks are computed simultaneously, and their outputs are added together.
• Reasoning: Parallel layers can improve parallelism and system efficiency.
• Examples: GPT-J, PaLM, Cohere Command A, Command R+, and Falcon 2 11B use parallel layers.

2. Position Embeddings

2.1. Evolution of Position Embeddings

• Early models used sine embeddings (original transformer), absolute embeddings (GPTs, OPT), or relative embeddings (T5, Gopher).
• Most modern models have converged to Rotary Position Embeddings (RoPE).

2.2. Rotary Position Embeddings (RoPE)

• Principle: Encode positional information using rotations, which are invariant to absolute position.
• Mathematical Formulation: The inner product of embeddings should be a function of the two words and the difference in their positions.
• Implementation: Divide the high-dimensional vector into blocks of two dimensions and rotate each pair by a different angle.
• Rotation Angle: Determined by the position of each word.
• Application: Applied at the attention layer, intervening on the queries and keys.
• Reasoning: RoPE allows for extrapolating context length and is empirically effective.

3. Hyperparameter Tuning

3.1. Feed-Forward Size

• dff: Feed-forward dimension (output hidden dimension of MLP).
• d model: Dimensionality of the input.
• Consensus: dff = 4 * d model for ReLU-style MLPs.
• GLU Variants: dff = (8/3) * d model ≈ 2.66 * d model.
• Examples: LLaMA-1, Qwen, DeepSeek, Yi, and T5 roughly follow the 2.6 rule.
• Exceptions: T5 uses a much larger multiplier (64x), with a hidden dim of 1,024 and a dff of 65,000. Gemma 2 follows in these footsteps and does a factor of 8.
• Evidence: Kaplan's scaling law paper shows that the optimal dff/d model ratio is around 4.

3.2. Head Dimensions

• Consensus: The dimension d (hidden dimension) should be equal to the head dimension times the number of heads.
• Examples: GPT-3, T5, LaMDA, PaLM, and LLaMA2 all have a ratio of one.
• Exception: T5 tried a big ratio of 16.

3.3. Aspect Ratio (Depth vs. Width)

• d model: Hidden dimension of the residual stream (controls width).
• n layer: Number of layers (controls depth).
• Consensus: About 128 hidden dimensions per layer.
• Reasoning: Aspect ratio affects parallelism and networking constraints.
• Evidence: Kaplan et al show that the optimal aspect ratio is around 100.
• Downstream Accuracy: E Tay et al found that deeper models might be better for downstream accuracy, even if they don't improve training loss.

3.4. Vocabulary Size

• Vocabulary sizes have been trending upwards.
• Early Models: 30,000 to 50,000 tokens.
• Multilingual/Production Systems: 100,000 to 250,000 tokens.
• Examples: Command A has very large vocab sizes. GPT-4 and others use around 100k tokens.

3.5. Regularization (Dropout and Weight Decay)

• Pre-training typically involves one epoch of training, so overfitting is not a major concern.
• Dropout has gone out of fashion.
• Weight decay is still commonly used.
• Reasoning: Weight decay interacts with learning rate schedules to improve training losses, especially at the tail end of training.

4. Stability Tricks

4.1. Stability Issues in Transformers

• Stability issues can arise in various parts of the transformer, but softmaxes are a common problem area.

4.2. Z-Loss

• Goal: Ensure that the softmax normalizer (z) is close to 1.
• Implementation: Add an auxiliary loss term to force log(z) to be close to 0.
• Formula: alpha * log(z)^2
• Examples: PaLM, Baichuan 2, DCLM, and OLMo 2 use z-loss.

4.3. QK Norm

• Goal: Control the inputs to the softmax to be bounded in size.
• Implementation: Pass the queries and keys through a LayerNorm layer before taking their inner product.
• Examples: Dehgani (2023), Chameleon, Idefcs, Gemma 2, DCLM, and OLMo 2 use QK norm.

4.4. Soft-Capping Logits

• Goal: Softly clip the logits that go into the softmax.
• Implementation: Pass the logits through a soft maximum operation (tanh).
• Examples: Gemma 2 and OLMo 2 use soft-capping logits.

5. Variations on Attention Heads

5.1. Grouped-Query Attention (GQA) and Multi-Query Attention (MQA)

• Motivation: Improve inference efficiency by reducing memory access patterns.
• MQA: Multiple query heads, but only one head for keys and values.
• GQA: Reduce the number of keys by some multiple.

5.2. Sparse Attention

• Motivation: Enable longer context lengths by reducing computational cost.
• Techniques:
  • Sparse attention patterns (OpenAI, 2019)
  • Sliding window attention
• Modern Implementation:
  • Transformer blocks with full self-attention (no position embedding) every few blocks.
  • Sliding window attention with RoPE in the other blocks.
  • Examples: LLaMA 4, Gemma, and Cohere command A.

6. Synthesis/Conclusion

This lecture provides a deep dive into the architecture and training of large language models, covering various design choices, hyperparameters, and stability tricks. It emphasizes learning from the experience of others by analyzing the evolution of transformer architectures and identifying key trends and consensus choices. The lecture also highlights the importance of considering system-level concerns, such as memory movement, in addition to computational complexity. By understanding these details, practitioners can make informed decisions when designing and training their own language models.