The Transformer Architecture Explained Step by Step

Every modern LLM, from GPT-style chatbots to Claude, Llama, and Gemini, runs on one core design: the Transformer. Introduced in the 2017 paper “Attention Is All You Need,” it replaced recurrence with self-attention, which lets a model process all tokens in a sequence in parallel. This guide walks through each component in order, with the math, the tensor shapes, and working PyTorch code.

Quick Takeaways

• Self-attention computes how much each token should “look at” every other token using Query, Key, and Value projections.
• Multi-head attention runs several attention operations in parallel (8 heads in the original model) so the model captures different relationships at once.
• Positional encoding injects word order, because attention alone is permutation-invariant.
• Modern LLMs are mostly decoder-only stacks with a causal mask. The original paper used an encoder-decoder design.

Step 1: Tokenization and Embeddings

A Transformer never sees raw text. A tokenizer (usually Byte-Pair Encoding or a SentencePiece variant) splits text into subword tokens and maps each to an integer ID.

An embedding matrix of shape [vocab_size, d_model] converts each ID into a dense vector. For a batch of B sequences with T tokens each, the input tensor becomes [B, T, d_model].

The original paper multiplies embeddings by √d_model to keep their scale compatible with the positional signal added next. Production models use vocabularies from roughly 32K to over 200K tokens, and embedding dimensions from 768 to well over 8,192 in frontier-scale models.

Step 2: Positional Encoding

Self-attention treats a sequence as a set. Without position information, “dog bites man” and “man bites dog” look identical. The original paper adds fixed sinusoidal encodings to the embeddings:

PE(pos, 2i)   = sin(pos / 10000^(2i / d_model))
PE(pos, 2i+1) = cos(pos / 10000^(2i / d_model))

Each dimension uses a different wavelength, so every position gets a unique signature, and relative offsets are easy to learn.

Modern models usually use other schemes:

Step 3: Scaled Dot-Product Self-Attention

This is the core mechanism. Each token vector x is projected by three learned matrices into:

• Query (Q): what this token is looking for.
• Key (K): what this token offers to others.
• Value (V): the information that gets passed along.

The attention formula:

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

Walk through it in four moves:

1. Score: Compute Q · Kᵀ, giving a [T, T] matrix of token-to-token similarity.
2. Scale: Divide by √d_k (8 when d_k = 64). Without this, large dot products push softmax into regions with vanishing gradients.
3. Normalize: Apply softmax per row so weights sum to 1.
4. Mix: Multiply by V to produce a weighted blend of value vectors.

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

def scaled_dot_product_attention(q, k, v, mask=None):
    # q, k, v: [batch, heads, seq_len, d_k]
    d_k = q.size(-1)
    scores = q @ k.transpose(-2, -1) / math.sqrt(d_k)   # [B, h, T, T]
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float("-inf"))
    weights = F.softmax(scores, dim=-1)
    return weights @ v, weights                          # [B, h, T, d_k]

Attention cost scales as O(T² · d) in compute and O(T²) in memory for the score matrix. That quadratic term is why long context windows are expensive, and why kernels like FlashAttention exist.

Step 4: Multi-Head Attention

One attention operation can only learn one weighting pattern per layer. Multi-head attention splits d_model into h smaller subspaces and runs attention in each one.

With d_model = 512 and h = 8, each head works in 64 dimensions. The outputs are concatenated and passed through a final linear projection W_O.

import torch.nn as nn

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model=512, n_heads=8):
        super().__init__()
        assert d_model % n_heads == 0
        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, mask=None):
        B, T, _ = x.shape
        split = lambda t: t.view(B, T, self.h, self.d_k).transpose(1, 2)
        q, k, v = split(self.w_q(x)), split(self.w_k(x)), split(self.w_v(x))
        out, _ = scaled_dot_product_attention(q, k, v, mask)
        out = out.transpose(1, 2).contiguous().view(B, T, -1)  # concat heads
        return self.w_o(out)

Attention Variants in Modern LLMs

Inference memory is dominated by the KV cache, so newer models share keys and values across heads.

Step 5: Masking (Padding and Causal)

Two mask types matter:

• Padding mask: hides <pad> tokens so batches of unequal length work.
• Causal (look-ahead) mask: a lower-triangular matrix that blocks token t from attending to tokens after t. This enforces autoregressive generation.

def causal_mask(T):
    return torch.tril(torch.ones(T, T)).unsqueeze(0).unsqueeze(0)  # [1,1,T,T]

Step 6: Position-wise Feed-Forward Network

After attention mixes information across tokens, a two-layer MLP processes each token independently:

FFN(x) = max(0, x·W₁ + b₁)·W₂ + b₂

