The transformer assembled: why Pre-LN wins, and what attention sinks actually are

Nearly every piece of a transformer has already been built somewhere in this series. The full architecture is token embeddings, then $N$ repetitions of [RMSNorm → multi-head attention → residual → RMSNorm → SwiGLU feed-forward → residual], then an output projection into logits. Attention’s forward and backward pass were derived and checked to $10^{-6}$. RMSNorm’s backward pass is a close relative of LayerNorm’s, simpler because it skips mean-centering and normalizes only by root-mean-square, the same Jacobian shape either way. The residual connection’s $1+\frac{dF}{dx}$ protection against vanishing gradients is exactly the mechanism this whole architecture leans on at every single layer. Cross-entropy’s clean $\hat y - y$ gradient is what actually kicks off backprop from the logits. Every matmul in the whole thing runs inside the roofline framework, and every one of those matmuls runs at whatever precision was chosen, with whatever stability tricks that choice requires. This post’s job isn’t to re-teach any of that. It’s to show how the pieces compose, then go past what any of them cover individually.

The pieces that are genuinely new here#

Multi-head attention splits $d_{\text{model}}$ into $H$ heads, each of dimension $d_k = d_{\text{model}}/H$, computing $H$ independent attention operations in parallel, each with its own $Q, K, V$ projections, then concatenating the $H$ outputs and applying one final projection. The reason for multiple heads instead of one large attention operation: each head can specialize in a different kind of relationship between tokens, one tracking syntactic dependency, another tracking coreference, and concatenation-then-projection lets the model combine these specialized views rather than being forced into a single, averaged one.

SwiGLU, the standard modern feed-forward block, is $(xW_1 \odot \text{SiLU}(xW_2))W_3$: two parallel matmuls, $xW_1$ and $xW_2$, one of which passes through SiLU ($\text{SiLU}(z) = z\cdot\sigma(z)$) and then gates the other via elementwise multiplication, before a final projection. The gating lets the network learn, per-dimension, how much of $xW_1$‘s signal to let through, a strictly more expressive block than a single matmul plus a fixed nonlinearity, at the cost of a third weight matrix.

Causal masking sets every entry in the attention score matrix’s upper triangle, future positions relative to each query, to $-\infty$ before softmax. Since $e^{-\infty}=0$, softmax assigns exactly zero probability to attending to future tokens, not approximately zero, exactly zero, guaranteed by the exponential function’s own mathematics rather than a separate rule bolted on afterward.

Why Pre-LN wins, derived rather than asserted#

This is stated everywhere as folklore, “pre-norm is more stable,” almost never connected to the specific Jacobian that makes it true, and it follows directly from the residual-gradient result already built.

Post-LN computes $x_{\ell+1}=\text{Norm}(x_\ell+\text{Sublayer}(x_\ell))$, normalization applied after the residual sum. Differentiating this, the gradient must pass through the normalization’s own Jacobian, the identity-minus-outer-product structure derived for LayerNorm, at every single layer, before it can reach the residual path’s protective $+1$ term.

Pre-LN computes $x_{\ell+1}=x_\ell+\text{Sublayer}(\text{Norm}(x_\ell))$, normalization applied only inside the sublayer, with the residual sum entirely outside it. The gradient with respect to $x_\ell$ has a completely unobstructed path:

$\frac{\partial x_{\ell+1}}{\partial x_\ell}=1+\frac{\partial\,\text{Sublayer}(\text{Norm}(x_\ell))}{\partial x_\ell}$

The $1$ term never has to pass through a normalization’s Jacobian at all. Post-LN forces every layer’s gradient through an extra Jacobian multiplication that Pre-LN’s architecture sidesteps entirely, which is precisely why 96-plus layer transformers use Pre-LN almost universally. It isn’t a stylistic preference, it’s the same $1+\frac{dF}{dx}$ argument from residual connections, applied one layer more carefully.

Attention sinks: a mystery that isn’t one once you know softmax’s constraint#

