The Universal
Architecture

Chapter 1: The Residual Stream

If you had to point to the single most important reason the transformer is universal, it wouldn't be attention. It would be something much simpler: the residual connection. That humble "add the input back to the output" after every sub-layer is what makes the entire architecture modular, composable, and trainable at depth. Without it, none of the retrofitting we saw in Chapter 0 would work.

Here's why. Think of a transformer as a highway — an information highway running through the model from input to output. Each layer (attention, FFN) is an off-ramp/on-ramp: it reads from the highway, computes something, and writes the result back onto the highway by adding it to the existing stream. The stream itself flows unimpeded from the first layer to the last.

The Math: What "Residual" Actually Means

Without residual connections, a two-layer network computes:

The output is the result of composing functions. The original input x is gone — it's been fully transformed. If f₁ loses some information, f₂ can never recover it.

With residual connections, the same network computes:

Expand this and you see something remarkable. The output is:

The original input x is always present. Layer f₁ adds its contribution. Layer f₂ adds another contribution. Neither layer needs to preserve information from the input — the residual connection does it automatically. Each layer just needs to compute what's missing or what needs correction.

Worked Example: Residual vs Non-Residual

Let's trace actual numbers. Suppose x = [1.0, 2.0] and we have two simple layers where f₁ doubles, f₂ halves:

Without residual:

We got back to where we started. The two layers cancelled. Worse: if f₁ mapped to [0, 0] (a common failure mode during training), f₂ sees nothing. Information is destroyed.

With residual:

Even if f₁ outputs zeros, h₁ = [1.0, 2.0] — the input survives. Even if f₂ outputs zeros, h₂ = [3.0, 6.0] — everything accumulated so far survives. No layer can destroy information. Each can only add to it.

python
import numpy as np

x = np.array([1.0, 2.0])

# Without residual — information can be destroyed
def no_residual(x, f1, f2):
    h1 = f1(x)         # if f1 → zeros, game over
    h2 = f2(h1)
    return h2

# With residual — input always survives
def with_residual(x, f1, f2):
    h1 = x + f1(x)     # even if f1 → zeros, h1 = x
    h2 = h1 + f2(h1)   # even if f2 → zeros, h2 = h1
    return h2

# Test with a "dead" layer
dead = lambda x: np.zeros_like(x)
double = lambda x: x * 2

print(no_residual(x, dead, double))   # [0. 0.] — destroyed!
print(with_residual(x, dead, double))  # [3. 6.] — input survived

Why This Makes the Architecture Universal

The residual stream has three consequences for universality:

1. Layers are optional. If a layer hasn't learned anything useful yet (as often happens early in training), it can output near-zeros and the stream flows through unharmed. This means you can add layers to a pretrained model and they start as no-ops — the model behaves as before while the new layers gradually learn to contribute.

2. Layers are modular. Each layer reads from and writes to the same shared representation. It doesn't matter whether the layer before it was attention or FFN or something entirely new — as long as it reads a vector and writes a vector of the same dimension, it plugs in. This is why you can insert cross-attention layers (Chapter 4) into an existing model without rewriting anything.

3. Gradients flow freely. During backpropagation, the gradient of the loss with respect to an early layer doesn't need to pass through every intervening layer — it has a direct path through the residual connections. This is what makes 100+ layer transformers trainable when a naive 100-layer network would have vanishing gradients.

The Gradient Highway — Why Deep Transformers Train

Let's trace the gradient. Consider a loss L at the output of a 4-layer residual network. The gradient with respect to the input x is:

That leading I (identity matrix) is the hero. It means the gradient always has a direct path back to the input, regardless of what the individual layers do. Even if ∂f_i/∂h is tiny (vanishing) or huge (exploding), the identity term guarantees a stable gradient path. This is why transformers can be 100+ layers deep.

Without residuals, the gradient would be:

A chain of multiplications. If each Jacobian has norm slightly less than 1 (say 0.9), after 100 layers the gradient magnitude is 0.9¹⁰⁰ ≈ 2.66 × 10^-5. Practically zero. The residual connection breaks this chain.

