Transformer Architecture (Full Stack) — Module 10.2

How This Lesson Fits the Module & Volume

You know the Transformer as a stacked block. This lecture zooms out to the end-to-end pipeline: token embeddings, positional signals, stacked encoder/decoder layers, and an output head.

It stitches Module 10.1 pieces—positional encoding, attention, FFN, LayerNorm—into one trainable system you can implement in PyTorch.

Learning Objectives

By the end of this lesson, students should be able to:

Draw the full Transformer pipeline from token IDs to logits.
Explain how embeddings + positional signals form the input representation.
Describe how stacked encoder and decoder blocks refine representations.
Identify the output projection (and optional tying with input embeddings).
Build a compact encoder–decoder sketch in PyTorch with correct tensor shapes.
Connect full-stack design choices to BERT/GPT/ViT variants later in the module.

Definition

The Transformer full stack is the complete model path: embedding lookup → add (or fuse) positional information → N encoder and/or decoder blocks (attention + FFN + residual + LayerNorm) → output head that maps final hidden states to task logits (e.g., vocabulary scores).

End-to-End Data Path

Token IDs

Integer sequence from tokenizer.

Embed + PE

Dense vectors + position.

Stacked Blocks

Encoder and/or decoder.

Output Head

Logits / predictions.

Stage	Typical shape	Module 10.1 / 10.2 link
Token embedding	(B, L, d_model)	Vol. 09 embedding layer; input to attention
Positional signal	(1 or B, L, d_model)	PE / learned PE
Self-attention block	(B, L, d_model)	Self-attn, MHSA
Cross-attention (dec)	(B, T, d_model)	Cross-attention
FFN	(B, L, d_model)	FFN
LM / class head	(B, L, V) or (B, C)	Task-specific linear map

Encoder Path vs Decoder Path

Encoder stack

Bidirectional self-attention.
Builds contextual source memory.
Used alone in BERT-like models.

Decoder stack

Causal self-attention + cross-attn.
Generates target tokens step by step.
Alone = GPT; with encoder = NMT.

Shared glue

Residuals / skips.
LayerNorm (pre or post).
Identical block depth often used.

Code: Tiny Full-Stack Encoder–Decoder

import math
import torch
from torch import nn

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=512):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        pos = torch.arange(max_len).unsqueeze(1).float()
        div = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(pos * div)
        pe[:, 1::2] = torch.cos(pos * div)
        self.register_buffer("pe", pe.unsqueeze(0))  # (1, max_len, d_model)

    def forward(self, x):
        return x + self.pe[:, :x.size(1)]

class TinySeq2SeqTransformer(nn.Module):
    def __init__(self, vocab=1000, d_model=64, nhead=4, num_layers=2):
        super().__init__()
        self.embed = nn.Embedding(vocab, d_model)
        self.pos = PositionalEncoding(d_model)
        layer_kwargs = dict(d_model=d_model, nhead=nhead, dim_feedforward=4*d_model, batch_first=True)
        self.encoder = nn.TransformerEncoder(nn.TransformerEncoderLayer(**layer_kwargs), num_layers)
        self.decoder = nn.TransformerDecoder(nn.TransformerDecoderLayer(**layer_kwargs), num_layers)
        self.head = nn.Linear(d_model, vocab)

    def forward(self, src, tgt):
        # src/tgt: (B, L) token ids
        memory = self.encoder(self.pos(self.embed(src)))
        out = self.decoder(self.pos(self.embed(tgt)), memory)
        return self.head(out)  # (B, T, vocab)

model = TinySeq2SeqTransformer()
src = torch.randint(0, 1000, (2, 8))
tgt = torch.randint(0, 1000, (2, 6))
print(model(src, tgt).shape)  # torch.Size([2, 6, 1000])

Strengths and Tradeoffs

Strengths

Clear modular stages—easy to swap heads or depths.
Same stack serves MT, LM, and classification with small changes.
Shapes stay interpretable: always track (B, L, d_model).

Tradeoffs

Full encoder–decoder is heavier than decoder-only for LM.
Masking (pad + causal) must be wired carefully.
Quadratic attention still dominates long contexts.

Common Misconception

“The output head is just another Transformer block.” The head is usually a linear projection (sometimes weight-tied with the input embedding). The deep work happens in the stacked attention/FFN blocks; the head maps final states to the vocabulary or label space.

Looking Ahead

Next lectures open the black boxes: Encoder Block and Decoder Block, then attention variants and positions.

Knowledge Check

Short Answer: List the four full-stack stages in order. Answer: Token embeddings, positional signal, stacked blocks, output head.
True/False: Positional information is optional for bag-of-tokens tasks but required for order-sensitive sequence modeling. Answer: True (attention alone has no inherent order).
Multiple Choice: After the encoder, memory typically has shape: (a) (B, V), (b) (B, S, d_model), (c) (d_model,). Answer: (b).
Short Answer: What does the LM head usually produce? Answer: Logits over the vocabulary for each target position.
True/False: Cross-attention lets decoder queries attend to encoder memory keys/values. Answer: True.
Multiple Choice: Weight tying often shares: (a) LayerNorm and dropout, (b) input embedding and output projection weights, (c) Adam and SGD. Answer: (b).
Short Answer: Name one Module 10.1 component inside every block. Answer: Attention, FFN, LayerNorm, or residual (any one).
Short Answer: Why track (B, L, d_model) at every stage? Answer: To verify the stack preserves width/length and catch shape bugs early.
Multiple Choice: A decoder-only full stack drops: (a) embeddings, (b) the encoder + cross-attention path, (c) LayerNorm. Answer: (b).
True/False: The full stack is the map; later lectures zoom into each box. Answer: True.

Key Takeaways

Full stack = embeddings + positions + stacked blocks + output head.
Encoder builds memory; decoder consumes it (or stands alone with causal masks).
Module 10.1 primitives appear at every layer; only wiring differs by variant.
Always verify tensor shapes through the pipeline in PyTorch.
Next: Encoder Block.

Trainer’s Guide

Hands-on idea: Print shapes after embed, after PE, after encoder, after decoder, and after the head for TinySeq2SeqTransformer.

Discussion prompt: For a chatbot, when is a full encoder–decoder worth the cost versus a decoder-only LM?

Recap: The Transformer full stack wires tokens to logits through embeddings, positions, and stacked blocks. Continue with the Encoder Block.