Transformer Architecture

"Attention Is All You Need"

The 2017 Revolution

The transformer, introduced by Vaswani et al. in "Attention Is All You Need," replaced recurrent and convolutional layers entirely with attention mechanisms. This enabled massive parallelization and unprecedented scale.

Why Transformers Won

RNNs process sequences step-by-step—slow and prone to forgetting. CNNs have fixed receptive fields. Transformers let every position attend to every other position directly, in parallel.

This direct global connectivity, combined with massive parallelization on GPUs, enabled training on internet-scale data—giving us GPT, BERT, and the current AI revolution.

The Inductive Bias Shift

Every architecture embeds an inductive bias—assumptions about data structure that guide learning:

CNNs: Assume spatial locality (nearby pixels are related) and translation invariance (a cat is a cat anywhere in the image).
RNNs: Assume sequential dependence (earlier tokens causally affect later ones) and recency (recent tokens matter most).
Transformers: Make minimal assumptions—every position can attend to every other. The model learns which positions matter.

This weaker inductive bias requires more data but enables greater flexibility. Given internet-scale training, transformers learn the "right" biases from data rather than having them hard-coded.

Architectural Simplicity

The transformer stacks self-attention and feed-forward layers, using residual connections and layer normalization for stable training. The architecture is remarkably simple—mostly attention and MLPs—yet scales to billions of parameters.

Transformer Block: The Building Block

A transformer is a stack of identical blocks. Each block has two main components: multi-head self-attention and a position-wise feed-forward network, each wrapped in residual connections and layer normalization.

Input

x

↓

Multi-Head Self-Attention

+

Layer Norm

↓

Feed-Forward Network

+

Layer Norm

↓

Output

y

🎯 Data Flow: Follow One Token Through

Let's trace the word "cat" through a transformer block:

"cat" enters as a 768-dimensional vector (its embedding)
Attention: "cat" looks at ALL other tokens → "What's relevant to me?" → Gets updated with context
Add & Normalize: Add original "cat" back (residual) + normalize values
Feed-Forward: "cat" gets transformed alone (no mixing with others) → Expand → Contract
Add & Normalize: Add previous version back + normalize
"cat" exits as a richer 768-dimensional vector, now carrying context from the sentence

Key insight: Attention = "mix information between tokens" • FFN = "process each token individually"

Multi-Head Self-Attention

Each position attends to all positions. Multiple heads capture different relationship types. This is where global context mixing happens.

MultiHead(Q, K, V) with Q=K=V=x

Feed-Forward Network (FFN)

Applied independently to each position. Typically expands dimension 4× then projects back. This is where per-position computation happens.

FFN(x) = GELU(xW₁ + b₁)W₂ + b₂

Layer Normalization

Normalizes activations across features (not batch). Stabilizes training and enables higher learning rates.

LN(x) = γ · (x - μ) / σ + β

Residual Connections

Adds input directly to output of each sub-layer. Enables gradient flow through very deep networks (50+ layers).

output = x + SubLayer(x)

Positional Encoding

◈

The Position Problem

Self-attention treats input as a set—it has no notion of order. "The cat sat" and "sat cat the" produce identical attention patterns. We must explicitly inject position information.

Sinusoidal Positional Encoding

The original transformer uses sine and cosine functions of different frequencies:

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

Each dimension has a different frequency, creating a unique "fingerprint" for each position. The sinusoidal form allows the model to learn relative positions through simple linear combinations.

Learned Positional Embeddings

Modern models (GPT, BERT) often learn position embeddings directly—a matrix of shape (max_seq_len, d_model) trained with the rest of the network.

Trade-off: Learned embeddings can capture position patterns specific to the training data but don't extrapolate to longer sequences.

Architecture Variants

Encoder-Only

BERT, RoBERTa

Encoder

→ [CLS] embedding

Bidirectional attention: Each token attends to all tokens in both directions. Best for understanding tasks (classification, NER, question answering).

Training: Masked language modeling—predict masked tokens from context.

Decoder-Only

GPT, LLaMA, Claude

Decoder

→ Next token

Causal attention: Each token attends only to previous tokens (left-to-right). Best for generation tasks.

Training: Next-token prediction—predict each token from all previous tokens.

Encoder-Decoder

T5, BART, Original Transformer

Enc

↓ Cross-Attention

Dec

Cross-attention: Decoder attends to encoder outputs. Best for sequence-to-sequence tasks (translation, summarization).

The Causal Mask

In decoder models, we prevent positions from attending to future tokens using a causal mask—ensuring the model can only "see" the past when predicting the next token.

The cat sat on

The ✓ ✗ ✗ ✗

cat ✓ ✓ ✗ ✗

sat ✓ ✓ ✓ ✗

on ✓ ✓ ✓ ✓

Causal mask: ✓ = can attend, ✗ = masked (set to -∞ before softmax)

Common Misconceptions

✗

"Transformers understand language like humans"

Transformers are pattern-matching machines optimized to predict tokens. They have no grounded understanding of meaning, only statistical correlations learned from text.

The accurate framing: Transformers capture useful patterns that often align with semantic meaning, but this emerges from prediction objectives, not understanding.

✗

"More layers always means better performance"

Beyond a certain depth, adding layers provides diminishing returns and increases training instability. The optimal depth depends on task and data.

The accurate framing: Modern scaling laws suggest balancing depth, width, and data. Very deep models need careful initialization and may still underperform wider, shallower alternatives.

⚗ Interactive Lab: Transformer Explorer

