How the Attention Mechanism Is Explained and Implemented in AI Engineering From Scratch

The AI-Engineering-From-Scratch curriculum teaches attention from first principles, starting with a NumPy-based scaled-dot-product self-attention implementation and progressing to a full PyTorch multi-head self-attention layer with causal masking and dropout.

The attention mechanism powers modern transformer architectures like GPT and BERT. This open-source curriculum by Rohit Ghumare provides a pedagogical journey from mathematical foundations to runnable code, with implementations split across two key phases of the repository.

Single-Head Scaled-Dot-Product Self-Attention

The foundational lesson introduces self-attention using pure NumPy. Given a sequence X of token embeddings with shape T × D (sequence length by model dimension), the mechanism derives queries (Q), keys (K), and values (V) through learned linear projections.

The core operation follows the scaled dot-product formula:

[ \text{weights}_{i,j}= \text{softmax}!\left(\frac{Q_i\cdot K_j^{\top}}{\sqrt{d_k}}\right) \qquad\text{output}= \text{weights},V ]

Scaling by √dₖ prevents dot-product magnitudes from exploding as dimensionality grows. Softmax normalizes scores into a probability distribution where each row sums to 1.

Implementation in self_attention.py

The file at phases/07-transformers-deep-dive/02-self-attention-from-scratch/code/self_attention.py contains:

• Lines 4-8: Numerically-stable softmax implementation (subtract max before exponentiating)
• Lines 10-15: scaled_dot_product_attention(Q, K, V) function computing scores, applying softmax, and returning weighted output plus attention matrix
• Lines 18-27: SelfAttention class with three projection matrices (Wq, Wk, Wv) initialized using He-style scaling
• Lines 28-32: forward() method building Q, K, V from input X

import numpy as np
from phases.07_transformers_deep_dive.02_self_attention_from_scratch.code.self_attention import SelfAttention

# Dummy token embeddings: 6 tokens, d_model=16

rng = np.random.default_rng(42)
X = rng.normal(0, 1, (6, 16))

attn = SelfAttention(d_model=16, dk=8, dv=8, seed=42)
output, weights = attn.forward(X)

print("Output shape:", output.shape)    # (6, 8)

print("Attention matrix shape:", weights.shape)  # (6, 6)

This implementation demonstrates that attention weights form an T × T matrix, where each token attends to every other token in the sequence.

Multi-Head Self-Attention with PyTorch

The advanced lesson in phases/19-capstone-projects/33-multihead-self-attention/code/main.py extends the concept to multi-head attention, allowing the model to learn multiple representation subspaces in parallel.

Key Architecture Decisions

Single QKV projection: Rather than separate projections per head, the implementation uses one nn.Linear(d_model, 3 * d_model) followed by head splitting—matching modern efficient transformer implementations.

Head splitting mechanics: The _split_heads method (lines 55-62) reshapes tensors from (B, T, D) to (B, n_head, T, d_head), enabling parallel attention computation across heads.

Causal Mask Implementation

The causal mask prevents autoregressive models from attending to future positions. Registered as a buffer (lines 52-54), it uses torch.tril to create a lower-triangular matrix:

mask_slice = self.causal_mask[:t, :t]  # shape (t, t)

scores = scores.masked_fill(mask_slice == 0, float("-inf"))

Positions above the diagonal are set to -inf, causing softmax to output zero probability for future tokens.

Full Forward Pipeline

The forward() method (lines 63-99) executes:

1. Single linear projection producing Q, K, V
2. Split into multiple heads
3. Scaled dot-product attention per head
4. Causal masking with masked_fill
5. Softmax normalization
6. Dropout on attention weights (attn_dropout)
7. Weighted sum with values
8. Head merging with _merge_heads
9. Final linear projection (out_proj)
10. Output dropout (out_dropout)

import torch
from phases.19_capstone_projects.33_multihead_self_attention.code.main import MultiHeadSelfAttention

torch.manual_seed(0)
batch, seq_len, d_model = 2, 7, 16
x = torch.randn(batch, seq_len, d_model)

