Lensa ML

The Transformer Architecture

Step 1 of 6: Why Replace Recurrence?

One hop between any two words, at an n² price

RNN: reads one token after anotherAttention: every token looks at every tokenfirst → last token: RNN n − 1 hops, attention 1
Sentence length8
Hops first → last (RNN)
7
Hops first → last (attention)
1
Steps that must wait (RNN)
8
Scores per layer (attention)
64

An RNN folds the sentence into one hidden state, token by token, so information from the first of 8 tokens needs 7 hops to reach the last, getting diluted on the way, and the 8 steps can't run in parallel. The Transformer (Vaswani et al., 2017) lets every token compare itself with every other token in one step, all at once on a GPU. The price is n² comparisons per layer; drag the length up to see it.

Drag the sentence length and compare the counts. A recurrent network processes tokens one after another, so training can't run in parallel across a sequence, and the signal between distant words passes through many steps. Self-attention connects every pair of positions directly. The price is quadratic: n² scores per layer and head, which is why long-context models rely on tricks such as FlashAttention, sliding windows, or sparse attention. Attention itself already existed alongside RNNs (Bahdanau et al., 2014); the 2017 paper dropped the recurrence.

Step 2 of 6: Scaled Dot-Product Attention

Query, keys, softmax, weighted sum of values

finance →↑ naturethebankbyriverqueryoutputscore = q·k / √d_k → softmax weightthe1.10.00bank5.70.41by1.20.00the1.10.00river6.00.58
Rotate query0°
"bank" attends most to
river 0.58
Output vector
(0.24, 0.78)
Weights add to
1.00

Each token makes a query ("what am I looking for?") and a key ("what do I contain?"). The query of "bank" is scored against every key with a dot product, divided by √d_k, and softmax turns the scores into weights that add to 1. The output is the weighted sum of the value vectors: here 0.58 of "river" plus 0.41 of itself. "bank" points halfway between finance and nature, so it locks onto whichever content word the sentence has: river. The vectors are hand-set so you can read them; a trained model learns them.

Pick a sentence and a query word, then rotate the query. Scaled dot-product attention computes softmax(QKᵀ/√d_k)·V for every query at once, as two matrix multiplications. Q, K, and V are three different linear projections of the same token vectors. Here W_q = W_k = 4·I and W_v = I so the arithmetic stays visible; a trained model learns all three, in hundreds of dimensions rather than two.

Step 3 of 6: Why Divide by √dₖ

Big dot products make softmax stop learning

one query vs 6 random keys, each with 64 dimensions0.004.10.004.90.00-2.71.0014.10.00-6.70.000.4

bars = softmax weights · numbers under them = scores

Dimensions d_k64
Score spread (rms)
6.99
Largest weight
1.000
Gradient through it, w(1−w)
0.0001

With 64 dimensions a dot product of random vectors adds up 64 terms, so scores spread out like √64 ≈ 8.0. Softmax then puts 1.000 on one key and almost nothing elsewhere, and the gradient through that weight is only 0.0001: the head can barely learn. Switch on the scaling.

Raise d_k with and without the scaling. If query and key entries have unit variance, their dot product over d_k dimensions has variance d_k, so scores grow like √d_k. Softmax over large scores becomes nearly one-hot, and the gradient through the winning weight, w(1 − w), vanishes. Dividing by √d_k keeps the variance near 1. The paper uses d_k = 64 per head, so the scores are divided by 8.

Step 4 of 6: Meaning Depends on Context

One word, two vectors, and the residual stream

finance →↑ naturebank beforebank · riverbank · moneyall vectors drawn at 0.6× to fit
bank before
45°
bank in the river sentence
60°
bank in the money sentence
30°
cos between the two banks
0.87

Before attention, "bank" has one vector at 45° (0° = pure finance, 90° = pure nature) wherever it appears. After one layer it has two: 60° next to "river" and 30° next to "money", with cosine 0.87 between them. Every layer adds what attention found to the token's own vector (x + attention(x)), a "residual stream" that layers keep writing into. That's how a Transformer builds context-dependent meaning, something a fixed word embedding can't do.

Switch sentences and toggle the residual. After attention, every token's vector depends on the whole sentence, which is why Transformer outputs are called contextual embeddings, unlike word2vec's single vector per word. The residual connection adds each sublayer's output to its input, so information flows through the network as a running sum, and layer normalization keeps its scale stable. Here one head already pulls the two banks 30° apart; real models refine this over many layers and heads.

Step 5 of 6: Position Without Recurrence

Attention alone ignores word order

a "previous word" head: row = query, column = keythebankbytheriver0: the0.200.200.200.200.201: bank0.200.200.200.200.202: by0.200.200.200.200.203: the0.200.200.200.200.204: river0.200.200.200.200.20dashed outline = the true previous word
Rows that find the previous word
0 / 4
"bank" output vs unshuffled
—

