Transformers, attention and language-model objectives

Why it matters

Applied AI work assumes you understand what a model computes when you send it a prompt: why context length costs what it costs, why a tokenizer mangles numbers, why a padding bug changes outputs. A standard research or ML engineering exercise is to write attention from scratch, with masks and shapes.

Core ideas

Attention. Each position produces a query, a key and a value by linear projections. With Q of shape (n, d_k), K of shape (m, d_k) and V of shape (m, d_v), Attention(Q, K, V) = softmax(Q K^T / sqrt(d_k)) V has shape (n, d_v): each output row is a convex combination of value rows, weighted by query-key similarity. If query and key components are roughly independent with unit variance, their dot product has variance d_k; dividing by sqrt(d_k) keeps the logits at unit scale so the softmax does not saturate and gradients do not vanish.

Multi-head attention. Split the model dimension into h heads, each with its own projections, attend in parallel, concatenate and project back. Heads can specialize (local syntax, copying, long-range links) at the same parameter cost as one full-width head. With batch B, tensors are typically shaped (B, h, n, d_head).

Masks. Masks act on the scores before the softmax by setting disallowed positions to a large negative value.

• Causal mask: position i attends only to positions up to i, so training on every position in parallel matches left-to-right generation.
• Padding mask: padded keys must be excluded, or real tokens attend to junk. A row with every key masked becomes NaN if negative infinity is used, which implementations must guard against.

The block and model families. A block applies attention and then a position-wise MLP (several times wider than the model dimension, with GELU or a gated variant such as SwiGLU), each wrapped in a residual connection with normalization. Most current LLMs normalize before each sublayer (pre-norm), which trains deep stacks more stably than the original post-norm layout, often with RMSNorm. Encoder-only models such as BERT see the whole input and suit classification, embeddings and reranking. Decoder-only models are causal and generate text. Encoder-decoder models (the original transformer, T5) encode an input and decode an output through cross-attention.

Tokenization. BPE starts from characters or bytes and repeatedly merges the most frequent adjacent pair in a training corpus, recording the merges in order; encoding replays them. Byte-level BPE starts from the 256 byte values, so any string can be encoded without unknown tokens. Consequences: common words become single tokens while rare words, identifiers, numbers and many non-English scripts split into several, which changes cost, usable context and arithmetic behavior. Leading spaces usually belong to tokens, so "cat" and " cat" differ. Vocabulary size trades sequence length against embedding and output-layer size.

Positional information. Without position information attention is permutation-equivariant: "dog bites man" and "man bites dog" contain the same tokens. Options: fixed sinusoidal encodings added to embeddings (the original paper); learned absolute position embeddings (BERT, GPT-2), which do not extend beyond the trained length; and rotary embeddings (RoPE), which rotate queries and keys by position-dependent angles so their dot product depends on relative offset. RoPE is common in recent open models, and context-extension methods rescale its frequencies.

Objectives. Causal language modeling minimizes the cross-entropy of each next token given the true prefix (teacher forcing). Perplexity, exp of the mean per-token loss, is comparable only across models with the same tokenizer and evaluation text. Masked language modeling (BERT) selects a fraction of input positions (15% in BERT), replaces most of them with a mask token, and predicts the originals from context on both sides; it yields strong encoders but no left-to-right generator.

Decoding and cost. Generation adds one token at a time: greedy, temperature-scaled sampling, top-k or top-p (nucleus). Materialized attention scores take O(n^2) memory per head and O(n^2 d) compute for sequence length n; IO-aware kernels avoid storing the full score matrix. During generation, a KV cache keeps past keys and values per layer, so each new token attends over the cache in time linear in context length, while cache memory grows with length and batch size.

Common misconceptions

• "The sqrt(d_k) factor normalizes the output." It scales the logits before the softmax to prevent saturation.
• "The causal mask matters only at inference." It is essential in training; without it the model sees the token it must predict.
• "Tokens are words." They are learned subword pieces whose count varies with language, formatting and digits.
• "Perplexity compares any two language models." Only with the same tokenizer and evaluation text.
• "Masked and causal LM differ only in the mask." They train different capabilities: bidirectional representations versus left-to-right generation.

How to approach it

Implement attention with masks in a notebook and test it against the framework's fused attention on random padded batches. Train a small BPE tokenizer and inspect its segmentations. Read the original paper with D2L chapter 11 open. In interviews, write shapes on every line and explain each design choice by the failure it prevents.