Transformer architecture explained, layer by layer

If you’re building on LLMs, you need the transformer architecture explained in terms of what actually happens to tensors at each step. This guide walks through the forward pass from embedding to logits, with minimal PyTorch implementations that mirror what runs in production. You’ll see where compute concentrates, where memory blows up, and which knobs matter when you’re optimizing inference or debugging training instability.

Token and position embeddings

The input sequence arrives as integer token IDs of shape (batch, seq_len). The first operation maps each ID to a learned vector of dimension d_model. Position information gets added because self-attention is permutation-invariant — without it, “dog bites man” and “man bites dog” produce identical representations.

import torch
import torch.nn as nn

class Embeddings(nn.Module):
    def __init__(self, vocab_size: int, d_model: int, max_seq_len: int, dropout: float = 0.1):
        super().__init__()
        self.token_emb = nn.Embedding(vocab_size, d_model)
        self.pos_emb = nn.Embedding(max_seq_len, d_model)
        self.dropout = nn.Dropout(dropout)
        self.register_buffer("positions", torch.arange(max_seq_len).unsqueeze(0))

    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        # input_ids: (batch, seq_len)
        seq_len = input_ids.size(1)
        token_vectors = self.token_emb(input_ids)          # (batch, seq_len, d_model)
        pos_vectors = self.pos_emb(self.positions[:, :seq_len])  # (1, seq_len, d_model)
        return self.dropout(token_vectors + pos_vectors)

Pitfall: Learned positional embeddings cap maximum sequence length at max_seq_len. RoPE (rotary position embeddings) and ALiBi avoid this by computing positions analytically — standard in modern LLMs. If you’re extending context at inference, learned embeddings require interpolation or retraining.

Tradeoff: d_model controls capacity but scales everything quadratically in attention. Typical values: 768 (BERT-base), 4096 (Llama-2-7B), 8192 (Llama-3-70B).

Multi-head self-attention

This is where the transformer architecture explained most tutorials get wrong: attention is not a single operation. It’s a composition of projections, scaled dot-products, and an output projection — all batched across heads.

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int, n_heads: int, dropout: float = 0.1, bias: bool = False):
        super().__init__()
        assert d_model % n_heads == 0
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_head = d_model // n_heads

        # Fused QKV projection is standard for memory bandwidth
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=bias)
        self.out_proj = nn.Linear(d_model, d_model, bias=bias)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor, attn_mask: torch.Tensor | None = None) -> torch.Tensor:
        batch, seq_len, _ = x.shape

        # Fused projection: (batch, seq_len, 3 * d_model)
        qkv = self.qkv(x)
        qkv = qkv.view(batch, seq_len, self.n_heads, 3 * self.d_head)
        qkv = qkv.transpose(1, 2)  # (batch, n_heads, seq_len, 3 * d_head)
        q, k, v = qkv.chunk(3, dim=-1)  # each: (batch, n_heads, seq_len, d_head)

        # Scaled dot-product attention
        # scores: (batch, n_heads, seq_len, seq_len)
        scores = torch.matmul(q, k.transpose(-2, -1)) / (self.d_head ** 0.5)

        if attn_mask is not None:
            # attn_mask: (batch, 1, 1, seq_len) or (batch, 1, seq_len, seq_len)
            scores = scores.masked_fill(attn_mask == 0, float("-inf"))

        attn_weights = torch.softmax(scores, dim=-1)
        attn_weights = self.dropout(attn_weights)

        # Weighted sum: (batch, n_heads, seq_len, d_head)
        context = torch.matmul(attn_weights, v)

        # Merge heads: (batch, seq_len, d_model)
        context = context.transpose(1, 2).contiguous().view(batch, seq_len, self.d_model)
        return self.out_proj(context)

Where compute lives: The q @ k.T matmul is O(batch * n_heads * seq_len^2 * d_head). For seq_len=4096, d_model=4096, n_heads=32, that’s ~2.7B FLOPs per layer just for attention scores. This is why flash attention (kernel-fused, memory-efficient) matters — it avoids materializing the full (seq_len, seq_len) matrix in HBM.

Common pitfall: Forgetting contiguous() before view() after transpose(). The tensor becomes non-contiguous; view() fails or silently corrupts data if you use reshape() instead.

