Inside the Transformer

The 2017 paper "Attention Is All You Need" by Vaswani et al. is probably the most consequential ML paper of the last decade. It introduced the Transformer, which threw out recurrent and convolutional sequence models in favor of a mechanism built entirely on attention. Every major language model since — GPT, BERT, PaLM, LLaMA, Claude — sits on this foundation.

If you work in AI, understanding the Transformer isn't optional. So I'll walk through it piece by piece.

The problem with recurrence

Before Transformers, sequence modeling meant RNNs and LSTMs. Those process tokens one at a time, carrying a hidden state that accumulates information left to right. That sequential design has two real problems:

Training won't parallelize. Each timestep depends on the previous one, so you can't shard the work across GPU cores in any useful way.
Long-range dependencies decay. Even LSTMs, which were specifically designed to keep information alive over long sequences, struggle when the relevant context is hundreds of tokens back. The hidden state is a bottleneck.

The Transformer drops recurrence entirely. Every token can attend to every other token directly, in parallel.

The architecture at a glance

The original Transformer is encoder-decoder. The encoder processes the input sequence and produces a set of representations. The decoder generates the output one token at a time, attending both to the encoder's output and to its own previously generated tokens.

For language models like GPT, only the decoder is used. For BERT-style models, only the encoder. The core mechanism — self-attention — is the same in both.

Self-attention: the core mechanism

Self-attention lets each token in a sequence look at every other token and decide how much to "attend" to it. Here's how it works.

Given an input sequence of n tokens, each represented as a d-dimensional vector, we compute three matrices:

Q (Queries): what is this token looking for?
K (Keys): what does this token offer?
V (Values): what information does this token carry?

Each one comes from multiplying the input by a learned weight matrix:

Q = XW_Q \qquad K = XW_K \qquad V = XW_V

with $Q, K \in \mathbb{R}^{n \times d_k}$ and $V \in \mathbb{R}^{n \times d_v}$ . Attention is computed as:

\mathrm{Attention}(Q, K, V) = \mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V

Walking through this:

$QK^\top$ produces an $n \times n$ matrix where entry $(i, j)$ is the dot product between query $i$ and key $j$ . That measures how much token $i$ should attend to token $j$ .
The $1/\sqrt{d_k}$ is a scaling factor. Without it, the dot products grow as $d_k$ grows, pushing softmax into regions with vanishing gradients. The square root keeps variance stable.
softmax normalizes each row to sum to 1, so the raw scores become a probability distribution over tokens.
The final multiply by $V$ gives the output: each token's representation is a weighted sum of all value vectors, weighted by attention.

In code:

import torch
import torch.nn.functional as F
import math

def scaled_dot_product_attention(Q, K, V, mask=None):
    d_k = Q.size(-1)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)

    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

The mask parameter shows up in the decoder, where you don't want tokens attending to future positions. That's causal masking: token i can only see tokens 0 through i.

Multi-head attention

A single attention head can only capture one type of relationship at a time — maybe "subject-verb agreement" or "coreference resolution." To capture several at once, the Transformer uses multi-head attention.

Instead of one attention op with d-dimensional Q, K, V, you split them into h heads, each with d/h dimensions. Every head computes attention independently, and the outputs are concatenated and projected:

