04 — Foundations

Self-Attention#

Scaled Dot-ProductMathematical 1 min read
◆ The PatternThe mechanism that lets every token look at every other token

Self-attention is the core operation of transformers. Each token produces a Query (what am I looking for?), a Key (what do I contain?), and a Value (what do I output?). Attention weights are the softmax of query-key dot products.

Attention(Q, K, V) = softmax(QKᵀ / √d_k) · V
Scale by √d_k to prevent softmax saturation as dimension grows. Output is a weighted mix of value vectors.
Q = XW_Q K = XW_K V = XW_V where W ∈ ℝ^(d_model × d_k)
Linear projections from the input. No bias in most modern architectures.

Complexity is O(n² · d) — quadratic in sequence length. This is the fundamental bottleneck that drives context window research. Each attention weight tells you how much token i "pays attention to" token j.

The √d_k scaling is crucial. Without it, dot products grow proportionally with dimension, pushing softmax into regions with tiny gradients.
Interactive — attention weight heatmap

Python — self-attention from scratch#

Pattern bridge: Every token attending to every other token — a complete correlation matrix computed at each layer. In markets, VWAP weights each price by volume — attention over the session.
Performance in practice
  • O(n²) is the bottleneck: at 128K tokens one attention matrix has 16 billion entries. FlashAttention avoids storing it by computing attention in tiles, so memory grows linearly with length (Dao et al. 2022)
  • FlashAttention-2 roughly doubles FlashAttention’s speed (Dao 2023); PyTorch 2 exposes fused attention kernels through torch.nn.functional.scaled_dot_product_attention
  • For inference: the KV-cache makes attention O(n) per new token, not O(n²). Cache size = 2 × n_layers × n_kv_heads × d_head × seq_len × bytes per value
When to use this
✓ Use when: Building or understanding any transformer model. Debugging attention patterns to understand model behavior. Designing custom architectures. Understanding why context length is limited.
✗ Skip when: Using LLMs via API — attention is handled for you. Working with very long sequences (>100K tokens) where sub-quadratic alternatives (state-space models such as Mamba, or RWKV) may be more efficient.
Worked example

The attention scores for one head in one layer form an n × n matrix. In fp16 that is 0.03 GB at 4k tokens, 2.1 GB at 32k and 34 GB at 128k — for a single head. FlashAttention never stores the matrix: it computes attention in tiles that fit in fast on-chip memory, which is what made long contexts practical.

Where it misleads
  • FlashAttention removes the memory cost, not the compute: the work still grows with n².
  • Attention weights are not explanations of the output (Jain & Wallace 2019).
  • Longer context is not the same as better use of context (see Context Windows).
In code
def score_matrix_gb(n, bytes_per=2):            # one head, one layer, fp16
    return n * n * bytes_per / 1e9

sizes = {n: round(score_matrix_gb(n), 3) for n in (4_096, 32_768, 131_072)}
Python, in your browser — edit the code and run it again

Assumes import numpy as np. Configurations named after real models use their published shapes; other numbers are made up or simulated, as the comments say.

The same idea with a library
import torch
import torch.nn.functional as F

def self_attention(x, W_q, W_k, W_v):
    """x: (batch, seq_len, d_model)"""
    Q = x @ W_q  # (batch, seq_len, d_k)
    K = x @ W_k
    V = x @ W_v
    d_k = Q.size(-1)
    scores = Q @ K.transpose(-2, -1) / d_k**0.5
    weights = F.softmax(scores, dim=-1)
    return weights @ V  # (batch, seq_len, d_v)
Sources
  • A. Vaswani et al., “Attention Is All You Need”, NeurIPS, 2017
  • T. Dao, D. Y. Fu, S. Ermon, A. Rudra & C. Ré, “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness”, NeurIPS, 2022

Last reviewed October 2026 · How this site checks its content · Report a mistake

Linked from 4
← Previous
Positional Encoding
Open in the full reader, with the topic sidebar