Attention and Transformers

LESSON 7.1

Queries, keys, and values

A query describes what a position seeks, a key provides a matching representation, and a value carries retrievable information. They are learned projections rather than hand-written tags. Query–key dot products become scores; row-wise softmax turns them into distributions over source keys.

The weighted value sum has a fixed feature width even when source length changes. Cross-attention can use different query and source lengths, allowing each decoder position to retrieve from all permitted encoder positions.

Q = Xq WQ; K = Xk WK; V = Xk WV
A = softmax(QKᵀ / √dk + mask), along keys
Output = A V

LESSON 7.2

Calculate an attention read

Subtracting the maximum score before exponentiation improves numerical stability without changing softmax. Scaling dot products by √dk moderates their magnitude under common variance assumptions. It does not replace all normalization or initialization choices.

The correct softmax of [2,0.5,−1] is approximately [0.7856,0.1753,0.0391]. Normalize over the keys for each query, and check that masked keys receive zero. The lab's adjustable temperature illustrates sharpness separately from the standard dimension scaling.

LESSON 7.3

Self-attention, cross-attention, and position

Self-attention derives Q, K, and V from the same sequence. Cross-attention gets Q from one sequence and K/V from another. The matrix operation is the same; the information sources differ.

Unmasked content-only self-attention is permutation equivariant: reordering the inputs reorders the outputs. Learned absolute embeddings, sinusoidal encodings, or relative-position mechanisms can supply order information. A causal mask also introduces an ordering constraint.

Original sinusoidal positions:
PE(pos,2i) = sin(pos / 10000^(2i/dmodel))
PE(pos,2i+1) = cos(pos / 10000^(2i/dmodel))

LESSON 7.4

Several retrieval spaces inside one block

Each attention head learns its own projections. Head outputs are concatenated and projected to model width. Heads can specialize, but they are not guaranteed to represent clean human categories.

The block also includes a position-wise feed-forward network, residual paths, and normalization. Attention mixes positions; the feed-forward network transforms each position's features. Pre-norm and post-norm orderings change the computation and its gradient routes.

headᵢ = Attention(XWQᵢ, XWKᵢ, XWVᵢ)
MultiHead(X) = concat(head₁,…,headₕ) WO
dmodel=64, heads=4 → 16 features per head

mha = nn.MultiheadAttention(64, 4, batch_first=True)
out, weights = mha(x,x,x, average_attn_weights=False)
# x, out: [N,T,64]; weights: [N,4,T,T]

LESSON 7.5

Choose which information is available

An encoder uses bidirectional context within an observed input. An encoder–decoder adds a causal output stream that cross-attends to the encoder. A decoder-only model predicts using earlier context without requiring a separate encoder.

The architecture establishes what a prediction may observe. Define the inference boundary before selecting masks. A model intended for online use cannot rely on future observations.

LESSON 7.6

Block future tokens before softmax

For next-token training, supervise each input position with the following token. A causal mask prevents the model from reading later input positions. Shifting targets without masking still permits leakage.

Boolean conventions differ by API. In MultiheadAttention and TransformerEncoder masks, True means blocked. In scaled_dot_product_attention, True means allowed. Padding masks address nonexistent positions; causal masks address future positions. Both may be needed.

inputs, targets = tokens[:, :-1], tokens[:, 1:]
T = inputs.size(1)
blocked = torch.triu(torch.ones(T,T,dtype=torch.bool), diagonal=1)
features = encoder(embedding(inputs) + positions[:T], mask=blocked)
logits = head(features)
loss = F.cross_entropy(logits.reshape(-1, vocab), targets.reshape(-1))

Keep this idea. Change a sequence suffix in evaluation mode. Earlier logits must remain unchanged in a causal model.

LESSON 7.7

Parallel training, sequential generation

Teacher-forced training can evaluate known positions in parallel behind a causal mask. Free generation remains sequential because each newly generated token becomes context for the next. Key/value caching reduces repeated work but does not remove that dependency.

Dense attention performs quadratic pairwise score work as sequence length grows. A naïve implementation also stores a quadratic score matrix; optimized kernels can avoid fully materializing it. Separate arithmetic complexity from actual memory behavior.

A fair recurrent-versus-transformer study fixes splits, input information, and tuning budgets, then reports accuracy and measured cost. Attention maps are diagnostics, not automatic causal explanations: values, residual streams, and later transformations also affect predictions.

WHEN THE MODEL SURPRISES YOU

Mistakes worth catching early

FROM EQUATIONS TO A RUNNING MODEL

Train a tiny causal transformer

Learn modular counting sequences using positional embeddings, two causal transformer blocks, shifted targets, and validation checkpoint selection. A future-token perturbation assertion runs before training. The final report includes reserved test loss and generated tokens.

Download the script, create a Python environment, and run the commands below in the folder containing the download. The example uses PyTorch on CPU and generates its own small dataset. No dataset or model download is required.

