Embeddings and Contrastive Learning, The Transformer’s Three Attention Patterns and Where Am I? Position in Attention, from Sinusoids to RoPE built the pieces of a Transformer: token vectors, attention, and position. This post puts them into one model. The thread through it is a single vector per token position, which starts as the token’s embedding, is added to by every layer, and ends as a prediction of the next token. That vector is the residual stream, and this blog is named after it.
The layers never overwrite the stream; each reads it and adds an update. This post follows what that design means, from the arithmetic of one block to measurements of the stream inside a real model, Qwen2.5-1.5B-Instruct: how large it grows, how much each layer adds, when the prediction appears in it, and what one layer is worth.
The residual stream: every sublayer reads a normalized copy of the stream and adds its output back. Nothing on the main line but additions.
- The stream is one vector per token, changed only by addition: in Qwen2.5-1.5B-Instruct its size grows from 0.03 at the embedding to 8.02 after 28 blocks.
- Each block adds a modest update: after the first block, attention adds 0.18 to 0.46 times the stream it reads and the MLP 0.26 to 0.54 times (0.73 in the last block).
- The prediction forms late: read through the unembedding, the stream’s top token matches the final one 2% of the time after 14 blocks, 38% after 24 and 74% after 27.
- On one passage, middle blocks are each dispensable: removing any one of blocks 3 to 27 raised the loss by 0.03 to 0.30 nats; removing block 1 raised it by 9.3.
- The MLPs hold most of the weights: 75% of Qwen2.5-1.5B-Instruct’s, against 10% for attention.
Table of contents
Open Table of contents
Tokens become vectors
A model never sees text. A tokenizer splits text into a sequence of IDs from a fixed vocabulary (Byte-Pair Encoding), and an embedding matrix , one learned row per vocabulary entry, turns each ID into a vector. With a toy vocabulary of four tokens and 2-dimensional vectors:
Embedding is a table lookup, no arithmetic. Position is not added here in most current models: they rotate the queries and keys inside attention instead (Where Am I? Position in Attention, from Sinusoids to RoPE). is the residual stream’s first value at this position.
One running vector
Every component of a Transformer works on the same vector at the same position. The embedding writes its first value, every sublayer reads it and writes back by addition, and the unembedding reads the result. With standing for everything block adds:
After blocks the stream is the embedding plus every block’s update. Elhage et al. (2021), who set out this view, put it as: “All components of a transformer (the token embedding, attention heads, MLP layers, and unembedding) communicate with each other by reading and writing to different subspaces of the residual stream.”
Nothing in the architecture assigns those subspaces. Every update is added to the same vector, so two layers’ contributions can overlap, and models pack more features into the stream than it has dimensions, at the cost of some interference between them; Elhage et al. (2022) call this superposition. The stream is a shared, additive workspace, not a set of labeled slots.
One block: attention, then the MLP
A Transformer stacks identical blocks. Each block makes two read-compute-add steps:
- Attention. The stream is normalized, the normalized copy goes into attention (The Transformer’s Three Attention Patterns derives what it computes), and attention’s output is added to the un-normalized stream: . Attention is the only place positions exchange information.
- MLP. The updated stream is normalized again, passed through a small feed-forward network at each position separately, and added back: .
The normalization rescales a vector to a fixed typical size so that training stays stable; Dropout and Normalization derives LayerNorm and RMSNorm, which Qwen uses. Both sublayers follow one contract: read a normalized copy, compute, add back. Nothing in that contract depends on what is computed, which is why attention or the MLP can be replaced in a Transformer variant without changing the stream.
Where the normalization sits has changed. The 2017 Transformer (Vaswani et al.) normalized after each addition, , called Post-LN. Wang et al. (2019) and GPT-2 (Radford et al., 2019) moved it inside the branch, Pre-LN, as above, and GPT-2 added “an additional layer normalization … after the final self-attention block”. The un-normalized stream keeps growing through the layers (measured below), so the final norm brings it back to a fixed size before the unembedding. Xiong et al. (2020) showed why the placement matters: with Post-LN the norm sits on the main path, the gradients near the output are large at initialization, and training needs a careful warmup. Pre-LN leaves an unbroken additive path from the embedding to the last block. Build a Tiny GPT, End to End, and Train It on This Blog measured the swap on a small model: post-norm cost 0.06 to 0.08 bits per byte at 4 layers, with settings tuned for pre-norm. Current models use Pre-LN, most with RMSNorm; some also normalize each sublayer’s output before the add (Gemma 2 does both, OLMo 2 only the output), and the main path stays free of normalization either way.
Inside the MLP sublayer
The MLP applies the same small network to each position’s vector on its own. Write for the stream’s width. In the 2017 Transformer it expands the vector to , applies a nonlinearity , and projects back:
with a ReLU in the original, and GELU, a smooth version of it (Hendrycks and Gimpel, 2016), from GPT and BERT on. Most current models (Llama, Mistral, Qwen, DeepSeek) use a gated version, SwiGLU (Shazeer, 2020): a second projection of the input multiplies the activated one element by element, so each hidden unit’s output is scaled by a learned, input-dependent gate:
where is the sigmoid. Three matrices instead of two, so at matched cost the hidden width shrinks from to : weights, as for . Shazeer’s comparison at matched cost, on a T5-base encoder-decoder, found the gated variants better: held-out log-perplexity 1.944 with SwiGLU against 1.997 with ReLU after 65,536 steps, and 1.636 against 1.677 after 524,288. The paper offers no explanation beyond the results. Counted:
class SwiGLU(torch.nn.Module):
def __init__(self, D, F):
super().__init__()
self.gate = torch.nn.Linear(D, F, bias=False)
self.up = torch.nn.Linear(D, F, bias=False)
self.down = torch.nn.Linear(F, D, bias=False)
def forward(self, x):
return self.down(torch.nn.functional.silu(self.gate(x)) * self.up(x))
count = lambda m: sum(p.numel() for p in m.parameters())
D = 4096
with torch.device("meta"): # shapes only, no memory
gelu_mlp = torch.nn.Sequential(torch.nn.Linear(D, 4 * D, bias=False), torch.nn.GELU(),
torch.nn.Linear(4 * D, D, bias=False))
swiglu = SwiGLU(D, 11008) # Llama 1 7B: 8/3 * 4096, rounded up to a multiple of 256
attention = torch.nn.ModuleList([torch.nn.Linear(D, D, bias=False) for _ in range(4)]) # Q, K, V, output
print(f"D = {D}: GELU MLP {count(gelu_mlp) / 1e6:.1f}M weights (8 D^2 = {8 * D * D / 1e6:.1f}M), "
f"SwiGLU {count(swiglu) / 1e6:.1f}M, attention {count(attention) / 1e6:.1f}M (4 D^2)")
D = 4096: GELU MLP 134.2M weights (8 D^2 = 134.2M), SwiGLU 135.3M, attention 67.1M (4 D^2)
context 2,048: MLP 64% of the block's FLOPs, attention's projections 32%, scores and weighted sums 4%
context 8,192: MLP 57% of the block's FLOPs, attention's projections 28%, scores and weighted sums 14%
context 32,768: MLP 40% of the block's FLOPs, attention's projections 20%, scores and weighted sums 40%
context 131,072: MLP 18% of the block's FLOPs, attention's projections 9%, scores and weighted sums 73%
Llama 3 8B block: MLP 176.2M, attention 41.9M, MLP share 81%
Qwen2.5-1.5B-Instruct: 1.54B weights, D = 1536, MLP width 8960, 28 blocks
embedding (tied) 233.4M 15%
attention 154.2M 10%
MLP 1156.1M 75%
norms 0.1M 0%
At matched cost a block holds about weights, in attention’s four projections and in the MLP, so two thirds of a block is MLP. Newer models tilt further. Llama 3 8B widens its SwiGLU to and shares key and value heads across queries (grouped-query attention), so its MLP holds 81% of each block. Qwen2.5-1.5B-Instruct widens to 8,960, about , and shares each key and value head among six query heads: its MLPs hold 75% of all its weights, its attention 10%, and its embedding 15%.
The MLPs also do most of the arithmetic at ordinary context lengths. Attention’s own work, the scores and weighted sums, grows with the context: about FLOPs per token against the MLP’s , so it overtakes the MLP at about , 32,768 tokens at Llama’s width. That counts only the half of the scores the causal mask keeps; Build a Tiny GPT, End to End, and Train It on This Blog counts all of them, as the PaLM paper does, which moves the crossover to . FLOPs are not time either: how fast attention runs depends on its memory traffic as much as its arithmetic (FlashAttention).
The stream in a real model
All of this can be watched. Hooks on Qwen2.5-1.5B-Instruct read the stream as it enters each of its 28 blocks, the stream each sublayer reads, and each sublayer’s update, on one passage of 512 tokens of this blog, in 16-bit on an Apple M4 laptop GPU. One model, one passage and one run: the exact numbers would move with the text, and other models were not measured. At each point the stream is also read the way the model reads it at the end, through the final norm and the unembedding, which shows what the model would predict if it stopped there; this is the logit lens (nostalgebraist, 2020):
lens = lambda x: model.lm_head(model.model.norm(x.to(model.dtype))).float()
loss = lambda logits: F.cross_entropy(logits[0, :-1], ids[0, 1:]).item()
rms = lambda x: x[0, 1:].pow(2).mean(-1).sqrt().mean().item() # positions 1 onward
512 tokens; final next-token loss 2.175 nats; stream width 1536
the next block adds, relative to
the stream each sublayer reads:
point in the stream size (RMS) attention MLP logit-lens loss top token = final
embedding 0.03 11.87x 0.92x 107.32 2%
after block 1 0.43 0.46x 0.41x 14.87 1%
after block 2 0.55 0.32x 0.54x 13.47 1%
after block 3 0.70 0.30x 0.47x 13.69 0%
after block 4 0.82 0.25x 0.49x 13.66 1%
after block 5 0.92 0.27x 0.42x 14.07 0%
after block 6 0.96 0.25x 0.43x 12.83 1%
after block 7 1.04 0.26x 0.40x 13.08 1%
after block 8 1.08 0.24x 0.39x 12.91 1%
after block 9 1.09 0.24x 0.36x 13.49 1%
after block 10 1.11 0.28x 0.36x 13.02 1%
after block 11 1.13 0.32x 0.33x 12.91 1%
after block 12 1.21 0.24x 0.31x 12.28 1%
after block 13 1.27 0.30x 0.26x 12.27 1%
after block 14 1.38 0.18x 0.28x 11.93 2%
after block 15 1.36 0.23x 0.27x 12.14 1%
after block 16 1.38 0.23x 0.28x 12.10 2%
after block 17 1.44 0.23x 0.29x 11.65 2%
after block 18 1.55 0.25x 0.30x 11.26 4%
after block 19 1.67 0.22x 0.39x 10.42 6%
after block 20 1.79 0.32x 0.34x 9.45 14%
after block 21 2.01 0.31x 0.35x 8.76 19%
after block 22 2.44 0.23x 0.37x 7.46 23%
after block 23 3.00 0.21x 0.35x 7.03 27%
after block 24 3.72 0.19x 0.33x 5.87 38%
after block 25 4.67 0.18x 0.29x 4.63 48%
after block 26 5.51 0.21x 0.28x 4.11 59%
after block 27 6.73 0.29x 0.73x 2.94 74%
after block 28 8.02 2.17 100%
first position's stream size: 0.02 at the embedding, up to 300 (after block 12); others' at most 8
Qwen2.5-1.5B-Instruct on 512 tokens. The stream grows as blocks add to it; its prediction, read through the unembedding, appears only in the last few blocks.
Data: residual-stream-explained/stream.pyThree things show. The stream grows: its size rises from 0.03 to 8.02, fastest in the last third, which is why Pre-LN models normalize once more before the unembedding. The first position is an outlier, its stream reaching 300 by block 12, while the other positions average at most 8; such massive activations on a few tokens are common in LLMs (Sun et al., 2024) and are left out of the sizes (the losses include it).
The prediction forms late. Read through the unembedding, the stream predicts the final top token for 2% of positions or fewer through block 17, then 14% after block 20, 38% after 24 and 74% after 27; the logit-lens loss falls from about 12 to 14 nats through the middle to 2.94 after block 27 and 2.17 at the end. The middle blocks write things the unembedding does not read directly. (At the embedding the loss is 107 nats: Qwen2.5-1.5B ties its unembedding to its embedding, so an unchanged token vector predicts the token itself, rarely the next one.)
What blocks 2 to 28 add, relative to the stream each sublayer reads. Block 1’s attention, reading the tiny embedding, adds about 12 times its size and is left off.
Data: residual-stream-explained/stream.pyEach update is modest. After the first block, attention adds 0.18 to 0.46 times the stream it reads and the MLP 0.26 to 0.54 times, except the last MLP at 0.73. No block replaces the stream; each nudges it. The first block is the exception: the embedding is tiny (0.03), and block 1’s attention adds 11.87 times its size, so the stream after block 1 is mostly what block 1 wrote.
Why “residual”
The name comes from residual connections, the shortcut , introduced by He et al. (2016) for image models. Plain deep networks, past some depth, got worse on both training and test data as layers were added. He et al. checked that it was not vanishing gradients (their networks used batch normalization and had healthy gradient norms): deep plain networks were harder to optimize. The residual add fixes this, and the reason is one derivative. For a block :
with the identity. Across many blocks the derivative is a product, , and expanding it always leaves a bare : part of the gradient at the top reaches block unchanged, whatever the branches do (He et al., 2016). A block can start out adding almost nothing and still pass the signal on, so near-identity is the easy default. The identity path stays clean only if nothing sits on it after the addition, which is the case for Pre-LN.
A consequence is that a trained model tolerates losing a block. Removing each of Qwen’s blocks in turn, the update simply not added:
def skip_block(i):
def hook(module, args, kwargs, out):
if i == skip:
return args[0] if args else kwargs["hidden_states"] # the block adds nothing
return hook
block removed: loss (increase)
block 1 11.476 (+9.301)
block 2 4.589 (+2.414)
block 3 2.462 (+0.288)
block 4 2.310 (+0.135)
block 5 2.347 (+0.173)
block 6 2.387 (+0.212)
block 7 2.379 (+0.205)
block 8 2.337 (+0.163)
block 9 2.296 (+0.121)
block 10 2.289 (+0.114)
block 11 2.273 (+0.098)
block 12 2.253 (+0.078)
block 13 2.266 (+0.091)
block 14 2.296 (+0.121)
block 15 2.260 (+0.085)
block 16 2.206 (+0.031)
block 17 2.234 (+0.059)
block 18 2.282 (+0.108)
block 19 2.288 (+0.114)
block 20 2.296 (+0.121)
block 21 2.283 (+0.108)
block 22 2.405 (+0.230)
block 23 2.378 (+0.203)
block 24 2.374 (+0.200)
block 25 2.405 (+0.230)
block 26 2.475 (+0.301)
block 27 2.399 (+0.224)
block 28 2.943 (+0.768)
Removing one block at a time. The first two are essential; any one of the middle blocks costs a fraction of a nat.
Data: residual-stream-explained/stream.pyWithout block 1 the loss rises from 2.175 to 11.476 nats and without block 2 to 4.589: the early blocks turn the tiny embedding into something the rest can use. Removing any one of blocks 3 to 27 costs 0.03 to 0.30 nats, and the last block 0.77. The other blocks still read a stream that is nearly what they expect, because each block only adds to it. Gromov et al. (2025) found something similar at larger scale: removing a contiguous run of deeper layers, then briefly fine-tuning to repair the gap, kept question-answering accuracy until roughly half the layers were gone. Build a Tiny GPT, End to End, and Train It on This Blog shows the other side: without residual connections at all, a 4-layer model trained at a normal learning rate collapses to predicting the same distribution of characters whatever its input.
From the stream to a prediction
After the last block, one more normalization and one matrix multiplication turn the stream into scores over the vocabulary. The unembedding matrix projects the final stream to one score per vocabulary entry, a logit, and a softmax turns the logits into probabilities. is sometimes the embedding matrix transposed, (weight tying, Press and Wolf, 2017). Qwen2.5-1.5B ties, as does Gemma at every size; Llama 3 8B and 70B and DeepSeek-V3 learn a separate . Embedding weights grow with vocabulary times width and the rest with layers times width squared, so at scale tying saves a small share.
Continuing the toy example, take the final stream to be and unembed:
What softmax actually computes
Given logits , softmax exponentiates each and divides by the sum:
Exponentiating makes every entry positive, and dividing by the total makes them sum to 1, so the output is a probability distribution that keeps the logits’ order. On the toy logits, then shifted and scaled:
def softmax(z):
e = np.exp(z - z.max()) # subtracting the max changes nothing, and avoids overflow
return e / e.sum()
exp(z): [1.492 3.004 0.670 0.333] sum 5.499
softmax(z): [0.271 0.546 0.122 0.061]
softmax(z + 10): [0.271 0.546 0.122 0.061]
softmax(0.5 * z): [0.281 0.398 0.188 0.133]
softmax(2 * z): [0.188 0.764 0.038 0.009]
softmax(10 * z): [0.001 0.999 0.000 0.000]
Scaling the logits sharpens or flattens the distribution; shifting them changes nothing.
Data: residual-stream-explained/softmax.pyTwo properties follow from the formula:
- Only differences between logits matter. Adding a constant to
every logit multiplies numerator and denominator by , which
cancels:
softmax(z + 10)issoftmax(z). Code subtracts the largest logit first for that reason; it changes nothing and keeps from overflowing. - Scale controls sharpness. Multiplying the logits by stretches their differences, so the largest takes more of the mass, 0.999 at ; as grows the output approaches one-hot and the derivative goes to zero. flattens it toward uniform. This is why attention divides its scores by (The Transformer’s Three Attention Patterns), and what sampling temperature controls (Decoding).
The logits can be dragged below, with the scale and shift controls, to check both properties:
Softmax: drag a logit, watch the distribution
| logit z | s·z + c | exp | p | ||
|---|---|---|---|---|---|
| the | 0.4 | 0.40 | 1.492 | 0.271 | |
| cat | 1.1 | 1.10 | 3.004 | 0.546 | |
| sat | -0.4 | -0.40 | 0.670 | 0.122 | |
| . | -1.1 | -1.10 | 0.333 | 0.061 | |
| sum of exp | 5.499 | 1.000 | |||
largest probability 0.546 · its gradient factor p(1−p) = 0.2479
Multiplying every logit by s sharpens (s > 1) or flattens (s < 1) the distribution without changing the argmax.
cat is the most likely next token here, 0.546, for no reason but the toy
numbers: a prediction is whatever accumulated on the stream, projected onto
the vocabulary and normalized. Cross-Entropy and the Softmax
Gradient derives the loss trained on
these probabilities.
Where this leads
- Build a Tiny GPT, End to End, and Train It on This Blog builds real blocks, attention and SwiGLU with RMSNorm before each, trains them on this blog, and removes the residual connection and moves the normalization to see what each costs.
- How LLM Architecture Evolved, 2017–2026 follows the substitutions inside the block since 2017, and Mixture of Experts replaces the MLP with many.
- The KV Cache stores what attention needs from past positions, so generation does not recompute the stream.
Takeaways
- A Transformer’s core object is the residual stream: one vector per token position, starting at the embedding and changed only by addition. In Qwen2.5-1.5B-Instruct it grows from 0.03 to 8.02 in size over 28 blocks, which is why a final norm precedes the unembedding.
- Every sublayer follows one contract: read a normalized copy, compute, add back. Modern models normalize before each sublayer (Pre-LN), keeping the main path free of normalization.
- Each sublayer adds a modest update, about 0.2 to 0.5 times the stream it reads after the first block, and the prediction forms late: the stream’s top token matches the final one 2% of the time after 14 blocks and 74% after 27.
- The residual add’s identity term passes signal and gradient through every block. On one passage, any one middle block of Qwen could be removed for 0.03 to 0.30 nats; the first two were essential.
- The MLP works on each position alone and holds most of the weights: two thirds of a matched-cost block, 75% of Qwen2.5-1.5B-Instruct. It does most of the arithmetic until the context reaches about tokens.
- Prediction is mechanical once the stream is built: normalize, unembed, softmax.
References
Start with Elhage et al.’s A Mathematical Framework for Transformer Circuits, where the residual-stream view comes from, then He et al. for the skip connection and Xiong et al. for Pre-LN against Post-LN.
- Elhage et al., A Mathematical Framework for Transformer Circuits (Anthropic, 2021): the residual-stream view.
- Elhage et al., Toy Models of Superposition (Anthropic, 2022): more features than dimensions.
- He, Zhang, Ren and Sun, Deep Residual Learning for Image Recognition (CVPR 2016): residual connections and the degradation problem they fix.
- He, Zhang, Ren and Sun, Identity Mappings in Deep Residual Networks (ECCV 2016): the identity path, and why nothing should sit on it after the addition.
- Vaswani et al., Attention Is All You Need (NeurIPS 2017): the original block, with Post-LN.
- Radford et al., Language Models are Unsupervised Multitask Learners (2019): GPT-2, Pre-LN with a final norm.
- Wang et al., Learning Deep Transformer Models for Machine Translation (ACL 2019): pre-norm for deep Transformers.
- Xiong et al., On Layer Normalization in the Transformer Architecture (ICML 2020): why Post-LN needs warmup and Pre-LN does not.
- Shazeer, GLU Variants Improve Transformer (2020): SwiGLU and the other gated MLPs, at matched cost.
- Hendrycks and Gimpel, Gaussian Error Linear Units (GELUs) (2016): GELU.
- Touvron et al., LLaMA: Open and Efficient Foundation Language Models (2023): SwiGLU at 8/3 D (11,008 for the 7B model).
- nostalgebraist, interpreting GPT: the logit lens (2020): reading the stream through the unembedding.
- Sun et al., Massive Activations in Large Language Models (COLM 2024): the few positions whose stream grows far larger than the rest.
- Gromov et al., The Unreasonable Ineffectiveness of the Deeper Layers (ICLR 2025): removing layers from trained LLMs.
- Press and Wolf, Using the Output Embedding to Improve Language Models (EACL 2017): weight tying.