The Transformer Architecture

Scaled Dot-Product Attention

The atomic operation of the Transformer is the scaled dot-product attention function. Given a sequence of $n$ tokens, we project each token into three vectors: a query $q$, a key $k$, and a value $v$. The full sequence is packed into matrices $Q \in \mathbb{R}^{n \times d_k}$, $K \in \mathbb{R}^{n \times d_k}$, $V \in \mathbb{R}^{n \times d_v}$.

Explanation: Each row of the output is a weighted sum of all value vectors, where the weights are the softmax-normalized dot-product similarities between that row's query and all keys.

Symbol glossary:

• $Q$ — query matrix; "what am I looking for?"
• $K$ — key matrix; "what do I advertise?"
• $V$ — value matrix; "what do I actually return if selected?"
• $d_k$ — dimension of the key/query vectors
• $\sqrt{d_k}$ — temperature scaling factor to prevent softmax saturation in high dimensions

The $\sqrt{d_k}$ scaling is critical: without it, dot products grow large with dimension, pushing softmax into regions with near-zero gradients. In PyTorch: ```python import torch import torch.nn.functional as F def scaled_dot_product_attention(Q, K, V, mask=None): d_k = Q.size(-1) scores = torch.matmul(Q, K.transpose(-2, -1)) / d_k**0.5 # (B, H, n, n) if mask is not None: scores = scores.masked_fill(mask == 0, float('-inf')) weights = F.softmax(scores, dim=-1) return torch.matmul(weights, V), weights ```

Multi-Head Attention

Rather than computing a single attention function, the Transformer runs $h$ attention heads in parallel, each operating in a lower-dimensional subspace, then concatenates and projects the results:

Explanation: Each head independently learns a different type of relationship (syntactic, semantic, positional), and their outputs are linearly combined into a single representation. where each head is:

Explanation: The projection matrices $W_i^Q, W_i^K, W_i^V$ linearly map the full-dimensional input into a lower-dimensional subspace where each head specializes.

Symbol glossary:

• $W_i^Q \in \mathbb{R}^{d_{\text{model}} \times d_k}$ — query projection for head $i$
• $W_i^K \in \mathbb{R}^{d_{\text{model}} \times d_k}$ — key projection for head $i$
• $W_i^V \in \mathbb{R}^{d_{\text{model}} \times d_v}$ — value projection for head $i$
• $W^O \in \mathbb{R}^{h d_v \times d_{\text{model}}}$ — output projection
• Typically $d_k = d_v = d_{\text{model}} / h$ so total compute matches single-head

```python import torch.nn as nn class MultiHeadAttention(nn.Module): def __init__(self, d_model, h): super().__init__() assert d_model % h == 0 self.d_k = d_model // h self.h = h self.W_q = nn.Linear(d_model, d_model) self.W_k = nn.Linear(d_model, d_model) self.W_v = nn.Linear(d_model, d_model) self.W_o = nn.Linear(d_model, d_model) def forward(self, Q, K, V, mask=None): B, n, _ = Q.shape

# Project and split into heads: (B, n, d_model) -> (B, h, n, d_k) Q = self.W_q(Q).view(B, n, self.h, self.d_k).transpose(1, 2) K = self.W_k(K).view(B, n, self.h, self.d_k).transpose(1, 2) V = self.W_v(V).view(B, n, self.h, self.d_k).transpose(1, 2) x, _ = scaled_dot_product_attention(Q, K, V, mask)

# Concat: (B, h, n, d_k) -> (B, n, d_model) x = x.transpose(1, 2).contiguous().view(B, n, self.h * self.d_k) return self.W_o(x) ```

Positional Encoding

Because self-attention is permutation-equivariant (shuffling tokens produces shuffled — not different — outputs), order must be injected explicitly. The original paper uses fixed sinusoidal encodings:

Explanation: Each position gets a unique vector of interleaved sines and cosines at geometrically-spaced frequencies, so relative offsets between positions correspond to fixed linear transformations — allowing the model to generalize to unseen lengths.

Symbol glossary:

• $pos$ — position index in the sequence
• $i$ — dimension index (ranges over $0 \ldots d_{\text{model}}/2$)
• $10000$ — base wavelength hyperparameter spanning short to very long periodicities

Feed-Forward Sublayer

After attention, each position passes independently through the same position-wise feed-forward network (FFN):

Explanation: A two-layer MLP with ReLU nonlinearity applied identically to each token — this is where the model stores and retrieves factual knowledge, as shown by mechanistic interpretability research.

Symbol glossary:

• $W_1 \in \mathbb{R}^{d_{\text{model}} \times d_{ff}}$, $W_2 \in \mathbb{R}^{d_{ff} \times d_{\text{model}}}$ — weight matrices; typically $d_{ff} = 4 \times d_{\text{model}}$
• $b_1, b_2$ — bias vectors

Each Transformer encoder layer = MultiHeadAttention → Add & LayerNorm → FFN → Add & LayerNorm. Residual connections wrap both sublayers to enable stable training at depth.

Attention Mechanisms · Positional Encoding and Rotary Embeddings · Layer Normalization · BERT and Masked Language Modeling · Mixture-of-Experts Models · FlashAttention · Autoregressive Language Models and GPT