This is genuinely underdiscussed relative to how reliably it shows up, and the mechanism only becomes visible once softmax’s Jacobian is understood precisely, which this series already built from scratch.

The phenomenon. A small number of tokens, most commonly the very first token in a sequence, attract a disproportionate share of attention weight across many heads and layers, despite carrying little or no semantic content. This is an attention sink. Separately but correlated, a small number of hidden-state channels at specific tokens take on extreme values far outside the normal activation range, called massive activations.

The mechanism, directly off softmax’s own sum-to-one constraint. Attention layers sometimes need to perform something close to a no-op, a position where the model deliberately doesn’t want attention to update the residual stream much. But softmax’s output must sum to exactly 1 across every position in the sequence. There’s no way to make every attention weight small simultaneously; if one weight goes down, the leftover probability mass has to go somewhere. The model’s actual solution: drive the attention logits for most tokens toward $-\infty$, forcing their softmax outputs toward zero, and whatever token absorbs the leftover mass, usually the first token, since it’s present in every sequence and therefore a reliable target, becomes the sink. Not because it’s meaningful. As an unavoidable mathematical consequence of a normalization constraint that has no other release valve.

Why this connects directly back to the very first post in this series. Massive activations are exactly the outlier problem that motivated LLM.int8()‘s mixed-precision decomposition, but with a distinction that’s easy to miss and documented as a real, historical confusion in the field itself: token-wise outliers (attention sinks) and channel-wise outliers (the kind LLM.int8() addresses) are mechanistically distinct phenomena requiring different mitigation strategies. A quantization scheme built to protect specific weight channels does not automatically protect specific tokens, and vice versa. Treating every large activation as the same kind of problem is a real mistake, clarified only as recently as 2024, not a pedantic distinction.

There’s a genuinely open research direction worth naming with an honest hedge rather than certainty. Recent work carries the title “Attention Sinks Induce Gradient Sinks: Massive Activations as Gradient Regulators in Transformers,” suggesting these tokens aren’t just an inference-time quantization nuisance but may serve an active, functional role in regulating gradient flow during training. I have the paper’s title and citation context, not its full argument, so this is flagged as an open question, not a settled fact, but the connection itself, a phenomenon forced by softmax’s constraint turning out to also regulate gradients during backprop, is exactly the kind of synthesis worth sitting with rather than resolving prematurely.

RoPE, and the failure that never throws an error#

RoPE (Rotary Position Embeddings) doesn’t add a fixed positional vector to each token’s embedding. It rotates the query and key vectors in 2D subspaces by an angle proportional to position, with rotation frequency for dimension $2i$ given by $\theta_{m,i}=m\cdot\theta_{\text{base}}^{-2i/d_h}$. The key property: after rotation, the dot product between a rotated query and a rotated key depends only on their relative position $m-n$, not their absolute positions, so the model learns relationships that generalize across position rather than memorizing what happens at a specific slot.

Here’s the failure worth naming precisely because it’s silent rather than loud. A model trained with positional encodings tuned for a 4K context, deployed at 8K, sees positions its encoding scheme was never tuned for during training. Quality degrades, but there’s no error, no crash, no NaN. The model simply becomes progressively less reliable as position grows, in a way that looks like ordinary quality variance rather than a specific, nameable bug. This is exactly why RoPE’s base frequency is chosen deliberately per model rather than left at a default: Llama 3’s base frequency of 500,000 was specifically increased over Llama 2’s default to extend reliable behavior toward a 128K context target.

Real architecture numbers, checked rather than assumed#

GPT-3: 96 layers, 96 attention heads, $d_{\text{model}}=12288$. Its Adam-state memory, using the exact formula derived earlier in this series, FP32 weights plus FP32 gradients plus FP32 first and second moments, 16 bytes per parameter: $175\times10^9 \times 16 = 2.8$TB. That figure confirms this was, at the time, a genuine model-parallelism-forcing number, no single machine’s memory came remotely close.

