Attention Variants, LM Heads & Tokenisation

Part 1 described the textbook version of self-attention, softmax, and next-token prediction. Real production models rarely use the textbook version unmodified. Three components in particular get swapped, compressed, or redesigned per-model: how attention is computed, how the final hidden state becomes a token (the lm_head), and how raw text is split into tokens in the first place. This page is a field guide to the variants.

1. Attention Variants

Standard (multi-head) attention gives every token its own Q, K, and V projection per head. It's expressive, but the key/value cache it needs during generation grows linearly with both sequence length and head count — that's usually the actual memory bottleneck, not the model weights.

Multi-Head Attention (MHA) — the Part 1 default. Each of the h heads gets its own K and V. Most expressive, largest KV cache.
Multi-Query Attention (MQA) — all heads share a single K and V, only Q stays per-head. Shrinks the KV cache by a factor of h. Faster inference, small quality cost.
Grouped-Query Attention (GQA) — a middle ground: heads are split into groups, each group shares one K/V pair. Used by LLaMA 2/3, Mistral, Gemini. Most of MHA's quality, most of MQA's speed.
Multi-Head Latent Attention (MLA) — compresses K and V into a small shared latent vector before caching, then reconstructs them per-head on the fly. DeepSeek-V2/V3's approach to shrinking the KV cache further than GQA without giving up per-head expressivity.
Sliding-window / local attention — each token only attends to a fixed-size window of nearby tokens instead of the whole sequence. O(n) instead of O(n²). Used by Mistral, Longformer for long-context efficiency.
Sparse attention — a fixed pattern of far-apart tokens (strided, "global" anchor tokens, block-local) attend to each other, approximating full attention at a fraction of the cost. BigBird, Longformer.
Linear attention — replaces the softmax similarity with a kernel trick so attention can be computed in O(n) instead of O(n²), at some cost to quality. Performer, linear Transformers.

Worth separating from the list above: FlashAttention isn't an approximation — it computes exactly the same Attention(Q,K,V) as the standard formula. It's a fused GPU kernel that avoids ever materialising the full n×n attention matrix in slow memory, tiling the computation to fit in fast on-chip SRAM instead. Same maths, same output, dramatically less memory traffic. It's an implementation detail, not a variant — but it's why most attention variants above are practical at all at long context lengths.

2. LM Head Methods

After the final layer, each token has a dense hidden vector (often 2048–8192 numbers). The lm_head is the linear layer that projects this vector into logits — one raw score per vocabulary entry, ready for softmax.

Weight tying (tied embeddings) — reuse the input embedding matrix, transposed, as the output projection. Halves the parameter count spent on vocabulary (this matrix can be huge: vocab_size × d_model). Standard since GPT-2.
Untied heads — learn a separate output matrix instead of reusing embeddings. More parameters, occasionally better quality — some large models decouple input and output representations deliberately.
Adaptive / hierarchical softmax — clusters the vocabulary (frequent words get cheap direct lookup, rare words go through an extra step) so training and inference don't pay full cost for a huge vocabulary. Common before large-scale subword tokenisation made huge vocabularies less necessary.
Mixture of Softmaxes (MoS) — computes several separate softmax distributions and blends them, instead of one. A single softmax over a linear projection is provably limited in the distributions it can represent; mixing several increases expressivity at some extra compute cost.
Speculative-decoding / multi-token heads — extra lm_head-like modules (e.g. Medusa heads) predict several future tokens at once from the same hidden state. A larger model verifies them in parallel instead of generating strictly one token at a time, trading a bit of extra compute for much lower latency per accepted token.

Everything downstream of the lm_head — temperature, top-k, top-p, repetition penalties, banned-token lists — operates on these logits before softmax ever runs. See Part 5: Determinism, Seeds & Sampling for how that stage introduces (and controls) randomness.

3. Tokenisation Methods

Tokenisation happens before the model ever sees a number — it decides what counts as one indivisible unit of text. The choice trades off vocabulary size, sequence length, and how gracefully the model handles rare words, typos, code, or other languages.

Word-level — one token per whitespace-separated word. Simple, but any word not seen during training is an unrepresentable "unknown" token. Mostly obsolete.
Character-level — one token per character. No out-of-vocabulary problem at all, but sequences get very long, which is expensive since attention cost grows with sequence length.
Byte-Pair Encoding (BPE) — start from individual characters, greedily merge the most frequent adjacent pair into a new token, repeat until the vocabulary hits a target size. Common words end up as single tokens; rare words fall back to smaller pieces. Used by GPT-2/3.
Byte-level BPE — the same merge algorithm, but starting from raw UTF-8 bytes instead of characters. Every possible input is representable — no "unknown" token can ever occur, even for emoji or malformed unicode. GPT-2/3/4's tiktoken uses this.
WordPiece — like BPE, but merges are chosen by which pair most increases the likelihood of the training corpus, not just raw frequency. Used by BERT.
Unigram LM tokenisation — the reverse approach: start from a large candidate vocabulary and iteratively remove the pieces that hurt overall corpus likelihood the least, until reaching the target size. One of the two algorithms offered by SentencePiece.
SentencePiece — not a merge algorithm itself, but a language-agnostic framework (supporting both BPE and Unigram) that treats input as a raw stream of unicode/bytes with no whitespace-based pre-tokenisation. This makes it work uniformly across languages that don't separate words with spaces. Used by LLaMA, T5.

Vocabulary size is a direct trade-off: a bigger vocabulary means more common words get their own single token (shorter sequences, cheaper attention), but it also means bigger embedding and lm_head matrices, since both scale with vocab_size. Most production models settle somewhere between 32k and 128k+ tokens (GPT-4's cl100k_base uses roughly 100k, LLaMA 3 moved to ~128k).

None of these choices change the core idea from Part 1 — Q, K, V, softmax, weighted blend. They change its cost curve: how it scales with context length, how big the model's parameter budget goes toward vocabulary versus reasoning, and how gracefully it handles text it wasn't trained on.