Transformer Architecture — Complete Introduction

🎯 Learning Objectives

• Understand why RNNs/LSTMs were replaced by Transformers for sequence modeling
• Explain the encoder-decoder architecture of the original Transformer
• Derive the scaled dot-product attention formula and implement it
• Describe multi-head attention and its advantages over single-head attention
• Implement positional encoding and understand why it is necessary
• Trace the forward pass through a complete Transformer block

📋 Prerequisites

• Deep Learning fundamentals (MLP, backpropagation) — needed to understand layer composition
• Sequence modeling basics (RNN, LSTM) — helpful for appreciating why Transformers were invented
• Linear algebra (matrix multiplication, vector spaces) — the attention mechanism is built on linear transformations
• Softmax and probability — attention weights are normalized using softmax

1. 📖 Core Content

1.1 Intuition: Why a New Architecture?

1.2 High-Level Architecture

1. Embedding layer: Converts input tokens to dense vectors
2. Positional encoding: Adds information about token position in the sequence
3. Multi-head attention: Allows the model to focus on different parts of the sequence
4. Add & Layer normalization: Residual connections + normalization for stable training
5. Feed-forward network: MLP applied independently at each position
6. Cross-attention (decoder only): Lets the decoder look at the encoder's output

1.3 Scaled Dot-Product Attention

1.3.1 Intuition

• You have a query: the topic you're interested in
• Each book has a key: its title/subject
• Each book has a value: its content
• The librarian computes how well your query matches each book's key, then retrieves the most relevant content In the Transformer, each token produces three vectors: Query (Q), Key (K), and Value (V). The attention mechanism computes the compatibility between Q and K to determine how much "attention" to pay to each V.

1.3.2 Formal Definition

• $Q \in \mathbb{R}^{n \times d_k}$: Query matrix (n tokens, each of dimension d_k)
• $K \in \mathbb{R}^{n \times d_k}$: Key matrix (n tokens, each of dimension d_k)
• $V \in \mathbb{R}^{n \times d_v}$: Value matrix (n tokens, each of dimension d_v)
• $d_k$: Dimension of keys/queries (the scaling factor)
• $\sqrt{d_k}$: Scaling factor to prevent large dot products from pushing softmax into regions with extremely small gradients

1.3.3 Step-by-Step Computation

1.3.4 Python Implementation

python
Run Inline
# runnable
import numpy as np
def scaled_dot_product_attention(Q, K, V):
    """
    Scaled Dot-Product Attention
    Args:
        Q: Query matrix (n_tokens, d_k)
        K: Key matrix (n_tokens, d_k)
        V: Value matrix (n_tokens, d_v)
    Returns:
        output: Contextualized representations (n_tokens, d_v)
        attention_weights: Attention distribution (n_tokens, n_tokens)
    """
    d_k = Q.shape[-1]
    # 1. Compute compatibility scores
    scores = np.dot(Q, K.T)  # (n_tokens, n_tokens)
    # 2. Scale to prevent vanishing gradients
    scores = scores / np.sqrt(d_k)
    # 3. Apply softmax row-wise for attention weights
    exp_scores = np.exp(scores - np.max(scores, axis=-1, keepdims=True))  # numerical stability
    attention_weights = exp_scores / np.sum(exp_scores, axis=-1, keepdims=True)
    # 4. Weighted sum of values
    output = np.dot(attention_weights, V)
    return output, attention_weights
# Test with our example
Q = np.array([[1, 0, 2, 1],
              [0, 2, 1, 0],
              [1, 1, 0, 2]], dtype=float)
K = np.array([[0, 1, 1, 0],
              [2, 0, 0, 1],
              [1, 2, 1, 0]], dtype=float)
V = np.array([[1, 0, 1],
              [1, 1, 0],
              [0, 1, 1]], dtype=float)
output, attn = scaled_dot_product_attention(Q, K, V)
print("Attention Weights:\n", np.round(attn, 3))
print("\nOutput:\n", np.round(output, 3))

1.3.5 Edge Cases & Gotchas

1. Very long sequences: The $O(n^2)$ complexity means attention becomes computationally prohibitive for long sequences. A 1000-token sequence produces a 1,000,000-element attention matrix.
2. Numerical instability: Large dot products (before scaling) can push softmax into regions where gradients vanish. The $\sqrt{d_k}$ scaling is specifically designed to prevent this.
3. Single token sequence: When there's only one token, attention is trivial — the output equals the value vector, and the attention weight is always 1.0.
4. Identical Q, K vectors: If all queries and keys are identical, attention becomes uniform across tokens, giving each token equal weight.

1.4 Multi-Head Attention

1.4.1 Intuition