In the base model, W₁ expands 512 → 2048 and W₂ projects back to 512. Most of a Transformer’s parameters sit in these FFN layers. Modern LLMs swap ReLU for SwiGLU or GELU variants and often set the hidden size to roughly 8/3 × d_model to compensate for the extra gate matrix.

Step 7: Residual Connections and Layer Normalization

Each sublayer (attention and FFN) is wrapped in a residual connection and normalization. These make deep stacks trainable.

• Post-LN (original): LayerNorm(x + Sublayer(x))
• Pre-LN (modern default): x + Sublayer(LayerNorm(x)). It trains more stably at depth and often removes the need for a long warmup.
• RMSNorm: a cheaper normalization (no mean subtraction) used in Llama-class models.

class TransformerBlock(nn.Module):
    def __init__(self, d_model=512, n_heads=8, d_ff=2048, p=0.1):
        super().__init__()
        self.attn = MultiHeadAttention(d_model, n_heads)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff), nn.GELU(), nn.Linear(d_ff, d_model)
        )
        self.ln1, self.ln2 = nn.LayerNorm(d_model), nn.LayerNorm(d_model)
        self.drop = nn.Dropout(p)

    def forward(self, x, mask=None):
        x = x + self.drop(self.attn(self.ln1(x), mask))   # Pre-LN attention
        x = x + self.drop(self.ffn(self.ln2(x)))          # Pre-LN FFN
        return x

Step 8: Stacking the Layers

Stack N blocks. The output of each block feeds the next, and the residual stream keeps a running representation of every token. Depth scales with model size:

Step 9: The Output Head and Decoding

The final hidden states pass through a linear layer of shape [d_model, vocab_size], producing logits. Softmax turns the last position’s logits into a probability distribution over the next token.

Decoding settings shape the output:

• Temperature: 0.0–0.3 for code and factual tasks, 0.7–1.0 for creative writing.
• Top-p (nucleus): commonly 0.9–0.95.
• Top-k: commonly 40–50.

Generation is autoregressive: sample a token, append it, and repeat. A KV cache stores past keys and values so each new token costs one forward step instead of recomputing the whole sequence.

Architecture Families: Encoder, Decoder, or Both

In encoder-decoder models, the decoder adds a cross-attention sublayer. Its queries come from the decoder, while its keys and values come from the encoder output.

Transformers vs. Alternatives

Practical Workflows and Real-World Use Cases

Build a Minimal GPT in Under 30 Lines

Combine the pieces above into a decoder-only model:

class MiniGPT(nn.Module):
    def __init__(self, vocab=32000, d_model=512, n_layers=6, n_heads=8, max_len=1024):
        super().__init__()
        self.tok = nn.Embedding(vocab, d_model)
        self.pos = nn.Embedding(max_len, d_model)
        self.blocks = nn.ModuleList(
            [TransformerBlock(d_model, n_heads) for _ in range(n_layers)]
        )
        self.ln_f = nn.LayerNorm(d_model)
        self.head = nn.Linear(d_model, vocab, bias=False)

    def forward(self, ids):
        B, T = ids.shape
        x = self.tok(ids) + self.pos(torch.arange(T, device=ids.device))
        mask = causal_mask(T).to(ids.device)
        for blk in self.blocks:
            x = blk(x, mask)
        return self.head(self.ln_f(x))   # [B, T, vocab] logits

Train it with cross-entropy loss on next-token prediction, AdamW (lr ≈ 3e-4, betas = (0.9, 0.95)), a warmup schedule, and gradient clipping at 1.0.

Where Teams Apply Each Design

• Semantic search and RAG: encoder models produce embeddings; a decoder model generates answers grounded in retrieved passages, which reduces hallucinations.
• Code assistants: decoder-only models with long context windows read whole repositories.
• Translation and summarization pipelines: encoder-decoder models remain strong where input and output are clearly separate.

Performance Levers for Engineers

PyTorch ships an optimized fused kernel, so in production replace the manual attention function with F.scaled_dot_product_attention(q, k, v, is_causal=True).

FAQ

What is the main idea behind the Transformer architecture?

The Transformer uses self-attention to let every token weigh its relationship to every other token in a single parallel step. It removes the sequential bottleneck of RNNs and trains efficiently on GPUs.

Why is attention divided by √d_k?

Dot products grow in magnitude with dimension size. Dividing by √d_k keeps scores in a stable range, so softmax does not saturate and gradients stay useful.

What is the difference between encoder and decoder in a Transformer?

An encoder reads the full input with bidirectional attention and builds contextual representations. A decoder generates output one token at a time using a causal mask, and in encoder-decoder models it also attends to the encoder through cross-attention.

Why do Transformers struggle with very long inputs?

Standard self-attention scales quadratically with sequence length, and the KV cache grows linearly with context during generation. Techniques like FlashAttention, GQA, sliding-window attention, and hybrid SSM layers reduce this cost.