python
# Gradient magnitude after N layers
import numpy as np

n_layers = 100
jacobian_norm = 0.9  # each layer slightly shrinks gradients

# Without residual: product of Jacobians
no_res_grad = jacobian_norm ** n_layers
print(f"Without residual: {no_res_grad:.2e}")  # 2.66e-05 — vanished

# With residual: each Jacobian is (I + df/dx), so product ≈ (1 + jac)^N
# The identity term dominates — gradient stays O(1)
with_res_grad = (1 + jacobian_norm) ** n_layers  # explodes, but LayerNorm tames it
print(f"With residual (raw): {with_res_grad:.2e}")  # 1.38e+27
# In practice, LayerNorm keeps this in check — the point is it doesn't vanish

Chapter 2: Tokenize Everything

The transformer doesn't know what a word is. It doesn't know what a pixel is. It doesn't know what a sound wave is. All it knows is: "I receive a sequence of vectors, each of dimension d_model. I process them with attention and FFN. I output a sequence of vectors." That's it. The entire architecture is built around this one abstraction: a sequence of d-dimensional vectors.

This means the entire burden of domain adaptation falls on the tokenizer — the component that converts raw domain data (text, images, audio, point clouds, robot states) into that universal format. Get the tokenizer right, and the transformer does the rest. This chapter is about how that conversion works for each major domain.

Pattern: Tokenize, Embed, Add Position

Every domain follows the same three-step recipe:

Text: Subword Tokenization

Text was the original domain. The tokenizer splits text into subword tokens using algorithms like BPE (Byte-Pair Encoding). Common words stay whole ("the", "and"), uncommon words get split ("unbelievable" → "un", "believ", "able"). Each token maps to a row in a learned embedding matrix.

Images: Patch Tokenization (ViT)

This was the big breakthrough. Instead of feeding individual pixels (which would create impossibly long sequences), ViT cuts the image into a grid of non-overlapping patches. Each patch is flattened and linearly projected to d_model dimensions.

That's it. A 224×224 image becomes 197 tokens of dimension 768 — the same shape a 197-word sentence would have. The transformer can't tell the difference.

Audio: Spectrogram Patches (AST)

Audio is first converted to a mel spectrogram — a 2D image where the x-axis is time and the y-axis is frequency. Then it's patched exactly like ViT.

Video: Spatiotemporal Tubes (ViViT)

Video adds a time dimension. ViViT extracts tubelet tokens — 3D patches spanning space AND time:

Point Clouds: Per-Point Features

3D point clouds are already discrete — each point has (x, y, z) coordinates plus optional features (color, normals). The tokenizer just embeds each point:

Robot States: Action Tokens (RT-2)

RT-2 does something clever: it discretizes continuous robot actions (joint angles, gripper open/close) into text tokens. A 7-DOF action becomes 7 integer tokens, concatenated to the language instruction and image tokens. The transformer processes all three modalities in a single sequence.

Notice the pattern: no matter the domain, the output is always [N, d_model]. The number N varies (197 for images, 512 for audio, 3136 for video), and longer sequences cost quadratically more in attention, but the format is identical. This is the abstraction barrier that makes the transformer universal.

python
import torch
import torch.nn as nn

class PatchTokenizer(nn.Module):
    """Universal pattern: chunk → flatten → project → add position"""
    def __init__(self, in_dim, d_model, n_tokens):
        super().__init__()
        self.proj = nn.Linear(in_dim, d_model)  # The only learned part
        self.pos  = nn.Parameter(torch.randn(n_tokens, d_model) * 0.02)

    def forward(self, patches):
        # patches: [batch, N, in_dim]
        tokens = self.proj(patches)  # [batch, N, d_model]
        tokens = tokens + self.pos   # add positional encoding
        return tokens               # ready for transformer

