Tier 3 · Deep Learning & Neural Networks

The Transformer
(Part 2) & Modern Architectures

Assemble attention, positional encoding, and feedforward layers into the complete transformer block. Then see how the same idea — applied with small variations — built BERT, GPT, and Vision Transformers.

📚 4 Lessons 🏗 Full Architecture 🖼 ViT Included ⏱ ~4 hours

Lesson 8.1

Full Transformer Architecture (Encoder-Decoder)

You now have all the pieces: multi-head attention, positional encoding. The transformer assembles these into encoder and decoder stacks, each made of repeated identical blocks, with two more critical ingredients: residual connections and layer normalization, plus a position-wise feedforward network.

The transformer encoder block

1Multi-head self-attention — tokens attend to each other within the sequence.
2Add & Norm — residual connection (x + Attention(x)) followed by LayerNorm. The residual lets gradients flow directly; LayerNorm stabilizes the scale of activations.
3Position-wise feedforward network — a 2-layer MLP applied independently to each position: FFN(x) = max(0, xW₁+b₁)W₂+b₂. Typically expands to 4× the model dimension then contracts back.
4Add & Norm — another residual + LayerNorm around the FFN.

This block is stacked N times (6 in the original paper, dozens in modern LLMs). The decoder block adds a third sub-layer: cross-attention, where decoder queries attend to encoder keys/values — this is how the decoder incorporates information from the source sequence.

🎭

Masked self-attention in the decoder The decoder must not attend to future tokens — otherwise it would "cheat" by seeing the answer during training. A causal mask sets attention scores to −∞ for any position j > i (future positions), so after softmax, those positions get zero attention weight. This is what makes language models autoregressive: each token can only see what came before it.

Pythoncomplete transformer encoder block

import torch
import torch.nn as nn

class TransformerEncoderBlock(nn.Module):
    def __init__(self, d_model=512, n_heads=8, d_ff=2048, dropout=0.1):
        super().__init__()
        self.attn  = MultiHeadAttention(d_model, n_heads)   # from Module 7
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),                  # GELU, not ReLU — smoother gradient
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model)
        )
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        # Sub-layer 1: multi-head attention + residual + norm
        attn_out, _ = self.attn(x, mask)
        x = self.norm1(x + self.dropout(attn_out))   # residual connection

        # Sub-layer 2: feedforward + residual + norm
        ffn_out = self.ffn(x)
        x = self.norm2(x + self.dropout(ffn_out))    # residual connection

        return x

class TransformerEncoder(nn.Module):
    def __init__(self, n_layers=6, d_model=512, n_heads=8, d_ff=2048):
        super().__init__()
        self.layers = nn.ModuleList([
            TransformerEncoderBlock(d_model, n_heads, d_ff)
            for _ in range(n_layers)
        ])

    def forward(self, x, mask=None):
        for layer in self.layers:
            x = layer(x, mask)
        return x

# Causal mask for decoder-only models (GPT-style)
def causal_mask(seq_len):
    # Lower triangular: position i can attend to positions 0..i
    mask = torch.tril(torch.ones(seq_len, seq_len))
    return mask   # (seq_len, seq_len), 1=visible, 0=masked

# Full model test
encoder = TransformerEncoder(n_layers=6, d_model=512)
x = torch.randn(2, 20, 512)
out = encoder(x)
print(out.shape)   # (2, 20, 512) — same shape, contextualized

📐

Pre-norm vs post-norm The original transformer applied LayerNorm after the residual addition (post-norm, shown above). Most modern LLMs (GPT-2 onward, LLaMA) apply LayerNorm before the sub-layer (pre-norm): x = x + Attention(Norm(x)). Pre-norm trains more stably at scale and doesn't need learning rate warmup as critically — it's the default choice for new large-scale models.

Lesson 8.2

BERT vs. GPT-Style Architectures — The Key Structural Difference

The original transformer has an encoder and decoder. Most modern large language models use only one half. The choice of which half — and how attention is masked — defines two fundamentally different model families with different capabilities.

	BERT (encoder-only)	GPT (decoder-only)
Attention	Bidirectional — every token sees every other token	Causal/masked — each token only sees previous tokens
Training objective	Masked Language Modeling — predict randomly masked tokens using full context	Next-token prediction — predict the next token given all previous ones
Best for	Understanding tasks: classification, NER, extractive QA	Generation tasks: completion, chat, creative writing, code
Can generate text?	Not naturally — no autoregressive structure	Yes — this is its core design
Examples	BERT, RoBERTa, DeBERTa	GPT-2/3/4, LLaMA, Claude, Gemini

