Understanding Transformers: Self-Attention from First Principles

By Prahlad Menon Published 2026-06-19 7 min read

Transformers power every major LLM — GPT, Claude, Gemini, LLaMA. Most explanations jump straight to architecture diagrams without building intuition. Here’s a ground-up explanation following the pedagogical approach from Jurafsky and Martin’s Speech and Language Processing Chapter 8.

The Architecture at a Glance

Before diving into details, here’s the big picture: a transformer processes each token through a “residual stream” — a d-dimensional vector that flows up through N stacked transformer blocks. Each block enriches this stream through attention (which mixes in information from other tokens) and feedforward layers.

The key insight: attention is the only component that looks at other tokens. Everything else processes each position independently.

Why Attention? The Contextual Meaning Problem

Static embeddings (Word2Vec, GloVe) assign each word a fixed vector. But word meaning depends heavily on context:

The chicken didn’t cross the road because it was too tired. The chicken didn’t cross the road because it was too wide.

In sentence one, “it” = chicken. In sentence two, “it” = road. Same word, completely different referent.

Even more interesting — reading left-to-right up to just “…because it”:

The chicken didn’t cross the road because it…

At this point, “it” could plausibly refer to either chicken or road! A good representation of “it” should encode information about both potential referents, weighted by plausibility.

More examples of distant contextual dependencies:

“The keys to the cabinet are on the table” — subject-verb agreement spans multiple words
“I walked along the pond… the bank” — “bank” means riverbank because of “pond” many words earlier

Attention is the mechanism that lets each token integrate information from relevant context words, no matter how far away.

Self-Attention Layer: Information Flow

In a self-attention layer, when processing token position i, the model attends to all prior positions (1 through i) to compute the output.

Figure: Self-attention information flow When processing each input xᵢ, the model attends to all inputs up to and including xᵢ.

This happens in parallel for all positions — that’s the computational advantage over sequential RNNs.

Attention: Start Simple, Then Add Complexity

The Simplest Version

At its core, attention is just a weighted sum of prior token representations:

aᵢ = Σⱼ≤ᵢ αᵢⱼ · xⱼ

Each αᵢⱼ is a weight indicating how much token j should contribute to the representation of token i.

How do we compute these weights? Through similarity. Tokens more similar to the current token contribute more:

score(xᵢ, xⱼ) = xᵢ · xⱼ           (dot product)
αᵢⱼ = softmax(score(xᵢ, xⱼ))     (normalize to sum to 1)

This simplified version captures the essence: compare current token to prior tokens, weight the sum by similarity.

The Problem: One Embedding, Three Roles

The simplified version compares raw embeddings directly. But during attention, each token actually plays three different roles:

Query: “What am I looking for?” — the current token’s request for relevant context
Key: “What do I have to offer?” — what each token advertises about its content
Value: “Here’s my actual contribution” — the content that gets weighted and summed

Query, Key, Value Projections

Transformers introduce three weight matrices — Wᴷ, Wᑫ, Wⱽ — to let each token have distinct representations for these roles:

qᵢ = xᵢ · Wᑫ    (query: what position i is looking for)
kⱼ = xⱼ · Wᴷ    (key: what position j offers)
vⱼ = xⱼ · Wⱽ    (value: what position j contributes)

Now similarity is computed between the query of the current token and the keys of prior tokens:

score(xᵢ, xⱼ) = (qᵢ · kⱼ) / √dₖ

The √dₖ scaling is critical: without it, dot products grow large as dimensions increase, causing softmax to produce extreme values (near 0 or 1), leading to vanishing gradients.

The final output uses values:

αᵢⱼ = softmax((qᵢ · kⱼ) / √dₖ)   for all j ≤ i
headᵢ = Σⱼ≤ᵢ αᵢⱼ · vⱼ
aᵢ = headᵢ · Wᴼ                   (reshape to output dimension)

The separation of Q/K/V is powerful: the model can learn that features to match on (keys) differ from content to retrieve (values). Like a database: search by key, retrieve the value.

Visualizing Attention Weights

Here’s what learned attention looks like in practice:

Figure: Attention weights for "it" Self-attention weight distribution for the word “it” — darker = higher attention. The model attends heavily to “chicken” and “road” since “it” could refer to either.

