Self-Attention Mechanism

AI & Machine Learning Foundations beginner How AI Works 11 min read

Overview

Self-attention (sometimes called intra-attention) is the mechanism at the heart of every Transformer model. It is the process by which each token in a sequence computes a weighted relationship with every other token, allowing the model to understand context, resolve ambiguity, and capture long-range dependencies in a single computational step.

Consider the sentence: “The animal didn’t cross the street because it was too tired.” What does “it” refer to? A human instantly knows “it” means “the animal.” Self-attention gives the model the same ability — when processing the token “it,” the attention mechanism assigns high weight to “animal” and low weight to “street,” effectively resolving the reference.

For AI engineers, understanding self-attention is essential because it directly governs the quality of model outputs, the computational cost of inference, and the limits of what context window sizes are practical. Every token of every API call you make passes through dozens of self-attention layers, and the mechanism’s properties determine how well the model can reason about the information you provide.

How It Works

Step 1: Creating Query, Key, and Value Vectors

For each token in the input sequence, the model creates three vectors by multiplying the token’s embedding by three learned weight matrices: W_Q, W_K, and W_V.

Q = X · W_Q    (Query: "What am I looking for?")
K = X · W_K    (Key: "What do I contain?")
V = X · W_V    (Value: "What information do I contribute?")

Think of it like a library lookup system:

These vectors are typically smaller than the full embedding dimension. If d_model = 512 and there are 8 attention heads, each head’s Q, K, V vectors have dimension d_k = 64.

Key Formula — Scaled Dot-Product Attention

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

This single equation is the core of all Transformer computation. Every modern LLM runs billions of instances of this formula per inference call.

Step 2: Computing Attention Scores

The attention score between two tokens is the dot product of one token’s Query vector and another token’s Key vector. For a sequence of n tokens, this produces an n × n score matrix where each cell [i, j] represents how much token i should attend to token j.

scores = Q · K^T

Example for 4 tokens:
         Key_1  Key_2  Key_3  Key_4
Query_1 [  12     3      7      1  ]
Query_2 [   5    15      2      8  ]
Query_3 [   2     6     11      4  ]
Query_4 [   9     1      3     14  ]

Higher scores mean stronger relationships. In the example above, token 1 attends most strongly to itself (score 12) and token 3 (score 7).

Step 3: Scaling

The raw dot products are divided by √d_k (the square root of the key dimension). Without this scaling, large dot products push the softmax function into regions where its gradient is extremely small, causing training to stall. For d_k = 64, we divide by √64 = 8.

scaled_scores = scores / √d_k

# Our example with d_k = 64:
scaled_scores = scores / 8

Step 4: Softmax Normalization

The softmax function converts the scaled scores into a probability distribution — each row sums to 1.0. This determines the attention weights: how much each token contributes to the output representation of the current token.

attention_weights = softmax(scaled_scores)

# Each row now sums to 1.0:
         Token_1  Token_2  Token_3  Token_4
Token_1 [ 0.52    0.05     0.38     0.05  ]
Token_2 [ 0.08    0.65     0.03     0.24  ]
...

Step 5: Weighted Sum of Values

Finally, each token’s output is computed as the weighted sum of all Value vectors, using the attention weights. If token 1 has attention weight 0.52 on itself and 0.38 on token 3, its output is heavily influenced by those two tokens’ Value vectors.

output = attention_weights · V

# Token 1's output:
output_1 = 0.52 * V_1 + 0.05 * V_2 + 0.38 * V_3 + 0.05 * V_4

Multi-Head Attention: Parallel Perspectives

A single attention computation captures one type of relationship. Multi-head attention runs h independent attention computations in parallel, each with its own learned W_Q, W_K, W_V matrices. Each head can specialize in different patterns:

The outputs of all heads are concatenated and passed through a final linear projection:

MultiHead(Q, K, V) = Concat(head_1, ..., head_h) · W_O

where head_i = Attention(Q · W_Qi, K · W_Ki, V · W_Vi)

Causal (Masked) Self-Attention

In decoder models (GPT, Claude), a causal mask prevents tokens from attending to future positions. This is implemented by setting future positions to -∞ before the softmax, which drives their attention weights to zero. This ensures autoregressive generation: each token can only see itself and previous tokens.

mask = [[ 0,  -∞, -∞, -∞],
        [ 0,   0, -∞, -∞],
        [ 0,   0,  0, -∞],
        [ 0,   0,  0,  0]]

masked_scores = scores + mask

Practical Applications

Why Self-Attention Quality Matters for RAG

When you inject retrieved context into a prompt, self-attention is what allows the model to find and reason over the relevant passages. If your retrieved chunks are buried among irrelevant text, the attention mechanism may fail to weight them properly — the model’s output degrades even though the right information is technically present. This is why context engineering (careful placement and formatting of retrieved data) matters so much.

Computational Cost: The O(n²) Reality

Self-attention computes an n × n score matrix for every head in every layer. For a 128K token context window with 96 attention heads across 80 layers, that’s 128,000² × 96 × 80 ≈ 100 trillion attention score computations per forward pass. This quadratic scaling is the fundamental reason context windows are expensive and why longer contexts cost more in API pricing.

Attention Patterns in Practice

Researchers have visualized attention patterns and found that models develop interpretable behaviors:

Practical Code Example (PyTorch)

import torch
import torch.nn.functional as F

def scaled_dot_product_attention(Q, K, V, mask=None):
    d_k = Q.size(-1)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / (d_k ** 0.5)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))
    weights = F.softmax(scores, dim=-1)
    return torch.matmul(weights, V), weights

Key Takeaways

Further Reading