Transformer Architecture

The Transformer Revolution

Before 2017, sequence models used RNNs and LSTMs — networks that processed tokens one at a time. This sequential nature made them:

• Slow to train (no parallelism)
• Prone to forgetting over long sequences (vanishing gradients)
• Hard to scale

The paper "Attention Is All You Need" (Vaswani et al., 2017) introduced the Transformer: a model built entirely on attention mechanisms, enabling full parallelism and capturing long-range dependencies effortlessly.

Today, virtually every state-of-the-art model — GPT-4, Claude, Gemini, T5, BERT — is a Transformer variant.

Self-Attention: The Core Mechanism

Self-attention computes a weighted sum of value vectors, where weights come from how well each query matches each key.

The Q, K, V Framework

For each token, we project its embedding into three vectors:

• Query (Q): "What am I looking for?"
• Key (K): "What do I contain?"
• Value (V): "What do I contribute if selected?"

Scaled Dot-Product Attention

$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$

1. Compute dot products of Q with all K: raw attention scores
2. Divide by √d_k to prevent vanishing gradients in softmax
3. Apply softmax to get a probability distribution (attention weights)
4. Multiply weights by V vectors and sum them

The result for each token is a context-aware representation informed by the entire sequence.

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)
    Returns: (batch, heads, seq_len, d_v)
    """
    d_k = Q.size(-1)

    # Step 1: raw attention scores
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
    # scores: (batch, heads, seq_len, seq_len)

    # Step 2: apply causal mask for autoregressive models
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))

    # Step 3: softmax over last dim → attention weights
    attn_weights = F.softmax(scores, dim=-1)

    # Step 4: weighted sum of values
    output = torch.matmul(attn_weights, V)
    return output, attn_weights

# Example
batch, heads, seq_len, d_k = 2, 8, 16, 64
Q = torch.randn(batch, heads, seq_len, d_k)
K = torch.randn(batch, heads, seq_len, d_k)
V = torch.randn(batch, heads, seq_len, d_k)

output, weights = scaled_dot_product_attention(Q, K, V)
print(f"Output shape: {output.shape}")       # (2, 8, 16, 64)
print(f"Attention weights: {weights.shape}") # (2, 8, 16, 16)

Multi-Head Attention

Instead of a single attention computation, transformers use h parallel attention heads, each learning different relationship patterns:

$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, ..., \text{head}_h)W^O$

where each head_i = Attention(QW_i^Q, KW_i^K, VW_i^V)

Why multiple heads?

• Head 1 might learn syntactic relationships (subject-verb agreement)
• Head 2 might learn coreference (pronoun → antecedent)
• Head 3 might learn positional proximity
• Together they capture richer representations than any single view

A typical GPT-3 class model uses 96 attention heads with d_model=12288.

Positional Encoding

Self-attention is permutation invariant — it sees the same regardless of token order. Positional encodings inject order information.

Sinusoidal (original)

$PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i/d_{model}}}\right)$ $PE_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{2i/d_{model}}}\right)$

Rotary Position Embedding (RoPE)

Used by LLaMA, GPT-NeoX, and Mistral. Instead of adding absolute position embeddings, RoPE rotates Q and K vectors by position-dependent angles before computing attention. This enables:

• Better generalization to sequence lengths beyond training
• Relative position information naturally encoded in dot products

ALiBi (Attention with Linear Biases)

Used by MPT, BLOOM. Adds a position-dependent bias to attention scores (no added parameters).

The Full Transformer Block

Each Transformer layer consists of:

Input → LayerNorm → Multi-Head Self-Attention → Residual Add
      → LayerNorm → Feed-Forward Network       → Residual Add → Output

Feed-Forward Network (FFN)

After attention, each position independently passes through a 2-layer MLP:

$\text{FFN}(x) = \text{GELU}(xW_1 + b_1)W_2 + b_2$

The FFN is 4× wider than the model dimension and is where most factual knowledge is stored (per the "memories in transformers" research).

Residual Connections

output = x + sublayer(LayerNorm(x)) — Pre-norm formulation (used in modern models) places LayerNorm before the sublayer for more stable training.

Layer Stacking

Models stack N identical blocks. Deeper networks learn increasingly abstract representations:

• Early layers: syntax, local patterns
• Middle layers: semantics, entity relationships
• Late layers: task-specific, abstract reasoning

Encoder vs Decoder vs Encoder-Decoder

Encoder-Only (BERT, RoBERTa)

• Bidirectional: each token attends to all other tokens (no causal mask)
• Best for: classification, NER, sentence embeddings
• Cannot generate sequences autoregressively

Decoder-Only (GPT, Claude, LLaMA, Mistral)

• Causal masking: token i can only attend to tokens 0..i (no future leakage)
• Best for: text generation, chat, reasoning
• The dominant architecture for modern LLMs

Encoder-Decoder (T5, BART, mT5)

• Encoder reads the input with full bidirectional attention
• Decoder generates output while cross-attending to encoder representations
• Best for: translation, summarization, structured prediction

Architecture	Attention	Use Case	Examples
Encoder-only	Bidirectional	Understanding	BERT, RoBERTa
Decoder-only	Causal	Generation	GPT, Claude, LLaMA
Enc-Dec	Cross-attention	Seq2Seq	T5, BART

import torch
import torch.nn as nn
import math

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        assert d_model % num_heads == 0
        self.d_k = d_model // num_heads
        self.num_heads = num_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, C = x.shape
        # Project and reshape to (B, heads, T, d_k)
        Q = self.W_q(x).view(B, T, self.num_heads, self.d_k).transpose(1, 2)
        K = self.W_k(x).view(B, T, self.num_heads, self.d_k).transpose(1, 2)
        V = self.W_v(x).view(B, T, self.num_heads, self.d_k).transpose(1, 2)

        # Attention
        scores = (Q @ K.transpose(-2, -1)) / math.sqrt(self.d_k)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        attn = torch.softmax(scores, dim=-1)
        out = (attn @ V).transpose(1, 2).contiguous().view(B, T, C)
        return self.W_o(out)

class TransformerBlock(nn.Module):
    def __init__(self, d_model=512, num_heads=8, ff_dim=2048, dropout=0.1):
        super().__init__()
        self.attn = MultiHeadAttention(d_model, num_heads)
        self.ff = nn.Sequential(
            nn.Linear(d_model, ff_dim), nn.GELU(),
            nn.Linear(ff_dim, d_model),
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.drop = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        # Pre-norm + residual
        x = x + self.drop(self.attn(self.norm1(x), mask))
        x = x + self.drop(self.ff(self.norm2(x)))
        return x

# Stack multiple blocks
class MiniGPT(nn.Module):
    def __init__(self, vocab_size=50257, d_model=512, n_layers=6, n_heads=8, max_seq=1024):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, d_model)
        self.pos_embed = nn.Embedding(max_seq, d_model)
        self.blocks = nn.ModuleList([TransformerBlock(d_model, n_heads) for _ in range(n_layers)])
        self.norm = nn.LayerNorm(d_model)
        self.head = nn.Linear(d_model, vocab_size, bias=False)

    def forward(self, tokens):
        B, T = tokens.shape
        pos = torch.arange(T, device=tokens.device)
        x = self.embed(tokens) + self.pos_embed(pos)
        # Causal mask
        mask = torch.tril(torch.ones(T, T, device=tokens.device))
        for block in self.blocks:
            x = block(x, mask)
        return self.head(self.norm(x))  # logits over vocabulary

Summary

• Self-attention enables every token to gather information from the full sequence in parallel
• Q, K, V decompose attention into what-to-look-for, what-I-have, and what-I-give
• Multi-head attention learns diverse relationship patterns simultaneously
• Positional encodings (sinusoidal, RoPE, ALiBi) inject order information
• Encoder-only models excel at understanding; decoder-only at generation; enc-dec at seq2seq
• Modern improvements: Flash Attention for memory efficiency, RoPE for better length generalization

Next chapter, we'll see how these building blocks scale into Large Language Models and what emerges at scale.