Self-Attention Engine
Implement causal scaled dot-product attention and multi-head attention in PyTorch
What You Will Build
Self-attention is the mathematical heart of the Transformer architecture. Unlike recurrent networks that compress history sequentially, self-attention allows every token position `t` in a sequence of length `T` to directly query and route information from any preceding token `j ≤ t` in parallel via matrix multiplication: `Attention(Q, K, V) = softmax(QKᵀ / √d_k + M)V`. In this project, you will build both single-head and multi-head causal self-attention modules from scratch in PyTorch using only basic linear layers and tensor operations (never calling `nn.MultiheadAttention` inside your module). You will trace every 4D tensor reshape `(B, T, H, d_k) -> (B, H, T, d_k)`, verify that future token perturbations have zero effect on past outputs (strict causality test), and confirm numerical parity against PyTorch's reference implementation.
Why This Project Matters in the AI/LLM Stack
Mastering the exact tensor shapes and causal masking mechanics of Multi-Head Attention (`(B, H, T, T)`) makes building the full Mini GPT architecture (Project 10) and KV-cached inference (Project 13) straightforward.
Learning Objectives
- Implement Q, K, and V projections, `1 / √d_k` scaling, causal `-inf` masking, and value aggregation from scratch in PyTorch.
- Reshape and transpose 3D sequence tensors `(B, T, d_model)` into 4D multi-head tensors `(B, H, T, d_k)` and recombine heads cleanly.
- Write a strict causality unit test proving that changing input tokens at positions `> t` alters outputs at positions `≤ t` by `0.0`.
- Verify forward and backward numerical equivalence against `torch.nn.functional.scaled_dot_product_attention(..., is_causal=True)`.
Prerequisites
- Phase 1 Linear Algebra (dot products, batched matrix multiplication) & Phase 3 PyTorch
- Phase 5 Module 5.1 & 5.2 (Scaled Dot-Product Attention & Multi-Head Causal Attention)
Required Mathematical Concepts
- Query, Key, Value linear projections: Q = XW_Q, K = XW_K, V = XW_V
- Scaled Dot-Product Attention: Attention(Q, K, V) = softmax((QKᵀ) / √d_k + M) V
- Variance scaling derivation: if q_i, k_i ~ N(0, 1), Var(q · k) = d_k ⇒ divide by √d_k
- Causal upper-triangular mask: M_{i,j} = 0 if j ≤ i else -∞
- Multi-head decomposition: d_model = H × d_k and output projection W_O ∈ ℝ^{d_model×d_model}
Step-by-Step Implementation Tasks (13 Steps)
Check off each engineering task as your code passes its acceptance criteria. Progress is saved automatically in your browser.
Embeddings, Q/K/V Projections & Scaled Attention Scores
Project input embeddings into Queries, Keys, and Values and compute scaled pairwise compatibility scores.
Self-attention operates on a batch of continuous token representations `X` of shape `(B, T, d_model)`.
Separate linear projections allow each token to ask a question (`Q`), advertise its content (`K`), and pass along payload information (`V`).
Batched matrix multiplication `q @ k.transpose(-2, -1)` computes the raw dot-product similarity between every query `i` and every key `j`.
Without dividing by `√d_k`, dot products grow in variance as `d_k` increases, pushing Softmax into flat saturation regions where gradients vanish.
Causal Masking, Softmax Probabilities & Value Mixing
Apply upper-triangular `-inf` masking, compute row-wise softmax weights, and aggregate Value vectors.
In autoregressive language modeling, position `t` must never attend to future tokens `j > t`. Setting upper-triangular entries to `-inf` before Softmax drives their attention probabilities to exact `0.0`.
Applying Softmax along the last dimension (`dim=-1`, the Key dimension) turns each Query row into a valid probability distribution summing to `1.0`.
Multiplying the `(..., T, T)` attention probabilities by `V` of shape `(..., T, d_k)` blends Value vectors according to their relevance to each Query.
Multi-Head Splitting, Concatenation & Output Projection
Generalize to `H` parallel attention heads with fused QKV projection and final output projection `W_O`.
Splitting `d_model` into `H` subspaces of size `d_k = d_model // H` lets different heads simultaneously attend to syntactic, positional, and semantic relationships.
After attention, the `H` head outputs `(B, H, T, d_k)` are transposed back to `(B, T, H, d_k)` and flattened into `(B, T, d_model)`.
Concatenating heads simply stacks them side-by-side; the output projection `W_O` (`nn.Linear(d_model, d_model)`) mixes information across all `H` heads.
Shape Verification, Causality Proof & Reference Comparison
Assert tensor shapes at every intermediate line, prove zero future leakage, and verify parity with PyTorch SDPA.
Explicit shape assertions catch subtle broadcasting bugs (like forgetting `.transpose(1, 2)`).
The definitive test of causal self-attention is prefix invariance: mutating tokens at positions `t+1 .. T-1` must have ZERO effect on the output at positions `0 .. t`.
PyTorch 2.0+ provides `F.scaled_dot_product_attention(q, k, v, is_causal=True)`. Comparing your manual attention math against it verifies 100% numerical accuracy.
Starter & Verification Code Examples
Multi-Head Causal Self-Attention in PyTorch
Shows the canonical 4D tensor reshape, `1/√d_k` scaling, `register_buffer` causal mask, and head recombination.
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
class CausalMultiHeadAttention(nn.Module):
def __init__(self, d_model: int, num_heads: int, max_seq_len: int = 1024) -> None:
super().__init__()
if d_model % num_heads != 0:
raise ValueError(f"d_model ({d_model}) must be divisible by num_heads ({num_heads})")
self.d_model = d_model
self.num_heads = num_heads
self.head_dim = d_model // num_heads
self.c_attn = nn.Linear(d_model, 3 * d_model)
self.c_proj = nn.Linear(d_model, d_model)
self.register_buffer(
"causal_mask",
torch.tril(torch.ones(max_seq_len, max_seq_len)).view(1, 1, max_seq_len, max_seq_len),
)
def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
B, T, C = x.size()
qkv = self.c_attn(x)
q, k, v = qkv.split(self.d_model, dim=2)
# Reshape (B, T, C) -> (B, H, T, d_k)
q = q.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
k = k.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
v = v.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
scores = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(self.head_dim))
scores = scores.masked_fill(self.causal_mask[:, :, :T, :T] == 0, float("-inf"))
attn_weights = F.softmax(scores, dim=-1)
out = attn_weights @ v # (B, H, T, d_k)
out = out.transpose(1, 2).contiguous().view(B, T, C)
return self.c_proj(out), attn_weightsHow to Test Your Implementation
Run tensor shape checks, attention weight row-sum & upper-triangle zero checks, prefix causality invariance tests, and PyTorch SDPA parity tests.
Prefix Causality Invariance Test
Prove that token position `t` cannot see or be influenced by any token at position `> t`.
Steps: Compare `attn(x)[:, :t, :]` before and after replacing `x[:, t:, :]` with random noise in `eval()` mode.
Pass Criteria: Max absolute difference on positions `0..t-1` is `< 1e-7`.
Attention Matrix Row-Stochastic & Upper-Triangle Zero Test
Verify that post-softmax attention weights sum to 1.0 per row and are strictly 0.0 for `j > i`.
Steps: Inspect `attn_weights` of shape `(B, H, T, T)`: check `torch.triu(attn_weights, diagonal=1).abs().max()` and `attn_weights.sum(dim=-1)`.
Pass Criteria: Upper triangle max is `0.0` and all row sums equal `1.0`.
Overall Project Acceptance Criteria
- Single-head and multi-head causal self-attention implemented from scratch in PyTorch without calling `nn.MultiheadAttention`.
- 4D tensor reshapes `(B, H, T, d_k)`, `1/√d_k` scaling, lower-triangular causal masking, and `W_O` output projection verified.
- Strict causality unit test proves zero information leakage from future positions `j > i`.
- Forward and backward parity verified against `F.scaled_dot_product_attention(..., is_causal=True)`.
Expandable Hints & Debugging Tips
Try implementing each stage on your own first. Expand a hint below if you want to verify tensor shapes, mathematical identities, or common edge cases.
Optional Stretch Challenges
Implement Rotary Positional Embeddings (RoPE)
Implement 2D complex-plane rotation of Query and Key vectors (`apply_rope(q, k, cos, sin)`) so attention scores depend purely on relative token distance `i - j`.
What You Learn: How modern open-weights LLMs (Llama, Mistral, Qwen, Gemma) encode relative position directly inside self-attention.
Grouped-Query Attention (GQA)
Extend your module to support `num_kv_heads < num_query_heads` (sharing Key and Value heads across groups of Query heads) to reduce KV-cache memory.
What You Learn: How Grouped-Query Attention reduces inference memory bandwidth while preserving model quality.
What Your Completed Repository Should Contain
A reusable causal attention module.
Deliverable Artifacts
- Standalone `SingleHeadCausalAttention` and `CausalMultiHeadAttention` PyTorch modules
- Strict prefix causality and PyTorch SDPA reference parity test suite
- Attention heatmap visualization utility
After Completing This Project, You Should Be Able To...
- Derive and implement scaled dot-product causal self-attention `softmax(QKᵀ/√d_k + M)V` in PyTorch.
- Track and transform 4D tensors `(B, H, T, d_k)` across multi-head splitting and concatenation without memory-layout bugs.
- Prove strict autoregressive causality using prefix perturbation unit tests.
- Explain why `1/√d_k` scaling and `-inf` pre-softmax masking are essential for stable Transformer training.
GitHub Repository Submission
Link your public or private GitHub repository URL to keep a structured portfolio of your from-scratch AI builds. Stored locally in your browser (self-reported; no GitHub tokens or remote code execution).
Suggested Repository: causal-self-attention-engine
Recommended Directory Layout:
attention/mha.py attention/visualize.py tests/test_self_attention.py README.md
Recommended README Documentation Checklist
- Step-by-step tensor shape trace from `(B, T, d_model)` through `(B, H, T, T)`
- Instructions to run causality and PyTorch SDPA parity tests
- Sample causal attention heatmap plot