Self-attention is a mechanism that lets every token in a sequence compute a weighted sum of all other tokens’ representations, where the weights come from pairwise compatibility scores between learned query and key vectors. Unlike recurrence or convolution, it captures global dependencies in a single layer with no sequential bottleneck. The Transformer architecture builds on this to process entire sequences in parallel while modeling arbitrary long-range interactions.
How self-attention works
At its core, self-attention takes a sequence of token embeddings and produces a new sequence where each position aggregates information from the entire input. The computation follows three linear projections per token: query (Q), key (K), and value (V). For a sequence of length n and model dimension d, you pack the embeddings into a matrix X ∈ ℝⁿˣᵈ and compute:
Q = X W_Q # (n, d_k)
K = X W_K # (n, d_k)
V = X W_V # (n, d_v)The attention scores come from scaled dot-product between queries and keys:
scores = Q Kᵀ / √d_k # (n, n)
weights = softmax(scores) # (n, n)
output = weights @ V # (n, d_v)Each row of the output is a convex combination of all value vectors, weighted by how much that position’s query attends to each key. The √d_k scaling prevents gradients from vanishing when d_k is large — without it, the dot products grow with dimension and push softmax into saturation.
The projection matrices are learned
W_Q, W_K, W_V are not fixed; they’re learned parameters. This means the model discovers what “compatibility” means for its task. In early layers, attention often tracks syntactic relationships (subject-verb agreement, adjective-noun binding). In deeper layers, it captures semantic coreference, discourse structure, and task-specific reasoning patterns.
# Minimal PyTorch implementation
import torch
import torch.nn as nn
import torch.nn.functional as F
class SelfAttention(nn.Module):
def __init__(self, d_model, d_k, d_v):
super().__init__()
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_v, bias=False)
self.scale = d_k ** -0.5
def forward(self, x, mask=None):
# x: (batch, seq_len, d_model)
q = self.W_q(x) # (batch, seq_len, d_k)
k = self.W_k(x) # (batch, seq_len, d_k)
v = self.W_v(x) # (batch, seq_len, d_v)
scores = torch.matmul(q, k.transpose(-2, -1)) * self.scale # (batch, seq_len, seq_len)
if mask is not None:
scores = scores.masked_fill(mask == 0, float('-inf'))
weights = F.softmax(scores, dim=-1)
return torch.matmul(weights, v) # (batch, seq_len, d_v)Why multi-head attention matters
Single-head attention forces all relationship types through one compatibility function. Multi-head attention splits d_model into h heads, each with its own W_Q, W_K, W_V projections of reduced dimension d_k = d_v = d_model / h. Each head learns a different “relation type” — one might track pronoun resolution, another verb-argument structure, another positional patterns.
class MultiHeadAttention(nn.Module):
def __init__(self, d_model, n_heads):
super().__init__()
assert d_model % n_heads == 0
self.d_k = d_model // n_heads
self.n_heads = n_heads
self.W_q = nn.Linear(d_model, d_model, bias=False)
self.W_k = nn.Linear(d_model, d_model, bias=False)
self.W_v = nn.Linear(d_model, d_model, bias=False)
self.W_o = nn.Linear(d_model, d_model, bias=False)
def forward(self, x, mask=None):
batch, seq_len, _ = x.shape
q = self.W_q(x).view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2)
k = self.W_k(x).view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2)
v = self.W_v(x).view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2)
scores = torch.matmul(q, k.transpose(-2, -1)) * (self.d_k ** -0.5)
if mask is not None:
scores = scores.masked_fill(mask == 0, float('-inf'))
weights = F.softmax(scores, dim=-1)
out = torch.matmul(weights, v) # (batch, n_heads, seq_len, d_k)
out = out.transpose(1, 2).contiguous().view(batch, seq_len, -1)
return self.W_o(out)The output projection W_O mixes information across heads. This design lets the model attend to different positions for different reasons simultaneously — something a single softmax distribution cannot express.
Positional information is not built in
Self-attention is permutation-equivariant: shuffle the input tokens and the output shuffles identically. The model has no inherent notion of order. Transformers inject positional information via learned or fixed positional encodings added to the input embeddings before the first attention layer.
# Sinusoidal positional encoding (Vaswani et al., 2017)
def sinusoidal_encoding(seq_len, d_model):
pe = torch.zeros(seq_len, d_model)
position = torch.arange(0, seq_len).unsqueeze(1).float()
div_term = torch.exp(torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model))
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
return pe.unsqueeze(0) # (1, seq_len, d_model)Learned positional embeddings work similarly but add parameters. Relative positional encodings (Shaw et al., 2018; T5, Transformer-XL) bias attention scores directly based on distance rather than absolute position, which generalizes better to longer sequences.
A concrete example: pronoun resolution
Consider the sentence: “The trophy doesn’t fit in the brown suitcase because it is too large.”
The word “it” could refer to the trophy or the suitcase. A self-attention head that learns to resolve this will assign high weight from the query at “it” to the key at “trophy” (or “suitcase”) based on semantic compatibility learned through W_Q and W_K. The value vector at that position then contributes to “it”’s updated representation, carrying the size property forward.
In practice, you can visualize this by extracting attention weights:
# Extract attention weights for analysis
def get_attention_weights(model, input_ids, layer_idx, head_idx):
# Assumes model returns (output, attention_weights) where
# attention_weights is list of (batch, n_heads, seq_len, seq_len)
_, attn_weights = model(input_ids, return_attention=True)
return attn_weights[layer_idx][0, head_idx].detach().cpu().numpy()Visualization often reveals heads that specialize: one attends to the previous token (local syntax), another to the sentence root (dependency), another to repeated entities (coreference). No single head does everything — the ensemble does.
Computational complexity and memory
Self-attention is O(n²d) in time and O(n²) in memory for the attention matrix. This quadratic scaling is the primary bottleneck for long sequences. A 32K context window with d_model=4096 and 32 heads requires ~8 GB just for the attention scores at FP16 (32K × 32K × 2 bytes × 32 heads ≈ 64 GB — actually distributed across layers, but still massive).
Several approaches mitigate this:
- FlashAttention (Dao et al., 2022): IO-aware kernel that fuses softmax and matmul, avoiding materializing the full n×n matrix in HBM. Standard in modern implementations.
- Sliding window attention: Restrict each token to attend to w neighbors (Longformer, BigBird). Reduces to O(nwd).
- Linear attention: Approximate softmax with kernel feature maps (Performer, Linformer) to get O(nd²).
- Block-sparse patterns: Fixed or learned sparse masks (BigBird, Sparse Transformer).
FlashAttention is now the default in PyTorch 2.0+ via torch.nn.functional.scaled_dot_product_attention — use it instead of manual matmul+softmax.
# Modern PyTorch: uses FlashAttention automatically when available
out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, is_causal=True)Common misconceptions
“Attention weights are interpretable”
Attention weights show where gradient flows during training, not necessarily what the model “looks at.” They can be high for positions that get suppressed later in the network, or low for positions that matter through residual connections. Treat them as hints, not explanations. For rigorous attribution, use integrated gradients or attention rollout.
“Self-attention replaces recurrence”
It replaces sequential recurrence with parallel pairwise interaction. But deep Transformers still compose information across layers — each layer’s output becomes the next layer’s input. The effective receptive field grows with depth, similar to recurrent steps. The difference: all positions update simultaneously per layer.
“More heads always helps”
Beyond a point, heads become redundant. Empirically, 8–32 heads works well for d_model 512–4096. Very large models (PaLM, GPT-3) use 64–96 heads but with d_k ≈ 128. The d_k per head matters more than head count — too small and each head lacks capacity; too large and you lose the multi-head benefit.
“Causal masking is only for decoding”
Causal masking (preventing attention to future tokens) is essential for autoregressive generation. But it also appears in bidirectional models for prefix LM objectives or when modeling streaming inputs. The mask is just a structural prior — use whatever matches your data’s temporal structure.
“Self-attention is the only attention”
Cross-attention (queries from one sequence, keys/values from another) powers encoder-decoder models (translation, summarization) and retrieval-augmented generation. The mechanics are identical — only the source of K and V changes. If you’re building RAG, you’re using cross-attention whether you call it that or not.
Practical considerations for engineers
Initialization matters. Scale W_Q, W_K, W_V with 1/√d_model (or use the default Xavier/He init in your framework). Poor initialization causes attention collapse — all weights uniform — and kills gradients.
Gradient checkpointing. For long sequences, recompute attention during backward pass instead of storing activations. torch.utils.checkpoint.checkpoint wraps any module; apply per Transformer block.
Mixed precision. FP16/BF16 works for attention scores but softmax can underflow. PyTorch’s scaled_dot_product_attention handles this with internal FP32 accumulation. Don’t write your own softmax in FP16.
KV caching. For autoregressive inference, cache K and V from previous steps. Each new token only computes its own q, k, v and attends to the cached history. This reduces per-step cost from O(n²) to O(n).
# Simplified KV cache update
class KVCache:
def __init__(self, max_seq_len, n_heads, d_k, device):
self.k_cache = torch.zeros(1, n_heads, max_seq_len, d_k, device=device)
self.v_cache = torch.zeros(1, n_heads, max_seq_len, d_k, device=device)
self.seq_len = 0
def update(self, k, v):
# k, v: (1, n_heads, 1, d_k)
self.k_cache[:, :, self.seq_len] = k.squeeze(2)
self.v_cache[:, :, self.seq_len] = v.squeeze(2)
self.seq_len += 1
return self.k_cache[:, :, :self.seq_len], self.v_cache[:, :, :self.seq_len]Sequence packing. For training on variable-length sequences, pack multiple examples into one batch with block-diagonal attention masks. Avoids padding waste. Hugging Face’s DataCollatorForLanguageModeling with pad_to_multiple_of helps.
When self-attention isn’t the answer
- Very long sequences (>16K) where quadratic cost dominates: consider state-space models (Mamba, RWKV), linear attention, or hybrid architectures.
- Strict latency budgets on short sequences: a small CNN or RNN can be faster due to lower constant factors.
- Strong local inductive bias (images, audio): convolution or locality-sensitive attention often matches Transformers with fewer parameters.
- Memory-constrained edge deployment: quantized linear attention or distilled smaller models work better than pruning attention heads.
Summary
Self-attention computes pairwise token affinities via learned query-key projections, then aggregates values with those affinities. Multi-head attention splits this into parallel relation-specific channels. Positional encodings break permutation symmetry. The O(n²) complexity drives most architectural innovation around long-context modeling.
If you’re implementing a Transformer from scratch, start with F.scaled_dot_product_attention and a standard pre-norm residual block. The rest — positional encodings, KV caching, gradient checkpointing, mixed precision — are engineering concerns that compound. Get the core attention math right and the system scales.