How a Transformer Works: A Decoder-Only Model, End to End

The Big Picture

A decoder-only transformer does one job: given a sequence of tokens, produce a probability for every possible next token. Everything else — chatbots, code completion, long essays — comes from running that one job in a loop. Here is the whole path for our example prompt, with the shape of the data at each stage. T is the number of tokens (5 here), d is the model width, and V is the vocabulary size.

1. Input text"The cat sat on the"string
2. Tokenizesplit into pieces a model knows (BPE)5 tokens
3. Token IDseach piece becomes an integer(T)
4. Token embeddingslook up one learned vector per ID(T, d)
5. + Position embeddingsadd a learned vector per position(T, d)
6. N transformer blocksattention + feed-forward, repeated; shape never changes(T, d)
7. Final LayerNormone last normalization(T, d)
8. LM headproject every position onto the vocabulary(T, V)
9. Softmaxturn the last position's scores into probabilities(V)
10. Samplepick one token, append it, repeat1 token

The single most useful fact to hold on to: inside the stack, the data is always a T × d matrix — one d-dimensional vector per token. Every block reads that matrix and writes a matrix of the same shape. What changes is what those vectors mean: at the bottom each vector describes one word in isolation; by the top, each vector encodes that word in the context of everything before it.

From Text to Vectors

Tokenization splits text into pieces from a fixed vocabulary — common words become one token, rare words break into several. Real models use byte-pair encoding, covered from scratch in Part 2 of the CS336 series. Each token then becomes an integer ID; GPT-2's tokenizer, for example, maps our five words to five IDs. No learning happens here — it's a lookup.

Token embeddings are the first learned component: a table E with one row of d numbers per vocabulary entry. Looking up five IDs gives a 5 × d matrix. Training moves these rows around so that tokens used in similar ways end up with similar vectors.

Position embeddings fix a real problem: attention, as we'll see, treats its input as an unordered set. Without position information, "the cat sat on the mat" and "the mat sat on the cat" would look identical. GPT-2's answer is a second learned table P with one row per position; row 0 is added to the first token, row 1 to the second, and so on. (Most newer models use RoPE instead, which rotates vectors by an angle that depends on position and has no learned table — see Part 4.)

Inside One Block

A block has two halves, and each half follows the same pattern: normalize, transform, add back. Attention lets each token gather information from earlier tokens; the feed-forward layer then processes each token on its own. The "add back" — the residual connection — means each half only has to learn a change to the vector, not rebuild it, which is a large part of why very deep stacks train at all.

In code, the whole block is two lines — the rest of this article is about what's inside them:

x = x + attention(layer_norm_1(x))  # mix across tokens
x = x + mlp(layer_norm_2(x))        # per-token

Attention, Step by Step

Attention is how the vector for "the" (the last word of our prompt) can pick up that the sentence is about a cat sitting on something. Each token asks a question, every earlier token offers an answer, and each token takes a weighted blend of the answers that match its question best. Concretely, with the shapes our script printed (T = 5, d = 64, h = 4 heads, so dk = 16 per head):

1. Project to queries, keys and values(5, 64) eachQ = X·W_Q ("what am I looking for"), K = X·W_K ("what do I contain"), V = X·W_V ("what I'll hand over").
2. Split into h heads(4, 5, 16)Reshape, don't recompute: 4 heads each get a 16-number slice, so they can learn to look for different things.
3. Score every pair: QKᵀ / √d_k(4, 5, 5)A 5 × 5 grid per head: how well each token's query matches each token's key. Dividing by √d_k keeps the numbers in a range where softmax still has useful gradients.
4. Apply the causal mask(4, 5, 5)Set every score above the diagonal to −∞, so no token can see a later one. This is what makes it a decoder.
5. Softmax each row(4, 5, 5)Each row becomes attention weights that are ≥ 0 and sum to 1; masked positions become exactly 0.
6. Weighted sum of values(4, 5, 16)Each token's output is its row of weights times V: a blend of what the earlier tokens offered.
7. Concatenate heads, project with W_O(5, 64)Glue the 4 heads back into 64 numbers per token, then mix them with one more linear layer. Back to (T, d).

Try it: who can see whom?

Click a token on the left to see which positions it's allowed to attend to. Each row is one token's view; the last row — the token whose output predicts the next word — sees the whole prompt.