# Text tokenizer: embedding lookup (in_dim = vocab_size, one-hot)
text_tok  = PatchTokenizer(50257, 768, 512)

# Image tokenizer: 16×16 RGB patch → 768
image_tok = PatchTokenizer(16*16*3, 768, 197)

# Audio tokenizer: 16×16 spectrogram patch → 768
audio_tok = PatchTokenizer(16*16*1, 768, 512)

# Same transformer processes all three — it sees [batch, N, 768] every time

Chapter 3: Agnostic Attention

You've tokenized your data into [N, d_model]. Now what? It enters the attention mechanism. And here's the crucial insight: attention doesn't know what the tokens represent. It's a pure set operation. It takes N vectors, computes pairwise similarity scores between all pairs, and produces N output vectors. Whether those vectors came from words, image patches, audio frames, or robot joint states — the computation is identical.

This isn't a bug. It's the design. Attention is an adaptive pooling operation: each token computes a weighted average of all other tokens, where the weights are learned from the data. The only inductive bias it has is "some tokens are more relevant to each other than others." It discovers which tokens are relevant from training data alone.

The Set Operation View

Mathematically, self-attention is a function on sets of vectors. Given a set X = {x₁, ..., x_N}, attention computes:

Let's trace what this does for a single query token x_i:

Notice: nothing in this computation references position, spatial structure, or domain. It's purely about vector similarity. If token 5 and token 42 have similar query-key dot products, they'll attend to each other — regardless of whether they're adjacent words, distant image patches, or one is a text token and the other is an image token.

Worked Example: Same Mechanism, Different Domains

Imagine 4 tokens, each 3-dimensional. We'll trace attention with the same weights, but different input domains:

python
import numpy as np

# Same attention weights for both domains
np.random.seed(42)
W_Q = np.random.randn(3, 3) * 0.5
W_K = np.random.randn(3, 3) * 0.5
W_V = np.random.randn(3, 3) * 0.5

def attention(X):
    Q = X @ W_Q                       # [4, 3]
    K = X @ W_K
    V = X @ W_V
    scores = Q @ K.T / np.sqrt(3)    # [4, 4]
    weights = np.exp(scores)
    weights /= weights.sum(axis=1, keepdims=True)  # softmax
    return weights @ V                 # [4, 3]

# "Text tokens" — embeddings for [The, cat, sat, down]
text = np.array([[0.2,0.8,0.1], [0.9,0.1,0.7],
                  [0.3,0.6,0.4], [0.5,0.3,0.2]])

# "Image patches" — embeddings for 4 image patches
image = np.array([[0.7,0.2,0.5], [0.1,0.9,0.3],
                   [0.4,0.4,0.8], [0.6,0.1,0.6]])

# Same function, same weights, different data
text_out  = attention(text)    # Works perfectly on text
image_out = attention(image)   # Works perfectly on images
# The attention function doesn't know — or care — what domain the tokens came from

Where Domain Knowledge Enters: Positional Encoding

If attention is domain-agnostic, how does the model know about spatial structure? Through positional encoding — and this is one of the few domain-specific components.

Different domains use different positional encodings because their data has different structure:

The positional encoding is the transformer's only inductive bias for spatial/temporal structure. Everything else — which patches relate to which, what spatial patterns matter — is learned from data through the attention weights.

Why No Inductive Bias is a Superpower

CNNs have strong inductive bias: local connectivity (a pixel relates most to its neighbors) and translation equivariance (a pattern at position A is the same pattern at position B). This makes CNNs data-efficient for images — they "know" about spatial locality from the start.

Transformers have almost no inductive bias. They don't assume locality, translation equivariance, or any spatial structure. This seems like a weakness, and it is — on small datasets. ViT trained on ImageNet-1K (1.3M images) underperforms ResNet. But ViT pretrained on JFT-300M (300M images) crushes ResNet.

Why? Because with enough data, the model discovers the right inductive bias from the data itself. And the bias it discovers might be better than what a human engineer would have hardcoded. Early ViT layers learn local patterns (like a CNN), but later layers learn long-range dependencies that CNNs fundamentally cannot represent.

