45 / 90

Day 45: Multi-Head Attention Mechanics and Numerical Stability

Deconstruct Multi-Head Attention into linear subspace projections W^Q, W^K, W^V, W^O. Master tensor reshaping across h parallel heads (B, h, T, d_k), verify exact parameter ledgers, and resolve edge-case numerical instabilities in attention implementation.

Models & Training 480 active minutes (8 hours focused study)
01 120 min

Multi-Head Subspaces: Why and How Heads Divide d_model

Agent-Assisted Engineering Check

Prerequisites

Day 43-44 scaled dot-product attention and causal masking.

Learning Objective

Explain the representational necessity of multi-head attention, define linear projections W^Q, W^K, W^V, W^O, and verify parameter ledgers for tiny and standard configurations.

Instructions

Required work: explain why a single attention head compresses multiple distinct relationship types (e.g., syntactic dependency vs semantic topic) into a single average. Define MultiHead(Q,K,V) = Concat(head_1, ..., head_h) W^O, where each head_i = Attention(Q W_i^Q, K W_i^K, V W_i^V). Let d_model = 4, h = 2, so d_k = d_v = 2. Derive the parameter ledger: W^Q, W^K, W^V, W^O each have shape (4, 4) = 16 parameters (total 64 without bias, or 80 with 4x4 biases). Then compute the ledger for GPT-2 Small: d_model = 768, h = 12, d_k = 64 (4 * 768*768 = 2,359,296 weights, + 4*768 = 3,072 biases, total 2,362,368).

Practice & Prediction Log

1. Why can a single attention head not simultaneously focus on positional proximity and long-range syntactic agreement?
2. Write the parameter ledger for d_model = 4, h = 2 with and without biases.
3. Write the exact parameter ledger for GPT-2 Small attention layer (d_model = 768, h = 12).

Independent Mastery Criterion

Precise conceptual explanation of subspace routing and error-free parameter accounting for both tiny and GPT-2 scales.

Authoritative Source Resources & Provenance

linking arXiv.org perpetual non-exclusive license / Open Access

Attention Is All You Need

Author: Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Łukasz Kaiser, Illia Polosukhin | Original Source Link ↗

Relevant Section / Range: Section 3.2.2: Multi-Head Attention architecture, linear projections, and parameter equations

Verified Version & Date: NeurIPS 2017 / arXiv:1706.03762v7 (Verified: 2026-10-30)

Selection Rationale:

The foundational primary literature establishing scaled dot-product attention, multi-head projection mechanisms, and autoregressive causal masking.

Source Prerequisites:

Linear algebra, matrix multiplication, softmax probability distributions

Safe Fallback (Offline Resilience):

Complete original algebraic derivations, matrix examples, variance proofs, and tensor diagrams are self-contained in the curriculum text without external document dependency.

Compared Candidate Sources & Selection Analysis (3)
Bahdanau et al.: Neural Machine Translation by Jointly Learning to Align and Translate (ICLR 2015) ↗ compared_alternative [arXiv Open Access]

Pioneered additive attention in RNNs, but Vaswani scaled dot-product attention replaced additive alignment in modern LLM architectures due to parallel matrix multiplication speed.

Shaw et al.: Self-Attention with Relative Position Representations (NAACL 2018) ↗ compared_deferred [arXiv Open Access]

Introduces relative positional biases into attention; deferred to subsequent positional embedding modules after absolute scaled attention is thoroughly mastered.

Karpathy nanoGPT (causal self-attention module) ↗ compared_selected_co_primary [MIT]

Selected jointly as the primary code reference: translates Section 3.2 of the Vaswani paper into minimal, idiomatic PyTorch tensor operations.

linking MIT

The Annotated Transformer

Author: v2022: Austin Huang, Suraj Subramanian, Jonathan Sum, Khalid Almubarak, Stella Biderman; original: Sasha Rush (Harvard NLP) | Original Source Link ↗

Relevant Section / Range: MultiHeadedAttention module implementation: linear projections and subspace splitting

