Skip to content

Attention and Transformers

Learning contract override: Prerequisite: sequence-model basics and PyTorch training loops. Time: 60–90 minutes with the encoder-decoder lab. Evidence: attention alignment, decoded examples, and one written shape check.

What This Is

Attention lets a model combine information across positions using data-dependent weights. Transformers combine attention, feed-forward layers, residual connections, and normalization without requiring recurrence or convolution in the core sequence block.

The core idea: instead of processing sequences left-to-right, let every position attend to every other position in parallel.

Two things make attention a workflow skill, not only a theory topic:

  • attention weights are an inspectable internal signal — useful for debugging, but not a causal explanation by themselves
  • the same block powers encoders (bidirectional) and decoders (causal), and the mask is what separates them

Most real mistakes with transformers come from misreading the mask, misreading the positional encoding, or training from scratch when fine-tuning was the right call.

When You Use It

  • processing text where long-range dependencies matter
  • building or fine-tuning language models (BERT, GPT, T5)
  • applying transformers to vision (ViT) or audio
  • tasks where relationships between elements matter and attention's compute cost is acceptable
  • multi-modal work where attention bridges modalities (CLIP, Flamingo)
  • you need to inspect attention routing as one debugging signal

The Attention Mechanism

Attention computes a weighted sum of values, where the weights come from comparing queries to keys:

Attention(Q, K, V) = softmax(QKᵀ / √d_k) × V

In plain terms:

  1. Each position creates a query ("what am I looking for?")
  2. Each position creates a key ("what do I contain?")
  3. Each position creates a value ("what information do I carry?")
  4. Queries compare against all keys to produce attention weights
  5. Values are combined using those weights

The scaling by √d_k is not cosmetic. Under common initialization assumptions, unscaled dot-product variance grows with key dimension, which can push softmax toward saturation and weaken its gradients.

Self-Attention Step By Step

import torch
import torch.nn.functional as F

def self_attention(X, W_q, W_k, W_v):
    Q = X @ W_q    # queries
    K = X @ W_k    # keys
    V = X @ W_v    # values

    d_k = K.shape[-1]
    scores = Q @ K.transpose(-2, -1) / (d_k ** 0.5)   # scaled dot-product
    weights = F.softmax(scores, dim=-1)                # attention weights
    output = weights @ V                                # weighted combination
    return output, weights

What to inspect at this step:

  • the weights tensor — every row should sum to 1.0 and be non-negative
  • how the pattern changes when you swap two rows of X (should permute rows and columns of weights consistently)
  • the scale of scores before softmax — if it is saturating the softmax, your scaling is off

Multi-Head Attention

Instead of one attention operation, split Q/K/V into multiple "heads" that each attend to different aspects:

import torch.nn as nn

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        assert d_model % n_heads == 0
        self.n_heads = n_heads
        self.d_head = d_model // n_heads
        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_out = nn.Linear(d_model, d_model)

    def forward(self, x, mask=None):
        B, T, C = x.shape
        q = self.W_q(x).view(B, T, self.n_heads, self.d_head).transpose(1, 2)
        k = self.W_k(x).view(B, T, self.n_heads, self.d_head).transpose(1, 2)
        v = self.W_v(x).view(B, T, self.n_heads, self.d_head).transpose(1, 2)
        scores = (q @ k.transpose(-2, -1)) / (self.d_head ** 0.5)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float("-inf"))
        attn = F.softmax(scores, dim=-1)
        out = (attn @ v).transpose(1, 2).contiguous().view(B, T, C)
        return self.W_out(out), attn

Why multiple heads: separate projections let the layer represent several attention patterns in parallel. Some heads may specialize, but human-readable roles are not guaranteed.

Rules of thumb:

  • d_head is an architecture choice constrained by quality, memory, and throughput; common model families often use values in the tens to low hundreds
  • n_heads = d_model // d_head — tune d_head and derive n_heads, not the other way around
  • number of heads is not a quality knob — more heads with the same d_model means each head is smaller

Causal Masking For Decoders

A decoder must not look at future tokens. The causal mask is a lower-triangular matrix that zeros out attention to future positions:

import torch

def causal_mask(T, device):
    return torch.tril(torch.ones(T, T, device=device)).bool()

# in the forward:
mask = causal_mask(T, x.device)       # (T, T)
scores = scores.masked_fill(~mask, float("-inf"))

If you forget the causal mask in a decoder, the model will "cheat" by looking at future tokens during training. Training loss will drop beautifully. Generation will then collapse, because at inference time the future tokens are gone.

Inspection habit: print attn[0, 0] and make sure all values above the diagonal are zero.

Transformer Block

A transformer block combines:

  1. Multi-head self-attention — look at all positions
  2. Add & Norm — residual connection + layer normalization
  3. Feed-forward network — two linear layers with a non-linearity
  4. Add & Norm — another residual connection
Input → LayerNorm → Self-Attention → + → LayerNorm → FFN → + → Output
          |                            ^            |        ^
          |____________________________|            |________|
                     (residual)                      (residual)

