Decoder-Only Transformers

Purpose

This note covers the decoder-only transformer as the standard autoregressive language-model architecture. The serving note Transformer Architecture and Implementation explains the same model from a systems angle. This note stays on the modeling side: what the block computes, how GPT-style models differ from the original transformer, and why causal self-attention became the default sequence primitive.

From the Original Transformer to GPT-Style Models

The original transformer paper is encoder-decoder. Decoder-only models keep the masked self-attention stack and language-model objective, then drop the encoder and cross-attention path.

The probability factorization is:

p (x_{1 : T}) = t = 1 \prod T p (x_{t} ∣ x_{< t})

That matches next-token prediction exactly.

Scaled Dot-Product Self-Attention

Given hidden states $X \in \mathbb{R}^{T \times d_{model}}$ :

Q = X W_{Q}, K = X W_{K}, V = X W_{V}

and

Attn (Q, K, V) = softmax (\frac{Q K ^{⊤}}{d _{k}} + M) V

where $M$ is the causal mask. Future positions get $-\infty$ before softmax, so position $t$ can only attend to $1,\dots,t$ .

Causal masking is what makes teacher forcing valid

Without the mask, position $t$ could attend to position $t+1$ and simply copy the token it is being trained to predict. The mask is not an optimization detail; it is the mechanism that lets every position train in parallel while still matching the autoregressive factorization $p(x_t \mid x_{<t})$ used at generation time.

The $1/\sqrt{d_k}$ scale from Vaswani et al. is there because raw dot products grow with dimension and can otherwise saturate the softmax.

Multi-Head Attention

Instead of one attention map, split the model into $h$ heads:

MHA (X) = Concat (head_{1}, \dots, head_{h}) W_{O}

with

head_{i} = Attn (X W_{Q}^{(i)}, X W_{K}^{(i)}, X W_{V}^{(i)})

In the original transformer base model:

$N = 6$ layers
$d_{model} = 512$
$d_{ff} = 2048$
$h = 8$ heads

Those numbers matter because they became the template many later variants scaled from.

Decoder Block

A modern GPT-style block is usually pre-norm:

H_{1} H_{2} = X + MHA (LN (X)) = H_{1} + MLP (LN (H_{1}))

The MLP is typically

MLP (x) = W_{2} ϕ (W_{1} x + b_{1}) + b_{2}

often with expansion ratio around $4\times$ .

The original transformer used post-norm. GPT-2 moves layer normalization to the input of each sub-block, which the paper explicitly notes, making it resemble a pre-activation residual network.

The full model is this block repeated $N$ times between an embedding and an unembedding:

flowchart TB
    tok["Token IDs"] --> emb["Token embedding + positional information"]
    emb --> ln1
    subgraph block["Decoder block, repeated N times"]
        ln1["LayerNorm"] --> attn["Causal self-attention"]
        attn --> add1["Add"]
        add1 --> ln2["LayerNorm"]
        ln2 --> mlp["MLP, ~4x expansion"]
        mlp --> add2["Add"]
    end
    emb -. "residual" .-> add1
    add1 -. "residual" .-> add2
    add2 --> lnf["Final LayerNorm"]
    lnf --> head["Unembedding to vocabulary logits"]

The dotted residual edges are the identity paths. They carry both activations forward and gradients backward untouched, which is what makes 48-layer stacks like GPT-2 1.5B trainable.

Positional Information

Self-attention alone is permutation-equivariant. Order has to be injected.

The original transformer uses sinusoidal positional encodings:

PE (p os, 2 i) = sin (p os /1000 0^{2 i / d_{m o d e l}})

PE (p os, 2 i + 1) = cos (p os /1000 0^{2 i / d_{m o d e l}})

GPT-style models often replace this with learned or rotary position schemes, but the need is the same.

Training Objective

For a tokenized sequence, maximize

t = 1 \sum T lo g p (x_{t} ∣ x_{< t})

Teacher forcing makes this parallel over positions during training. Every position predicts the next token simultaneously, with the mask enforcing causality.

GPT-2 Architectural Details