Verified Version & Date: Updated 2022 edition (Verified: 2026-10-30)

Selection Rationale:

Widely used pedagogical commentary that pairs every paragraph of the original Vaswani paper with an executable, testable code cell.

Source Prerequisites:

Intermediate Python, PyTorch broadcasting and tensor slicing semantics

Safe Fallback (Offline Resilience):

The curriculum provides standalone hand-worked numerical tables, explicit causal mask verification scripts, and offline analytical proofs.

Compared Candidate Sources & Selection Analysis (3)
Jay Alammar: The Illustrated Transformer ↗ compared_alternative [CC BY-NC 4.0]

Spectacular visual diagrams and intuitive color-coding, but lacks executable tensor code and formal numerical shape validation.

Stanford CS224N: Natural Language Processing with Deep Learning (Lecture 10) ↗ compared_alternative [Stanford University Copyright]

Rigorous academic slides and video lectures, but access to assignment grading is restricted to Stanford enrolled students.

Vaswani et al.: Attention Is All You Need ↗ compared_selected_co_primary [arXiv Open Access]

The original primary paper; Annotated Transformer annotates it directly.

02 120 min

Tensor Reshaping, Transposition, and Batched Attention (B, h, T, d_k)

Independent Mastery Check

Prerequisites

Complete Day 45 Activity 1.

Learning Objective

Trace tensor rank and dimension transformations across projections, head splitting (B, T, h, d_k) -> (B, h, T, d_k), batched matrix multiplication, and concatenation.

Instructions

Required work: trace an input batch X of shape (B, T, d_model) = (1, 2, 4) through Multi-Head Attention with h=2, d_k=2. Step 1: linear projection Q = X W^Q, K = X W^K, V = X W^V of shape (1, 2, 4). Step 2: reshape to (B, T, h, d_k) = (1, 2, 2, 2). Step 3: transpose dims 1 and 2 to (B, h, T, d_k) = (1, 2, 2, 2). Step 4: batched score matmul Q @ K^T of shape (1, 2, 2, 2) @ (1, 2, 2, 2) -> (1, 2, 2, 2) via transpose(-2, -1) yielding (B, h, T, T) = (1, 2, 2, 2). Step 5: scale by sqrt(2) and apply causal mask. Step 6: matmul with V yielding (1, 2, 2, 2). Step 7: transpose back to (1, 2, 2, 2) and flatten to (1, 2, 4). Step 8: project through W^O of shape (4, 4) -> (1, 2, 4). Trace with X = [[1, 0, 1, 0], [0, 2, 0, 1]] and identity projection weights.

Practice & Prediction Log

1. List the tensor shapes at every step from input X to output Out.
2. For X = [[1, 0, 1, 0], [0, 2, 0, 1]] with identity projections, compute the causal attention weights for Head 0 and Head 1.
3. Compute the final concatenated output after W^O.

Independent Mastery Criterion

Complete step-by-step tensor shape derivation and correct numerical evaluation of head-split attention and concatenation.

Authoritative Source Resources & Provenance

linking MIT

nanoGPT: Simplest, Fastest Repository for Training and Finetuning Medium GPTs

Author: Andrej Karpathy | Original Source Link ↗

Relevant Section / Range: model.py: view and transpose operations: B, T, nh, hs -> B, nh, T, hs

Verified Version & Date: GitHub master commit (Verified: 2026-10-30)

Selection Rationale:

Detailed modern reference implementation of clean, idiomatic batched causal self-attention in fewer than 60 lines of code.

Source Prerequisites:

Python, NumPy/PyTorch tensor fundamentals, and feed-forward network basics

Safe Fallback (Offline Resilience):

Self-contained step-by-step tensor derivations, numeric matrix traces, and pure Python/NumPy implementations are directly provided in curriculum exercises.

Compared Candidate Sources & Selection Analysis (3)
Hugging Face transformers (GPT2Attention) ↗ compared_rejected_for_foundations [Apache-2.0]

Includes extensive production abstraction layers, configuration polymorphism, and cache tracking that obscure direct mathematical tensor operations for beginner learners.