Two conventions exist:

  • post-norm (original paper): LN(x + Attn(x)) then LN(x + FFN(x))
  • pre-norm: x + Attn(LN(x)) then x + FFN(LN(x)) — often easier to optimize at depth

Both conventions remain in use. The original Transformer used post-norm; later work analyzes why pre-norm can improve optimization stability (Xiong et al., 2020). Check the target architecture rather than inferring the convention from depth.

Full pre-norm block in PyTorch:

class TransformerBlock(nn.Module):
    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = MultiHeadAttention(d_model, n_heads)
        self.ln2 = nn.LayerNorm(d_model)
        self.ff = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model),
        )
        self.drop = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        h, _ = self.attn(self.ln1(x), mask=mask)
        x = x + self.drop(h)
        x = x + self.drop(self.ff(self.ln2(x)))
        return x

Positional Encoding

Attention has no notion of position — "the cat sat" and "sat the cat" look the same. Positional encoding injects order.

Sinusoidal (original)

import math, torch

def sinusoidal_pe(seq_len, d_model):
    pe = torch.zeros(seq_len, d_model)
    position = torch.arange(seq_len).unsqueeze(1)
    div_term = torch.exp(torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model))
    pe[:, 0::2] = torch.sin(position * div_term)
    pe[:, 1::2] = torch.cos(position * div_term)
    return pe

Fixed and non-learned. The formula can be evaluated beyond the training length, but useful length extrapolation is not guaranteed.

Learned positional embeddings

self.pos_emb = nn.Embedding(max_len, d_model)
x = token_emb + self.pos_emb(torch.arange(T, device=x.device))

Used in BERT, GPT-2. Does not extrapolate past max_len.

Rotary position embedding (RoPE)

Used in LLaMA, Mistral, and many other decoder models. Instead of adding a positional vector, RoPE rotates query/key vectors in 2D subspaces by an angle that depends on position. In the attention dot product, the positional rotation depends on the relative offset between query and key positions.

Key property: RoPE encodes relative offsets inside the query-key interaction without a learned position table. Quality beyond the trained context still requires measurement and may need scaling or interpolation methods.

ALiBi

A bias added directly to the attention scores that penalizes distance. Cheap, extrapolates well, used in BLOOM and some MPT models.

Which to use:

  • small from-scratch project: sinusoidal or learned — fastest to implement
  • reproducing or extending an existing LLM: use that architecture's positional scheme; RoPE is a common current choice
  • if extrapolation to long contexts at inference is a hard requirement: ALiBi or RoPE with position interpolation

Encoder vs Decoder vs Encoder-Decoder

Component Sees Mask Used By
Encoder all positions none (bidirectional) BERT, RoBERTa, ViT, sentence embeddings, classification
Decoder only past positions causal GPT family, LLaMA, text generation
Encoder-Decoder encoder: all; decoder: past + cross-attention to encoder causal on decoder; cross-attn on output T5, BART, translation, summarization

Cross-attention is the same operation as self-attention, but queries come from the decoder and keys/values come from the encoder output.

Decision tree for which to build / fine-tune:

  • classification, embedding extraction, token-level tagging → encoder-only (BERT-family)
  • open-ended generation, chat, code completion → decoder-only (GPT-family)
  • input-to-output text tasks with a strong structural contract (translation, summarization with tight constraints) → encoder-decoder (T5-family)

Most modern LLM work uses decoder-only because the architecture is simpler to scale and the same model handles classification via prompting.

KV-Cache For Inference

At inference time, decoder attention recomputes keys and values for every past token on every new token. The fix is to cache them:

# pseudocode
kv_cache = []    # list of (k, v) per layer

for step in range(max_new_tokens):
    x = embed(last_token)
    for layer_idx, block in enumerate(blocks):
        k_new, v_new = block.attn.project_kv(x)
        if step > 0:
            k_prev, v_prev = kv_cache[layer_idx]
            k_new = torch.cat([k_prev, k_new], dim=-2)
            v_new = torch.cat([v_prev, v_new], dim=-2)
        kv_cache[layer_idx] = (k_new, v_new)
        x = block.forward_with_kv(x, k_new, v_new)
    next_token = sample(lm_head(x))

At prefix length T, recomputing the full prefix costs O(T²) attention work for that step; reusing cached keys and values reduces the attention work for the new query to O(T). Across a growing generated sequence, this changes the attention component from cubic to quadratic total work (ignoring constants and other layers).

What to inspect:

  • memory usage — KV-cache can dominate GPU memory for long contexts
  • whether positions and masks account for the cached prefix; in one-token decoding, only the new query is evaluated against cached and new keys

Vision Transformer (ViT)

A transformer works on images by treating the image as a sequence of patches:

class ViTPatchEmbed(nn.Module):
    def __init__(self, img_size=224, patch_size=16, in_ch=3, d_model=768):
        super().__init__()
        self.proj = nn.Conv2d(in_ch, d_model, kernel_size=patch_size, stride=patch_size)
        n_patches = (img_size // patch_size) ** 2
        self.cls_token = nn.Parameter(torch.zeros(1, 1, d_model))
        self.pos_emb = nn.Parameter(torch.zeros(1, n_patches + 1, d_model))

    def forward(self, x):
        x = self.proj(x).flatten(2).transpose(1, 2)     # (B, n_patches, d_model)
        cls = self.cls_token.expand(x.size(0), -1, -1)
        x = torch.cat([cls, x], dim=1)
        return x + self.pos_emb

Key design calls:

  • patch size 16 → 196 patches for a 224×224 image — enough for most tasks
  • the [CLS] token can be read out for classification after the transformer stack; some variants use pooled patch tokens instead
  • compare fine-tuning a pretrained ViT with from-scratch training; the result depends on dataset size, augmentation, regularization, and model scale

Efficiency Trade-Offs

Self-attention is O(n²) in sequence length. Three scaling responses exist:

  • FlashAttention — exact attention with an I/O-aware implementation; speedup and memory savings depend on hardware, shapes, dtype, and backend (Dao et al., 2022)
  • Sliding-window attention — each token only attends to the nearest w tokens (used in Longformer, Mistral)
  • Linear attention — use a kernel or recurrence-compatible reformulation to avoid materializing the full score matrix (for example, Performer); approximation and task trade-offs vary

Decision order:

  1. does your framework, hardware, dtype, and mask support an optimized scaled-dot-product-attention backend? benchmark torch.nn.functional.scaled_dot_product_attention
  2. do you actually need longer context? if no, stop
  3. can sliding window cover your task? try that before linear attention

Common Transformer Models To Know

Model Architecture Main Use
BERT encoder-only classification, NER, embeddings
RoBERTa encoder-only (revised BERT training recipe) classification, NER, embeddings
GPT-2 / GPT-3 decoder-only text generation, in-context learning
LLaMA / Mistral decoder-only, RoPE, RMSNorm modern open-weight LLMs
T5 / Flan-T5 encoder-decoder text-to-text tasks
BART encoder-decoder summarization, translation
ViT encoder-only on patches image classification
CLIP dual encoder image-text retrieval

What To Inspect

  • attention weights on a real sample — useful for catching mask or indexing bugs; do not require a human-interpretable pattern from every head
  • causal mask shape for decoders — usually lower-triangular including the current-token diagonal
  • the scale of pre-softmax scores — saturating scores point to missing √d_k scaling or bad initialization
  • gradient norms by block — locate where gradients grow or vanish instead of assuming attention fails first
  • KV-cache correctness — cached and uncached logits should agree within an explicit numerical tolerance on the same prefix

Failure Pattern

Training a transformer from scratch without comparing a pretrained baseline. The required data depends on model size, objective, tokenizer, and domain; for modest task datasets, fine-tuning a suitable pretrained model is usually the stronger first experiment.

Another failure: ignoring the quadratic score matrix in full self-attention. When profiling shows attention memory or runtime is the bottleneck, benchmark an optimized exact-attention backend, then consider local or sliding-window patterns if the task permits them.

A third: treating attention weights as explanation. They are internal routing coefficients, not causal proof of which input determined the output. Attribution methods answer different questions and also require validation against perturbation or task-specific tests.

Common Mistakes

  • forgetting the causal mask in decoder models, which lets the model "cheat" by looking at future tokens
  • not scaling the dot products by √d_k, leading to saturating softmax and vanishing gradients
  • using too few attention heads for the model dimension (or too many — the right question is d_head)
  • ignoring the positional encoding and wondering why word order does not matter
  • mixing post-norm and pre-norm in the same stack by accident
  • training with batch_first=False on a new model and getting confused by tensor shapes
  • forgetting to adjust the causal mask length when the sequence length changes batch to batch
  • using torch.nn.MultiheadAttention with wrong key_padding_mask vs attn_mask — one is per-key, the other is per-query-key pair

Practice

  1. Implement single-head self-attention from scratch and inspect the attention weights. Show that rows of the weight matrix sum to 1.
  2. Show what happens when you remove positional encoding from a sequence-classification task. Compare accuracy.
  3. Compare a 2-head and 8-head attention with the same d_model. Explain whether more heads help, in terms of d_head.
  4. Explain the difference between encoder attention (bidirectional) and decoder attention (causal) by printing both masks.
  5. Fine-tune a pretrained BERT model and compare against training from scratch on the same dataset.
  6. Implement a KV-cache for a small decoder and verify that cached generation matches uncached generation on the same prefix.
  7. Switch a ViT block from post-norm to pre-norm at 12 layers and report which one trains more stably.
  8. Replace standard attention with torch.nn.functional.scaled_dot_product_attention and measure the speedup.

Runnable Example

Run the offline attention workflow from the repository root:

.venv/bin/python labs/encoder-decoder-translation/src/encoder_decoder_workflow.py

Inspect attention_example.png and the decoded examples. Verify that attention weights have the expected source-by-target shape before interpreting the alignment.

Longer Connection

Continue with Transfer and Fine-Tuning for practical fine-tuning of transformer backbones, Text Representations and Order for the simpler text representations that serve as baselines before transformers, and Text Generation and Language Models for the workflow side of using these architectures.