Llama 3 405B: 126 layers, 128 attention heads, $d_{\text{model}}=16384$, RoPE base frequency 500,000, confirmed directly against Meta’s own paper. Its specific grouped-query-attention KV-group count and SwiGLU intermediate size are consistent with what’s publicly known about the model’s design, but not independently confirmed to the same standard as the numbers above, worth treating as plausible rather than verified until checked directly.

Failure modes, each traced to its mechanism#

Failure	Mechanism	Why it’s dangerous
Missing causal mask	Model attends to future tokens during training	Training loss looks perfect, the model has learned to cheat using information it won’t have at generation time, and the failure only appears at inference as incoherent generation, with no warning during training
Wrong RoPE scaling	Positional encodings out-of-distribution beyond the trained context length	Silent, gradual quality degradation, easily mistaken for ordinary variance
Missing residual connections	Vanishing gradient through many layers	The model simply never converges, diagnosable but easy to misattribute to a learning-rate or data problem instead
Post-LN instead of Pre-LN	Every layer’s gradient forced through an extra normalization Jacobian	Less catastrophic than missing residuals entirely, but measurably less stable at extreme depth

Common mistakes#

Treating “pre-norm is more stable” as folklore to accept rather than a direct consequence of where the normalization sits relative to the residual sum: the derivation above is short, and it’s worth being able to reproduce rather than just cite.

Assuming an attention sink is a bug or an artifact of bad training: it’s a forced consequence of softmax needing every row to sum to exactly 1, with no other way to express “attend to almost nothing here.”

Applying a channel-wise outlier fix (built for weights) to a token-wise outlier problem (attention sinks) and expecting it to work: they’re mechanistically distinct, and conflating them is a documented, real mistake in quantization research, not a hypothetical one.

Assuming a model degrading at long context is a data or capacity problem: if the degradation is smooth and silent rather than a hard failure, check the positional encoding scheme’s trained range before assuming anything else.

What this takes to be frontier-job-ready#

The technical axis: given a quantization scheme that works well on one model but degrades badly on another, know to ask whether the outliers involved are token-wise or channel-wise before assuming the scheme itself is broken. That distinction is exactly the diagnostic this scenario requires, and conflating the two wastes real debugging time chasing the wrong fix.

The operational axis: a model passing every training-time check while failing incoherently at deployment, the missing-causal-mask failure above, is exactly the kind of gap between “the loss curve looked fine” and “the actual behavior is wrong” that separates recognizing a real production incident from trusting a dashboard that has no way to see the actual failure mode.

The autonomy axis: the attention-sink gradient-regulation question above is presently unsettled. Whether massive activations are purely a quantization nuisance or an active part of how gradients flow during training is an open question, not a solved one, and engaging with it seriously requires being willing to hold a genuinely uncertain technical question rather than waiting for a textbook to resolve it first.

Try it yourself#

Beginner. Write out the Post-LN and Pre-LN gradient expressions for a 3-layer toy network by hand, and count how many times the gradient passes through a normalization Jacobian in each case before reaching the input layer.

Intermediate. Using the RoPE relative-position property, show explicitly that the dot product between a query at position $m$ and a key at position $n$, both rotated by RoPE, depends only on $m-n$, not on $m$ and $n$ individually.

Advanced. Using softmax’s sum-to-one constraint, construct a minimal toy example (3 to 4 tokens) where forcing one attention weight toward zero necessarily increases at least one other weight, and use it to explain concretely why a transformer cannot express “attend to nothing” without some token absorbing the leftover mass.

---

The one-sentence version: a transformer is almost entirely pieces already derived elsewhere in this series composed into one architecture, Pre-LN beats Post-LN for the same reason residual connections work at all, an unobstructed $+1$ gradient path, and attention sinks aren’t a mystery or a bug, they’re the one predictable place softmax’s own sum-to-one constraint has to send probability mass it isn’t allowed to simply discard. Every one of those tokens the model attends over started as a decision made before training even began, about how to cut text into pieces in the first place.