A tiny example with real numbers

Three tokens, one head, dk = 2, with small hand-picked vectors so you can check the arithmetic: Q = [[1,0], [0,1], [1,1]], K = [[1,0], [1,1], [0,1]], V = [[1,0], [0,1], [2,2]]. Computed with PyTorch:

Read the rows. Token 1 can only see itself, so its weight is 1.00 on itself and its output is simply its own value, [1, 0]. Token 2 matches token 2's key more strongly (0.71 vs 0.00), so it takes two thirds of its answer from token 2. Token 3's query [1, 1] lines up best with key [1, 1], giving weights 0.25 / 0.50 / 0.25 and output 0.25·[1,0] + 0.50·[0,1] + 0.25·[2,2] ≈ [0.75, 1.00] (0.744 exactly, before rounding the weights). That blend of earlier information is what attention adds to a token.

The Feed-Forward Layer

After attention has mixed information across positions, the feed-forward network (MLP) processes each token independently, with the same weights at every position:

The expansion to 4d gives the layer room to compute many intermediate features, and GELU — x · Φ(x), where Φ is the standard normal CDF — is the nonlinearity that makes those features more than a linear reshuffle. A useful mental model, backed by interpretability research, is that attention moves information between tokens while MLPs store and apply much of what the model knows.

From Vectors to a Prediction

Stack N blocks (12 in GPT-2 small, 2 in our toy model) and the shape is still (T, d). Three steps turn it into a prediction:

Try it: temperature and top-k

The softmax output is a probability for every token; how you pick from it changes the text a lot. The scores below are illustrative (not from a real model) for continuing "The cat sat on the". Move the sliders:

Low temperature sharpens the distribution toward the single most likely token (greedy decoding, the limit as temperature → 0, always picks "mat"). High temperature flattens it, making unusual continuations more likely. Top-k keeps only the k highest-scoring tokens and renormalizes; nucleus (top-p) sampling instead keeps the smallest set whose probabilities sum to p, so the cutoff adapts to how confident the model is — the method introduced by Holtzman et al. to avoid the bland, repetitive text that pure likelihood-maximizing decoding produces.

The Generation Loop

A transformer predicts one token. Text comes from a loop: predict, sample, append, run the model again on the longer sequence. That's what "autoregressive" means.

Two practical consequences follow. First, generation is sequential — token 100 can't be computed before token 99 exists — which is why output tokens are slower and cost more than input tokens. Second, running the full model over the whole sequence at every step would repeat almost all of the work; real inference engines cache each layer's keys and values (the KV cache) so each new step only computes attention for the one new token.

How Many Parameters?

Only a few stages have learned weights: the embedding tables, the four attention projections and two MLP matrices in each block, the norms, and the LM head. Ignoring biases and norm scales (a tiny fraction), the count is:

Worked example at the size of a typical "7B" model — d = 4,096, d_ff = 16,384, N = 32, V = 50,000:

For every term, including biases and norm scales, and a PyTorch script that counts them from a real model, see Part 4's parameter count.

How It Learns: Backpropagation

Everything so far was the forward pass: run the model and get a prediction. Training adds a backward pass that works out, for every one of those parameters, which direction to nudge it so the prediction gets better. The training loop in our script is three lines:

logits = model(data[:, :-1])                                            # forward: predict each next token
loss = F.cross_entropy(logits.reshape(-1, V), data[:, 1:].reshape(-1))  # how wrong were we?
opt.zero_grad(); loss.backward(); opt.step()                            # backward, then update

loss.backward() computes a gradient for every parameter; opt.step() uses them to update the weights. The rest of this section is what those two calls do.

1. The loss

Training text supplies the right answer at every position: after "The cat sat on the" the next token really was "mat". The loss is cross-entropy — the negative log of the probability the model gave the correct token:

2. The first gradient: p − y

Backpropagation starts at the loss and applies the chain rule backwards. The first step has a famously clean answer: for softmax followed by cross-entropy, the gradient of the loss with respect to each logit is simply predicted probability minus target, where the target y is 1 for the correct token and 0 for everything else. Checked with PyTorch's autograd:

Read it as instructions. Gradient descent moves opposite the gradient, so "mat"'s score gets pushed up (its gradient is negative), and every other token's score gets pushed down in proportion to how much probability it wrongly took. The rows sum to zero: probability can only move between tokens.