Explore the transformer architecture interactively. See how data flows through layers and understand the role of each component.

Architecture

Number of Layers 3

Hover over components to learn more.

Visualization of sinusoidal positional encodings. Each row is a position, each column is a dimension. Notice the different frequencies across dimensions.

Understanding tensor shapes is crucial for implementing transformers. Follow the data through each operation with concrete dimensions.

Batch Size (B):

Sequence Length (T):

Model Dimension (d_model):

Heads (h):

Input Tokens

[32, 128]

Integer token IDs

↓ Embedding Lookup

Embeddings + Position

[32, 128, 512]

(B, T, d_model)

↓ Linear Projections (Q, K, V)

Q, K, V (each)

[32, 128, 512]

Each is (B, T, d_model)

↓ Reshape for Multi-Head

Multi-Head Q, K, V

[32, 8, 128, 64]

(B, h, T, d_head) where d_head = d_model/h

↓ Q @ K.T (Attention Scores)

Attention Weights

[32, 8, 128, 128]

(B, h, T, T) — Every position attends to every other

↓ Softmax, then @ V

Attention Output

[32, 8, 128, 64]

(B, h, T, d_head)

↓ Concat Heads + Linear

Projected Output

[32, 128, 512]

(B, T, d_model) — Back to original shape

↓ FFN (expand → contract)

FFN Intermediate

[32, 128, 2048]

(B, T, 4×d_model) — Expansion layer

↓ Project back

Layer Output

[32, 128, 512]

(B, T, d_model) — Ready for next layer

Memory bottleneck: The attention matrix [B, h, T, T] is why long sequences are expensive. For T=4096, this is 4096² ≈ 16M elements per head per batch.

Here's how to implement a transformer from scratch in PyTorch. This is production-quality code used as a foundation for real LLMs.

                                    Multi-Head Attention
                                    import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        assert d_model % num_heads == 0
        
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_head = d_model // num_heads
        
        # Q, K, V projections (often combined into one)
        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
        
        # Project to Q, K, V
        Q = self.W_q(x)  # [B, T, d_model]
        K = self.W_k(x)
        V = self.W_v(x)
        
        # Reshape for multi-head: [B, T, d_model] → [B, h, T, d_head]
        Q = Q.view(B, T, self.num_heads, self.d_head).transpose(1, 2)
        K = K.view(B, T, self.num_heads, self.d_head).transpose(1, 2)
        V = V.view(B, T, self.num_heads, self.d_head).transpose(1, 2)
        
        # Scaled dot-product attention
        scores = (Q @ K.transpose(-2, -1)) / math.sqrt(self.d_head)
        # scores: [B, h, T, T]
        
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        
        attn = F.softmax(scores, dim=-1)
        out = attn @ V  # [B, h, T, d_head]
        
        # Concat heads: [B, h, T, d_head] → [B, T, d_model]
        out = out.transpose(1, 2).contiguous().view(B, T, self.d_model)
        return self.W_o(out)
                                

                                    Transformer Block
                                    class TransformerBlock(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super().__init__()
        self.attn = MultiHeadAttention(d_model, num_heads)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        
        # Feed-forward network (expand → contract)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),  # d_model → 4×d_model
            nn.GELU(),
            nn.Linear(d_ff, d_model),   # 4×d_model → d_model
            nn.Dropout(dropout)
        )
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x, mask=None):
        # Attention with residual connection
        x = x + self.dropout(self.attn(self.norm1(x), mask))
        # FFN with residual connection  
        x = x + self.dropout(self.ffn(self.norm2(x)))
        return x
                                

                                    Full GPT-style Model
                                    class GPTModel(nn.Module):
    def __init__(self, vocab_size, d_model, num_heads, num_layers, max_len):
        super().__init__()
        self.token_emb = nn.Embedding(vocab_size, d_model)
        self.pos_emb = nn.Embedding(max_len, d_model)
        
        self.blocks = nn.ModuleList([
            TransformerBlock(d_model, num_heads, 4*d_model)
            for _ in range(num_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)
        
        # Embed tokens + positions
        x = self.token_emb(idx) + self.pos_emb(pos)
        
        # Causal mask for autoregressive generation
        mask = torch.tril(torch.ones(T, T)).unsqueeze(0).unsqueeze(0)
        
        for block in self.blocks:
            x = block(x, mask)
        
        x = self.ln_f(x)
        logits = self.head(x)  # [B, T, vocab_size]
        return logits
                                

This is real code: With minor modifications, this is the foundation for GPT-2, LLaMA, and most modern LLMs. The magic is in the scale, not the architecture.

Key Observations

Residual connections create "highways" for gradients, enabling deep networks.
Layer norm appears after each sub-layer, stabilizing training.
FFN expands and contracts: d_model → 4×d_model → d_model.

Learning Objectives

"Attention Is All You Need"

Transformer Block: The Building Block

🎯 Data Flow: Follow One Token Through

Multi-Head Self-Attention

Feed-Forward Network (FFN)

Layer Normalization

Residual Connections

Positional Encoding

The Position Problem

Sinusoidal Positional Encoding

Learned Positional Embeddings

Architecture Variants

Encoder-Only

Decoder-Only

Encoder-Decoder

The Causal Mask

Common Misconceptions

"Transformers understand language like humans"

"More layers always means better performance"

⚗ Interactive Lab: Transformer Explorer

Multi-Head Attention

Transformer Block

Full GPT-style Model

Key Observations

✓ Check Your Understanding