Masking: Causal mask (lower-triangular) for autoregressive decoding. Padding mask for variable-length batches. Combine them: causal_mask & padding_mask[:, None, None, :].

Feed-forward network (MLP)

Each position processes independently through a two-layer MLP with a non-linearity. This is where the model “thinks” — attention moves information, the MLP transforms it.

class FeedForward(nn.Module):
    def __init__(self, d_model: int, d_ff: int | None = None, dropout: float = 0.1, bias: bool = False):
        super().__init__()
        d_ff = d_ff or 4 * d_model  # Standard expansion factor
        self.gate_proj = nn.Linear(d_model, d_ff, bias=bias)
        self.up_proj = nn.Linear(d_model, d_ff, bias=bias)
        self.down_proj = nn.Linear(d_ff, d_model, bias=bias)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # SwiGLU: gate * silu(up)  — standard in Llama, Mistral, Gemma
        gate = self.gate_proj(x)
        up = self.up_proj(x)
        return self.down_proj(self.dropout(torch.nn.functional.silu(gate) * up))

Why SwiGLU? Gated linear units (GLU variants) consistently outperform ReLU/GELU at equal parameter count. The gate learns which features to amplify per token. Expansion factor 4x is standard; some models (Phi-3) use 3.5x, others (Nemotron) push to 8x for quality.

Memory note: The intermediate d_ff activation is batch * seq_len * d_ff * 2 bytes (bf16). At seq_len=8192, d_model=8192, d_ff=32768, that’s ~4 GB per layer just for the MLP activations. Gradient checkpointing recomputes this on backward pass to trade compute for memory.

Layer normalization and residual connections

Pre-norm (LayerNorm before sublayer) is now universal. Post-norm (original Transformer paper) causes gradient instability at depth.

class TransformerBlock(nn.Module):
    def __init__(self, d_model: int, n_heads: int, d_ff: int | None = None, dropout: float = 0.1):
        super().__init__()
        self.attn_norm = nn.LayerNorm(d_model, eps=1e-5)
        self.attn = MultiHeadAttention(d_model, n_heads, dropout)
        self.ffn_norm = nn.LayerNorm(d_model, eps=1e-5)
        self.ffn = FeedForward(d_model, d_ff, dropout)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor, attn_mask: torch.Tensor | None = None) -> torch.Tensor:
        # Pre-norm attention with residual
        h = self.attn_norm(x)
        h = self.attn(h, attn_mask)
        x = x + self.dropout(h)

        # Pre-norm FFN with residual
        h = self.ffn_norm(x)
        h = self.ffn(h)
        x = x + self.dropout(h)
        return x

Why pre-norm? The residual path x + sublayer(norm(x)) keeps gradient norm stable. With post-norm norm(x + sublayer(x)), the sublayer output must stay small or gradients explode — requiring careful initialization and learning rate warmup.

RMSNorm variant: Llama and derivatives use RMSNorm (no mean centering, no bias). Slightly faster, marginally better stability. Drop-in replacement: nn.RMSNorm in PyTorch 2.4+.

The full stack: stacking blocks

class Transformer(nn.Module):
    def __init__(
        self,
        vocab_size: int,
        d_model: int,
        n_layers: int,
        n_heads: int,
        max_seq_len: int,
        d_ff: int | None = None,
        dropout: float = 0.1,
        tie_weights: bool = True,
    ):
        super().__init__()
        self.embeddings = Embeddings(vocab_size, d_model, max_seq_len, dropout)
        self.blocks = nn.ModuleList([
            TransformerBlock(d_model, n_heads, d_ff, dropout)
            for _ in range(n_layers)
        ])
        self.final_norm = nn.LayerNorm(d_model, eps=1e-5)
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)

        if tie_weights:
            # Share input/output embeddings — saves ~vocab_size * d_model params
            self.lm_head.weight = self.embeddings.token_emb.weight

    def forward(self, input_ids: torch.Tensor, attn_mask: torch.Tensor | None = None) -> torch.Tensor:
        x = self.embeddings(input_ids)
        for block in self.blocks:
            x = block(x, attn_mask)
        x = self.final_norm(x)
        logits = self.lm_head(x)  # (batch, seq_len, vocab_size)
        return logits

Weight tying: The tie_weights flag shares the token embedding matrix with the output projection. This cuts ~150M params for a 50k vocab / 4096 d_model model. It also acts as a regularizer — the model learns consistent representations for input and output.