• One attention head might focus on syntactic relationships ("cat" ↔ "sat")
• Another head might focus on coreference ("it" ↔ "mat")
• Another might focus on adjective-noun relationships ("warm" ↔ "mat") Multi-head attention runs multiple attention mechanisms in parallel, each with different learned projection matrices, allowing the model to capture different types of relationships.

1.4.2 Formal Definition

• $h$: Number of heads (typically 8 or 16)
• $W_i^Q \in \mathbb{R}^{d_{model} \times d_k}$, $W_i^K \in \mathbb{R}^{d_{model} \times d_k}$, $W_i^V \in \mathbb{R}^{d_{model} \times d_v}$: Learned projection matrices for head $i$
• $W^O \in \mathbb{R}^{h d_v \times d_{model}}$: Output projection matrix
• $d_k = d_v = d_{model} / h$: Dimension per head (e.g., if d_model=512, h=8, then d_k=d_v=64)

1.4.3 Worked Example: Multi-Head Attention

• Head 1 might learn to focus on semantic relationships — connecting nouns with their modifiers
• After attention: Head 1 output is ℝ^{3×4} For Head 2: Using $W_2^Q$, $W_2^K$, $W_2^V$ (each 8×4):
• Head 2 might learn to focus on positional relationships — connecting verbs with their subjects
• After attention: Head 2 output is ℝ^{3×4} Concatenation: Concat(Head1, Head2) = ℝ^{3×8} Output projection: Multiply by $W^O$ (8×8) → ℝ^{3×8} The final output contains information from both attention patterns.

1.4.4 Python Implementation

python
Run Inline
# runnable
import numpy as np
class MultiHeadAttention:
    def __init__(self, d_model, num_heads):
        """
        Multi-Head Attention
        Args:
            d_model: Model dimension (e.g., 512)
            num_heads: Number of attention heads (e.g., 8)
        """
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
        # Initialize weight matrices (in practice, these are learned)
        np.random.seed(42)
        self.W_Q = np.random.randn(d_model, d_model) * 0.1
        self.W_K = np.random.randn(d_model, d_model) * 0.1
        self.W_V = np.random.randn(d_model, d_model) * 0.1
        self.W_O = np.random.randn(d_model, d_model) * 0.1
    def split_heads(self, x):
        """Split last dimension into (num_heads, d_k)"""
        batch_size, seq_len, _ = x.shape
        x = x.reshape(batch_size, seq_len, self.num_heads, self.d_k)
        return x.transpose(0, 2, 1, 3)  # (batch, heads, seq, d_k)
    def combine_heads(self, x):
        """Inverse of split_heads"""
        batch_size, heads, seq_len, d_k = x.shape
        x = x.transpose(0, 2, 1, 3)  # (batch, seq, heads, d_k)
        return x.reshape(batch_size, seq_len, self.d_model)
    def forward(self, Q, K, V):
        batch_size = Q.shape[0]
        # 1. Linear projections
        Q_proj = Q @ self.W_Q  # (batch, seq, d_model)
        K_proj = K @ self.W_K
        V_proj = V @ self.W_V
        # 2. Split into heads
        Q_split = self.split_heads(Q_proj)  # (batch, heads, seq, d_k)
        K_split = self.split_heads(K_proj)
        V_split = self.split_heads(V_proj)
        # 3. Apply scaled dot-product attention per head
        scores = np.matmul(Q_split, K_split.transpose(0, 1, 3, 2))  # (batch, heads, seq, seq)
        scores = scores / np.sqrt(self.d_k)
        # Softmax
        exp_scores = np.exp(scores - np.max(scores, axis=-1, keepdims=True))
        attn_weights = exp_scores / np.sum(exp_scores, axis=-1, keepdims=True)
        # Weighted sum
        head_outputs = np.matmul(attn_weights, V_split)  # (batch, heads, seq, d_k)
        # 4. Combine heads
        combined = self.combine_heads(head_outputs)
        # 5. Output projection
        output = combined @ self.W_O
        return output, attn_weights
# Test with random data
mha = MultiHeadAttention(d_model=512, num_heads=8)
x = np.random.randn(2, 10, 512)  # batch=2, seq=10, d_model=512
output, attn = mha.forward(x, x, x)  # Self-attention: Q=K=V=x
print(f"Input shape: {x.shape}")
print(f"Output shape: {output.shape}")
print(f"Attention weights shape: {attn.shape}")  # (batch, heads, seq, seq)

1.5 Positional Encoding

1.5.1 Intuition

1.5.2 Formal Definition

