Self-Attention: The Heart of Transformers¶

Prerequisites¶

Lesson 01: Introduction to Attention — QKV abstraction, scaled dot-product formula
Module 05 L02: Activation Functions — softmax

What You'll Learn¶

Concept	Why it matters
Self-attention vs cross-attention	Self-attention operates on one sequence; cross-attention relates two
Q, K, V projection matrices	How the model learns what to look for vs what to advertise
Scaled dot-product attention	The exact formula used in every Transformer
Causal masking	Why decoder models cannot look at future tokens
O(n²) complexity	Why long sequences are expensive

Intuition: A Sequence Attending to Itself¶

Regular (cross) attention: the decoder asks questions of the encoder.

Self-attention: each position asks questions of every other position in the same sequence.

Think of a sentence as a roundtable discussion. Each word gets to say: "What context do I need from the rest of the sentence to understand my own meaning?"

"The animal didn't cross the street because it was too tired."

When processing "it":
  "animal" answers strongly — "it" refers to me
  "street" answers weakly   — probably not me
  "tired"  answers weakly   — I'm an attribute, not the referent

→ Self-attention assigns high weight to "animal" when encoding "it"

flowchart TD
    classDef proj fill:#eef2ff,stroke:#6366f1,stroke-width:2px;
    classDef dot fill:#f0fdf4,stroke:#10b981,stroke-width:2px;
    classDef out fill:#f5f3ff,stroke:#8b5cf6,stroke-width:2px;

    X["Input Embeddings X (N x d_model)"]:::proj --> ProjQ["Query Projection Q = X * W_Q"]:::proj
    X --> ProjK["Key Projection K = X * W_K"]:::proj
    X --> ProjV["Value Projection V = X * W_V"]:::proj

    ProjQ & ProjK --> DotMat["Dot Product Scores S = Q * K^T (N x N)"]:::dot
    DotMat --> Scale["Scale by 1/sqrt(d_k)"]:::dot
    Scale --> Mask["Causal Masking (Optional)"]:::dot
    Mask --> Softmax["Softmax Row-Wise (Attention Weights A)"]:::dot

    Softmax & ProjV --> Output["Context Output Matrix Output = A * V (N x d_v)"]:::out

Self-Attention Architecture — adapted from "Attention Is All You Need" (Vaswani et al., 2017).

The Three Projection Matrices¶

For each token embedding x_i ∈ ℝ^{d_model}, self-attention creates three vectors:

q_i = x_i · W_Q    (query  — "what am I looking for?")  shape: (d_k,)
k_i = x_i · W_K    (key    — "what do I advertise?")    shape: (d_k,)
v_i = x_i · W_V    (value  — "what do I actually give?") shape: (d_v,)

Where: - W_Q ∈ ℝ^{d_model × d_k} — query projection - W_K ∈ ℝ^{d_model × d_k} — key projection - W_V ∈ ℝ^{d_model × d_v} — value projection

Why separate Q and K?

If we used the same matrix for both, every token would have the same query and key. The separation lets the model learn different representations for "what I seek" vs "what I offer." This is crucial for heads that specialize (e.g., one head tracks syntactic subject, another tracks semantic theme).

In matrix form for a sequence of n tokens:

Q = X · W_Q    shape: (n, d_k)   — all queries stacked
K = X · W_K    shape: (n, d_k)   — all keys stacked
V = X · W_V    shape: (n, d_v)   — all values stacked

The Full Formula¶

Attention(Q, K, V) = softmax(Q K^T / √d_k) · V

Step by step:

Step	Operation	Input shape	Output shape
1. Scores	`Q @ K.T`	(n, d_k) × (d_k, n)	(n, n)
2. Scale	`/ sqrt(d_k)`	(n, n)	(n, n)
3. (Optional mask)	add -∞ to masked positions	(n, n)	(n, n)
4. Softmax	along last dim	(n, n)	(n, n)
5. Context	`@ V`	(n, n) × (n, d_v)	(n, d_v)

The (n, n) matrix is the attention matrix A. Entry A[i, j] says how much token i attends to token j.

Worked Numerical Example¶

Sentence: "The cat sat" — 3 tokens, d_model = 4, d_k = d_v = 4.

Embeddings X (3, 4):
  The = [1, 0, 1, 0]
  cat = [0, 1, 0, 1]
  sat = [1, 1, 0, 0]

Let W_Q = W_K = W_V = I (identity for this toy example).

So Q = K = V = X.

Step 1 — Raw scores (Q @ K^T):

Q @ K^T =
  [1,0,1,0]   [1,0,1,0]^T       row 0 (The):  [2, 0, 1]
  [0,1,0,1] · [0,1,0,1]^T  =    row 1 (cat):  [0, 2, 1]
  [1,1,0,0]   [1,1,0,0]^T       row 2 (sat):  [1, 1, 2]