3. The gradient flows backward through the model

Every component does the same two things on the way back: compute the gradient for its own weights (to be updated), and the gradient for its input (to pass further back). Four pieces are worth understanding:

Two things never receive a gradient: tokenization (a discrete lookup, not a differentiable function), and the embedding rows of tokens that didn't appear in the batch — their gradient is exactly zero that step.

4. The update: Adam

With a gradient g for every parameter θ, the simplest update is θ ← θ − η·g (plain gradient descent, learning rate η). Language models almost always use Adam, which keeps two running averages per parameter and scales each step by them:

Dividing by √v̂ gives every parameter its own step size. It shows up clearly on the very first step: our check applied gradients of 0.5 and −1.5 to two parameters with η = 0.1, and both moved by exactly 0.1, because m̂/√v̂ reduces to the sign of the gradient when there's only one step of history. Our training script uses AdamW, which applies weight decay directly to the weights instead of mixing it into the gradient — the standard choice for transformers. Adam also stores m and v for every parameter, which is why optimizer state takes more memory than the model itself (see Part 3).

Check every claim in this section

This script verifies each rule above with PyTorch's autograd: p − y, both linear-layer rules, the residual gradient, and one Adam step by hand versus torch.optim.Adam.

import torch
import torch.nn.functional as F

torch.manual_seed(0)

# 1. Softmax + cross-entropy: the gradient of the loss w.r.t. the logits is p - y
vocab = ["mat", "floor", "sofa", "bed", "dog", "roof"]
z = torch.tensor([3.2, 2.4, 2.1, 1.8, 0.9, 0.6], requires_grad=True)   # logits
target = torch.tensor(0)                                                # true next token: "mat"
loss = F.cross_entropy(z[None], target[None])
loss.backward()
p = z.softmax(-1).detach()
y = F.one_hot(target, len(vocab)).float()
print(f"loss = -log p(mat) = {loss.item():.3f}   (p(mat) = {p[0]:.3f})")
print("autograd dL/dz:", [round(g, 3) for g in z.grad.tolist()])
print("p - y        :", [round(g, 3) for g in (p - y).tolist()])

# 2. Linear layer Y = X W: dL/dX = dL/dY W^T and dL/dW = X^T dL/dY
X = torch.randn(5, 8, requires_grad=True)       # 5 tokens, d_in = 8
W = torch.randn(8, 3, requires_grad=True)       # d_out = 3
Y = X @ W
G = torch.randn_like(Y)                         # pretend upstream gradient dL/dY
Y.backward(G)
print("dL/dX == dL/dY @ W^T :", torch.allclose(X.grad, G @ W.T))
print("dL/dW == X^T @ dL/dY :", torch.allclose(W.grad, X.T @ G))

# 3. Residual y = x + f(x): the gradient reaching x is dL/dy PLUS the branch's contribution
x = torch.randn(4, requires_grad=True)
f = lambda t: 0.001 * torch.tanh(t)              # a branch that has barely learned anything
y = x + f(x)
y.sum().backward()
print("residual grad at x (~1.0 even though f is tiny):", [round(g, 4) for g in x.grad.tolist()])

# 4. One Adam step (bias-corrected), by hand vs torch.optim.Adam
theta = torch.tensor([1.0, -2.0], requires_grad=True)
opt = torch.optim.Adam([theta], lr=0.1, betas=(0.9, 0.999), eps=1e-8)
g = torch.tensor([0.5, -1.5])
theta.grad = g.clone()
m = 0.1 * g                      # m = b1*0 + (1-b1)*g
v = 0.001 * g ** 2               # v = b2*0 + (1-b2)*g^2
m_hat, v_hat = m / (1 - 0.9), v / (1 - 0.999)
by_hand = theta.detach() - 0.1 * m_hat / (v_hat.sqrt() + 1e-8)
opt.step()
print("Adam step by hand:", [round(t, 4) for t in by_hand.tolist()], " torch:", [round(t, 4) for t in theta.tolist()])

Output:

loss = -log p(mat) = 0.790   (p(mat) = 0.454)
autograd dL/dz: [-0.546, 0.204, 0.151, 0.112, 0.046, 0.034]
p - y        : [-0.546, 0.204, 0.151, 0.112, 0.046, 0.034]
dL/dX == dL/dY @ W^T : True
dL/dW == X^T @ dL/dY : True
residual grad at x (~1.0 even though f is tiny): [1.0, 1.0008, 1.0009, 1.0002]
Adam step by hand: [0.9, -1.9]  torch: [0.9, -1.9]

Build It in PyTorch

Everything above in one runnable script: a 2-block, 64-wide GPT with 103,296 parameters, a toy word-level vocabulary standing in for BPE, a shape trace at every stage, a 200-step training run on a 20-word corpus, and the generation loop with temperature and top-k. Install PyTorch (pip install torch), save as tiny_gpt.py, and run python tiny_gpt.py. It finishes in seconds on a laptop CPU.

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

torch.manual_seed(0)

# --- 1-3. Text -> tokens -> IDs (a toy word-level vocabulary instead of BPE) ---
vocab = ["<unk>", "The", "cat", "sat", "on", "the", "mat", "dog", "."]
stoi = {w: i for i, w in enumerate(vocab)}
prompt = "The cat sat on the"
ids = torch.tensor([[stoi.get(w, 0) for w in prompt.split()]])   # (B=1, T=5)

V, d_model, n_heads, d_ff, n_layers, T_max = len(vocab), 64, 4, 256, 2, 32
TRACE = True
def show(name, x):
    if TRACE:
        print(f"{name:<34} {tuple(x.shape)}")


class CausalSelfAttention(nn.Module):
    def __init__(self):
        super().__init__()
        self.h, self.d_k = n_heads, d_model // n_heads
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)

    def forward(self, x):
        B, T, _ = x.shape
        q, k, v = self.W_q(x), self.W_k(x), self.W_v(x)                  # (B, T, d_model) each
        show("  Q = X W_q", q)
        # split into h heads: (B, T, d_model) -> (B, h, T, d_k)
        q, k, v = (t.view(B, T, self.h, self.d_k).transpose(1, 2) for t in (q, k, v))
        show("  split heads", q)
        scores = q @ k.transpose(-2, -1) / math.sqrt(self.d_k)           # (B, h, T, T)
        show("  QK^T / sqrt(d_k)", scores)
        mask = torch.triu(torch.ones(T, T, dtype=torch.bool), diagonal=1)
        scores = scores.masked_fill(mask, float("-inf"))                  # hide future tokens
        weights = scores.softmax(dim=-1)                                  # each row sums to 1
        show("  softmax weights", weights)
        out = weights @ v                                                 # (B, h, T, d_k)
        show("  weights x V (per head)", out)
        out = out.transpose(1, 2).reshape(B, T, d_model)                  # concat heads
        show("  concat heads", out)
        return self.W_o(out)


class Block(nn.Module):
    def __init__(self):
        super().__init__()
        self.ln1, self.attn = nn.LayerNorm(d_model), CausalSelfAttention()
        self.ln2 = nn.LayerNorm(d_model)
        self.mlp = nn.Sequential(nn.Linear(d_model, d_ff), nn.GELU(), nn.Linear(d_ff, d_model))

    def forward(self, x):
        x = x + self.attn(self.ln1(x))     # pre-norm attention + residual
        x = x + self.mlp(self.ln2(x))      # pre-norm MLP + residual
        show("  block output", x)
        return x


class TinyGPT(nn.Module):
    def __init__(self):
        super().__init__()
        self.tok_emb = nn.Embedding(V, d_model)
        self.pos_emb = nn.Embedding(T_max, d_model)
        self.blocks = nn.ModuleList(Block() for _ in range(n_layers))
        self.ln_f = nn.LayerNorm(d_model)
        self.lm_head = nn.Linear(d_model, V, bias=False)

    def forward(self, ids):
        B, T = ids.shape
        show("token IDs", ids)
        x = self.tok_emb(ids) + self.pos_emb(torch.arange(T))            # (B, T, d_model)
        show("embeddings + positions", x)
        for i, block in enumerate(self.blocks):
            if TRACE: print(f"block {i}")
            x = block(x)
        logits = self.lm_head(self.ln_f(x))                              # (B, T, V)
        show("logits", logits)
        return logits


