What this page is for
The previous article, Neural Language Models Before Transformers, ends with the problem Transformers were built to solve: recurrent models carry the past through a sequential hidden state, while causal CNNs parallelize training but communicate through a fixed receptive field.
A Transformer takes a different route. Every token builds a query about what it needs, a key describing what it offers, and a value containing the information that can be retrieved. Attention compares queries with keys and uses those scores to mix values. The rest of the architecture makes that operation trainable at scale: multiple heads, residual paths, normalization, feed-forward networks, positional information, and repeated blocks.
This page reconstructs that system from the tensor shapes upward. The main reference is Vaswani et al. (2017), Attention Is All You Need. For the two major language-model branches, it also points to BERT for encoder-only modelling and OpenAI's Generative Pre-Training paper for the decoder-only GPT line.
The implementation below intentionally uses a small decoder-only language model because causal next-token prediction makes every moving part visible in one compact program.
1. Start with the tensor contract: (B, T, C)
Most confusion disappears once every operation is attached to a shape.
B — batch size: independent sequences processed together.
T — sequence length: the number of token positions in this forward pass.
C — channel or model dimension: usually called dmodel or nembd.
If token IDs have shape (B, T), an embedding table maps each integer token to a vector:
$ (B,T) \rightarrow (B,T,C) $
The Transformer stack preserves (B, T, C). Attention mixes information across the T dimension. The feed-forward network transforms each token across the C dimension. Residual connections require the block to return the same shape it received.
A useful mental split is therefore:
Operation Main job Shape in → out Token embedding token ID → vector (B,T) → (B,T,C) Self-attention communication between positions (B,T,C) → (B,T,C) Feed-forward network computation inside each position (B,T,C) → (B,T,C) Vocabulary projection hidden state → token logits (B,T,C) → (B,T,V)
V here is vocabulary size, not the value matrix used inside attention.
2. Self-attention: queries, keys, values
For one attention head, three learned linear projections turn the same input X into queries, keys, and values:
$ Q=XWQ, \qquad K=XWK, \qquad V=XWV $
If one head has dimension D, then:
$ X:(B,T,C), \qquad Q,K,V:(B,T,D) $
The query at position i is compared with every permitted key. The raw compatibility matrix is:
$ S = \frac{QK^T}{\sqrt{D}} $
The last two dimensions are what multiply:
$ (B,T,D) @ (B,D,T) \rightarrow (B,T,T) $
So the attention-score matrix is not a hidden-state tensor. It is a token-to-token matrix. Row i contains the scores produced by query token i against the keys at all positions.
The division by $\sqrt{D}$ is the scaling introduced in the original Transformer. Without it, dot products grow in magnitude as the head dimension increases, pushing softmax toward very sharp distributions and weaker gradients.
After masking, softmax converts each row to weights:
$ A = \operatorname{softmax}(S) $
and the output is the weighted sum of values:
$ Y = AV $
with shapes:
$ (B,T,T) @ (B,T,D) \rightarrow (B,T,D) $
A single causal head in PyTorch
This is the entire mathematical core of causal self-attention.
3. Why decoder attention needs a causal mask
A next-token language model must not use future tokens to predict an earlier token. During training we already know the whole target sequence, so without a mask the model could cheat.
For a sequence of length four, the permitted attention pattern is lower triangular:
Future score entries are replaced with -inf before softmax. Their softmax probability then becomes zero.
An encoder such as BERT normally does not use this causal mask: a token may use context on both sides. A GPT-style decoder-only model does use it because position t is trained to predict the next token from the prefix available up to t.
4. Multi-head attention: split the channel dimension
One head creates one learned similarity space. Multi-head attention runs several such spaces in parallel.
If the model dimension is C and there are H heads, the usual head dimension is:
$ D = \frac{C}{H} $
For example, C=384 and H=6 gives D=64. After projection, it is convenient to expose the head axis explicitly:
$ (B,T,C) \rightarrow (B,H,T,D) $
Each head independently produces an attention matrix (B, H, T, T). The resulting head outputs (B, H, T, D) are concatenated back to (B, T, C) and passed through an output projection.
The heads are not assigned fixed linguistic jobs. They are simply separate learned projections, giving the model several subspaces in which different token relationships can become useful.
5. Attention alone does not know token order
A bare self-attention operation has no intrinsic notion that token 2 came after token 1. The model therefore needs positional information.
The original Transformer adds deterministic sinusoidal position vectors to token embeddings. Learned absolute position embeddings are another common choice. GPT-style models later moved toward relative schemes, and many modern LLMs use Rotary Position Embedding (RoPE), which rotates query and key coordinates as a function of position rather than simply adding a position vector once to the input.
This supplied SVG depicts the additive-position convention: token and position vectors are combined before the encoder stack. That is appropriate for sinusoidal or learned absolute embeddings. RoPE is different: positional information is applied to queries and keys inside attention.
So the statement “positional encoding is added once before the Transformer” is not universal. It describes additive positional embeddings, not every modern Transformer.
6. The encoder block: communication, then per-token computation
An encoder block has two major learned sublayers:
multi-head self-attention, where tokens exchange information;
position-wise feed-forward network, where each token is transformed independently with the same MLP.
Residual connections route the original representation around each sublayer. Layer normalization stabilizes the scale and optimization dynamics; the normalization operation itself comes from Ba, Kiros, and Hinton (2016).
The supplied encoder SVG uses the modern pre-norm arrangement: normalize, apply a sublayer, then add the residual. The original 2017 Transformer used post-norm (“Add & Norm”). Both are Transformers; the ordering is an implementation choice with important optimization consequences.
A pre-norm block can be written compactly as:
For a bidirectional encoder, the structure is the same but self-attention does not apply the causal future mask.
7. Encoder, decoder, and decoder-only are three different layouts
“Transformer” describes a family of block structures, not one single input/output topology.
Architecture Attention pattern Typical objective Canonical example Encoder-only bidirectional self-attention masked-token / representation learning BERT Decoder-only causal self-attention next-token prediction Generative Pre-Training Encoder-decoder encoder self-attention + decoder causal self-attention + cross-attention sequence-to-sequence prediction Attention Is All You Need
The original Transformer decoder has a third sublayer: cross-attention over encoder outputs. A modern GPT-style decoder-only block does not, because there is no separate encoder stream.
This supplied SVG is best read as a decoder-only language-model block : masked self-attention followed by an FFN. A sequence-to-sequence Transformer decoder would also contain cross-attention to encoder states.
That distinction is why “BERT is the encoder and GPT is the decoder” is a useful shortcut, but not a literal claim that they are two halves of one model.
8. Assemble a tiny GPT-style language model
A decoder-only language model adds four pieces around the repeated blocks:
token embeddings;
positional information;
a stack of causal Transformer blocks;
a final projection from C hidden channels to vocabsize logits.
The supplied full-model SVG shows the high-level decoder-only path. Its additive positional-encoding and embedding-scaling details follow the original Transformer style; a GPT implementation can instead use learned positions or RoPE.
Here is a complete small model using learned absolute positions because they keep the implementation explicit:
The model never outputs words directly. It outputs logits, one score per vocabulary item at every sequence position.
9. Training: one shifted sequence creates many next-token examples
For causal language modelling, targets are the input shifted by one token.
If tokenids and targets are both (B, T), the model produces logits (B, T, V). Cross-entropy compares each position's vocabulary distribution with the correct next-token ID.
Training is parallel across all T positions because the complete target sequence is already available and the causal mask enforces what each position is allowed to see. This is the crucial training advantage over recurrence: token t does not have to wait for a recurrent hidden state produced at token t-1.
The computational trade-off is attention's token-to-token interaction. Standard full attention forms a T × T score structure per head, so its arithmetic scales quadratically with sequence length.
10. Generation: training is parallel, decoding is autoregressive
At inference time, a decoder-only language model predicts one new token, appends it, then repeats.
The important asymmetry is:
training: all target positions can be evaluated together under a causal mask;
generation: the next token does not exist yet, so decoding remains sequential across newly generated positions.
This is where inference systems introduce a KV cache.
11. KV cache: do not recompute the past at every decoding step
In naive generation, the model reprocesses the whole prefix every time a token is appended. But in causal attention, the key and value vectors for old tokens do not change merely because one new token was added.
A KV cache stores the past keys and values for every decoder layer. On the next step:
compute the new token's query, key, and value;
append the new key/value to the cache;
compare the new query with all cached keys;
mix the cached values;
predict the next token.
Conceptually, without a cache:
With a cache:
The cache trades memory for much less repeated computation during autoregressive decoding. It does not make long context free: the new query still has to interact with the permitted past keys, and the cache itself grows with sequence length.
12. FlashAttention changes the execution, not the attention equation
A common mistake is to describe FlashAttention as a different approximation to attention. The original FlashAttention paper is explicitly about computing exact attention with an IO-aware algorithm.
The mathematical result remains:
$ \operatorname{softmax}\left(\frac{QK^T}{\sqrt D}\right)V $
The implementation changes how tiles of Q, K, and V move between high-bandwidth memory and fast on-chip memory, avoiding materializing the full attention matrix in slow memory. That can dramatically reduce memory traffic and practical memory use.
FlashAttention therefore does not turn ordinary dense attention's arithmetic into a linear-time algorithm. It makes the same attention computation much more hardware-efficient.
Modern PyTorch can dispatch to optimized scaled-dot-product attention kernels:
For learning, it is still worth implementing the score matrix manually once. For production, use the optimized primitive unless you have a reason not to.
13. Quantization is a model-storage/computation technique, not an attention mechanism
Quantization reduces the precision used to represent weights and sometimes activations. For example, LLM.int8() studies 8-bit matrix multiplication for large Transformer inference, while QLoRA uses a frozen 4-bit quantized base model to make fine-tuning far more memory-efficient.
It helps to keep three optimizations conceptually separate:
KV cache: reuse past keys and values during autoregressive generation.
FlashAttention: compute attention with better memory movement.
Quantization: represent/model matrix operations at lower precision to reduce memory and often improve throughput.
They can be used together because they attack different bottlenecks.
14. RoPE: the modern positional idea worth understanding next
Additive positional embeddings give each token representation a position signal. RoPE instead transforms query and key coordinates with position-dependent rotations. The dot product then naturally depends on relative position.
That is why RoPE belongs inside the attention story, not simply in an “embedding plus position vector” box. The original RoFormer paper is the primary reference.
This site has a separate implementation-level article for it:
ROTARY POSITION EMBEDDING →
15. Common Transformer misconceptions
“Every input must have exactly the same fixed number of tokens”
No. A model has a maximum context it supports under a given configuration, but individual inputs can be shorter. Padding is commonly used to batch different sequence lengths efficiently; it is not a mathematical requirement that every example always contain exactly, say, 512 tokens.
“Attention is the whole Transformer”
No. Attention handles token-to-token communication. The FFN supplies substantial per-token nonlinear computation, residual paths preserve and update representations, normalization stabilizes optimization, and positional mechanisms encode order.
“The attention matrix is the model's memory”
Not in the persistent sense. The T × T attention weights are computed for a forward pass from the current activations. A KV cache preserves keys and values across generation steps, but that is an inference cache, not learned long-term memory.
“BERT and GPT use different attention equations”
The scaled dot-product operation is the same core mechanism. The crucial difference is what positions each token is allowed to attend to and the training objective. BERT uses bidirectional encoder attention; GPT uses causal decoder-only attention.
“FlashAttention is approximate or changes model behaviour”
The original FlashAttention algorithm computes exact attention while reducing memory traffic. Numerical implementations can differ at floating-point precision, but the intended attention operation is unchanged.
“A head corresponds to one human-interpretable concept”
Not necessarily. Heads are learned subspaces. Some develop interpretable patterns, but the architecture does not assign a semantic job such as “syntax head” or “name head” in advance.
16. Shape reference
For a decoder-only model with batch B, sequence length T, model width C, number of heads H, head width D=C/H, and vocabulary size V:
Quantity Shape token IDs (B, T) token representations (B, T, C) Q, K, V after splitting heads (B, H, T, D) attention scores (B, H, T, T) attention weights (B, H, T, T) per-head weighted values (B, H, T, D) concatenated attention output (B, T, C) FFN output (B, T, C) vocabulary logits (B, T, V)
The two formulas worth being able to recover without notes are:
$ D = \frac{C}{H} $
and
$ \operatorname{Attention}(Q,K,V)= \operatorname{softmax}\left(\frac{QK^T}{\sqrt D}+M\right)V $
where M is the mask: zero/no restriction for permitted positions and effectively -∞ for forbidden positions before softmax.
17. Papers to read in order
Vaswani et al. (2017), Attention Is All You Need. The encoder-decoder Transformer, scaled dot-product attention, multi-head attention, sinusoidal positions, residual connections, and position-wise FFNs.
Devlin et al. (2018), BERT. The encoder-only branch: deep bidirectional Transformer representations and masked-language-model pre-training.
Radford et al. (2018), Improving Language Understanding by Generative Pre-Training. The early GPT decoder-only language-model path.
Su et al. (2021), RoFormer / RoPE. Rotary positional information applied to query/key geometry.
Dao et al. (2022), FlashAttention. Exact attention reorganized around GPU memory hierarchy and IO cost.
Dettmers et al. (2022), LLM.int8(). Lower-precision Transformer matrix multiplication for large-model inference.
Dettmers et al. (2023), QLoRA. 4-bit base-model quantization combined with parameter-efficient fine-tuning.
The first three explain the architecture family. RoPE and FlashAttention explain two ideas that are central to modern implementations without changing the basic Transformer abstraction.
What to retain
A Transformer block alternates communication (attention) with per-token computation (FFN).
The basic state tensor is (B, T, C).
One head projects to Q, K, and V; QKᵀ creates a (T,T) token-relation matrix.
Multi-head attention splits C into H heads of width D=C/H, then concatenates them back to C.
Encoder attention is usually bidirectional; decoder-only language modelling uses a causal mask.
Residual connections preserve a direct representation path around attention and FFN sublayers.
Positional information is essential, but additive embeddings and RoPE insert it differently.
Training evaluates many positions in parallel; generation is still autoregressive.
KV caching, FlashAttention, and quantization solve different inference bottlenecks.
If these points and the shape table are clear, the larger Transformer variants stop being separate architectures to memorize. They become modifications of one computational skeleton.