(dot products of each pair of embeddings)

Step 2 — Scale by √d_k = √4 = 2:

[[2, 0, 1],    /2 =    [[1.00, 0.00, 0.50],
 [0, 2, 1],             [0.00, 1.00, 0.50],
 [1, 1, 2]]             [0.50, 0.50, 1.00]]

Step 3 — Softmax (row-wise):

Row 0: softmax([1.00, 0.00, 0.50]) = [0.462, 0.170, 0.368]
Row 1: softmax([0.00, 1.00, 0.50]) = [0.170, 0.462, 0.368]
Row 2: softmax([0.50, 0.50, 1.00]) = [0.244, 0.244, 0.512]

Interpretation: "The" attends mostly to itself (0.46), then "sat" (0.37), less to "cat" (0.17).

Step 4 — Output = A @ V (V = X here):

Output[0] ("The") = 0.462×[1,0,1,0] + 0.170×[0,1,0,1] + 0.368×[1,1,0,0]
                  = [0.462,0,0.462,0] + [0,0.170,0,0.170] + [0.368,0.368,0,0]
                  = [0.830, 0.538, 0.462, 0.170]

The output for "The" is a mixture of all token embeddings, weighted by their relevance. This contextualizes each token with information from the whole sequence.

Implementation: Pure NumPy¶

import numpy as np


def softmax(x: np.ndarray) -> np.ndarray:
    """Numerically stable softmax along last axis."""
    x = x - x.max(axis=-1, keepdims=True)
    return np.exp(x) / np.exp(x).sum(axis=-1, keepdims=True)


def scaled_dot_product_attention(
    Q: np.ndarray,             # (B, n, d_k)
    K: np.ndarray,             # (B, n, d_k)
    V: np.ndarray,             # (B, n, d_v)
    mask: np.ndarray | None = None,  # (B, n, n)  True = mask out
) -> tuple[np.ndarray, np.ndarray]:
    """
    Batched scaled dot-product attention.

    Returns
    -------
    output  : (B, n, d_v)
    weights : (B, n, n)   — attention distribution, rows sum to 1
    """
    d_k = Q.shape[-1]
    # (B, n, d_k) × (B, d_k, n) → (B, n, n)
    scores = Q @ K.transpose(0, 2, 1) / np.sqrt(d_k)

    if mask is not None:
        scores = np.where(mask, -1e9, scores)

    weights = softmax(scores)          # (B, n, n)
    output  = weights @ V              # (B, n, d_v)
    return output, weights


class SelfAttention:
    """Single-head self-attention layer."""

    def __init__(self, d_model: int, d_k: int | None = None, d_v: int | None = None):
        self.d_model = d_model
        self.d_k = d_k or d_model
        self.d_v = d_v or d_model

        # Learnable projections (initialized with small random values)
        scale = 1.0 / np.sqrt(d_model)
        self.W_Q = np.random.randn(d_model, self.d_k) * scale
        self.W_K = np.random.randn(d_model, self.d_k) * scale
        self.W_V = np.random.randn(d_model, self.d_v) * scale

    def forward(
        self,
        x: np.ndarray,              # (B, n, d_model)
        mask: np.ndarray | None = None,  # (B, n, n)
    ) -> tuple[np.ndarray, np.ndarray]:
        """
        Returns
        -------
        output  : (B, n, d_v)
        weights : (B, n, n)
        """
        Q = x @ self.W_Q   # (B, n, d_k)
        K = x @ self.W_K   # (B, n, d_k)
        V = x @ self.W_V   # (B, n, d_v)
        return scaled_dot_product_attention(Q, K, V, mask)


# ── Smoke test ───────────────────────────────────────────────────────────────
B, n, d_model = 2, 6, 64
x = np.random.randn(B, n, d_model)

sa = SelfAttention(d_model=64, d_k=32, d_v=32)
out, attn = sa.forward(x)

print(f"Input:   {x.shape}")    # (2, 6, 64)
print(f"Output:  {out.shape}")  # (2, 6, 32)
print(f"Weights: {attn.shape}") # (2, 6, 6)
print(f"Row sum (should be 1): {attn[0].sum(axis=-1).round(4)}")

Causal (Masked) Self-Attention¶

Decoder models (GPT family) must not allow token i to see token j > i — that would be cheating during training.

We enforce this with a causal mask (upper-triangular):

def causal_mask(n: int) -> np.ndarray:
    """
    Returns boolean mask (n, n):
      True  = this position should be masked (future token)
      False = this position is visible (current or past)
    """
    return np.triu(np.ones((n, n), dtype=bool), k=1)