model = TinyGPT()
logits = model(ids)
probs = logits[0, -1].softmax(dim=-1)            # only the LAST position predicts the next token
print("next-token probabilities (untrained, so essentially random):")
for p, i in sorted(zip(probs.tolist(), range(V)), reverse=True)[:4]:
    print(f"  {vocab[i]!r:>8}: {p:.3f}")

# --- Train for a few seconds on a tiny corpus (next-token prediction) ---
TRACE = False
corpus = ("The cat sat on the mat . The dog sat on the mat . "
          "The cat sat on the mat . The dog sat on the cat .").split()
data = torch.tensor([[stoi[w] for w in corpus]])
opt = torch.optim.AdamW(model.parameters(), lr=3e-3)
for step in range(200):
    logits = model(data[:, :-1])                                         # predict token t+1 from tokens <= t
    loss = F.cross_entropy(logits.reshape(-1, V), data[:, 1:].reshape(-1))
    opt.zero_grad(); loss.backward(); opt.step()
    if step in (0, 199):
        print(f"step {step:>3}  loss {loss.item():.3f}")

probs = model(ids)[0, -1].softmax(dim=-1)
print("next-token probabilities after training:")
for p, i in sorted(zip(probs.tolist(), range(V)), reverse=True)[:4]:
    print(f"  {vocab[i]!r:>8}: {p:.3f}")

# --- Autoregressive generation loop ---
@torch.no_grad()
def generate(ids, steps, temperature=1.0, top_k=None):
    for _ in range(steps):
        logits = model(ids)[:, -1, :] / temperature                      # (B, V)
        if top_k is not None:
            kth = torch.topk(logits, top_k).values[:, -1, None]
            logits = logits.masked_fill(logits < kth, float("-inf"))
        next_id = torch.multinomial(logits.softmax(-1), num_samples=1)   # sample 1 token
        ids = torch.cat([ids, next_id], dim=1)                           # append and repeat
    return ids

out = generate(ids, steps=4, top_k=3)
print("generated:", " ".join(vocab[i] for i in out[0].tolist()))
print(f"parameters: {sum(p.numel() for p in model.parameters()):,}")

Output (PyTorch 2.14, CPU; the seed makes it reproducible, but exact numbers can differ on other versions or hardware):

token IDs                          (1, 5)
embeddings + positions             (1, 5, 64)
block 0
  Q = X W_q                        (1, 5, 64)
  split heads                      (1, 4, 5, 16)
  QK^T / sqrt(d_k)                 (1, 4, 5, 5)
  softmax weights                  (1, 4, 5, 5)
  weights x V (per head)           (1, 4, 5, 16)
  concat heads                     (1, 5, 64)
  block output                     (1, 5, 64)
block 1
  Q = X W_q                        (1, 5, 64)
  split heads                      (1, 4, 5, 16)
  QK^T / sqrt(d_k)                 (1, 4, 5, 5)
  softmax weights                  (1, 4, 5, 5)
  weights x V (per head)           (1, 4, 5, 16)
  concat heads                     (1, 5, 64)
  block output                     (1, 5, 64)
logits                             (1, 5, 9)
next-token probabilities (untrained, so essentially random):
     'cat': 0.256
     'mat': 0.175
     'sat': 0.127
     'The': 0.090
step   0  loss 2.511
step 199  loss 0.001
next-token probabilities after training:
     'mat': 0.999
     'cat': 0.000
      'on': 0.000
   '<unk>': 0.000
generated: The cat sat on the mat . The dog
parameters: 103,296

Three things to notice in that output. The shape trace matches the diagram exactly: (1, 5, 64) in, split into (1, 4, 5, 16) heads, a (1, 4, 5, 5) attention grid, back to (1, 5, 64), and (1, 5, 9) logits over the 9-word vocabulary. Before training, the next-token guess is essentially random. After 200 steps the loss falls from 2.51 to 0.001 and the model puts 0.999 on "mat". Be clear about what that means: with 103K parameters and a 20-word corpus, the model has memorized the text rather than learned English. The mechanism is identical to a frontier model's; the difference is scale and data.

Shape Cheat Sheet

Add a leading batch dimension B to every row for batched inputs, as in the script's (1, 5, 64). The (T, T) attention grid is why attention's cost grows with the square of the context length — the reason long-context models need tricks like FlashAttention and sparse or linear attention.

⚠️ Notes

🔗 References