Why BERT uses masked language modeling

If BERT tried next-token prediction with bidirectional attention, it would be trivial — the model could just "look ahead" to see the answer. Masked Language Modeling (MLM) solves this: randomly mask ~15% of input tokens, train the model to predict them using context from both directions. This forces genuinely bidirectional understanding without the trivial shortcut.

Why GPT uses causal masking

GPT's objective — predict the next token — requires that position i cannot see positions > i, or the task becomes trivial (just copy the answer). The causal mask enforces this. The upside: at inference, you can generate text one token at a time, feeding each prediction back in — exactly how ChatGPT-style generation works.

🔍

BERT: understanding specialist

Excellent at tasks requiring full-sentence understanding: sentiment classification, named entity recognition, extractive question answering (find the span of text that answers a question).

✍️

GPT: generation specialist

Naturally generates fluent, coherent text one token at a time. Scales remarkably well — larger GPT-style models show emergent capabilities (few-shot learning, reasoning) that smaller models lack.

🔀

Encoder-decoder (T5, BART)

Uses both halves — full transformer. Good for sequence-to-sequence tasks: translation, summarization. Encoder builds bidirectional understanding of input; decoder generates output autoregressively, attending to the encoder via cross-attention.

📈

Why decoder-only won for LLMs

Decoder-only models unify training and generation into one simple objective (next-token prediction), scale predictably with data and compute, and can be prompted to do almost any task via in-context learning. This simplicity is why GPT-style architectures dominate modern LLMs.

Pythoncomparing attention masks: BERT vs GPT

import torch

seq_len = 5

# BERT-style: bidirectional — full attention, no masking
bert_mask = torch.ones(seq_len, seq_len)
print("BERT mask (1=can attend):")
print(bert_mask)
# [[1 1 1 1 1]
#  [1 1 1 1 1]
#  [1 1 1 1 1]
#  [1 1 1 1 1]
#  [1 1 1 1 1]]  — every position sees every position

# GPT-style: causal — lower triangular mask
gpt_mask = torch.tril(torch.ones(seq_len, seq_len))
print("\nGPT mask (causal):")
print(gpt_mask)
# [[1 0 0 0 0]   ← token 0 sees only itself
#  [1 1 0 0 0]   ← token 1 sees tokens 0,1
#  [1 1 1 0 0]   ← token 2 sees tokens 0,1,2
#  [1 1 1 1 0]
#  [1 1 1 1 1]]  ← token 4 sees everything before + itself

# MLM masking for BERT training (separate from attention mask)
def mask_tokens(input_ids, mask_token_id, vocab_size, mlm_prob=0.15):
    labels = input_ids.clone()
    prob_matrix = torch.full(labels.shape, mlm_prob)
    masked_indices = torch.bernoulli(prob_matrix).bool()
    labels[~masked_indices] = -100   # -100 = ignored in loss computation
    input_ids[masked_indices] = mask_token_id   # replace with [MASK]
    return input_ids, labels

Lesson 8.3

Why Transformers Scaled So Well

RNNs hit a wall — bigger RNNs didn't reliably get better with more data and compute, and they were slow to train due to sequential processing. Transformers broke through this wall. Understanding why explains the entire trajectory from BERT to GPT-4 and beyond.

⚡

Parallelization

Unlike RNNs, every position in a sequence is processed simultaneously during training (not generation). This means transformers fully utilize GPU/TPU parallelism — training is dramatically faster for the same compute budget, which means you can train bigger models in the same wall-clock time.

📏

Constant path length

Any two tokens, regardless of distance, are connected by a single attention operation (path length O(1), vs O(T) for RNNs). This means gradients for long-range dependencies don't have to survive T sequential multiplications — they flow directly.

📈

Predictable scaling laws

Kaplan et al. (2020) and later Chinchilla (2022) showed transformer loss follows smooth power-law curves as you scale parameters, data, and compute together. This predictability let researchers confidently invest in ever-larger training runs, knowing performance would improve.

🧩

Architectural simplicity

The transformer block is the same building block repeated N times. This uniformity makes the architecture easy to scale — just add more layers, more heads, wider dimensions — without redesigning the model for each size.

📊

The scaling law finding Loss decreases predictably as a power law in three variables: model size (N), dataset size (D), and compute (C). Crucially, the Chinchilla paper found earlier models (GPT-3) were undertrained relative to their size — for a fixed compute budget, training a smaller model on more data outperforms training a larger model on less data. This reshaped how labs allocate compute budgets.

The cost: quadratic attention

