🧠 Transformers & Attention

To interpret “the animal did not cross the street because it was tired,” a model needs to connect a token with useful information elsewhere in the sentence. Attention supplies a learned read operation: a query scores available keys, and the resulting weights mix their values. A transformer repeatedly combines these reads with local feature transformations.

Before you start

Read Linear Algebra for dot products and matrix shapes, then Deep Learning Basics for softmax, backpropagation, and residual paths. Pre-Transformer Architectures explains how attention differs from recurrent state. Numerical Computing supports the precision and memory sections.

The main route is one attention read → batched shapes → masks and position → a full block → training and serving costs. By the end, you should be able to calculate a small attention output and explain which costs grow with context length.

1. Compute one attention read by hand

Take query q=[1,0], keys k1=[1,0], k2=[0,1], and values v1=[2,0], v2=[0,4]. The key/query width d_k is 2.

1. Dot products are [q·k1,q·k2]=[1,0].
2. Divide by √2 to get scores approximately [0.7071,0].
3. Softmax gives weights [0.6698,0.3302], which sum to one.
4. The output is .6698[2,0]+.3302[0,4] ≈ [1.3395,1.3210].

Keys determine where the query reads; values determine what it receives. Here the output is a convex combination of the value vectors before any output projection. Attention does not simply return the best-matching token or its probability.

If this query is at the first position of a causal sequence, it cannot read position two. Masking the second score to negative infinity before softmax gives weights [1,0] and output [2,0].

Check yourself: if both allowed scores are equal, what is the output? Solution: weights are [.5,.5], so output is [1,2]. Multiplying both values by two doubles the output without changing the attention weights.

Use the self-attention playground to inspect a query, change score scaling, and toggle the causal mask. Its small vectors illustrate the operation; they are not learned semantic evidence about an actual language model.

Why divide by √d_k?

If query and key coordinates are approximately independent, zero-mean, and unit-variance, their dot product is a sum with variance about d_k. Scaling by 1/√d_k keeps that variance near one under those assumptions. Very large score gaps can make softmax concentrate and yield small derivatives for many alternatives. Learned vectors need not satisfy those assumptions exactly, so the scale is a useful design choice rather than a stability guarantee. Temperature, normalized queries/keys, additive scoring, and other attention functions change this calculation.

2. Lift the example to tokens, batches, and heads

Let B be batch size, N the query length, M the key/value length, d the model width, h the number of query heads, and d_k/d_v the per-head key/value widths. In self-attention M=N and queries, keys, and values come from the same input X:B×N×d.

Tensor	Shape	Meaning
W_Q	d×(h d_k)	Query projection
W_K	d×(h d_k)	Key projection for ordinary multi-head attention
W_V	d×(h d_v)	Value projection
Q	B×h×N×d_k	Projected queries after reshape and transpose
K	B×h×M×d_k	Projected keys
V	B×h×M×d_v	Projected values
Scores / weights	B×h×N×M	One distribution over keys per query/head
Per-head output	B×h×N×d_v	Weighted value reads

The batched computation is A=softmax(QKᵀ/√d_k + mask, axis=keys) and O=AV, where the transpose swaps only the final two K axes. Concatenate head outputs into B×N×(h d_v) and apply W_O:(h d_v)×d. A common design uses d_k=d_v=d/h, but that equality is a design convention rather than part of the abstract attention definition.

In cross-attention, Q comes from the decoder state and K,V from encoder states; N and M can differ. Multi-head attention gives separately parameterized read distributions. Heads can learn different relationships, redundant ones, or mixtures; they are not assigned guaranteed “syntax” or “meaning” roles.

A careful implementation outline

project input to Q,K,V
reshape and transpose into batch, head, token, feature axes
scores = Q @ transpose_last_two_axes(K) / sqrt(d_k)
scores = scores + additive_mask    # 0 allowed, −infinity disallowed
weights = stable_softmax(scores, axis=last)
output = weights @ V
transpose/concatenate heads, then apply output projection

Combine causal, padding, and any segment masks correctly. An all-masked row needs an explicit implementation convention: naive softmax of all negative infinities produces an undefined subtraction. Avoid it or define an appropriate zero-output behavior. Masking after softmax without correct renormalization is not equivalent, because disallowed logits affected the denominator. Check boolean-mask polarity and causal-mask alignment in the chosen API, especially when query and cache lengths differ. Attention dropout is normally disabled at evaluation; some functional APIs require explicitly passing probability zero.