Without positions, every token looks identical to this head, so each row is a flat 0.20 everywhere and it can't find the previous word. Shuffle the words: the outputs just move around with them, unchanged.

Toggle the position vectors and shuffle the sentence. Self-attention is permutation-equivariant: reorder the inputs and the outputs are reordered the same way, with nothing else changing. The paper adds sinusoidal encodings of many frequencies to the embeddings. Because a shift in position is a rotation of each sine–cosine pair, a head can learn to attend to relative offsets, as the previous-word head does here. Many newer models use rotary embeddings (RoPE), which apply the rotation inside the query and key instead.

Step 6 of 6: The Full Architecture

Encoder, decoder, and where 65M parameters live

Encoder × 6self-attentionadd & normfeed-forwardadd & normDecoder × 6masked self-attentioncross-attentionfeed-forwardadd & norm (×3)where the parameters live
■ embeddings 18.9M■ encoder attention 6.3M■ encoder FFN 12.6M■ decoder attention (self + cross) 12.6M■ decoder FFN 12.6M
d_model512
Layers N6
Heads h8
Total parameters
63.1M
Paper reports
65M
Per-head size d_k = d/h
64

Each encoder layer is attention plus a per-token feed-forward network (d_ff = 2048), each wrapped in a residual add and layer normalization. Each decoder layer adds cross-attention, where the output text's queries read the encoder's keys and values, and masks its self-attention so it can't peek at future words. These settings are the paper's base model: the count comes to 63.1M, close to the reported 65M, with a shared 37,000-token vocabulary. Note that changing the heads leaves the total unchanged: 8 heads of size 64 use the same matrices as one head of size 512, split into slices. Keep only the encoder and you have BERT; keep only the decoder and you have GPT.

Load the two paper presets, then change the width, depth, and number of heads. The base model has d_model = 512, 8 heads, 6 layers, and d_ff = 2048; the big model has 1024, 16 heads, and d_ff = 4096. With a shared 37,000-token vocabulary this count lands within a few percent of the published 65M and 213M. The base model trained for 12 hours on 8 GPUs and reached 27.3 BLEU on English–German translation. Encoder-only models such as BERT and decoder-only models such as GPT reuse the same layers.

What is the Transformer architecture?

The Transformer is a neural network architecture that processes a whole sequence at once, using attention to let every token look directly at every other token. It was introduced in the 2017 paper "Attention Is All You Need" (Vaswani et al.) for machine translation, and it is the foundation of almost every modern language model, including GPT, BERT, and T5.

Before Transformers, sequence models were mostly recurrent (RNNs and LSTMs): they read one token at a time and squeezed everything they had seen into a single hidden state. The Transformer drops recurrence entirely. That makes training highly parallel on GPUs, and it means information between two distant words travels in one step instead of many.

Key ideas

  • Self-attention lets every token read every other token directly, in parallel.
  • Attention(Q, K, V) = softmax(QKᵀ / √d_k) · V: scores from queries and keys, a weighted sum of values.
  • Multiple heads learn different relationships at about the cost of one.
  • Positional encodings put word order back in, because attention alone ignores it.
  • Residual connections and layer normalization make deep stacks trainable.
  • Encoder-only (BERT) and decoder-only (GPT) models are both built from these blocks.

Common questions

What does the feed-forward layer do?

After attention mixes information between positions, each block applies the same small two-layer network to every position separately: expand from 512 to 2,048 dimensions, apply ReLU, and project back to 512. In each encoder block this feed-forward layer holds about two-thirds of the weights.

How many parameters does the original Transformer have?

The base model in the paper has about 65 million parameters and the big model about 213 million. Modern large language models use the same building blocks at a much larger scale, with billions of parameters.

Is GPT a Transformer?

Yes. GPT keeps only the decoder half: masked self-attention and feed-forward blocks, without cross-attention, trained to predict the next token. BERT keeps only the encoder half and is trained to fill in masked words.

Why is attention O(n²)?

Every token computes a score against every other token, so a sequence of n tokens needs an n × n score matrix in each head and layer. This is why long contexts are expensive and why efficient-attention variants exist.

Related explainers

References

  1. Vaswani et al. (2017), "Attention Is All You Need", NeurIPS
  2. Bahdanau, Cho & Bengio (2014), "Neural Machine Translation by Jointly Learning to Align and Translate"
  3. Ba, Kiros & Hinton (2016), "Layer Normalization"
  4. Su et al. (2021), "RoFormer: Enhanced Transformer with Rotary Position Embedding"
  5. Ainslie et al. (2023), "GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints"
  6. He et al. (2016), "Deep Residual Learning for Image Recognition", CVPR

Last updated