The Full Transformer Architecture

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. The Encoder Block
5. The Decoder Block
6. Cross-Attention Mechanism
7. The Feed-Forward Network
8. The Residual Stream
9. Encoder-Only Models (BERT)
10. Decoder-Only Models (GPT)
11. Parameter Count Analysis
12. Common Pitfalls
13. Summary
14. Exercises

Learning Objectives

1. Describe the full encoder-decoder transformer architecture from Vaswani et al. (2017).
2. Explain the function of each sublayer and the order of operations.
3. Derive the causal mask for autoregressive decoding.
4. Compute the total parameter count as a function of $d$, $H$, $d_{ff}$, $L$.
5. Compare encoder-only, decoder-only, and encoder-decoder variants.

Notation

• $d$ — model dimension
• $d_{ff}$ — feed-forward hidden dimension (typically $4d$)
• $H$ — number of attention heads
• $L$ — number of layers (blocks)
• $T$ — sequence length
• $V$ — vocabulary size

Core Intuition

The transformer processes sequences through alternating attention (mixing information across positions) and feed-forward (transforming each position independently) layers. Residual connections and layer normalization enable training very deep stacks. The architecture's power comes from attention's ability to model arbitrary pairwise interactions in $O(1)$ depth.

The Encoder Block

Each encoder layer applies two sublayers:

The encoder uses bidirectional self-attention — every position can attend to every other position. After $L$ layers, the encoder output $\mathbf{Z} \in \mathbb{R}^{T_{\text{src}} \times d}$ encodes the full input context.

The Decoder Block

Each decoder layer applies three sublayers:

Masked self-attention prevents position $t$ from attending to future positions $t' > t$.

Causal mask. The attention scores are masked:

Adding $-\infty$ before softmax ensures those positions receive zero attention weight.

Cross-Attention Mechanism

In cross-attention, queries come from the decoder and keys/values from the encoder:

This allows each decoder position to attend to all encoder positions — the mechanism by which the decoder "reads" the input.

The Feed-Forward Network

Applied identically and independently to each position:

where $\mathbf{W}_1 \in \mathbb{R}^{d_{ff} \times d}$, $\mathbf{W}_2 \in \mathbb{R}^{d \times d_{ff}}$, and $\phi$ is the activation (ReLU in original, GELU in GPT/BERT).

Parameters per FFN layer: $2d \cdot d_{ff} + d_{ff} + d \approx 8d^2$ (for $d_{ff} = 4d$).

Interpretation: The FFN acts as a position-wise "memory lookup" — it stores and retrieves factual knowledge. Attention routes information; FFN processes it.

The Residual Stream

The residual stream interpretation (Elhage et al., 2021): The hidden state $\mathbf{h}$ is a "residual stream" that accumulates contributions from each sublayer:

Each attention and FFN layer reads from and writes to this stream. This view makes gradient flow transparent: the gradient flows directly from loss to any layer through the identity path.

Encoder-Only Models (BERT)

• Remove the decoder entirely.
• Bidirectional attention over the full input.
• Pre-trained with Masked Language Modeling (MLM): predict randomly masked tokens.
• Used for: classification, NER, sentence similarity, retrieval.

Decoder-Only Models (GPT)

• Remove the encoder and cross-attention.
• Causal (masked) self-attention only.
• Pre-trained with next-token prediction: $p(x_t \mid x_1, \ldots, x_{t-1})$.
• Used for: text generation, in-context learning, reasoning.
• Most modern LLMs (GPT-4, LLaMA, Claude) are decoder-only.

Why decoder-only dominates: Simpler architecture; scales better; naturally supports generation; in-context learning emerges at scale.

Parameter Count Analysis

For a decoder-only transformer:

Per layer:

• Self-attention: $4d^2$ (Q, K, V, O projections)
• FFN: $2d \cdot d_{ff} \approx 8d^2$
• LayerNorm: $4d$ (negligible)
• Total per layer: $\approx 12d^2$

Full model:

• Embedding: $V \cdot d$
• $L$ layers: $12Ld^2$
• Final LN + output head: $\approx V \cdot d$ (often tied with embedding)

Example (GPT-3 175B): $L=96, d=12288, H=96, d_{ff}=49152, V=50257$. This gives $\approx 12 \times 96 \times 12288^2 \approx 174\text{B}$.

Common Pitfalls

Pitfall 1. Thinking the FFN is unimportant. It contains $\sim 2/3$ of all parameters and stores most learned "knowledge." Attention just routes information.

Pitfall 2. Not using the causal mask for generation. Without it, the model sees future tokens during training — the loss becomes trivially low but the model cannot generate.

Pitfall 3. Confusing Pre-Norm with Post-Norm. Modern models use Pre-Norm ($\text{LN}$ before attention/FFN) for training stability. The original paper used Post-Norm.

Summary

• The transformer alternates multi-head attention (inter-position mixing) and FFN (per-position transformation).
• Residual connections + LayerNorm enable training deep stacks.
• Causal masking enforces autoregressive structure for generation.
• Cross-attention connects encoder outputs to decoder queries.
• Decoder-only is the dominant architecture for modern LLMs.
• Total parameters scale as $\sim 12Ld^2$.

Exercises

Exercise 1. For a transformer with $L=12, d=768, d_{ff}=3072, V=30000$, compute the total parameter count.

Exercise 2. Derive the causal mask and show that position $t$'s output depends only on positions $1, \ldots, t$.

Exercise 3. Prove that the encoder-decoder cross-attention computes the same function as a retrieval operation: the decoder "queries" the encoder "database."

Exercise 4. Compute the FLOPs for one forward pass through a single transformer layer with sequence length $T$ and model dimension $d$.

Exercise 5. Explain why weight tying (sharing embedding and output projection matrices) is valid and compute the parameter savings.