This lack of hardcoded bias is precisely what makes the transformer universal. A CNN can only process grid-structured data (images). An RNN can only process sequential data (text). The transformer can process anything that can be expressed as a set of vectors — because it makes no assumptions about the structure.

Chapter 4: Cross-Attention — The Universal Glue

Self-attention lets tokens within a single sequence talk to each other. But what if you need two different representations to interact? A diffusion model needs to condition on a text prompt. A VLM needs image features to inform text generation. A robot policy needs language instructions to guide motor outputs. In every case, you need information from one modality to influence another.

This is where cross-attention comes in — and it's arguably the most important design pattern in modern AI. Cross-attention is identical to self-attention with one change: the queries come from one representation, while the keys and values come from another.

The Mechanics

In self-attention, all three projections come from the same input X:

In cross-attention, queries come from the target (the thing being updated) and keys/values come from the source (the conditioning signal):

The attention score is still the dot product between query and key, but now the query asks "what information do I need?" and the key/value from the source answers "here's what I have." The result is a representation of the target that's been conditioned on the source.

Worked Example: Image Conditioned on Text

Suppose you're building Stable Diffusion. You have noisy image features (the target) and a text prompt embedding (the source). Let's trace the shapes:

python
import torch
import torch.nn as nn

# Dimensions
d_model = 768          # transformer width
n_image_tokens = 256   # 16×16 latent patches
n_text_tokens = 77     # CLIP max sequence length

# Input representations
image_features = torch.randn(1, n_image_tokens, d_model)  # [1, 256, 768]
text_features  = torch.randn(1, n_text_tokens, d_model)   # [1, 77, 768]

# Cross-attention projections
W_Q = nn.Linear(d_model, d_model)  # queries from IMAGE
W_K = nn.Linear(d_model, d_model)  # keys from TEXT
W_V = nn.Linear(d_model, d_model)  # values from TEXT

# Compute cross-attention
Q = W_Q(image_features)   # [1, 256, 768] — "what does each patch need?"
K = W_K(text_features)    # [1, 77, 768]  — "what does each word offer?"
V = W_V(text_features)    # [1, 77, 768]  — "what info does each word carry?"

# Attention weights: [1, 256, 768] × [1, 768, 77] → [1, 256, 77]
scores = Q @ K.transpose(-2, -1) / (d_model ** 0.5)
weights = torch.softmax(scores, dim=-1)   # [1, 256, 77]

# Each image patch gets a weighted average of text token values
output = weights @ V  # [1, 256, 77] × [1, 77, 768] → [1, 256, 768]

# weights[0, 42, :] tells us: for image patch 42,
# how much does it attend to each of the 77 text tokens?
# If the prompt is "a red car", patch 42 (in the car region)
# will attend strongly to "car" and "red".

The critical shape to remember: the attention matrix is [n_target, n_source]. Each target token has a distribution over source tokens. This is a soft lookup: each image patch retrieves the most relevant text information.

Cross-Attention as "Soft Database Query"

Here's the analogy that makes cross-attention click. Think of it as a database:

The difference from a real database: it's soft (retrieves a weighted combination, not a single exact match) and learned (the W matrices are trained to define what "relevant" means).

Where Cross-Attention Appears

Cross-attention is everywhere in modern AI:

Chapter 5: The Conditioning Zoo

Cross-attention is one way to inject conditioning information into a transformer. But it's not the only way — and it's not always the best way. Over the past few years, researchers have discovered a whole zoo of conditioning mechanisms, each with different trade-offs in compute cost, expressiveness, and architectural complexity.

The fundamental question is always the same: how do I get information from signal C into representation X? The answer depends on what C looks like (scalar? vector? sequence?), how much compute you can afford, and whether C should influence the content or the statistics of X.

The Five Major Mechanisms

1. Cross-Attention