Harvard NLP: The Annotated Transformer ↗ compared_selected_co_primary [MIT]

Selected jointly: provides line-by-line literary commentary mapping equations directly to PyTorch tensors.

Karpathy: build-nanogpt (YouTube & GitHub 2024) ↗ compared_alternative [MIT]

Expanded video walkthrough; software is MIT-licensed while video audio/subtitles remain linking-only without transcript import.

linking MIT

The Annotated Transformer

Author: v2022: Austin Huang, Suraj Subramanian, Jonathan Sum, Khalid Almubarak, Stella Biderman; original: Sasha Rush (Harvard NLP) | Original Source Link ↗

Relevant Section / Range: MultiHeadedAttention forward pass tensor reshaping and transposition trace

Verified Version & Date: Updated 2022 edition (Verified: 2026-10-30)

Selection Rationale:

Widely used pedagogical commentary that pairs every paragraph of the original Vaswani paper with an executable, testable code cell.

Source Prerequisites:

Intermediate Python, PyTorch broadcasting and tensor slicing semantics

Safe Fallback (Offline Resilience):

The curriculum provides standalone hand-worked numerical tables, explicit causal mask verification scripts, and offline analytical proofs.

Compared Candidate Sources & Selection Analysis (3)
Jay Alammar: The Illustrated Transformer ↗ compared_alternative [CC BY-NC 4.0]

Spectacular visual diagrams and intuitive color-coding, but lacks executable tensor code and formal numerical shape validation.

Stanford CS224N: Natural Language Processing with Deep Learning (Lecture 10) ↗ compared_alternative [Stanford University Copyright]

Rigorous academic slides and video lectures, but access to assignment grading is restricted to Stanford enrolled students.

Vaswani et al.: Attention Is All You Need ↗ compared_selected_co_primary [arXiv Open Access]

The original primary paper; Annotated Transformer annotates it directly.

03 120 min

Numerical Stability in Attention: Max-Shift, Precision Limits, and Masking Hazards

Agent-Assisted Engineering Check

Prerequisites

Complete Day 45 Activity 2.

Learning Objective

Identify and mitigate floating-point vulnerabilities in attention: max-shift invariant softmax, all-masked row handling, and fp16 vs fp32 dynamic range limits.

Instructions

Required work: analyze three critical numerical edge cases. Case 1: Exponential overflow without max-shift. In float32, exp(z) overflows to inf at z >= 88.7; in float16, overflow occurs at z >= 11.1. Prove the shift identity softmax(z)_i = softmax(z - c)_i for c = max(z). Case 2: All-masked row hazard. If an entire row is masked by -inf (e.g., cross-attention with padded context), max(z) = -inf, producing -inf - (-inf) = NaN. Show how safe implementations replace NaN rows with 0 or uniform distribution. Case 3: Mask representation in mixed precision. Explain why -1e9 is unsafe in float16: it is outside the float16 range, so PyTorch's masked_fill raises an overflow error, and an explicit cast turns it into -inf, which brings back the all-masked-row NaN hazard above. Explain why -1e4 is preferred for half-precision safety.

Practice & Prediction Log

1. Prove algebraically that softmax(z - c) = softmax(z) for any scalar c.
2. In Python/PyTorch float16, what happens when computing torch.exp(torch.tensor(15.0, dtype=torch.float16)) vs float32?
3. How should an attention implementation handle an input sequence with empty valid keys without crashing into NaNs?

Independent Mastery Criterion

Rigorous mathematical proof of shift invariance and correct diagnostic handling of float16 overflow and all-masked row edge cases.

Authoritative Source Resources & Provenance

linking MIT

nanoGPT: Simplest, Fastest Repository for Training and Finetuning Medium GPTs

Author: Andrej Karpathy | Original Source Link ↗

Relevant Section / Range: model.py: FlashAttention / scaled_dot_product_attention numerical stability and fallback path

Verified Version & Date: GitHub master commit (Verified: 2026-10-30)

Selection Rationale:

