Attention And Transformers
Self-Attention & Multi-Head Attention
Sentence: "The animal didn't cross the street because it
Jr Codex Deep Learning Notes
Level: Advanced Prerequisites: Chapter 1: The Attention Mechanism Time to complete: ~25 minutes
Table of Contents
- Self-Attention — Attending Within One Sequence
- Why Self-Attention Replaces Recurrence
- The Limitation of a Single Attention "View"
- Multi-Head Attention
- Implementing Multi-Head Attention
- Masked Self-Attention
- Multi-Head Attention in PyTorch
- Summary & Next Steps
1. Self-Attention — Attending Within One Sequence
Self-attention is the specific case where the queries, keys, and values (Chapter 1, Section 3) all come from the same sequence — every position attends to every other position (including itself) within one input, rather than one sequence (a decoder) attending to a different sequence (an encoder).
Self-Attention, Concretely
─────────────────────────────────────────
Sentence: "The animal didn't cross the street because it
was too tired"
When computing a NEW representation for the word "it,"
self-attention lets the model directly compute how relevant
EVERY other word is — and, crucially, LEARNS (via training)
that "it" should attend strongly to "animal," resolving the
reference correctly, REGARDLESS of the distance between them
in the sentence.
─────────────────────────────────────────
2. Why Self-Attention Replaces Recurrence
Self-Attention vs an RNN's Hidden State (Module 5)
─────────────────────────────────────────
RNN/LSTM (Module 5): information about word 1 reaches
word 10 only by passing THROUGH
every intermediate hidden state,
step by step, SEQUENTIALLY —
and can decay along the way
(Ch.3-5 of Module 5)
Self-attention: information about word 1 reaches
word 10 in EXACTLY ONE step —
a DIRECT computation between
the two positions, computed IN
PARALLEL with every other
pairwise computation
─────────────────────────────────────────
This directly resolves all three of Module 5, Chapter 5's limitations: no fixed-size bottleneck (every position is directly available), full parallelization (every pairwise score can be computed simultaneously), and no long-range decay (distance between positions doesn't affect how directly they can interact).
3. The Limitation of a Single Attention "View"
A single attention computation (Chapter 1) can only capture one kind of relationship pattern at a time — but language (and sequences generally) often has multiple, simultaneous kinds of relevant relationships.
Multiple Relationship Types, Simultaneously
─────────────────────────────────────────
"The cat that chased the mouse was hungry"
One USEFUL attention pattern: "was" should attend
strongly to "cat"
(subject-verb agreement)
A DIFFERENT useful pattern: "chased" should attend
to BOTH "cat" (who
did the chasing) and
"mouse" (who was
chased)
A SINGLE attention computation, with ONE set of Q/K/V
weight matrices, is forced to find ONE compromise pattern
that partially captures BOTH needs — not ideal.
─────────────────────────────────────────
4. Multi-Head Attention
Multi-head attention runs several independent attention computations ("heads") in parallel, each with its own learned Q/K/V projection weights, letting each head specialize in a different kind of relationship — then combines all heads' outputs together.
Multi-Head Attention, Structurally
─────────────────────────────────────────
Input
│
├──► Head 1: own W_q, W_k, W_v ──► Attention ──► output_1
├──► Head 2: own W_q, W_k, W_v ──► Attention ──► output_2
├──► Head 3: own W_q, W_k, W_v ──► Attention ──► output_3
└──► Head 4: own W_q, W_k, W_v ──► Attention ──► output_4
│
CONCATENATE all head outputs
│
Final LINEAR projection
│
output
─────────────────────────────────────────
Each head typically operates on a smaller dimension than the full model (e.g. 8 heads of dimension 64 each, for a 512-dimensional model) — so multi-head attention's total compute cost is comparable to a single, full-dimension attention computation, while gaining the ability to capture multiple relationship types simultaneously.
5. Implementing Multi-Head Attention
import numpy as np
def softmax(x, axis=-1):
exp_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
return exp_x / np.sum(exp_x, axis=axis, keepdims=True)
def scaled_dot_product_attention(Q, K, V):
d_k = Q.shape[-1]
scores = Q @ K.T / np.sqrt(d_k)
weights = softmax(scores)
return weights @ V
def multi_head_attention(x, num_heads, d_model):
d_k = d_model // num_heads # each head gets a SLICE of the full dimension
head_outputs = []
for head in range(num_heads):
W_q = np.random.randn(d_model, d_k) * 0.1 # EACH head has its OWN weights
W_k = np.random.randn(d_model, d_k) * 0.1
W_v = np.random.randn(d_model, d_k) * 0.1
Q, K, V = x @ W_q, x @ W_k, x @ W_v
head_output = scaled_dot_product_attention(Q, K, V)
head_outputs.append(head_output)
concatenated = np.concatenate(head_outputs, axis=-1) # combine ALL heads
W_o = np.random.randn(d_model, d_model) * 0.1 # final output projection
return concatenated @ W_o
x = np.random.randn(5, 512) # a 5-token sequence, 512-dimensional representations
output = multi_head_attention(x, num_heads=8, d_model=512)
print(f"Output shape: {output.shape}") # [5, 512] — same shape as input6. Masked Self-Attention
For tasks like language generation (Module 5, Chapter 4's autoregressive generation), a position must not be allowed to attend to future positions — that would let the model "see the answer" during training. Masking enforces this by setting future positions' attention scores to negative infinity before the softmax, so they receive a weight of exactly zero.
def masked_scaled_dot_product_attention(Q, K, V):
d_k = Q.shape[-1]
seq_len = Q.shape[0]
scores = Q @ K.T / np.sqrt(d_k)
mask = np.triu(np.ones((seq_len, seq_len)), k=1).astype(bool) # upper triangle = FUTURE positions
scores = np.where(mask, -np.inf, scores) # block attending to the future
weights = softmax(scores)
return weights @ VMasked Attention's Effect on the Weight Matrix
─────────────────────────────────────────
Position 1 can attend to: position 1 only
Position 2 can attend to: positions 1-2
Position 3 can attend to: positions 1-3
...
Position N can attend to: positions 1-N (everything)
This is exactly the causal structure required for
autoregressive generation (Module 5, Ch.4) — used in the
DECODER of a Transformer, and throughout every modern large
language model covered in the NLP & LLM Notes.
─────────────────────────────────────────
7. Multi-Head Attention in PyTorch
import torch
import torch.nn as nn
multihead_attn = nn.MultiheadAttention(embed_dim=512, num_heads=8, batch_first=True)
sequence = torch.randn(1, 5, 512) # (batch, seq_len, embed_dim)
# Self-attention: query, key, AND value are all the SAME sequence (Section 1)
output, attention_weights = multihead_attn(query=sequence, key=sequence, value=sequence)
print(f"Output shape: {output.shape}") # [1, 5, 512]
print(f"Attention weights shape: {attention_weights.shape}") # [1, 5, 5] — averaged across heads
# Masked self-attention (Section 6) for a decoder
seq_len = 5
causal_mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool()
masked_output, _ = multihead_attn(query=sequence, key=sequence, value=sequence, attn_mask=causal_mask)8. Summary & Next Steps
Key Takeaways
- Self-attention lets every position in a sequence directly attend to every other position, replacing the sequential dependency chain of RNNs entirely.
- A single attention computation can only capture one kind of relationship pattern; multi-head attention runs several independent attention computations in parallel, each specializing in different relationships, then combines them.
- Masked self-attention prevents a position from attending to future positions, which is essential for autoregressive generation tasks and Transformer decoders.
- PyTorch's
nn.MultiheadAttentionimplements this entire mechanism directly, including optional masking.
Concept Check
- Why does self-attention allow full parallelization across a sequence, unlike an RNN?
- What problem does multi-head attention solve that a single attention computation cannot?
- Why is masking necessary in a Transformer decoder, but not in a Transformer encoder?
Next Chapter
→ Chapter 3: Positional Encoding
Jr Codex — 1-on-1 Personalized Coaching | Back to Module Index | Back to DL Index