class MultiHeadAttention(torch.nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.n_heads = n_heads
        self.d_k = d_model // n_heads

        self.W_Q = torch.nn.Linear(d_model, d_model)
        self.W_K = torch.nn.Linear(d_model, d_model)
        self.W_V = torch.nn.Linear(d_model, d_model)
        self.W_O = torch.nn.Linear(d_model, d_model)

    def forward(self, x, mask=None):
        batch_size, seq_len, _ = x.size()

        # Project and reshape into (batch, heads, seq_len, d_k)
        Q = self.W_Q(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        K = self.W_K(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        V = self.W_V(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)

        # Attention per head
        attn_output, _ = scaled_dot_product_attention(Q, K, V, mask)

        # Concatenate heads and project
        attn_output = attn_output.transpose(1, 2).contiguous().view(
            batch_size, seq_len, -1
        )
        return self.W_O(attn_output)

In practice, different heads learn to attend to different things. Interpretability work has shown that some heads specialize in syntactic relationships (attending to the previous token or the start of the sentence), while others pick up semantic ones (attending to the subject of a relative clause, for instance).

Positional encoding

Self-attention is permutation-invariant. Shuffle the input tokens and the attention scores change, but the mechanism itself has no built-in notion that position matters. Since word order obviously matters for language, the Transformer adds positional encodings to the input embeddings.

The original paper uses sinusoidal functions:

PE_{(pos,\,2i)} = \sin\!\left(\frac{pos}{10000^{2i/d_{model}}}\right) \qquad PE_{(pos,\,2i+1)} = \cos\!\left(\frac{pos}{10000^{2i/d_{model}}}\right)

Each position gets a unique pattern of sines and cosines at different frequencies. The intuition is that each dimension oscillates at a different rate, so every position gets a distinct "fingerprint." A handy property is that the encoding for position $p+k$ can be written as a linear transformation of the encoding for position $p$ , which may help the model learn relative positions.

Modern models have moved on. BERT and GPT-2 use learned positional embeddings. LLaMA and many current LLMs use rotary position embeddings (RoPE), which encodes position by rotating the query and key vectors. RoPE naturally decays attention to distant tokens and generalizes better to sequence lengths the model didn't see during training.

The full Transformer block

A single Transformer block has:

Multi-head self-attention
Residual connection + layer normalization
Feed-forward network (two linear layers with a nonlinearity)
Residual connection + layer normalization

\text{h} = \mathrm{LayerNorm}\big(x + \mathrm{MultiHeadAttention}(x)\big)

\text{output} = \mathrm{LayerNorm}\big(\text{h} + \mathrm{FFN}(\text{h})\big)

The feed-forward network (FFN) runs independently on each position. It typically expands the dimension by 4x, applies GELU or ReLU, then projects back down. Recent research suggests that the FFN layers act as key-value memories, storing factual knowledge learned during training — which is why model editing techniques target them.

Residual connections aren't optional. Without them, gradients vanish in deep models. The residual stream acts as a highway that carries information from early layers to later ones, with each layer adding its own contribution.

Encoder-decoder vs decoder-only

The original Transformer has both stacks. The encoder processes the full input with bidirectional attention — every token sees every other token. The decoder generates output autoregressively with causal masking, plus cross-attention layers that attend to the encoder's output.

The field has mostly converged on decoder-only for generation. GPT, LLaMA, and most modern LLMs are decoder-only with causal masking. BERT and its variants are encoder-only with bidirectional attention, which makes them strong for classification and extraction but unusable for generation. The encoder-decoder structure still shows up in T5 and in machine translation, where input and output are clearly distinct sequences.

Why Transformers won

The Transformer's dominance comes down to three properties.

Parallelism. Every token is processed at the same time during training. On modern GPU hardware that translates to massive speedups over recurrent models. Training GPT-3 with LSTMs would have been computationally infeasible.

Flexible context. Self-attention gives every token a direct path to every other token regardless of distance. An LSTM needs information to survive passage through every intermediate hidden state. A Transformer just needs the attention weight to be nonzero.

Scaling. Transformers scale predictably with compute, data, and parameters. The scaling laws (Kaplan et al., Hoffmann et al.) were discovered specifically for this architecture. More parameters, more data, more compute, and you reliably get a better model.

The big limitation is the quadratic cost of self-attention — $O(n^2)$ in sequence length. A 100K-token sequence has an attention matrix with 10 billion entries. That has pushed a lot of research into efficient variants: sparse attention, linear attention, FlashAttention (which is really a memory-access optimization), and sliding window approaches. But the core mechanism is still the scaled dot-product attention from the original paper.

Practical implications

If you're building systems on top of LLMs, understanding the Transformer helps you reason about how they behave. Long-range dependencies should work in theory, but attention dilutes over very long contexts in practice. The model's knowledge lives in the FFN layers and can be localized, which is how model editing works. Different attention heads capture different linguistic phenomena, which you can exploit for interpretability.

The Transformer isn't just an architecture. It's the computational substrate the current era of AI runs on.