Production scaled-dot-product-attention APIs may dispatch to fused kernels when device, dtype, shapes, masks, and configuration support them. Dispatch is not a guarantee that a particular kernel will run. Verify small outputs against a clear reference implementation and compare with tolerances appropriate to precision.

3. Masks and positional information answer different questions

A causal mask allows key position j only when j≤i for query i under ordinary left-to-right indexing. With a shifted language-model target, the representation at i predicts token i+1 without seeing it. This lets training evaluate known input positions together. At generation time, the next input token is not known until it is selected; ordinary autoregressive generation therefore remains sequential across newly generated tokens.

Bidirectional encoder attention can read all nonpadding positions. An encoder–decoder combines a bidirectional source encoder, causal target self-attention, and cross-attention to the source. These are different information-access patterns, not universal rankings of model quality. Encoder-only models commonly support classification, tagging, retrieval, and reranking; decoder-only models naturally support continuation; encoder–decoder models separate source understanding from generated output. Objectives or masks can be adapted, so an encoder architecture is not mathematically forbidden from every generative use.

Without positional information or an order-dependent mask, self-attention is permutation-equivariant: permuting input tokens permutes output vectors in the same way. It is not invariant in the sense of leaving the ordered output unchanged. A causal mask supplies an ordering constraint; positional features additionally encode positions or offsets.

Absolute positions and rotary positions

Learned absolute embeddings allocate learned vectors to positions and need a strategy beyond their trained table. Sinusoidal embeddings use a formula with multiple frequencies and can be computed at unseen positions, but successful extrapolation does not follow from computability alone.

RoPE rotates pairs of query/key coordinates using angles that depend on token position and frequency. For fixed content vectors q,k and orthogonal rotations R_m,R_n:

(R_m q)ᵀ(R_n k) = qᵀ R_(n−m) k.

The explicit rotation term depends on relative offset. In two dimensions, let q=k=[1,0] and the rotation frequency be π/4 per position. At m=1,n=2 the dot product is cos(π/4)≈.7071; shifting both to m=5,n=6 leaves that dot product unchanged. Different coordinate pairs use different frequencies in actual RoPE.

This identity holds for fixed content q,k. In a trained multilayer network, those content vectors depend on context, masks, and preceding computations, so it does not prove the whole model is invariant to every position shift. Extending context through frequency scaling, interpolation, or additional training changes the model's operating conditions and requires evaluation. RoPE does not guarantee retrieval or quality beyond trained lengths.

The RoPE playground holds the content vectors fixed to isolate this rotation identity. Move both positions together, then change their gap, and connect the change in score to the formula.

4. A transformer block alternates communication and local computation

One common pre-norm block is:

X: B×N×d
U = X + Attention(Norm1(X))
Y = U + FFN(Norm2(U))

Attention mixes information across allowed token positions. The feed-forward network applies the same learned feature transformation separately at each token. A two-projection FFN uses activation(XW1+b1)W2+b2, with W1:d×f and W2:f×d. The hidden width f is a design choice; f=4d is a familiar example. A SwiGLU FFN uses a gated product of two projections followed by a third. Matching about the same parameter count gives gated f≈8d/3 instead of 4d, ignoring biases and rounding.

Residual additions let a branch learn a change to a running representation. Their derivative includes an identity route but can still shrink, grow, or cancel with branch derivatives. Post-norm instead normalizes after addition, such as Norm(X+Attention(X)); it does normalize that residual result. Pre-norm normalizes the branch input and leaves the additive residual stream unnormalized at that point. Many designs add a final normalization. Placement changes gradient and scale behavior; neither placement guarantees stable training or eliminates the need to tune warmup.

For B×N×d, LayerNorm usually centers and rescales over d independently at each token. RMSNorm rescales by root mean square without centering; affine options vary. Its simpler arithmetic does not guarantee a fixed speedup or identical quality. See Deep Learning Basics for explicit normalization examples.

5. Separate attention arithmetic, working memory, and decoding state

For dense self-attention at fixed width, QKᵀ and AV together require work proportional to N²d. Projections and FFNs contribute work proportional to Nd² for common fixed width ratios. Doubling N quadruples the token-pair attention term, but only doubles those linear-in-N terms; total latency need not quadruple.

Naive attention materializes B h N² scores or weights. For one head with N=4096 and two bytes per stored element, one matrix takes 32 MiB; 32 heads take 1 GiB for that single matrix at batch one. This is not the entire training-memory budget. Activations, gradients, parameters, optimizer state, and other temporary buffers also matter.

