Self-attention is the specific case of attention where the queries, keys, and values all come from the same sequence — every token attends to every other token (including itself) within one single input, rather than one sequence attending to a different one.
The Defining Feature
Compare this to the Q/K/V setup in Query, Key, Value — the formulas are identical, but here \(\mathbf{X}\) is one single input sequence, used to generate all three projections. Every token's query gets compared against every token's key (including its own), letting the model directly relate any two positions in the sequence, however far apart they are.
Why This Matters — Direct, Position-Independent Relationships
Recall from RNN Vanishing Gradient that an RNN/LSTM/GRU must relate distant positions indirectly, by carrying information forward through many intermediate time steps — a process vulnerable to gradient shrinkage over long distances. Self-attention instead computes a direct connection between any two positions, regardless of how far apart they are in the sequence — token 1 and token 100 are exactly as easy to relate as token 1 and token 2, since the computation doesn't route through any intermediate steps at all.
Diagram — Every Token Attending to Every Token
Every token can directly attend to every other token in a single computation — no intermediate steps, no distance penalty.
Numerical/Practical Example — What a Token "Learns" to Attend To
In the sentence "The animal didn't cross the street because it was too tired," self-attention lets the model directly compute a strong connection between "it" and "animal" (resolving the pronoun's reference) — regardless of how many words separate them — by learning query/key projections that make "it"'s query vector align well with "animal"'s key vector. This kind of resolved reference is a genuinely famous illustrative example from the original Transformer research.
Code
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
class SelfAttention(nn.Module):
def __init__(self, d_model, d_k):
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_k, bias=False)
self.d_k = d_k
def forward(self, x): # x: (batch, seq_len, d_model) -- ONE sequence, used for Q, K AND V
Q, K, V = self.W_Q(x), self.W_K(x), self.W_V(x)
scores = Q @ K.transpose(-2, -1) / math.sqrt(self.d_k)
weights = F.softmax(scores, dim=-1)
return weights @ V
self_attn = SelfAttention(d_model=16, d_k=8)
x = torch.randn(1, 6, 16) # a sequence of 6 tokens
output = self_attn(x)
print(output.shape) # (1, 6, 8)Common Mistakes
- Confusing self-attention with cross-attention — self-attention's Q, K, and V all originate from the exact same sequence; cross-attention (next note) deliberately draws Q from one sequence and K/V from a different one.
- Assuming self-attention has some notion of sequence order built in — the raw mechanism, as described here, is actually order-agnostic (permuting the input tokens would just permute the output correspondingly, with the same relationships) — this is exactly why Transformers need a separate positional encoding mechanism, covered in the next category.
Interview Relevance
Q: "Why can self-attention relate two distant tokens in a sequence more effectively than an RNN can?" Self-attention computes a direct connection between every pair of positions in a single operation — there's no intermediate chain of hidden states the information has to survive passing through, unlike an RNN, where relating distant positions requires that information to be carried, step by step, across every intervening time step, making it vulnerable to the vanishing gradient problem over long distances.
Practice Question
Why is self-attention, as described in this note, considered "order-agnostic" — and what problem might that create for tasks like language modeling, where word order clearly matters?