# Visualize for n=5
mask = causal_mask(5)
print(mask.astype(int))
# [[0, 1, 1, 1, 1],   token 0 can only see itself
#  [0, 0, 1, 1, 1],   token 1 can see 0 and 1
#  [0, 0, 0, 1, 1],   token 2 can see 0,1,2
#  [0, 0, 0, 0, 1],
#  [0, 0, 0, 0, 0]]   token 4 can see everything before it

Adding the mask to attention:

B, n = 1, 5
x = np.random.randn(B, n, 64)
sa = SelfAttention(d_model=64)
m = causal_mask(n)[np.newaxis, :, :]   # (1, 5, 5) for broadcasting

out, weights = sa.forward(x, mask=m)

# Verify: upper triangle of weights[0] should be ~0
print("Upper triangle (future tokens):", weights[0][0, 1:].round(4))
# All near zero because -1e9 → softmax → ~0

PyTorch Implementation¶

import torch
import torch.nn as nn
import torch.nn.functional as F


class SelfAttentionTorch(nn.Module):
    def __init__(self, d_model: int, d_k: int):
        super().__init__()
        self.d_k = d_k
        self.W_Q = nn.Linear(d_model, d_k, bias=False)
        self.W_K = nn.Linear(d_model, d_k, bias=False)
        self.W_V = nn.Linear(d_model, d_k, bias=False)

    def forward(
        self,
        x: torch.Tensor,                     # (B, n, d_model)
        mask: torch.BoolTensor | None = None, # (B, n, n)
    ) -> tuple[torch.Tensor, torch.Tensor]:
        Q = self.W_Q(x)   # (B, n, d_k)
        K = self.W_K(x)   # (B, n, d_k)
        V = self.W_V(x)   # (B, n, d_k)

        # (B, n, d_k) × (B, d_k, n) → (B, n, n)
        scores = Q @ K.transpose(-2, -1) / (self.d_k ** 0.5)

        if mask is not None:
            scores = scores.masked_fill(mask, float("-inf"))

        weights = F.softmax(scores, dim=-1)  # (B, n, n)
        output  = weights @ V                # (B, n, d_k)
        return output, weights


# Test
B, n, d_model, d_k = 2, 10, 128, 64
x = torch.randn(B, n, d_model)

sa = SelfAttentionTorch(d_model=d_model, d_k=d_k)
out, weights = sa(x)

print(f"Output:  {out.shape}")   # torch.Size([2, 10, 64])
print(f"Weights: {weights.shape}") # torch.Size([2, 10, 10])

Diagram: Self-Attention Flow¶

flowchart LR
    X["X\n(B, n, d_model)"]

    subgraph Projections
        Q["Q = X·W_Q\n(B, n, d_k)"]
        K["K = X·W_K\n(B, n, d_k)"]
        V["V = X·W_V\n(B, n, d_v)"]
    end

    subgraph Score
        S["Q·Kᵀ / √d_k\n(B, n, n)"]
        M["+ Mask (optional)"]
        A["Softmax → Weights\n(B, n, n)"]
    end

    Out["Output = Weights·V\n(B, n, d_v)"]

    X --> Q & K & V
    Q & K --> S --> M --> A
    A & V --> Out

Complexity Analysis¶

Resource	Self-Attention	RNN
Time	O(n² · d)	O(n · d²)
Memory	O(n²)	O(n · d)
Parallelizable	Yes (all tokens at once)	No (sequential)
Long-range dependencies	O(1) path length	O(n) path length

For short sequences (n < 512) and large models (d = 4096+), the d² term in RNNs dominates. Self-attention wins on GPUs. For very long sequences (n > 32K), O(n²) memory becomes the bottleneck, motivating sparse and linear attention variants.

Edge Cases & Misconceptions¶

Misconception: Self-attention is permutation-invariant by default

Without positional encodings, self-attention treats the input as a set, not a sequence. "The cat sat" and "sat cat The" produce the same set of attention interactions (though different scores due to different token identities). Positional encodings break this symmetry — covered in Lesson 04.

Misconception: Attention weights = importance scores

High attention weight on a token does not mean the model finds it semantically important. The value vector v_i could be small in magnitude, so a high weight still contributes little to the output. Gradient-based attribution methods are more reliable than raw attention for interpretation.

Edge case: Attention sink tokens

LLMs learn to dump excess attention onto the first token (often <BOS>) or punctuation. Anthropic observed this in Claude; it is called an "attention sink." StreamingLLM exploits this to enable infinite-context streaming generation.

Production Connection¶

KV cache invalidation: In generation, the K and V matrices grow with each new token. Systems like vLLM use PagedAttention — borrowed from OS virtual memory — to manage fragmented KV caches across requests, dramatically improving GPU utilization.