Detailed modern reference implementation of clean, idiomatic batched causal self-attention in fewer than 60 lines of code.

Source Prerequisites:

Python, NumPy/PyTorch tensor fundamentals, and feed-forward network basics

Safe Fallback (Offline Resilience):

Self-contained step-by-step tensor derivations, numeric matrix traces, and pure Python/NumPy implementations are directly provided in curriculum exercises.

Compared Candidate Sources & Selection Analysis (3)
Hugging Face transformers (GPT2Attention) ↗ compared_rejected_for_foundations [Apache-2.0]

Includes extensive production abstraction layers, configuration polymorphism, and cache tracking that obscure direct mathematical tensor operations for beginner learners.

Harvard NLP: The Annotated Transformer ↗ compared_selected_co_primary [MIT]

Selected jointly: provides line-by-line literary commentary mapping equations directly to PyTorch tensors.

Karpathy: build-nanogpt (YouTube & GitHub 2024) ↗ compared_alternative [MIT]

Expanded video walkthrough; software is MIT-licensed while video audio/subtitles remain linking-only without transcript import.

linking arXiv.org perpetual non-exclusive license / Open Access

Attention Is All You Need

Author: Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Łukasz Kaiser, Illia Polosukhin | Original Source Link ↗

Relevant Section / Range: Section 3.2.1: Softmax numerical behavior and score magnitude scaling

Verified Version & Date: NeurIPS 2017 / arXiv:1706.03762v7 (Verified: 2026-10-30)

Selection Rationale:

The foundational primary literature establishing scaled dot-product attention, multi-head projection mechanisms, and autoregressive causal masking.

Source Prerequisites:

Linear algebra, matrix multiplication, softmax probability distributions

Safe Fallback (Offline Resilience):

Complete original algebraic derivations, matrix examples, variance proofs, and tensor diagrams are self-contained in the curriculum text without external document dependency.

Compared Candidate Sources & Selection Analysis (3)
Bahdanau et al.: Neural Machine Translation by Jointly Learning to Align and Translate (ICLR 2015) ↗ compared_alternative [arXiv Open Access]

Pioneered additive attention in RNNs, but Vaswani scaled dot-product attention replaced additive alignment in modern LLM architectures due to parallel matrix multiplication speed.

Shaw et al.: Self-Attention with Relative Position Representations (NAACL 2018) ↗ compared_deferred [arXiv Open Access]

Introduces relative positional biases into attention; deferred to subsequent positional embedding modules after absolute scaled attention is thoroughly mastered.

Karpathy nanoGPT (causal self-attention module) ↗ compared_selected_co_primary [MIT]

Selected jointly as the primary code reference: translates Section 3.2 of the Vaswani paper into minimal, idiomatic PyTorch tensor operations.

04 120 min

Independent Mastery Check: Multi-Head Accounting, Causal Invariance, and Stability

Independent Mastery Check

Prerequisites

Activities 45-01 through 45-03.

Learning Objective

Demonstrate unassisted mastery of multi-head attention parameter accounting, end-to-end forward computation trace, and numerical stability validation under declared AI-free conditions.

Instructions

Required work: work independently without AI assistants or external calculators. 1. Compute the exact parameter ledger for a Medium GPT attention block: d_model = 1024, h = 16 (weights and biases for W^Q, W^K, W^V, W^O). 2. Explain why causal invariance is strictly preserved across all h heads and through the linear output projection W^O. 3. Given logits [500.0, 502.0], compute the exact max-shifted softmax probabilities. 4. Complete the objective knowledge check quiz.

Practice & Prediction Log

1. Total parameters for d_model = 1024, h = 16 attention layer (with and without biases)?
2. Prove why W^O cannot leak future information across tokens if each individual head is causal.
3. Compute exact softmax on logits [500.0, 502.0].
4. Complete the 3-question quiz.

Objective Knowledge Check: Multi-Head Attention and Stability

Answer the following diagnostic questions to verify your understanding. Select your declared assistance condition before submitting.

Declared Assistance Condition

