04 Foundations 10 min read 2,026 words

Transformer Architecture

This chapter provides a comprehensive view of the complete transformer architecture, bringing together the components from previous chapters into a unified understanding.

fundamentalstransformersattentioncore
01overview

Architecture Overview

A decoder-only transformer (the architecture used by GPT, Claude, Llama) consists of:

02processing

Input Processing

2.1

Token Embedding

Convert token IDs to dense vectors:

Pythonpython · 7 lines
1234567
class TokenEmbedding(nn.Module):
    def __init__(self, vocab_size, d_model):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
    
    def forward(self, token_ids):
        return self.embedding(token_ids)

Dimensions:

  • Input: [batch_size, seq_len] token IDs
  • Output: [batch_size, seq_len, d_model] embeddings
2.2

Position Information

Position is incorporated via one of:

1. Rotary Position Embedding (RoPE): Applied within attention, not added to embeddings:

Pythonpython · 6 lines
123456
def apply_rope(q, k, positions):
    # Rotate q and k vectors based on position
    freqs = compute_frequencies(positions)
    q_rotated = rotate_embeddings(q, freqs)
    k_rotated = rotate_embeddings(k, freqs)
    return q_rotated, k_rotated

2. Learned Position Embeddings: Added directly to token embeddings:

Pythonpython · 2 lines
12
position_embeddings = nn.Embedding(max_seq_len, d_model)
x = token_embeddings + position_embeddings(positions)

Modern models (Llama, Mistral, GPT-4) use RoPE for better length generalization.

03block

The Transformer Block

3.1

Pre-Norm Structure

Modern transformers use pre-normalization:

Pythonpython · 23 lines
1234567891011121314151617181920212223
class TransformerBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.attn_norm = RMSNorm(config.d_model)
        self.attn = GroupedQueryAttention(
            d_model=config.d_model,
            n_heads=config.n_heads,
            n_kv_heads=config.n_kv_heads
        )
        self.ff_norm = RMSNorm(config.d_model)
        self.ff = SwiGLUFFN(
            d_model=config.d_model,
            d_ff=config.d_ff
        )
    
    def forward(self, x, mask=None, kv_cache=None):
        # Attention with residual
        h = x + self.attn(self.attn_norm(x), mask, kv_cache)
        
        # FFN with residual
        out = h + self.ff(self.ff_norm(h))
        
        return out
3.2

Attention Component