• $pos$: Position in the sequence (0, 1, 2, ...)
• $i$: Dimension index (0, 1, ..., d_model/2 - 1)
• $d_{model}$: Model dimension (e.g., 512) Why sinusoidal? The sinusoidal functions have a useful property: $PE_{pos+k}$ can be represented as a linear function of $PE_{pos}$. This means the model can easily learn to attend to relative positions.

1.5.3 Worked Example

• i=0 (even): $PE(0, 0) = \sin(0 / 10000^{0}) = \sin(0) = 0$
• i=0 (odd): $PE(0, 1) = \cos(0 / 10000^{0}) = \cos(0) = 1$
• i=1 (even): $PE(0, 2) = \sin(0 / 10000^{2/6}) = \sin(0) = 0$
• i=1 (odd): $PE(0, 3) = \cos(0 / 10000^{2/6}) = \cos(0) = 1$
• i=2 (even): $PE(0, 4) = \sin(0 / 10000^{4/6}) = \sin(0) = 0$
• i=2 (odd): $PE(0, 5) = \cos(0 / 10000^{4/6}) = \cos(0) = 1$ So $PE(0) = [0, 1, 0, 1, 0, 1]$ For position 1:
• $PE(1, 0) = \sin(1/10000^{0}) = \sin(1) = 0.841$
• $PE(1, 1) = \cos(1/10000^{0}) = \cos(1) = 0.540$
• $PE(1, 2) = \sin(1/10000^{2/6}) = \sin(1/10000^{0.333}) = \sin(1/10) = \sin(0.1) = 0.100$
• $PE(1, 3) = \cos(1/10000^{2/6}) = \cos(0.1) = 0.995$
• $PE(1, 4) = \sin(1/10000^{4/6}) = \sin(1/10000^{0.667}) = \sin(1/1000) = \sin(0.001) = 0.001$
• $PE(1, 5) = \cos(1/10000^{4/6}) = \cos(0.001) = 1.000$ (approx) So $PE(1) \approx [0.841, 0.540, 0.100, 0.995, 0.001, 1.000]$

1.5.4 Python Implementation

python
Run Inline
# runnable
import numpy as np
import matplotlib.pyplot as plt
def sinusoidal_positional_encoding(max_pos, d_model):
    """
    Compute sinusoidal positional encodings
    Args:
        max_pos: Maximum sequence length
        d_model: Model dimension
    Returns:
        PE: (max_pos, d_model) array of positional encodings
    """
    PE = np.zeros((max_pos, d_model))
    for pos in range(max_pos):
        for i in range(0, d_model, 2):
            # Even indices: sine
            PE[pos, i] = np.sin(pos / (10000 ** (i / d_model)))
            # Odd indices: cosine
            PE[pos, i + 1] = np.cos(pos / (10000 ** (i / d_model)))
    return PE
# Compute for 100 positions, d_model=512
PE = sinusoidal_positional_encoding(100, 512)
print(f"PE shape: {PE.shape}")
print(f"PE at position 0: {PE[0, :8]} ... {PE[0, -8:]}")
print(f"PE at position 1: {PE[1, :8]} ... {PE[1, -8:]}")
print(f"PE at position 50: {PE[50, :8]} ... {PE[50, -8:]}")

1.5.5 Edge Cases

1. Very long sequences: Sinusoidal encodings can theoretically handle any sequence length, but the frequencies become extremely low for high dimensions, making fine position discrimination difficult.
2. Position 0: All sine terms are 0 and all cosine terms are 1 at position 0, which could create a "bias" toward the first position.
3. Differentiability: The sinusoidal functions are smooth and differentiable, which is essential for gradient-based learning.

1.6 Feed-Forward Network (Position-wise FFN)

1.7 Residual Connections and Layer Normalization

• Residual connection ($x + \text{Sublayer}(x)$): Allows gradients to flow directly through the network, enabling training of deep models
• Layer normalization: Normalizes across the feature dimension, stabilizing training (Diagram)

4. 📐 Key Formulas / Concepts

5. ⚠️ Common Pitfalls

Pitfall 1: Confusing "attention" with "self-attention"

• Self-attention: Q, K, V all come from the same sequence (encoder input). Used in encoder layers and decoder's first attention layer.
• Cross-attention: Q comes from the decoder, K and V come from the encoder. Used in decoder's second attention layer.
• General attention: Q and K can come from different sources (e.g., text-to-image).

Pitfall 2: Forgetting the scaling factor $\sqrt{d_k}$

Pitfall 3: Assuming multi-head attention always improves performance

6. 📝 Practice Questions

Q1: Compute attention for Q=[1,0], K=[[1,0],[0,1],[1,1]], V=[[1],[0],[1]] with scaling

Step 1: Compute $QK^T = [1·1+0·0, 1·0+0·1, 1·1+0·1] = [1, 0, 1]$