Final norm: Critical for training stability. Without it, the last block’s residual stream can drift in scale, making the logits poorly calibrated.

KV cache for autoregressive decoding

During generation, you recompute attention over the full prefix at every step. KV cache stores K and V for past tokens so each step only computes attention for the new token.

class KVCache:
    def __init__(self, n_layers: int, batch: int, n_heads: int, d_head: int, max_seq_len: int, device: torch.device, dtype: torch.dtype):
        self.k_cache = torch.zeros(n_layers, batch, n_heads, max_seq_len, d_head, device=device, dtype=dtype)
        self.v_cache = torch.zeros(n_layers, batch, n_heads, max_seq_len, d_head, device=device, dtype=dtype)
        self.seq_len = 0

    def update(self, layer_idx: int, k: torch.Tensor, v: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        # k, v: (batch, n_heads, 1, d_head) for single-step decode
        seq_start = self.seq_len
        self.k_cache[layer_idx, :, :, seq_start:seq_start+1] = k
        self.v_cache[layer_idx, :, :, seq_start:seq_start+1] = v
        self.seq_len += 1
        return self.k_cache[layer_idx, :, :, :self.seq_len], self.v_cache[layer_idx, :, :, :self.seq_len]

Integration: Modify MultiHeadAttention.forward to accept optional kv_cache and layer_idx. During prefill, populate cache. During decode, read full K, V from cache, compute Q only for new token.

Memory math: For n_layers=32, batch=1, n_heads=32, d_head=128, max_seq_len=8192, bf16: 2 * 32 * 1 * 32 * 8192 * 128 * 2 bytes ≈ 1 GB. This is why multi-GPU tensor parallelism or paged attention (vLLM) matters for long contexts.

Common pitfalls checklist

Scaling laws and architectural choices

The transformer architecture explained here scales predictably. Kaplan et al. (2020) showed loss follows a power law in compute, parameters, and data — but the constants depend on architecture:

• Depth vs width: Deeper/narrower often beats shallower/wider at equal params (better gradient flow, more composition). Llama-2-7B: 32 layers × 4096. Mistral-7B: 32 layers × 4096. Same.
• Head dimension: d_head = 128 is a sweet spot. Smaller hurts quality; larger wastes compute on redundant heads.
• Expansion factor: 4x is standard. 8x helps quality but increases activation memory 2x.
• Normalization: RMSNorm > LayerNorm for stability at scale. QK-Norm (normalizing Q and K per head) helps very deep models (>80 layers).

Inference optimization priorities

If you’re serving this model, optimize in order:

1. Flash attention / SDPA — kernel-fused attention avoids HBM round-trips for QK^T. PyTorch 2.0+ F.scaled_dot_product_attention does this automatically on Hopper/Ampere.
2. KV cache + paged attention — vLLM’s block manager eliminates fragmentation. Essential for batch > 1.
3. Quantization — AWQ/GPTQ (4-bit weights, 16-bit activations) cuts memory 4x with <1% quality loss. bitsandbytes for 8-bit is easier but less compression.
4. Speculative decoding — Small draft model proposes tokens, large model verifies. 2-3x throughput gain for same quality.
5. Continuous batching — Don’t pad to max length. Append new requests as slots free up.

When to modify the architecture

You rarely need to invent new layers. But these modifications have proven value:

• Sliding window attention (Mistral): Local attention for most layers, global every N layers. Linear scaling with context.
• Grouped-query attention (Llama-2-70B, Mistral): n_kv_heads < n_heads. Reduces KV cache memory, minimal quality loss. n_kv_heads=8 for 32 heads is standard.
• Parallel attention + FFN (PaLM, GPT-NeoX): Compute attention and FFN simultaneously on split input, sum outputs. Saves one all-reduce in tensor parallelism.
• Mixture of experts (Mixtral, DeepSeekMoE): Sparse FFN with router. 8 experts, top-2 activation gives 4x compute efficiency at same params.

The transformer architecture explained through code reveals that most “magic” is careful composition of linear projections, normalization, and residuals. The hard problems — memory bandwidth, numerical stability, KV cache management — are systems problems, not math problems. If you’re building a gateway that routes across 240+ models, you’ll see every variant of these patterns in production. The code above compiles; the rest is engineering.