Pythonpython · 39 lines
123456789101112131415161718192021222324252627282930313233343536373839
class GroupedQueryAttention(nn.Module):
    def __init__(self, d_model, n_heads, n_kv_heads):
        super().__init__()
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.head_dim = d_model // n_heads
        
        self.q_proj = nn.Linear(d_model, n_heads * self.head_dim)
        self.k_proj = nn.Linear(d_model, n_kv_heads * self.head_dim)
        self.v_proj = nn.Linear(d_model, n_kv_heads * self.head_dim)
        self.o_proj = nn.Linear(n_heads * self.head_dim, d_model)
    
    def forward(self, x, mask, kv_cache):
        B, T, D = x.shape
        
        # Project
        q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim)
        k = self.k_proj(x).view(B, T, self.n_kv_heads, self.head_dim)
        v = self.v_proj(x).view(B, T, self.n_kv_heads, self.head_dim)
        
        # Apply RoPE
        q, k = apply_rope(q, k, positions)
        
        # Update KV cache
        if kv_cache is not None:
            k = torch.cat([kv_cache.k, k], dim=1)
            v = torch.cat([kv_cache.v, v], dim=1)
            kv_cache.update(k, v)
        
        # Repeat KV heads for GQA
        k = k.repeat_interleave(self.n_heads // self.n_kv_heads, dim=2)
        v = v.repeat_interleave(self.n_heads // self.n_kv_heads, dim=2)
        
        # Attention (using Flash Attention in practice)
        attn_out = flash_attention(q, k, v, mask)
        
        # Output projection
        out = self.o_proj(attn_out.view(B, T, -1))
        return out
3.3

Feed-Forward Network

Pythonpython · 12 lines
123456789101112
class SwiGLUFFN(nn.Module):
    def __init__(self, d_model, d_ff):
        super().__init__()
        # SwiGLU has 3 projections instead of 2
        self.gate_proj = nn.Linear(d_model, d_ff, bias=False)
        self.up_proj = nn.Linear(d_model, d_ff, bias=False)
        self.down_proj = nn.Linear(d_ff, d_model, bias=False)
    
    def forward(self, x):
        gate = F.silu(self.gate_proj(x))  # SiLU = Swish
        up = self.up_proj(x)
        return self.down_proj(gate * up)

FFN hidden dimension is typically 2.7x the model dimension for SwiGLU (vs 4x for standard FFN with GELU).

3.4

RMSNorm

Pythonpython · 9 lines
123456789
class RMSNorm(nn.Module):
    def __init__(self, d_model, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(d_model))
        self.eps = eps
    
    def forward(self, x):
        rms = torch.sqrt(torch.mean(x ** 2, dim=-1, keepdim=True) + self.eps)
        return self.weight * (x / rms)

Simpler and faster than LayerNorm since it skips mean centering.

04processing

Output Processing

4.1

Final Normalization

Apply RMSNorm after the last transformer block:

Pythonpython · 1 line
1
hidden_states = self.output_norm(hidden_states)
4.2

Language Model Head

Project to vocabulary size:

Pythonpython · 7 lines
1234567
class LMHead(nn.Module):
    def __init__(self, d_model, vocab_size):
        super().__init__()
        self.linear = nn.Linear(d_model, vocab_size, bias=False)
    
    def forward(self, x):
        return self.linear(x)  # Returns logits
05tied embeddings

Untied vs. Tied Embeddings

Standard Pattern (GPT-3, Llama 2): Weight Tying

  • Output head shares weights with input embeddings.
  • Pro: Saves memory (vocab_size * hidden_dim).
  • Con: Forces input and output latent spaces to be identical, which can be suboptimal.

2025 Frontier Pattern (Llama 3/4, GPT-5.2): Untied Embeddings

  • Output head has its own weights.
  • Why?: Larger vocabularies (128k+) make the embedding table a significant portion of the model. Untying allows the output head to specialize in "predictive logic" while input embeddings focus on "semantic understanding."
  • System Impact: Increases parameter count but often improves perplexity for multilingual and code tasks.
5.1

Getting Predictions

Pythonpython · 7 lines
1234567
# During generation
logits = lm_head(hidden_states[:, -1, :])  # Last position only
next_token = sample(logits)

# During training
logits = lm_head(hidden_states)  # All positions
loss = cross_entropy(logits, targets)
06architecture variations

Modern Architecture Variations

6.1

Llama 2/3 Architecture

ComponentImplementation
AttentionGrouped Query Attention (GQA)
PositionRotary Position Embedding (RoPE)
NormalizationRMSNorm (pre-norm)
ActivationSwiGLU
BiasNo bias in linear layers
6.2

Mistral Architecture

Same as Llama but adds:

  • Sliding Window Attention: Each layer only attends to 4K tokens
  • Still achieves effective 32K+ context via stacking
6.3

Mixture of Experts (MoE) & Hybrid Architectures

State-of-the-art models in 2025 often use Hybrid MoE/Dense blocks:

  • Periodic Dense Layers: Every few MoE layers, a dense layer is added to ensure "global" knowledge is shared across all experts.
  • Expert Parallelism: Distributing different experts across different GPUs. This makes inter-node bandwidth (NVLink/InfiniBand) a primary architecture bottleneck.
6.4

Multi-head Latent Attention (MLA) Integration

The standard attention block in DeepSeek-V3 and equivalent 2025 architectures replaces the standard Q/K/V projects with low-rank latent compressions.

  • Architectural Shift: The "KV Cache" is now a compressed latent representation, changing the memory/compute ratio of the entire transformer block.
6.5

Comparison of Choices

ChoiceOld ApproachModern ApproachBenefit
NormPost-LNPre-LN / RMSNormTraining stability, speed
PositionSinusoidal/LearnedRoPEBetter extrapolation
ActivationGELUSwiGLUQuality (+1% on benchmarks)
AttentionMHAGQA8x smaller KV cache
BiasWith biasNo biasFewer parameters, similar quality
07properties

Scaling Properties

7.1

Parameter Counts

ComponentParameters
Token embeddingvocab_size * d_model
Per layer Q/K/V3 * d_model * d_model (for MHA)
Per layer O projd_model * d_model
Per layer FFN3 * d_model * d_ff (for SwiGLU)
LM headd_model * vocab_size (often tied)

Approximation for decoder-only:

Texttext · 1 line
1
Total ≈ 12 * n_layers * d_model^2 (for d_ff = 4 * d_model, MHA)
7.2

Compute Requirements

Training: FLOPs per token ≈ 6 * parameters (forward + backward)

Inference: FLOPs per token ≈ 2 * parameters (forward only)

7.3

Scaling Laws

The Chinchilla scaling law suggests optimal allocation:

Texttext · 1 line
1
D (data tokens) ≈ 20 * N (parameters)

For a 70B model, train on ~1.4T tokens for compute-optimal training.

But: Many modern models overtrain relative to Chinchilla for better inference efficiency. Llama was trained on 2T+ tokens.

08comparison table

Architecture Comparison Table

ModelParamsLayersd_modelHeadsKV HeadsFFNContext
GPT-3175B96122889696GELU2K
Llama 2 70B70B808192648SwiGLU4K
Llama 3 405B405B1261638412816SwiGLU128K
DeepSeek V3671B1287168128MLAMoE128K
Llama 4 (spec)1T+140+1843219224MoE/H1M+

*Mistral uses sliding window attention for effective long context.

09questions

Interview Questions

Q: Walk me through the forward pass of a transformer.

Strong answer: For a decoder-only model generating text:

  1. Tokenization: Convert input text to token IDs
  2. Embedding: Look up token embeddings from the embedding table

  3. For each transformer layer:

    • Apply RMSNorm to input
    • Compute Q, K, V projections
    • Apply RoPE to Q and K for position
    • For generation: append new K, V to KV cache
    • Compute attention (masked, so each position only sees previous)
    • Project attention output and add residual
    • Apply RMSNorm
    • Pass through SwiGLU feed-forward network
    • Add residual
  4. Output norm: Apply final RMSNorm
  5. LM head: Project to vocabulary size to get logits
  6. Sample: Select next token from logits using temperature/top-p

For generation, repeat steps 3-6 for each new token, reusing the KV cache from previous positions.

Q: What is the difference between pre-norm and post-norm?

Strong answer: The difference is where layer normalization is applied relative to sublayers (attention, FFN):

Post-norm (original transformer):

Texttext · 1 line
1
x = LayerNorm(x + Sublayer(x))

Normalize after adding residual.

Pre-norm (modern transformers):

Texttext · 1 line
1
x = x + Sublayer(LayerNorm(x))

Normalize before the sublayer.

Pre-norm is preferred because:

  1. Gradients flow more directly through residual connections
  2. Training is more stable, especially for deep models
  3. Less sensitive to initialization and learning rate
  4. No need for learning rate warmup

The cost is slightly lower final performance in some benchmarks, but the training stability is worth it for large models.

Q: Explain GQA and why it matters for serving.

Strong answer: Grouped Query Attention (GQA) shares Key and Value heads across groups of Query heads.

Standard Multi-Head Attention: 64 query heads, 64 KV heads (1:1) GQA: 64 query heads, 8 KV heads (8:1)

Implementation: Each KV head is used by 8 query heads via repetition.

Why it matters: The KV cache stores K and V for all positions during generation. For Llama 70B at 8K context:

  • MHA: 2.6 MB/token * 8K = 21 GB per request
  • GQA (8:1): ~2.6 GB per request

8x reduction enables:

  • Larger batch sizes (more concurrent users)
  • Longer contexts
  • Lower GPU memory requirements

Quality impact: Minimal. Research shows GQA achieves 99%+ of MHA quality.

Q: What changed between GPT-2 and Llama 2?

Strong answer: Key architecture improvements:

ComponentGPT-2Llama 2
NormPost-LayerNormPre-RMSNorm
PositionLearned absoluteRoPE (rotary)
ActivationGELUSwiGLU
AttentionMHAGQA (for 70B)
BiasPresentRemoved

Impact:

  • RMSNorm: Faster and equally effective
  • RoPE: Better length extrapolation
  • SwiGLU: ~1% quality improvement
  • GQA: 8x smaller KV cache for serving
  • No bias: Fewer parameters, no quality loss

These changes enable training larger models more stably and serving them more efficiently.

10references

References

  • Vaswani et al. "Attention Is All You Need" (2017)
  • Touvron et al. "Llama: Open and Efficient Foundation Language Models" (2023)
  • Touvron et al. "Llama 2: Open Foundation and Fine-Tuned Chat Models" (2023)
  • Zhang and Sennrich. "Root Mean Square Layer Normalization" (2019)
  • Shazeer. "GLU Variants Improve Transformer" (2020)
  • Su et al. "RoFormer: Enhanced Transformer with Rotary Position Embedding" (2021)
  • Jiang et al. "Mistral 7B" (2023)

Previous: Attention Mechanisms | Next: Embeddings and Vector Spaces

• summary · added by this rebuild

Key takeaways

01

The modern block is five settled choices

Pre-norm RMSNorm, RoPE, grouped-query attention, SwiGLU feed-forward and no bias in linear layers — the comparison table pairs each with the older approach it replaced.

02

SwiGLU costs a third projection

Gate, up and down projections replace two, so the FFN hidden width drops to roughly 2.7x d_model instead of 4x to keep the parameter count comparable.

03

GQA turns 21GB into 2.6GB

Llama 70B needs about 2.6MB of KV cache per token under MHA, so an 8K-context request costs 21GB; 8:1 grouping brings that to roughly 2.6GB.

04

Frontier models are untying their embeddings

At 128K-plus vocabularies a shared input/output matrix is costly and constraining; separate weights let the LM head specialise, raising parameter count but improving multilingual and code perplexity.

05

Expert parallelism makes bandwidth the bottleneck

Hybrid MoE blocks spread experts across GPUs, making NVLink and InfiniBand bandwidth the architectural constraint; periodic dense layers are inserted so global knowledge stays shared.