The Transformer paper replaced recurrence with attention for sequence transduction. Its lasting contribution is not merely parallel training. It introduced a reusable computation pattern in which every token can construct a data-dependent view of other tokens.

These notes focus on the mechanics and the engineering consequences rather than a historical summary.

The 2017 model is an encoder–decoder Transformer with post-layer normalization. Decoder-only models, pre-norm blocks, grouped-query attention, and rotary embeddings are later choices. Collapsing all of them into “the Transformer architecture” hides exactly the implementation differences an engineer must track.

Start with a shape ledger

Let:

  • B be batch size;
  • Tq and Tk be query and key sequence lengths;
  • H be the number of query heads;
  • Dh be head dimension.

After projection and head splitting:

Q: [B, H, Tq, Dh]
K: [B, H, Tk, Dh]
V: [B, H, Tk, Dh]
scores = Q @ Kᵀ: [B, H, Tq, Tk]
output: [B, H, Tq, Dh]

Writing these shapes before code catches a large class of silent broadcasting and transpose bugs.

Scaled dot-product attention

Given input representations, learned projections produce queries Q, keys K, and values V:

Attention(Q, K, V) = softmax(QKᵀ / √dₖ) V

Each query-key dot product measures compatibility. Softmax turns those scores into weights, and the weighted value sum produces the output for each query position.

The division by √dₖ matters because dot-product variance grows with dimension. Large unscaled scores push softmax into saturated regions with very small gradients.

The mask is applied to scores before softmax. A causal mask disallows key_position > query_position; a padding mask excludes absent tokens. Boolean mask polarity differs between framework APIs, so test the exact function rather than assuming true always means “blocked.”

Stable softmax subtracts the row maximum:

m = max(scores)
p = exp(scores - m) / sum(exp(scores - m))

An all-masked row needs an explicit policy. Otherwise -∞ - (-∞) can produce NaN.

Self-attention is contextual routing

In self-attention, queries, keys, and values come from the same sequence. A token’s output is no longer its isolated embedding; it is a learned mixture of values from positions relevant to that token.

Masks constrain this routing:

  • padding masks prevent attention to nonexistent tokens;
  • causal masks prevent a decoder position from using future tokens;
  • application-level masks can represent structural constraints, although modern implementations must preserve efficient kernel paths.

Why multiple heads?

One attention operation creates one compatibility space. Multi-head attention projects the input into several lower-dimensional spaces, applies attention independently, concatenates the results, and projects again.

headᵢ = Attention(QWᵢQ, KWᵢK, VWᵢV)
MultiHead = Concat(head₁ ... headₕ)Wᴼ

Heads can represent different relationships, but avoid treating every visualization as a clean human-interpretable feature. The architectural benefit is multiple learned routing subspaces, not guaranteed semantic labels.

Position has to enter somewhere

Attention alone is permutation-equivariant: without position information, reordering tokens reorders outputs but does not otherwise change the computation.

The original model added sinusoidal positional encodings. Their frequencies provide a deterministic position signal and allow relative offsets to be represented through linear relationships. Later systems use learned absolute positions, relative biases, or rotary position embeddings, each with different extrapolation and implementation properties.

Complexity is both strength and limit

For sequence length n and representation dimension d, dense attention requires an n × n score matrix. Its sequence-length cost is O(n²d), while recurrence processes positions sequentially.

For moderate lengths, attention’s parallelism is highly effective on accelerators. For long contexts, the quadratic score matrix drives memory and compute pressure. This motivates sparse patterns, linear approximations, recurrence or state-space alternatives, and optimized exact-attention kernels.

An optimization such as FlashAttention does not change the mathematical result. It changes how the computation moves through memory, using tiling and recomputation to avoid materializing the entire score matrix in slow memory.

Online softmax across tiles

If keys are processed in blocks, each query row carries a running maximum m, normalization sum l, and partial output o. When a new block has maximum m_new, rescale the old accumulator into the new numerical frame:

m_next = max(m, m_new)
l_next = exp(m - m_next) * l + sum(exp(scores_new - m_next))
o_next = exp(m - m_next) * o + exp(scores_new - m_next) @ V_new

After the final tile, divide o by l. This is exact up to floating-point behavior; it is not sparse or approximate attention. FlashAttention still performs quadratic dense-attention FLOPs while reducing expensive high-bandwidth-memory traffic.

Residual streams and normalization

Each attention and feed-forward sublayer is wrapped with residual connections and normalization. The feed-forward network applies the same nonlinear transformation independently to each position.

Modern models often change normalization placement and activation functions, but the separation remains useful:

  • attention exchanges information across positions;
  • the feed-forward block transforms features within each position;
  • the residual stream carries and combines updates across layers.

Training and decoding are different workloads

Training and prompt prefill process many query positions in parallel. Autoregressive decoding produces one new query position per step. A KV cache avoids recomputing prior key/value projections, but cache memory grows with sequence length, layers, KV heads, and head dimension.

For one layer, ignoring metadata and alignment:

KV bytes ≈ 2 × B × Tk × Hkv × Dh × bytes_per_element

The leading 2 is for keys and values. Multi-query and grouped-query attention reduce Hkv relative to query heads, trading architectural behavior for lower cache bandwidth and memory.

Framework dispatch is part of correctness

PyTorch’s scaled_dot_product_attention can dispatch to different CUDA or math backends based on device, dtype, shape, mask, and other constraints. Do not claim FlashAttention is active because the code calls one API.

Two production footguns from the official API:

  • dropout follows dropout_p; callers must pass 0.0 during evaluation explicitly;
  • fused backends have input limitations and can produce small numerical differences from the math backend.

Benchmark the actual shapes, masks, dtype, training/inference mode, device, and selected backend. Include warmup and device synchronization; separate forward, backward, prefill, and decode measurements.

Engineering implications today

The paper suggests several durable lessons:

  1. Architecture follows hardware. Parallelizable computation helped make the design practical, not only mathematically elegant.
  2. Memory movement matters. A mathematically identical kernel can change feasible context length and serving cost.
  3. Masks define information flow. Small masking mistakes can become correctness or data-leakage failures.
  4. Sequence length is a capacity dimension. Context has latency and memory cost even when an API makes it look like a text field.
  5. A paper is a starting point. Production systems also need batching, caching, quantization, evaluation, observability, and safety boundaries.

Reading research well means separating the mathematical contribution, experimental evidence, implementation assumptions, and later improvements. That makes the paper useful as engineering knowledge instead of a citation to recognize.

References