We covered this in Chapter 4. Each target token dynamically selects which source tokens to attend to. Best when the conditioning signal is a rich sequence (text prompts, image features).

2. AdaLN-Zero (Adaptive Layer Normalization)

Instead of cross-attending to a sequence, AdaLN converts the conditioning signal into scale (γ) and shift (β) parameters for layer normalization. The conditioning signal (timestep, class label) is projected through an MLP to produce per-layer γ and β.

The "Zero" in AdaLN-Zero: the gate α is initialized to zero, so the conditioning layer starts as a no-op and gradually learns to contribute. This is the same "initialize as identity" trick that makes residual connections work.

3. FiLM (Feature-wise Linear Modulation)

FiLM is the predecessor to AdaLN. It applies a learned affine transformation to each feature channel: scale and shift, but applied to the features directly, not to a normalization layer.

The difference from AdaLN: FiLM applies scale/shift to raw features. AdaLN applies them to normalized features. In practice, AdaLN works better because LayerNorm stabilizes the features before modulation.

4. Concatenation

The simplest approach: just concatenate the conditioning tokens to the input sequence and let self-attention figure it out.

Used in: LLaVA (image tokens concatenated to text), RT-2 (action tokens concatenated to perception), many VLMs. Simple, but expensive when M is large.

5. Prefix Tuning / Prompt Tuning

Add learnable "virtual tokens" to the beginning of the sequence. These tokens carry the conditioning information and influence subsequent tokens through attention.

Side-by-Side Comparison

Worked Example: DiT's Choice of AdaLN-Zero

When the DiT paper (Peebles & Xie, 2022) designed a transformer for diffusion, they compared cross-attention, AdaLN, and in-context conditioning. The conditioning signal was simple: a class label (integer 0-999) plus a diffusion timestep (integer 0-999). Both are single vectors, not sequences.

Cross-attention would create Q/K/V projections and attention weights for what is essentially a 1-token source sequence. That's a lot of machinery for a single vector. AdaLN converts that vector into scale/shift parameters — much more efficient.

The results: AdaLN-Zero achieved FID 2.27 on ImageNet 256×256, beating cross-attention (FID 3.75) and in-context conditioning (FID 5.38). Simpler was better because the conditioning signal was simple.

python
import torch
import torch.nn as nn

class AdaLNZero(nn.Module):
    """DiT's conditioning mechanism."""
    def __init__(self, d_model, cond_dim):
        super().__init__()
        # One MLP produces 6 modulation parameters per layer:
        # gamma1, beta1, alpha1 (for attention)
        # gamma2, beta2, alpha2 (for FFN)
        self.mlp = nn.Sequential(
            nn.SiLU(),
            nn.Linear(cond_dim, 6 * d_model)
        )
        # Initialize output to zero → layer starts as no-op
        nn.init.zeros_(self.mlp[1].weight)
        nn.init.zeros_(self.mlp[1].bias)

    def forward(self, x, c):
        # c: [batch, cond_dim] — e.g., timestep + class embedding
        params = self.mlp(c)  # [batch, 6*d_model]
        g1, b1, a1, g2, b2, a2 = params.chunk(6, dim=-1)
        # Each is [batch, d_model]

        # Modulate attention sub-layer
        h = a1.unsqueeze(1) * self.attn(g1.unsqueeze(1) * self.norm1(x) + b1.unsqueeze(1))
        x = x + h  # residual

        # Modulate FFN sub-layer
        h = a2.unsqueeze(1) * self.ffn(g2.unsqueeze(1) * self.norm2(x) + b2.unsqueeze(1))
        x = x + h  # residual
        return x

Chapter 6: The Retrofitting Playbook

Now we have all the pieces: the residual stream (Chapter 1), tokenization (Chapter 2), domain-agnostic attention (Chapter 3), cross-attention (Chapter 4), and the conditioning zoo (Chapter 5). It's time to see how they come together. This chapter is the payoff — a complete field guide to how each major domain adapted the transformer.

The playbook has exactly four steps, and every successful adaptation follows them:

Case Study 1: ViT (Vision, 2020)

The simplest and most influential adaptation. Dosovitskiy et al. asked: what's the minimum change needed to make a transformer process images?

The total novelty: a patch embedding layer and 2D positional embeddings. Everything else is copy-paste from BERT.

Case Study 2: DiT (Diffusion, 2022)

DiT replaced the U-Net in diffusion models with a transformer. The key challenge: diffusion models need to condition on a timestep (how noisy is the current image) and a class label (what to generate).

DiT's novelty: AdaLN-Zero conditioning and operating on latent space patches instead of pixel patches. The transformer itself? Unchanged.

Case Study 3: ViViT (Video, 2021)

Video is images plus time. The challenge: a 32-frame video at ViT resolution creates 32 × 196 = 6,272 tokens. That's quadratic attention cost of O(6272²) ≈ 39M operations per attention layer. The solution: factored attention.

Factored attention: instead of one global attention over 6,272 tokens, do spatial attention (196 tokens within each frame) then temporal attention (32 tokens across frames for each spatial position). Cost drops from O(6272²) to O(196² × 32 + 32² × 196) — a ~30× reduction.

Case Study 4: RT-2 (Robotics, 2023)

RT-2 is perhaps the most elegant adaptation. Instead of designing a new architecture for robot control, the team took a pretrained Vision-Language Model (PaLM-E) and tokenized robot actions as text. The model generates action tokens the same way it generates word tokens.

RT-2 didn't modify the transformer AT ALL. It just added new tokens to the vocabulary. This is the purest example of the transformer's universality — the architecture doesn't even know it's controlling a robot.

The Universal Recipe

After reviewing every major adaptation, the recipe crystallizes:

If you can answer these four questions for your domain, you can build a transformer for it. The core — attention + FFN + residual — stays identical. The engineering decisions are ALL in the adapter layers.

Chapter 7: Composition Patterns

So far we've talked about adapting a single transformer to a new domain. But the real power emerges when you compose multiple pretrained models. You've trained a great vision encoder (DINOv2) and a great language model (LLaMA). How do you combine them into a VLM without retraining either from scratch?

This is the domain of composition patterns — the architectural strategies for connecting pretrained modules. Each pattern makes a different trade-off between flexibility, compute cost, and how much of the pretrained knowledge you preserve.

Pattern 1: Frozen Backbone + Trainable Adapter

The most common pattern. You freeze the backbone (keep its weights fixed) and train a small adapter module that translates between representations.

Why freeze? Two reasons. First, the backbone already encodes enormously valuable knowledge from pretraining (often on billions of examples). Fine-tuning risks catastrophic forgetting — the model unlearns its general capabilities while learning the new task. Second, freezing is cheap — you only need gradients through the adapter, not the backbone.