Attention's O(T²) cost in sequence length is the main scaling bottleneck for long contexts. A 100K-token context requires computing 10 billion attention score pairs. This motivated significant research: FlashAttention (IO-aware exact attention, much faster in practice), sparse attention patterns, linear attention approximations, and state-space models (Mamba) as alternatives.

Lesson 8.4

Vision Transformers (ViT) — Applying the Same Idea to Images

The Vision Transformer (Dosovitskiy et al., 2020) made a striking claim: you don't need convolutions for image classification at all. Split an image into fixed-size patches, treat each patch as a "token" the same way a word is a token in NLP, and feed the sequence of patch embeddings into a standard transformer encoder.

How ViT turns an image into a sequence

1Patch extraction: Split a 224×224 image into 16×16 patches → 14×14 = 196 patches. Each patch is 16×16×3 = 768 raw pixel values.
2Linear projection: Each flattened patch is projected through a learned linear layer into the model dimension (e.g. 768) — this is the "patch embedding," analogous to a word embedding.
3[CLS] token: A learnable classification token is prepended to the patch sequence, borrowed directly from BERT. After the transformer processes everything, this token's final representation is used for classification.
4Positional embeddings: Since patches lose their spatial arrangement once flattened into a sequence, learned positional embeddings are added — exactly as in NLP transformers.
5Standard transformer encoder: The sequence of patch embeddings + [CLS] token flows through standard multi-head self-attention and feedforward blocks, identical to BERT.

🖼️

No convolutional inductive bias Unlike CNNs, ViT has no built-in assumption that nearby pixels are related — every patch can attend to every other patch immediately, regardless of distance. This means ViT must learn spatial relationships from data rather than having them architecturally enforced. The tradeoff: ViT needs more training data than CNNs to reach comparable performance, but surpasses CNNs when trained on very large datasets.

PythonVision Transformer patch embedding

import torch
import torch.nn as nn

class PatchEmbedding(nn.Module):
    def __init__(self, img_size=224, patch_size=16, in_channels=3, embed_dim=768):
        super().__init__()
        self.n_patches = (img_size // patch_size) ** 2   # 14×14 = 196
        # A single conv layer does patch extraction + linear projection at once:
        # kernel_size = stride = patch_size means non-overlapping patches
        self.proj = nn.Conv2d(in_channels, embed_dim,
                              kernel_size=patch_size, stride=patch_size)

    def forward(self, x):
        # x: (batch, 3, 224, 224)
        x = self.proj(x)              # (batch, embed_dim, 14, 14)
        x = x.flatten(2)             # (batch, embed_dim, 196)
        x = x.transpose(1, 2)         # (batch, 196, embed_dim) — sequence of patches
        return x

class VisionTransformer(nn.Module):
    def __init__(self, img_size=224, patch_size=16, embed_dim=768,
                 n_layers=12, n_heads=12, n_classes=1000):
        super().__init__()
        self.patch_embed = PatchEmbedding(img_size, patch_size, embed_dim=embed_dim)
        n_patches = self.patch_embed.n_patches

        # Learnable [CLS] token — prepended to patch sequence
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        # Learnable positional embeddings — for [CLS] + all patches
        self.pos_embed = nn.Parameter(torch.zeros(1, n_patches + 1, embed_dim))

        self.encoder = TransformerEncoder(n_layers, embed_dim, n_heads)  # from Lesson 8.1
        self.norm = nn.LayerNorm(embed_dim)
        self.head = nn.Linear(embed_dim, n_classes)

    def forward(self, x):
        B = x.shape[0]
        x = self.patch_embed(x)                          # (B, 196, embed_dim)

        cls_tokens = self.cls_token.expand(B, -1, -1)        # (B, 1, embed_dim)
        x = torch.cat([cls_tokens, x], dim=1)             # (B, 197, embed_dim)
        x = x + self.pos_embed                            # add positional info

        x = self.encoder(x)                              # standard transformer
        x = self.norm(x)
        cls_output = x[:, 0]                              # take [CLS] token's final state
        return self.head(cls_output)                     # (B, n_classes)

# Test
vit = VisionTransformer(n_classes=10)
x = torch.randn(4, 3, 224, 224)
out = vit(x)
print(out.shape)   # (4, 10)

✅

CNN vs ViT: when to use which Small dataset (<100K images): CNN, or a pretrained ViT fine-tuned (don't train ViT from scratch on small data). Large dataset (>1M images) or pretrained ViT available: ViT often outperforms CNNs. Need interpretability via attention maps: ViT's attention weights are directly visualizable. Compute-constrained edge deployment: CNNs are generally more efficient at smaller scales.