The visualization shows attention weights αᵢⱼ for computing the representation of “it.” Darker shading = higher weight. Notice the model attends most to “chicken” and “road” — exactly the tokens that could be the referent of “it.”

Parallel Computation: Matrix Form

For efficiency, we pack all N tokens into matrices and compute everything at once:

Q = X · Wᑫ    [N × dₖ]
K = X · Wᴷ    [N × dₖ]
V = X · Wⱽ    [N × dᵥ]

All attention scores become a single matrix multiplication:

QKᵀ    [N × N matrix]

Entry (i, j) = dot product of query i with key j.

Causal Masking: Can’t See the Future

Problem: QKᵀ computes scores for all pairs, including future tokens. That’s cheating in language modeling!

Solution: mask the upper triangle to -∞ before softmax:

QKᵀ after masking:

[ q₁·k₁   -∞      -∞      -∞   ]
[ q₂·k₁  q₂·k₂   -∞      -∞   ]
[ q₃·k₁  q₃·k₂  q₃·k₃   -∞   ]
[ q₄·k₁  q₄·k₂  q₄·k₃  q₄·k₄ ]

Since softmax(-∞) = 0, future positions get zero weight. Each position only attends to itself and earlier positions.

Multi-Head Attention

One attention head learns one type of relationship. But language has many:

Syntactic (subject-verb agreement)
Semantic (word meaning disambiguation)
Coreference (what “it” refers to)

Multi-head attention runs A separate attention operations in parallel, each with its own Wᑫ, Wᴷ, Wⱽ. Each head can specialize.

MultiHead(X) = Concat(head₁, head₂, ..., headₐ) · Wᴼ

Typical configuration: d=512, A=8 heads, each head operates in dₖ=dᵥ=64 dimensions.

The Residual Stream

A key insight from transformer interpretability (Elhage et al., 2021): think of each token as having a residual stream — a d-dimensional vector that flows up through all layers.

The original embedding enters at the bottom
Attention reads from neighboring streams and writes back to the current stream
Feedforward layers process the stream independently at each position
Residual connections add outputs back in

Attention literally moves information between token streams. When computing “it,” attention reads from the “chicken” and “road” streams and adds that information into the “it” stream.

Figure: Residual stream information flow An attention head moves information from Token A’s residual stream into Token B’s residual stream.

The Full Transformer Block

Each block combines:

Layer Norm (stabilizes values)
Multi-Head Attention (token mixing — only component that sees other positions)
Residual connection (add attention output back)
Layer Norm
Feedforward Network (2-layer MLP, same for each position)
Residual connection

In equations (prenorm architecture):

t¹ = LayerNorm(x)
t² = MultiHeadAttention(t¹)
t³ = t² + x                    (residual)
t⁴ = LayerNorm(t³)
t⁵ = FFN(t⁴)
h = t⁵ + t³                    (residual)

Blocks stack because input and output both have dimension d. Large models use 12–96+ blocks.

Input: Token + Position Embeddings

The input to the first block combines:

Token embedding: Look up token ID in embedding matrix E [|V| × d]
Position embedding: Add position-specific vector [N × d]

Why position? Attention is permutation-invariant — it doesn’t inherently know position. We must encode it explicitly.

X = TokenEmbedding + PositionEmbedding

The Language Modeling Head

After the final block, we predict the next token:

Take final embedding at position N: shape [1 × d]
Multiply by unembedding matrix Eᵀ: shape [d × |V|]
Apply softmax → probability distribution over vocabulary

The unembedding matrix is often the transpose of the input embedding matrix (weight tying) — the same weights are optimized for both embedding and unembedding.

Computational Complexity

Attention is O(N²) in sequence length — we compute dot products between all pairs. This is why very long contexts are expensive. Modern architectures use various tricks (sparse attention, sliding windows, etc.) for longer contexts.

Why This Architecture Works

Parallelism: All positions compute simultaneously (unlike sequential RNNs)
Direct connections: Any token attends directly to any other — long-range dependencies are cheap
Learned relevance: Q/K/V matrices discover what to attend to
Modularity: Same dimensionality throughout → components are interchangeable
Residual streams: Information accumulates progressively through layers