Radford et al. train four decoder-only transformers with the following sizes:

Parameters	Layers	$d_{model}$
117M	12	768
345M	24	1024
762M	36	1280
1542M	48	1600

The paper also calls out several design changes relative to GPT:

layer normalization moved to the input of each sub-block
an extra layer norm after the final self-attention block
vocabulary expanded to 50,257
context length increased from 512 to 1024
residual-layer weights scaled at initialization by $1/\sqrt{N}$ where $N$ is the number of residual layers

These are the kinds of details that matter. A decoder-only transformer is not just “masked attention plus an MLP.” It is a very specific residual stack whose optimization behavior depends on these choices.

Why Decoder-Only Won for Language Modeling

The paper case for transformers already existed in Vaswani et al.: self-attention connects any two positions with constant sequential path length and parallelizes training. GPT-2 adds the empirical argument that simply scaling this setup on WebText produces strong zero-shot transfer.

The GPT-2 paper reports that the 1.5B parameter model achieves state-of-the-art results on 7 of 8 evaluated language-modeling datasets in a zero-shot setting while still underfitting WebText.

That is the basic scaling story of modern language models.

NumPy Self-Attention

import numpy as np
 
def causal_mask(T):
    mask = np.full((T, T), -1e9, dtype=np.float32)
    return np.triu(mask, k=1)
 
def softmax(x, axis=-1):
    x = x - x.max(axis=axis, keepdims=True)
    e = np.exp(x)
    return e / e.sum(axis=axis, keepdims=True)
 
def self_attention(x, Wq, Wk, Wv):
    Q = x @ Wq
    K = x @ Wk
    V = x @ Wv
    scores = Q @ np.swapaxes(K, -1, -2) / np.sqrt(Q.shape[-1])
    scores = scores + causal_mask(x.shape[1])[None, :, :]
    probs = softmax(scores, axis=-1)
    return probs @ V

PyTorch Block with `einops`

import torch
import torch.nn as nn
from einops import rearrange
 
class CausalSelfAttention(nn.Module):
    def __init__(self, d_model: int, n_heads: int):
        super().__init__()
        assert d_model % n_heads == 0
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.out = nn.Linear(d_model, d_model, bias=False)
 
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, D = x.shape
        qkv = self.qkv(x)
        q, k, v = qkv.chunk(3, dim=-1)
 
        q = rearrange(q, "b t (h d) -> b h t d", h=self.n_heads)
        k = rearrange(k, "b t (h d) -> b h t d", h=self.n_heads)
        v = rearrange(v, "b t (h d) -> b h t d", h=self.n_heads)
 
        scores = torch.matmul(q, k.transpose(-1, -2)) / (self.head_dim ** 0.5)
        mask = torch.triu(torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1)
        scores = scores.masked_fill(mask, float("-inf"))
        attn = scores.softmax(dim=-1)
        y = torch.matmul(attn, v)
        y = rearrange(y, "b h t d -> b t (h d)")
        return self.out(y)
 
class DecoderBlock(nn.Module):
    def __init__(self, d_model: int, n_heads: int, mlp_mult: int = 4):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = CausalSelfAttention(d_model, n_heads)
        self.ln2 = nn.LayerNorm(d_model)
        self.mlp = nn.Sequential(
            nn.Linear(d_model, mlp_mult * d_model),
            nn.GELU(),
            nn.Linear(mlp_mult * d_model, d_model),
        )
 
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.attn(self.ln1(x))
        x = x + self.mlp(self.ln2(x))
        return x

Executable Experiment

The companion notebook trains a small character-level decoder on WikiText-2, plots optimization progress, samples text, and visualizes learned attention maps. Its scale is small enough for a laptop while preserving the causal mask, residual blocks, and next-token objective used by larger models.

decoder-only-transformer-on-wikitext-2.ipynb

Goal

Build a tiny decoder-only Transformer on WikiText-2. We will use a cropped character-level corpus so the model still shows real attention behavior without turning the notebook into an hours-long training run.

Causal self-attention computes