Why adapter? The vision encoder and LLM typically have different embedding dimensions and different "languages" (the feature spaces don't align). The adapter bridges this gap. Different adapter designs have different expressiveness:

Worked Example: LLaVA's Composition

python
import torch
import torch.nn as nn

class LLaVA(nn.Module):
    def __init__(self, vision_encoder, llm, d_vision=1024, d_llm=4096):
        super().__init__()
        self.vision = vision_encoder  # CLIP ViT-L/14 — FROZEN
        self.llm = llm                # Vicuna-7B — initially frozen, then unfrozen

        # The only new thing: a 2-layer MLP adapter
        self.adapter = nn.Sequential(
            nn.Linear(d_vision, d_llm),
            nn.GELU(),
            nn.Linear(d_llm, d_llm)
        )
        # Adapter params: 1024×4096 + 4096 + 4096×4096 + 4096 ≈ 21M
        # vs Vision encoder: ~300M, LLM: ~7B
        # Adapter is 0.3% of total model!

    def forward(self, image, text_tokens):
        # Step 1: Extract vision features (frozen)
        with torch.no_grad():
            vis_features = self.vision(image)  # [1, 256, 1024]

        # Step 2: Adapt to LLM space (trainable)
        vis_tokens = self.adapter(vis_features)  # [1, 256, 4096]

        # Step 3: Concatenate with text and run LLM
        combined = torch.cat([vis_tokens, text_tokens], dim=1)
        output = self.llm(combined)
        return output

Notice: the adapter is 0.3% of total parameters. Yet it bridges a 300M-parameter vision encoder with a 7B-parameter LLM. This ratio — tiny adapter, massive pretrained backbone — is the hallmark of efficient composition.

Pattern 2: Dual Encoder

Two separate encoders process two modalities independently, producing embeddings in a shared space. The training objective aligns the spaces (e.g., contrastive loss pushes matching image-text pairs together and non-matching pairs apart).

Used by: CLIP, SigLIP, ALIGN. The encoders never directly interact — they communicate only through the shared embedding space. This makes dual encoders extremely efficient for retrieval (precompute all image embeddings, search by text) but limited for generation (no token-level cross-modal interaction).

Pattern 3: Mixture of Experts (MoE)

Instead of one FFN per layer, use N expert FFNs and a router that selects which expert(s) process each token. This is a composition pattern because it combines multiple specialized sub-networks within a single architecture.

Why it works: Different tokens route to different experts, creating implicit specialization. In multilingual models, different languages naturally cluster to different experts. In multimodal models, image tokens and text tokens may use different experts. The model gets 8× more parameters but only activates 2× the compute (if using 8 experts with top-2 routing).

python
class MoELayer(nn.Module):
    def __init__(self, d_model, n_experts=8, top_k=2):
        super().__init__()
        self.experts = nn.ModuleList([
            nn.Sequential(nn.Linear(d_model, 4*d_model), nn.GELU(),
                          nn.Linear(4*d_model, d_model))
            for _ in range(n_experts)
        ])
        self.router = nn.Linear(d_model, n_experts)
        self.top_k = top_k

    def forward(self, x):
        # x: [batch, seq, d_model]
        gates = torch.softmax(self.router(x), dim=-1)  # [B, S, N]
        top_vals, top_idx = gates.topk(self.top_k, dim=-1)
        # Only compute the top-k experts per token
        output = torch.zeros_like(x)
        for i in range(self.top_k):
            expert_idx = top_idx[..., i]  # which expert for this slot
            weight = top_vals[..., i]     # gate value
            for e in range(len(self.experts)):
                mask = (expert_idx == e)
                if mask.any():
                    output[mask] += weight[mask].unsqueeze(-1) * self.experts[e](x[mask])
        return output

Pattern 4: Backbone + Task Head

The simplest composition: a shared backbone (pretrained transformer) with task-specific heads (small networks appended to the output). The backbone extracts general features; the head adapts to the task.

When to Freeze vs Fine-Tune

The billion-dollar question. Here are the actual decision factors:

Chapter 8: Why Depth Works

We've established that the transformer is universal because its core is domain-agnostic and its adapter layers are thin. But there's one more mystery: why do deeper transformers consistently outperform wider ones at the same parameter count? GPT-3 has 96 layers. GPT-4 is rumored to have even more. Why not use 10 very wide layers instead?

The answer involves three interrelated ideas: feature hierarchies, the residual stream view, and the lottery ticket hypothesis. Together, they explain why stacking transformer layers works — and predict when it stops working.

Idea 1: Feature Hierarchies

Early layers learn simple patterns. Middle layers compose them into complex ones. Late layers build task-specific representations. This holds across every domain:

This hierarchy emerges naturally from training — nobody programs it. Depth creates the representational capacity for this hierarchy. A shallow network (2-3 layers) can't build the compositional features that a deep network can.

Here's a concrete trace. In a 12-layer ViT classifying "golden retriever":

Each layer adds one level of abstraction. You can't jump from "brown pixels" to "golden retriever" in one layer — the gap is too large. You need intermediate representations.

Idea 2: The Residual Stream View

From Chapter 1, we know each layer reads from and writes to a shared residual stream. This gives us a powerful way to think about depth: each layer makes a small edit to the stream. More layers = more edits = richer final representation.

Anthropic's research (Elhage et al., 2021) formalized this as the "residual stream" view of transformers. They showed that:

Attention heads READ from and WRITE to the stream independently. A head in layer 5 might read information written by a head in layer 2, even though layers 3 and 4 are in between. The residual connection enables this long-range communication.

The stream accumulates features, it doesn't transform them. After layer 1, the stream contains: original input + layer 1's contribution. After layer 12: original input + all 12 layers' contributions. Nothing is lost.

Idea 3: Lottery Tickets and Sparse Circuits

The lottery ticket hypothesis (Frankle & Carlin, 2019) suggests that large networks work because they contain many "lottery tickets" — small sub-networks that, if trained in isolation, would achieve good performance. Deeper networks contain exponentially more potential sub-networks because depth creates combinatorial diversity.

Think of it this way: a 12-layer network with 12 attention heads per layer has 144 heads total. But the number of circuits — paths through specific heads across layers — grows exponentially with depth. A 2-layer network with 12 heads per layer has at most 12 × 12 = 144 circuits. A 12-layer network has 12¹² ≈ 8.9 × 10¹² potential circuits. More depth = more lottery tickets = higher chance of finding a good solution.

Worked Example: Depth vs Width at Constant Parameters

Let's compare two models with the same parameter count (~85M):

python
# Model A: Deep and narrow
d_model_A = 512
n_layers_A = 24
params_per_layer_A = 4 * d_model_A**2 + 8 * d_model_A**2  # attn + FFN
total_A = params_per_layer_A * n_layers_A
print(f"Model A (24 layers, d=512): {total_A/1e6:.1f}M")
# 75.5M

# Model B: Shallow and wide
d_model_B = 1536
n_layers_B = 3
params_per_layer_B = 4 * d_model_B**2 + 8 * d_model_B**2
total_B = params_per_layer_B * n_layers_B
print(f"Model B (3 layers, d=1536): {total_B/1e6:.1f}M")
# 84.9M

# Similar parameter count, but:
# - Model A: 24 levels of abstraction, 12^24 possible circuits
# - Model B: 3 levels of abstraction, 12^3 = 1,728 circuits
# Model A consistently wins on benchmarks (Kaplan et al., 2020)

When Depth Stops Helping

Depth isn't free. Three failure modes:

1. Diminishing returns. Each additional layer adds less new information. Going from 12 to 24 layers helps a lot. Going from 96 to 192 helps very little. The scaling law (Kaplan et al., 2020) shows performance improves as a power law with depth: L(D) ∝ D^-α where α ≈ 0.076 for transformers. This means doubling depth gives ~5% improvement — less and less as you go deeper.

2. Training instability. Very deep networks (100+ layers) become harder to train. Gradients, despite residual connections, can still accumulate numerical errors. This is why techniques like pre-norm (LayerNorm before attention, not after) became standard for deep transformers.

3. Inference latency. Layers execute sequentially — you can't parallelize depth. A 96-layer model takes 96 sequential forward passes. Width, by contrast, parallelizes across GPU cores. For real-time applications, a shallower, wider model might be faster even if slightly less accurate.

Chapter 9: Connections

You've just learned the architectural design patterns that make the transformer a universal backbone. Let's map what we covered to where you can go deeper.

Cheat Sheet: The Universal Architecture Playbook

Related Lessons on Engineermaxxing

The Big Picture

The transformer's universality isn't an accident. It's the result of four deliberate design decisions that, together, create a maximally reusable architecture:

We are living through a remarkable convergence in AI architecture. For the first time in the field's history, the same design is state-of-the-art across nearly every modality and task. Understanding the design patterns behind that universality — which is what this lesson taught — is arguably the single most important architectural insight in modern AI.

"The transformer is not the final architecture. But it is the first universal one."