FlashAttention computes exact dense softmax attention algorithmically using tiles and online normalization, avoiding storage of the full score matrix in external memory. “Exact” contrasts with an attention approximation; floating-point output differences can remain. It keeps a running row maximum m, normalization sum l, and weighted value accumulator. When a tile changes the maximum, rescale previous accumulators before adding the new tile. This reproduces one global softmax rather than separately normalizing each tile.

Its benefit comes from reducing expensive data movement and improving kernel execution. It still has quadratic attention arithmetic, and recomputation in backward can add work. Kernel generation, hardware, shapes, sequence length, masks, and whether the workload is compute- or bandwidth-limited determine speedup. IO costs depend on head width and available fast memory as well as N; dividing N² by a byte count is not a complete complexity formula.

Windowed attention restricts each query to a neighborhood, giving roughly O(Nwd) attention work for fixed window w. Sparse/global patterns depend on the actual number of allowed edges; kernel overhead matters. Linear-attention methods change or approximate the computation and have feature-rank/width-dependent costs. They are distinct from an exact tiled implementation. Use the attention-complexity playground to compare modeled edge counts, then remember that those counts are not measured device latency.

KV caching changes the cost of generation

During prefill, process the prompt and compute its keys/values. During decode, append keys/values for a new token and let its query read cached entries. Caching avoids repeatedly recomputing the whole prefix, but each new full-attention query still reads a growing context. Position indices and masks must account for the existing cache.

For L layers, batch B, cache length N, h_kv key/value heads, common head width d_h, and s bytes per element:

KV bytes = 2 L B N h_kv d_h s.

The factor two is for keys and values. For L=32, B=1, N=4096, h_kv=8, d_h=128, s=2, the cache is 512 MiB, before paging/alignment/metadata overhead. With 32 KV heads it is 2 GiB. Grouped-query attention shares each KV head among multiple query heads; multi-query attention uses one KV head. Savings are the head-count ratio under matching precision and dimensions, with quality and kernel tradeoffs. Architectural modifications require trained or adapted weights; changing a configuration field alone is not equivalent.

Check yourself: double both cache length and active batch size. Solution: KV storage grows fourfold under the formula. This is separate from the quadratic prefill attention term and the approximately linear per-token attention read during decode.

For a hypothetical 70-billion-parameter model, two-byte raw weights require about 140 GB decimal. Quantized formats add scales and other metadata, and total serving memory includes cache and runtime buffers. A first-token latency target requires prompt-length and concurrency distributions, hardware, kernels, queueing, batching, and measured quality. Prefix caching can reduce repeated prefill work; speculative decoding mainly targets subsequent generation. Neither a GPU count nor a quantization label alone guarantees an SLO. Continue in Model Serving.

6. Tokens can represent more than English text

For text, tokenization chooses a mapping from strings to discrete IDs. BPE learns frequent merges from an initial symbol inventory; a unigram tokenizer selects segmentations under a learned vocabulary model. Byte coverage or byte fallback can avoid unknown symbols if configured properly, but normalization may alter exact text reversibility. Vocabulary size trades fewer tokens for larger embeddings/output layers. Measure lengths on the actual languages, code, and scripts; words are not a comparable counting unit for every writing system. A larger vocabulary does not guarantee uniform fertility or lower end-to-end latency.

For a multilingual moderation model, evaluate language, script, code-mixing, transliteration, and task-policy slices. Compare suitable pretrained backbones rather than assuming one model dominates every language. Normalization or transliteration can remove meaning and should be validated. Cascades need measured routing errors as well as speed; keyword rules are not automatically correct labels. Human-reviewed feedback needs an evaluation holdout and controls for selection bias.

A basic ViT converts a 224×224 RGB image into 196 patches of size 16×16. Each flattened patch has 768 entries and is projected to width d, with positional information and optionally a class token. The patch projection is equivalent to a stride-16, kernel-16 convolution. Patching, position, and attention are architectural choices with priors; ViT is not prior-free. Smaller patches increase token count and can preserve detail at higher cost. CNN, ViT, and hybrid performance depends on data, pretraining, augmentation, resolution, and compute. Multimodal VLMs explains how vision representations connect to language.

Sources and next lessons

• Attention Is All You Need — scaled attention, heads, masks, FFNs, and encoder–decoder structure.
• RoFormer: Enhanced Transformer with Rotary Position Embedding — rotary position construction and its relative-position identity.
• FlashAttention — exact tiled attention and IO-aware analysis.
• An Image is Worth 16×16 Words — ViT patch encoding and the paper's experimental comparisons.

Next: Embeddings & Retrieval for useful representations, Mixture of Experts for conditional feed-forward computation, and Model Serving for deployment budgets.