attn = MultiHeadSelfAttention(d_model=d_model, n_heads=4, max_context_length=10)
out, weights = attn(x, return_weights=True)

print("Input shape:", x.shape)          # (2, 7, 16)

print("Output shape:", out.shape)       # (2, 7, 16)

print("Weights shape:", weights.shape)  # (2, 4, 7, 7) — per-head attention

The return_weights=True flag (lines 96-98) enables inspection of attention patterns for interpretability.

Integration: Tiny Language Model

The curriculum culminates in TinyAttentionLM, a complete autoregressive language model demonstrating how attention layers integrate into larger architectures. Located in the same main.py file, this model stacks:

• Token embeddings: TokenEmbedding maps vocabulary IDs to vectors
• Positional encoding: SinusoidalPositionalEmbedding injects sequence position information
• Multi-head attention: The full attention layer described above
• Language modeling head: Linear projection to vocabulary logits

from phases.19_capstone_projects.33_multihead_self_attention.code.main import TinyAttentionLM, DemoConfig

cfg = DemoConfig()
model = TinyAttentionLM(
    vocab_size=cfg.vocab_size,
    d_model=cfg.d_model,
    n_heads=cfg.n_heads,
    max_context_length=cfg.seq_len,
)

ids = model.token_emb.weight.new_zeros(
    (cfg.batch_size, cfg.seq_len + 1)
).long().random_(0, cfg.vocab_size).split(1, dim=1)[0]
logits, attn_weights = model(ids, return_weights=True)

print("Logits shape:", logits.shape)  # (batch, seq_len, vocab_size)

The included training loop solves a repeat task, verifying that a single attention head attending to the previous token can learn the pattern.

Source Files and Testing

File	Purpose	Location
self_attention.py	NumPy single-head implementation	phases/07-transformers-deep-dive/02-self-attention-from-scratch/code/
main.py	PyTorch multi-head attention, embeddings, and tiny LM	phases/19-capstone-projects/33-multihead-self-attention/code/
test_attention.py	Unit tests for shape contracts, masking, and gradients	phases/19-capstone-projects/33-multihead-self-attention/code/tests/

The test suite validates that:

• Output shapes match input shapes (residual-compatible)
• Causal masking correctly zeros future attention
• Gradients flow through all parameters

Summary

• Foundational approach: The attention mechanism is taught through progressive complexity—pure NumPy first, then optimized PyTorch
• Scaled dot-product: Core operation uses √dₖ scaling and softmax normalization for stable training
• Multi-head design: Parallel attention heads with single QKV projection and explicit head splitting
• Causal masking: Lower-triangular mask prevents future-token leakage in autoregressive models
• Production features: Dropout, proper initialization, and optional attention weight extraction
• End-to-end validation: Tiny language model proves the implementation learns sequence patterns

Frequently Asked Questions

Why scale the dot-product attention by √dₖ?

Scaling by the square root of key dimension prevents dot-product magnitudes from growing too large as dimensionality increases. Without this scaling, softmax gradients become vanishingly small, slowing or preventing learning. The ai-engineering-from-scratch implementation follows the original "Attention Is All You Need" paper's approach.

What is the difference between single-head and multi-head attention?

Single-head attention computes one attention distribution over the input sequence. Multi-head attention runs h parallel attention operations with different learned projections, allowing the model to attend to information from different representation subspaces simultaneously. The repository's MultiHeadSelfAttention concatenates head outputs and applies a final linear mix.

How does the causal mask work in autoregressive models?

The causal mask (implemented with torch.tril in lines 52-54) creates a matrix where positions i < j are masked. Before softmax, these positions are set to -inf, ensuring the model assigns zero probability to attending to future tokens. This enforces that prediction for position i depends only on positions 0...i-1.

Why use a single QKV projection instead of separate per-head projections?

A single nn.Linear(d_model, 3 * d_model) followed by splitting is mathematically equivalent to separate projections but more computationally efficient on modern hardware. It reduces kernel launch overhead and enables fused operations. The repository follows this pattern (lines 20-30) to match production transformer implementations.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}