python -m venv .venv
# macOS / Linux: source .venv/bin/activate
# Windows PowerShell: .venv\Scripts\Activate.ps1
python -m pip install torch
python dl-week-7.py
# Fast execution check (not a convergence result):
python dl-week-7.py --smoke

"""Week 7: causal transformer on modular counting. CPU; no downloads.
python -m pip install torch
python dl-week-7.py [--smoke]
"""
import argparse
import copy
import math
import torch
from torch import nn
from torch.nn import functional as F

torch.set_num_threads(1)
VOCAB, WIDTH, LENGTH = 8, 32, 12


def sequences(n, seed):
    g = torch.Generator().manual_seed(seed)
    starts = torch.randint(VOCAB, (n, 1), generator=g)
    return (starts + torch.arange(LENGTH + 1)) % VOCAB


class TinyTransformer(nn.Module):
    def __init__(self):
        super().__init__()
        self.embedding = nn.Embedding(VOCAB, WIDTH)
        block = nn.TransformerEncoderLayer(WIDTH, 4, 64, dropout=0., batch_first=True)
        self.encoder = nn.TransformerEncoder(block, 2, enable_nested_tensor=False)
        # Encoder clones begin with equal weights; initialize layers independently.
        for layer in self.encoder.layers:
            for parameter in layer.parameters():
                if parameter.dim() > 1:
                    nn.init.xavier_uniform_(parameter)
        self.head = nn.Linear(WIDTH, VOCAB)
        pos = torch.arange(64)[:, None]
        freq = torch.exp(torch.arange(0, WIDTH, 2) * (-math.log(10000.) / WIDTH))
        pe = torch.zeros(64, WIDTH)
        pe[:, 0::2], pe[:, 1::2] = torch.sin(pos * freq), torch.cos(pos * freq)
        self.register_buffer('positions', pe)

    def forward(self, tokens):
        t = tokens.size(1)
        blocked = torch.triu(torch.ones(t, t, dtype=torch.bool), diagonal=1)
        x = self.embedding(tokens) * math.sqrt(WIDTH) + self.positions[:t]
        return self.head(self.encoder(x, mask=blocked))


def loss_on(model, tokens):
    return F.cross_entropy(model(tokens[:, :-1]).reshape(-1, VOCAB), tokens[:, 1:].reshape(-1))


def main():
    p = argparse.ArgumentParser()
    p.add_argument('--smoke', action='store_true')
    args = p.parse_args()
    torch.manual_seed(7)
    print('PyTorch', torch.__version__)
    train, val, test = sequences(256, 1), sequences(64, 2), sequences(128, 3)
    model = TinyTransformer()
    model.eval()
    with torch.no_grad():
        original = train[:4, :-1]
        changed = original.clone()
        changed[:, 5:] = (changed[:, 5:] + 3) % VOCAB
        error = (model(original)[:, :5] - model(changed)[:, :5]).abs().max().item()
        assert error < 1e-5, f'Future-token leakage: {error}'
        best_loss = loss_on(model, val).item()
    print('Causality check max prefix difference:', error)
    best = copy.deepcopy(model.state_dict())
    opt = torch.optim.AdamW(model.parameters(), lr=.003)
    for _ in range(2 if args.smoke else 25):
        model.train()
        for indices in torch.randperm(len(train)).split(64):
            opt.zero_grad()
            loss = loss_on(model, train[indices])
            assert torch.isfinite(loss)
            loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), 1.)
            opt.step()
        model.eval()
        with torch.no_grad():
            value = loss_on(model, val).item()
        if value < best_loss:
            best_loss, best = value, copy.deepcopy(model.state_dict())
    model.load_state_dict(best)
    model.eval()
    with torch.no_grad():
        print(f'Validation-selected loss={best_loss:.4f}; test={loss_on(model,test).item():.4f}; '
              f'uniform baseline={math.log(VOCAB):.4f}')
        generated = torch.tensor([[2]])
        for _ in range(11):
            token = model(generated)[:, -1].argmax(1, keepdim=True)
            generated = torch.cat([generated, token], dim=1)
        print('Generated:', generated[0].tolist())
    print('There are only eight distinct patterns. Held-out draws test the toy rule, not linguistic generalization.')
    if args.smoke:
        print('Smoke run checks execution, not convergence.')


if __name__ == '__main__':
    main()

Read the output, then challenge it

• The causality assertion must pass.
• Compare test cross-entropy with the uniform baseline log(vocabulary size). This is a toy sequence task, not a language benchmark.
• Check generated tokens for the learned increment pattern.
• Ablate position information while keeping the rest of the experiment fixed.

These are small mechanism studies, not benchmark claims. Seeds make runs easier to compare; versions and hardware can still change exact results. Decide your validation protocol before tuning and reserve final test results for reporting.

PREDICT → CALCULATE → EXPLAIN

Practice with worked answers

EXPLAIN IT WITHOUT THE CODE

Week 7 review questions

RESEARCH CONNECTIONS

Follow the original ideas

Course alignment: AI Algorithms 2, lecture Week 8. This site counts from Week 0. Explanations, browser labs, and runnable experiments are authored study companions to the supplied lecture material.