The 4 Lines of PyTorch That Power Self-Attention

When you strip away all the multi-head scaling and architectural boilerplate, the actual "attention" mechanism that powers GPT models is shockingly simple.

It's just this:

# 1. Compute attention scores: Q dot K^T, scaled by sqrt of dimension
scores = query @ key.transpose(-2, -1) * (1.0 / math.sqrt(k.size(-1)))

# 2. Prevent looking into the future (Causal Masking)
scores = scores.masked_fill(future_mask == 0, float('-inf'))

# 3. Convert raw scores to probabilities that sum to 1.0
weights = F.softmax(scores, dim=-1)

# 4. Multiply weights against Values to get the final context
out = weights @ value

• Query (Q): What information this token is looking for.
• Key (K): What information this token contains.
• Value (V): The actual information this token will share if selected.

Everything else in the Transformer architecture — Feed-Forward networks, LayerNorm, Residual connections — exists just to stabilize and process the output of these four lines.