Attn (Q, K, V) = softmax (\frac{Q K ^{⊤}}{d _{k}} + M) V,

where the mask (M) is (-\infty) above the diagonal so that token (t) cannot read from the future.

In [1]:

## %pip install -q datasets torch matplotlib einops
 
import math
import random
 
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.nn.functional as F
from datasets import load_dataset

In [2]:

ds = load_dataset("Salesforce/wikitext", "wikitext-2-raw-v1")
text = "\n".join(ds["train"]["text"])[:300000]
 
vocab = sorted(set(text))
stoi = {ch: i for i, ch in enumerate(vocab)}
itos = {i: ch for ch, i in stoi.items()}
data = torch.tensor([stoi[ch] for ch in text], dtype=torch.long)
 
split = int(0.9 * len(data))
train_ids = data[:split]
valid_ids = data[split:]
 
print(f"chars={len(data):,} vocab={len(vocab)}")

Out[2]:

chars=300,000 vocab=154

In [3]:

def get_batch(data: torch.Tensor, batch_size: int, block_size: int, device: str):
    starts = torch.randint(0, len(data) - block_size - 1, (batch_size,))
    x = torch.stack([data[s : s + block_size] for s in starts]).to(device)
    y = torch.stack([data[s + 1 : s + block_size + 1] for s in starts]).to(device)
    return x, y
 
 
class CausalSelfAttention(nn.Module):
    def __init__(self, d_model: int, n_heads: int, block_size: int):
        super().__init__()
        assert d_model % n_heads == 0
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads
        self.qkv = nn.Linear(d_model, 3 * d_model)
        self.proj = nn.Linear(d_model, d_model)
        mask = torch.tril(torch.ones(block_size, block_size))
        self.register_buffer("mask", mask.view(1, 1, block_size, block_size))
 
    def forward(self, x):
        B, T, C = x.shape
        qkv = self.qkv(x)
        q, k, v = qkv.chunk(3, dim=-1)
 
        q = q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
 
        scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        scores = scores.masked_fill(self.mask[:, :, :T, :T] == 0, float("-inf"))
        attn = scores.softmax(dim=-1)
        out = attn @ v
        out = out.transpose(1, 2).contiguous().view(B, T, C)
        return self.proj(out), attn
 
 
class Block(nn.Module):
    def __init__(self, d_model: int, n_heads: int, block_size: int, mlp_mult: int = 4):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = CausalSelfAttention(d_model, n_heads, block_size)
        self.ln2 = nn.LayerNorm(d_model)
        self.mlp = nn.Sequential(
            nn.Linear(d_model, mlp_mult * d_model),
            nn.GELU(),
            nn.Linear(mlp_mult * d_model, d_model),
        )
 
    def forward(self, x):
        attn_out, attn = self.attn(self.ln1(x))
        x = x + attn_out
        x = x + self.mlp(self.ln2(x))
        return x, attn
 
 
class TinyGPT(nn.Module):
    def __init__(self, vocab_size: int, block_size: int, d_model: int = 128, n_heads: int = 4, n_layers: int = 3):
        super().__init__()
        self.block_size = block_size
        self.token = nn.Embedding(vocab_size, d_model)
        self.pos = nn.Embedding(block_size, d_model)
        self.blocks = nn.ModuleList([Block(d_model, n_heads, block_size) for _ in range(n_layers)])
        self.ln_f = nn.LayerNorm(d_model)
        self.head = nn.Linear(d_model, vocab_size)
 
    def forward(self, idx):
        B, T = idx.shape
        pos = torch.arange(T, device=idx.device)
        x = self.token(idx) + self.pos(pos)[None, :, :]
        attn_maps = []
        for block in self.blocks:
            x, attn = block(x)
            attn_maps.append(attn)
        logits = self.head(self.ln_f(x))
        return logits, attn_maps
 
    @torch.no_grad()
    def generate(self, idx, max_new_tokens: int):
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -self.block_size :]
            logits, _ = self(idx_cond)
            probs = logits[:, -1].softmax(dim=-1)
            next_idx = torch.multinomial(probs, num_samples=1)
            idx = torch.cat([idx, next_idx], dim=1)
        return idx