Step 2: Scale: d_k=2, $\sqrt{2}=1.414$, so scores = [0.707, 0, 0.707]

Step 3: Softmax: $e^{0.707}=2.028, e^0=1, e^{0.707}=2.028$ Sum = 5.056 Weights = [2.028/5.056, 1/5.056, 2.028/5.056] = [0.401, 0.198, 0.401]

Step 4: Output = 0.401×1 + 0.198×0 + 0.401×1 = 0.802 Q2: Why must d_model be divisible by num_heads in multi-head attention?

To have equal split across heads. Each head gets dimension d_k = d_model / num_heads. If they're not divisible, some heads would have different dimensions, making concatenation and output projection mismatched. The output of all heads must be concatenated to exactly d_model dimensions. Q3: What happens to attention weights if we remove the scaling factor for d_k=1024?

Without scaling, the dot product values would have variance ≈ 1024, meaning typical values would be ~32. Softmax of large values produces near-one-hot distributions — the model focuses almost entirely on one token and ignores all others. The gradients through softmax become extremely small (vanishing gradient problem), making learning very slow or impossible. Q4: For a sequence of length 1000 with d_model=512 and 8 heads, what is the memory cost of storing the attention matrix in float32?

The attention matrix has shape (batch, heads, seq_len, seq_len) = (1, 8, 1000, 1000) = 8,000,000 elements. At 4 bytes per float32: 32 MB. For a batch of 32: 1 GB just for attention matrices. This quadratic scaling is why long sequences (e.g., 100K tokens) require optimized attention like Flash Attention. Q5: In the formula PE(pos, 2i) = sin(pos/10000^{2i/d_model}), what happens as i increases?

As i increases, the denominator $10000^{2i/d_{model}}$ grows exponentially. For d_model=512:

• i=0: frequency = 1/1 = 1 (very fast oscillation)
• i=128: frequency = 1/10000^{256/512} = 1/10000^{0.5} = 1/100 (medium oscillation)
• i=256: frequency = 1/10000 (very slow oscillation)

Low dimensions capture high-frequency (fine-grained) position information, high dimensions capture low-frequency (coarse) position information. Q6: Write the mathematical expression for attention when we have a mask M that prevents positions from attending to future positions.

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

Where M is a mask matrix with 0 for allowed positions and $-\infty$ (or a very large negative number) for masked positions. After adding M to scores, the softmax of masked positions becomes 0 (since $e^{-\infty} = 0$). Q7: In a Transformer with d_model=512 and 8 heads, what are the dimensions of W_i^Q?

$W_i^Q \in \mathbb{R}^{512 \times 64}$ because:

• Input dimension = d_model = 512
• Output dimension per head = d_k = d_model / num_heads = 512 / 8 = 64

Each head projects the 512-dimensional input into a 64-dimensional query space. Q8: Compare the information flow in RNN vs Transformer for a sentence of length 20.

RNN: Information flows sequentially. Token 1 → hidden → token 2 → hidden → ... → token 20. To connect token 20 with token 1, the signal must travel through 19 time steps (suffering from vanishing gradients). Path length = 19.

Transformer: Every token directly connects to every other token in one attention step. Token 20 can directly attend to token 1. Path length = 1.

This direct connectivity is a key advantage of Transformers for capturing long-range dependencies. Q9: What is the time complexity of self-attention for a sequence of length n?

$O(n^2 \cdot d_k)$ where n is sequence length and d_k is the dimension per head. The quadratic term comes from computing the $n \times n$ attention matrix. This is both the key strength (full connectivity) and key weakness (computational cost) of Transformers.

Compare to RNN: $O(n)$ sequential (but can't parallelize). Compare to CNN: $O(k \cdot n)$ where k is kernel size (limited receptive field). Q10: In the Transformer encoder, why do we use residual connections around each sublayer?

Residual connections (skip connections) serve two purposes:

1. Gradient flow: During backpropagation, gradients can bypass the attention and FFN sublayers and flow directly through the residual path. This prevents vanishing gradients in deep (6+ layer) Transformers.
2. Information preservation: The residual connection $x + \text{Sublayer}(x)$ ensures that the original input information is always available to later layers. The sublayer only needs to learn a "residual" (the change/addition to the input), which is typically easier than learning a complete transformation from scratch.

7. 🔗 Cross-References

• Next topic: Positional Encoding Deep Dive (Week 2)
• Related topic: Self-Attention and QKV Computation (Week 1)
• External: "Attention Is All You Need" — Vaswani et al., 2017
• Video: Week 1 lectures in BSDA5004 transcripts Join Discord NextSelf-Attention & QKV