Attention and Transformers

Recurrent neural networks process sequences one token at a time, passing a hidden state $h_t = f(h_{t-1}, x_t)$ forward. Two fundamental problems:

• Sequential bottleneck: each step depends on the previous one — you cannot parallelize across time. Training on long sequences is painfully slow on modern GPUs built for parallelism.
• Long-range forgetting: information from step 1 must survive through hundreds of hidden-state updates to influence step 500. Even LSTMs with their gating mechanisms lose distant context in practice — the hidden state is a fixed-size vector being rewritten at every step.

The 2017 paper Attention Is All You Need (Vaswani et al.) replaced recurrence entirely. Instead of forcing all context through a sequential bottleneck, attention lets every token directly attend to every other token in one parallel operation. Training became orders of magnitude faster; long-range dependencies became trivial.

Attention is a soft, differentiable lookup. Imagine you're translating the word "bank" and need to decide whether it means a financial institution or a riverbank. You want to look up other words in the sentence and weight their influence by relevance.

Formally, given a query $q$ (what you're looking for), a set of keys $K$ (what each candidate offers), and values $V$ (what each candidate contributes), attention computes:

1. Match query against every key: score $s_i = q \cdot k_i$
2. Normalize scores to weights: $\alpha_i = \text{softmax}(s / \sqrt{d_k})_i$
3. Weighted sum of values: $\text{output} = \sum_i \alpha_i v_i$

The scores tell you how relevant each position is. Softmax turns them into a probability distribution. The output is a context-weighted blend of values.

The $\sqrt{d_k}$ scaling prevents dot products from growing large with dimension, which would push softmax into saturation.

In self-attention, queries, keys, and values all come from the same sequence. For an input sequence $X \in \mathbb{R}^{T \times d_{\text{model}}}$, three learned linear projections produce:

$Q = XW^Q, \quad K = XW^K, \quad V = XW^V$

where $W^Q, W^K, W^V \in \mathbb{R}^{d_{\text{model}} \times d_k}$. The full attention computation:

$\text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right) V$

The matrix product $QK^\top \in \mathbb{R}^{T \times T}$ produces one score per pair of positions — every token attends to every other token simultaneously. This is $O(T^2 d_k)$ in time and space, the famous quadratic complexity in sequence length that motivates efficient attention research (Linformer, Flash Attention, etc.).

import torch, torch.nn.functional as F

def attention(Q, K, V):
    d_k = Q.shape[-1]
    scores = Q @ K.transpose(-2, -1) / d_k**0.5  # (B, T, T)
    weights = F.softmax(scores, dim=-1)
    return weights @ V  # (B, T, d_k)

A single attention head learns one type of relationship (e.g., subject–verb agreement). Multi-head attention runs $h$ attention heads in parallel with different projections, then concatenates and projects:

$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h)\, W^O$ $\text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)$

Each head attends to different parts of the representation space simultaneously. In BERT's analysis, different heads specialize: some track syntactic dependencies, others track coreference, others attend to the previous/next tokens.

# PyTorch built-in multi-head attention
attn = torch.nn.MultiheadAttention(
    embed_dim=512, num_heads=8, batch_first=True
)
x = torch.randn(2, 64, 512)  # (batch, seq_len, d_model)
out, weights = attn(x, x, x)  # self-attention: Q=K=V=x
# out shape: (2, 64, 512)

With $d_{\text{model}} = 512$ and $h = 8$ heads, each head uses $d_k = d_v = 64$. Total parameter count stays the same as a single head with $d_k = 512$ — the expressiveness gain is free.

Attention is permutation-equivariant: shuffle the input tokens and you get the same output, shuffled. For language, word order is everything — "dog bites man" ≠ "man bites dog". We must inject position information.

The original transformer adds fixed sinusoidal encodings to the input embeddings:

$PE_{(pos, 2i)} = \sin\!\left(\frac{pos}{10000^{2i/d_{\text{model}}}}\right), \quad PE_{(pos, 2i+1)} = \cos\!\left(\frac{pos}{10000^{2i/d_{\text{model}}}}\right)$

Properties: different frequencies encode different position ranges; the model can extrapolate to longer sequences than seen in training; relative position $pos_1 - pos_2$ can be expressed as a linear function of $PE_{pos_1}$ and $PE_{pos_2}$.

Modern models use learned positional embeddings (BERT, GPT) or rotary position embeddings — RoPE (LLaMA, GPT-NeoX), which encode position in the attention scores rather than the embeddings themselves, giving better length generalization.

A transformer block combines multi-head self-attention with a position-wise feedforward network, both wrapped in residual connections and layer normalization:

x' = LayerNorm(x + MultiHeadAttention(x, x, x))  # attention sublayer
x'' = LayerNorm(x' + FFN(x'))                     # feedforward sublayer

The FFN is just two linear layers with a nonlinearity: $\text{FFN}(x) = W_2 \cdot \text{ReLU}(W_1 x + b_1) + b_2$, with $d_{\text{ff}} = 4 \times d_{\text{model}}$. Surprisingly, this FF component stores most of the factual knowledge in GPT-style models.

The two sublayer pattern obeys Pre-LN (norm before attention) in modern practice rather than the original Post-LN: x' = x + MultiHeadAttention(LayerNorm(x)). Pre-LN trains more stably and doesn't require the careful learning-rate warmup that Post-LN demands.

One transformer block (Pre-LN variant). Both sublayers use residual connections.

Three transformer variants serve different tasks:

Causal masking in decoders sets attention scores for future positions to $-\infty$ before softmax, ensuring position $t$ can only attend to positions $\leq t$. This preserves the autoregressive property needed for generation: predict one token at a time, left to right.

Cross-attention in encoder-decoder models lets decoder queries attend to encoder keys and values — this is how the decoder "reads" the source sequence while generating the target.

Transformers are the dominant architecture today not because they are cleverly designed for language — they are general-purpose. They scale predictably with compute.

Chinchilla scaling laws (Hoffmann et al., 2022): for a given compute budget $C$ (in FLOPs), optimal model size $N$ and training tokens $D$ satisfy $N \propto C^{0.5}$ and $D \propto C^{0.5}$. Bigger model + more data beats just a bigger model.

Why transformers beat CNNs and RNNs at scale:

• Parallelism: self-attention on all tokens at once — GPU utilization is high.
• Capacity: FC layers in the FFN store information densely.
• Expressiveness: attention weights are input-dependent (unlike convolution weights, which are fixed after training).

The cost: $O(T^2)$ memory in attention limits context. Flash Attention (2022) reorders computation to use $O(T)$ memory via hardware-aware tiling — enabling 100k-token contexts on modern hardware.