Note: AI-free work is an honest declared condition, not claimed surveillance.

1.

In a Multi-Head Attention block with model dimension d_model = 768 and h = 12 heads, what is the head dimension d_k for each individual head?

2.

For a standard Multi-Head Attention layer with projection matrices W^Q, W^K, W^V, W^O in R^{d_model x d_model} without bias terms, how many weight parameters are present when d_model = 768?

3.

In the max-shifted softmax implementation softmax(z)_i = exp(z_i - max(z)) / sum_j exp(z_j - max(z)), why is subtracting max(z) mathematically valid and computationally essential?

Attempt History & Evidence

    Independent Mastery Criterion

    Exact parameter ledgers (4,194,304 / 4,198,400), rigorous proof of position-wise preservation through W^O, and 100% quiz score.

    Authoritative Source Resources & Provenance

    linking arXiv.org perpetual non-exclusive license / Open Access

    Attention Is All You Need

    Author: Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Łukasz Kaiser, Illia Polosukhin | Original Source Link ↗

    Relevant Section / Range: Section 3.2.2: Multi-Head Attention parameter count equations and projection matrices

    Verified Version & Date: NeurIPS 2017 / arXiv:1706.03762v7 (Verified: 2026-10-30)

    Selection Rationale:

    The foundational primary literature establishing scaled dot-product attention, multi-head projection mechanisms, and autoregressive causal masking.

    Source Prerequisites:

    Linear algebra, matrix multiplication, softmax probability distributions

    Safe Fallback (Offline Resilience):

    Complete original algebraic derivations, matrix examples, variance proofs, and tensor diagrams are self-contained in the curriculum text without external document dependency.

    Compared Candidate Sources & Selection Analysis (3)
    Bahdanau et al.: Neural Machine Translation by Jointly Learning to Align and Translate (ICLR 2015) ↗ compared_alternative [arXiv Open Access]

    Pioneered additive attention in RNNs, but Vaswani scaled dot-product attention replaced additive alignment in modern LLM architectures due to parallel matrix multiplication speed.

    Shaw et al.: Self-Attention with Relative Position Representations (NAACL 2018) ↗ compared_deferred [arXiv Open Access]

    Introduces relative positional biases into attention; deferred to subsequent positional embedding modules after absolute scaled attention is thoroughly mastered.

    Karpathy nanoGPT (causal self-attention module) ↗ compared_selected_co_primary [MIT]

    Selected jointly as the primary code reference: translates Section 3.2 of the Vaswani paper into minimal, idiomatic PyTorch tensor operations.

    linking MIT

    nanoGPT: Simplest, Fastest Repository for Training and Finetuning Medium GPTs

    Author: Andrej Karpathy | Original Source Link ↗

    Relevant Section / Range: model.py: CausalSelfAttention parameter accounting and projection weights verification

    Verified Version & Date: GitHub master commit (Verified: 2026-10-30)

    Selection Rationale:

    Detailed modern reference implementation of clean, idiomatic batched causal self-attention in fewer than 60 lines of code.

    Source Prerequisites:

    Python, NumPy/PyTorch tensor fundamentals, and feed-forward network basics

    Safe Fallback (Offline Resilience):

    Self-contained step-by-step tensor derivations, numeric matrix traces, and pure Python/NumPy implementations are directly provided in curriculum exercises.

    Compared Candidate Sources & Selection Analysis (3)
    Hugging Face transformers (GPT2Attention) ↗ compared_rejected_for_foundations [Apache-2.0]

    Includes extensive production abstraction layers, configuration polymorphism, and cache tracking that obscure direct mathematical tensor operations for beginner learners.

    Harvard NLP: The Annotated Transformer ↗ compared_selected_co_primary [MIT]

    Selected jointly: provides line-by-line literary commentary mapping equations directly to PyTorch tensors.

    Karpathy: build-nanogpt (YouTube & GitHub 2024) ↗ compared_alternative [MIT]

    Expanded video walkthrough; software is MIT-licensed while video audio/subtitles remain linking-only without transcript import.