In [4]:

device = "cuda" if torch.cuda.is_available() else "cpu"
block_size = 96
model = TinyGPT(len(vocab), block_size=block_size).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-4, weight_decay=1e-2)
 
losses = []
for step in range(400):
    xb, yb = get_batch(train_ids, batch_size=32, block_size=block_size, device=device)
    logits, _ = model(xb)
    loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), yb.reshape(-1))
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    losses.append(float(loss))
 
    if step % 100 == 0:
        print(f"step={step:04d} loss={float(loss):.4f}")

Out[4]:

step=0000 loss=5.2090

step=0100 loss=2.8399

step=0200 loss=2.6231

step=0300 loss=2.5646

In [5]:

plt.figure(figsize=(8, 4))
plt.plot(losses)
plt.title("Decoder-only Transformer training loss")
plt.xlabel("step")
plt.ylabel("cross entropy")
plt.show()

Out[5]:

<Figure size 800x400 with 1 Axes>

In [6]:

prompt = torch.tensor([[stoi["T"]]], dtype=torch.long, device=device)
sample = model.generate(prompt, max_new_tokens=600)[0].tolist()
print("".join(itos[i] for i in sample))

Out[6]:

The ceg frsthegou@ f pof nd presthan ) aseinxbus aZe r ttonilァ pey tonchewh Sope d fd GNs . thacond waleter axilanghietl Secor s きathest pbore pat lerle d ds achifinghat . pone . inia , sinee ped tho bithessin ebek on tan pHiove emsese th 100 ć 11 s ; omof s ced t Gathean Rhe his~o A c/ ben Thaid " atod Hed This iuro saleciechythe therand pesud aana ermug . obes acaste rens . Ked ess Crme ind tam rofo Hawens ds , aKinof . th acamed d . , y The trand " cndzeate tthe , ind , jt , thangfons . @ 'se Jentas ckeare 2 . . , the wro Hede tang opan ses g 
 cere 'sin Hedal Hengrllefolapo. ous layechiov0

In [7]:

xb, _ = get_batch(valid_ids, batch_size=1, block_size=64, device=device)
_, attn_maps = model(xb)
head0 = attn_maps[0][0, 0].detach().cpu()
 
plt.figure(figsize=(6, 5))
plt.imshow(head0, cmap="magma")
plt.title("Layer 0, head 0 attention map")
plt.xlabel("key position")
plt.ylabel("query position")
plt.colorbar()
plt.show()

Out[7]:

<Figure size 600x500 with 2 Axes>

Why This Notebook Matters

A character-level model is toy-sized, but the mechanics are real. The same masking rule, residual shape discipline, and next-token objective scale up to practical decoder-only language models.

Complexity

For sequence length $T$ and width $d$ , full self-attention is

O (T^{2} d)

in both compute and attention-score storage. That quadratic cost is the main reason long-context variants and inference-side memory work keep recurring.

Elijah's Notes

Explorer

Decoder-Only Transformers

Purpose

From the Original Transformer to GPT-Style Models

Scaled Dot-Product Self-Attention

Multi-Head Attention

Decoder Block

Positional Information

Training Objective

GPT-2 Architectural Details

Why Decoder-Only Won for Language Modeling

NumPy Self-Attention

PyTorch Block with `einops`

Executable Experiment

Goal

Why This Notebook Matters

Complexity

Sources

Graph View

Table of Contents

Elijah's Notes

Explorer

Decoder-Only Transformers

Purpose

From the Original Transformer to GPT-Style Models

Scaled Dot-Product Self-Attention

Multi-Head Attention

Decoder Block

Positional Information

Training Objective

GPT-2 Architectural Details

Why Decoder-Only Won for Language Modeling

NumPy Self-Attention

PyTorch Block with einops

Executable Experiment

Goal

Why This Notebook Matters

Complexity

Related Notes

Sources

Graph View

Table of Contents

PyTorch Block with `einops`