Grouped-Query Attention (GQA): LLaMA-3 and Mistral use GQA, where multiple query heads share a single key/value head. This reduces KV cache size by 8× with minimal quality loss. Understanding single-head self-attention is the prerequisite.

Tensor parallelism: At scale (GPT-4, 8 GPUs), the Q, K, V projection matrices are split across GPUs. Each GPU computes attention for a subset of heads. This is only possible because the heads are mathematically independent.

Flash Attention: Solving the O(n²) Memory Problem¶

The naive self-attention implementation stores the full attention matrix A ∈ ℝ^{n×n} in GPU high-bandwidth memory (HBM). At n=8192, this is 8192² × 4 bytes = 256 MB per head, per layer — before considering the batch dimension. For a 32-head model with batch=4, that's 32 GB for attention matrices alone.

Flash Attention (Dao et al., 2022) avoids materializing this matrix:

import numpy as np

def flash_attention_conceptual(
    Q: np.ndarray,   # (n, d_k)
    K: np.ndarray,   # (n, d_k)
    V: np.ndarray,   # (n, d_v)
    block_size: int = 64,
) -> np.ndarray:
    """
    Conceptual Flash Attention via tiling.

    Key idea: process attention in blocks that fit in fast SRAM.
    Never materialize the full (n, n) attention matrix.

    Algorithm:
      1. Divide Q into row-blocks of size B_r
      2. For each Q-block, iterate over all K, V blocks
      3. Compute partial attention scores, track running (max, sum) for numerical stability
      4. Accumulate output incrementally

    Complexity:
      Memory: O(n) instead of O(n²)
      Compute: same O(n²) — tiles the same flops but with better cache efficiency
    """
    n, d_k = Q.shape
    d_v = V.shape[-1]
    O = np.zeros((n, d_v))       # output accumulator

    scale = d_k ** -0.5

    for i in range(0, n, block_size):
        # Q block: rows i to i+block_size
        Qi = Q[i:i+block_size]              # (B_r, d_k)
        Oi = np.zeros((len(Qi), d_v))
        m_i = np.full(len(Qi), -np.inf)    # running max for log-sum-exp
        l_i = np.zeros(len(Qi))            # running denominator

        for j in range(0, n, block_size):
            # K, V block: rows j to j+block_size
            Kj = K[j:j+block_size]         # (B_c, d_k)
            Vj = V[j:j+block_size]         # (B_c, d_v)

            # Partial scores
            Sij = Qi @ Kj.T * scale        # (B_r, B_c)

            # Update running max for numerical stability
            m_ij = Sij.max(axis=-1)        # (B_r,)
            m_new = np.maximum(m_i, m_ij)  # (B_r,) — new global max

            # Rescale previous accumulator to new max
            Oi = Oi * np.exp(m_i - m_new)[:, None]
            l_i = l_i * np.exp(m_i - m_new)

            # Add contribution of this K, V block
            exp_Sij = np.exp(Sij - m_new[:, None])  # (B_r, B_c)
            Oi += exp_Sij @ Vj             # (B_r, d_v)
            l_i += exp_Sij.sum(axis=-1)    # (B_r,)

            m_i = m_new

        # Normalize: divide by running denominator
        O[i:i+block_size] = Oi / l_i[:, None]

    return O   # (n, d_v)

Memory comparison at n=4096, d_k=64, one head:

Implementation	Memory	Time
Naive (store full A)	64 MB	Baseline
Flash Attention v1	4 MB (8× reduction)	~2× faster
Flash Attention v2	4 MB	~4× faster (better parallelism)

Flash Attention v2 achieves ~70% of theoretical GPU FLOPs utilization, compared to ~30% for naive attention. This is the single biggest inference speedup in the Transformer stack.

Key Takeaways¶

Self-attention lets every token attend to every other token in the same sequence — O(1) path length between any two positions.
Q, K, V projections separate the roles of seeking, advertising, and carrying information; the model learns all three.
Scaling by √d_k is required to keep softmax gradients non-vanishing when d_k is large.
Causal masking enforces autoregressive generation: token i cannot see j > i.
O(n²) time and memory: the price of full pairwise attention; optimizations like Flash Attention and sparse attention mitigate this.
Without positional encoding, self-attention is permutation-invariant — word order doesn't exist yet.

📹 Recommended Videos¶

3Blue1Brown: Attention in transformers, visually explained — beautiful geometric intuition
Andrej Karpathy: Let's build GPT from scratch — full causal self-attention implementation
Sebastian Raschka: Self-Attention from scratch

🚀 Next Lesson¶

Lesson 3: Multi-Head Attention — running multiple self-attention operations in parallel, each learning a different relational pattern, and concatenating their outputs.