Continuous Function

Attention & Transformers

Pre-Norm Transformer Block

Where normalization sits inside a Transformer block, why attention and MLP branches must preserve residual-stream shape, and how pre-norm changes the local gradient route.

status: reviewimportance: criticaldifficulty 4/5math: graduateread: 15mlive demo
Back to Attention & TransformersNext: Decoder-Only Transformer Forward Pass (review)

Concept Structure

Pre-Norm Transformer Block

01Intuition

Start with the picture, metaphor, or geometric mechanism.

02Math

Make the objects explicit and connect them with notation.

03Code

Mirror the equations with runnable implementation details.

04Interactive Demo

Manipulate the mechanism and watch the idea respond.

5prerequisites
3next concepts
4related links

Learner Contract

What this page should let you do.

You are here becauseWhere normalization sits inside a Transformer block, why attention and MLP branches must preserve residual-stream shape, and how pre-norm changes the local gradient route.

This Attention & Transformers concept is the current idea: keep the same invariant visible across Intuition, Math, Code, Interactive Demo.

By the end4/4 sections ready | runnable code expected | live demo

Explain the mechanism, trace the main notation, and test one prediction in the live demo.

Do this firstIntuition

Read the intuition before the notation; the math should name a mechanism you already felt.

Test the linkManipulate one control and predict the visible change.Then continue to Decoder-Only Transformer Forward Pass (review)

Claim/source review status

Substantive review recorded

1/1 claims have bounded review metadata; still check caveats and source scope.Metadata-derived; review may be AI-assisted. Not a human certification.
Claims1/1 reviewed
Sources4 cited
Codeattached
Demolive
Reviewed2026-07-02
Updatedpage 2026-07-02

Learning item flow

4/4 sections readyAsk about thisResearch room
ConceptPre-Norm Transformer BlockAttention & TransformersEquationPre-Norm Transformer Block equation 1Exact equationCodePre-Norm Transformer Block runnable code 1Exact runnable codeDemoPre-Norm Transformer Block interactive demoInteractive demoClaimPre-norm and post-norm Transformer blocks preserve `(B,T,d)` residual...Exact claim checkSourceAttention Is All You NeedExact source
ConceptPre-Norm Transformer BlockAttention & Transformers
6 sources attachedLocal snapshot ready
concept:attention-transformers/pre-norm-transformer-block
CodenearbyPredictbefore revealRoomitem handoff
01IntuitionStart with the picture, metaphor, or geometric mechanism.02MathMake the objects explicit and connect them with notation.03CodeMirror the equations with runnable implementation details.04Interactive DemoManipulate the mechanism and watch the idea respond.
01

01

Intuition

Build the mental picture first so the rest of the page has something to attach to.

Section prompt

A Transformer block is easiest to understand as a residual-stream update.

The residual stream carries one vector per token. Attention writes one update into that stream. The MLP writes another update. Both updates must have the same shape as the stream, otherwise the residual addition is not legal.

The slippery question is not whether we use normalization. The question is:

Do we normalize before the branch reads from the stream, or after the branch has been added back?

Post-norm, the original Transformer layout, does this for a sublayer:

y=Norm⁡(x+F(x)).y = \operatorname{Norm}(x + F(x)).

The residual addition happens, then the sum is normalized.

Pre-norm moves normalization inside the branch:

y=x+F(Norm⁡(x)).y = x + F(\operatorname{Norm}(x)).

Now the branch still reads a normalized vector, but the residual stream has a more direct identity route through the addition. That is why pre-norm is often introduced when discussing deeper Transformer stacks and training stability.

Keep the caveat sharp: pre-norm is not a universal magic switch. It changes the local computation graph. Full training behavior also depends on depth, initialization, learning-rate schedule, normalization variant, residual scaling, optimizer, dropout, data, and task.

Vaswani et al. define the original post-norm-style sublayer output as LayerNorm(x + Sublayer(x)) and require sublayers to output d_model so residual addition works.

D2L's Transformer chapter teaches the same AddNorm idea, the positionwise feed-forward branch, and the requirement that residual inputs have matching shapes.

Xiong et al. analyze layer-normalization placement and argue that the original Post-LN placement can create large expected gradients near the output layer at initialization, while Pre-LN places normalization inside residual blocks and has better-behaved initialization gradients in their setting.

02

02

Math

Translate the story into symbols, assumptions, and a derivation you can inspect.

Section prompt
Equation 1A(⋅)∈RB×T×d,M(⋅)∈RB×T×d.A(\cdot) \in \mathbb R^{B \times T \times d}, \qquad M(\cdot) \in \mathbb R^{B \times T \times d}.Equation 2Ul=Xl+A(N(Xl)),U_l = X_l + A(N(X_l)),

Let:

  • B be batch size.
  • T be sequence length.
  • d be residual-stream width.
  • X_l in R^{B x T x d} be the residual stream entering block l.
  • N be LayerNorm or RMSNorm applied per token over the feature axis.
  • A be the attention branch.
  • M be the MLP branch.

Both branches must return residual-stream width:

A(⋅)∈RB×T×d,M(⋅)∈RB×T×d.A(\cdot) \in \mathbb R^{B \times T \times d}, \qquad M(\cdot) \in \mathbb R^{B \times T \times d}.

That shape invariant is what lets a block update the stream rather than replace the interface.

For a two-sublayer pre-norm block:

Ul=Xl+A(N(Xl)),U_l = X_l + A(N(X_l)), Xl+1=Ul+M(N(Ul)).X_{l+1} = U_l + M(N(U_l)).

The attention branch and MLP branch each read a normalized stream, but each residual addition has an explicit identity path.

For the corresponding post-norm layout:

Ul=N(Xl+A(Xl)),U_l = N(X_l + A(X_l)), Xl+1=N(Ul+M(Ul)).X_{l+1} = N(U_l + M(U_l)).

The same shapes appear, but the norm sits after each residual addition.

The local Jacobian picture explains the demo. For one sublayer F:

pre-norm:y=x+F(N(x)),∂y∂x=I+JFJN.\text{pre-norm:}\quad y = x + F(N(x)), \qquad \frac{\partial y}{\partial x} = I + J_F J_N.

There is an explicit identity term.

For post-norm:

post-norm:y=N(x+F(x)),∂y∂x=JN(I+JF).\text{post-norm:}\quad y = N(x + F(x)), \qquad \frac{\partial y}{\partial x} = J_N (I + J_F).

The residual addition is still present, but the path out of the sublayer passes through the normalization Jacobian. This local calculation is an intuition for placement, not a full convergence theorem.

03

03

Code

Keep the implementation aligned with the notation so the algorithm is legible.

Section prompt
Runnable code 1import numpy as np rng = np.random.default_rng(4) B, T, d = 1, 4, 6 x = rng.normal(size=(B, T...python
import numpy as np

rng = np.random.default_rng(4)
B, T, d = 1, 4, 6
x = rng.normal(size=(B, T, d))
Wa = rng.normal(size=(d, d)) * 0.08
W1 = rng.normal(size=(d, 2 * d)) * 0.08
W2 = rng.normal(size=(2 * d, d)) * 0.08

def ln(u, eps=1e-5):
    return (u - u.mean(-1, keepdims=True)) / np.sqrt(u.var(-1, keepdims=True) + eps)

def attn(u):
    return u @ Wa                 # toy same-shape attention branch

def mlp(u):
    return np.tanh(u @ W1) @ W2   # expand then project back to d

def pre_block(u):
    v = u + attn(ln(u))
    return v + mlp(ln(v))

def post_block(u):
    v = ln(u + attn(u))
    return ln(v + mlp(v))

def jac(fn, flat):
    eps, y0 = 1e-4, fn(flat)
    J = np.zeros((y0.size, flat.size))
    for i in range(flat.size):
        step = np.zeros_like(flat); step[i] = eps
        J[:, i] = (fn(flat + step) - y0) / eps
    return J

flat = x.reshape(-1)
Jpre = jac(lambda z: pre_block(z.reshape(B, T, d)).reshape(-1), flat)
Jpost = jac(lambda z: post_block(z.reshape(B, T, d)).reshape(-1), flat)

print("pre/post shapes:", pre_block(x).shape, post_block(x).shape)
print("trace per dim:", round(np.trace(Jpre) / flat.size, 3), round(np.trace(Jpost) / flat.size, 3))
04

04

Interactive Demo

Use direct manipulation to connect the explanation to a moving system.

Section prompt

The lab below compares one branch at a time.

Before reveal, choose which norm placement gives the residual stream a more direct identity route for gradients. After reveal, inspect the formulas, same-shape ledger, and local gradient-path summary.

The key invariant is: attention and MLP write same-shape updates into the residual stream; pre-norm moves normalization before the branch, while post-norm normalizes after the residual addition.

The demo is deliberately narrow. It does not say every model must use pre-norm, and it does not turn a local Jacobian sketch into a full training theorem.

Live Concept Demo

Explore Pre-Norm Transformer Block

The stage is code-native and interactive. Use it to test the explanation against the mechanism.

difficulty 4/5graduatecode-aligned
Demo Prediction Checkpoint

Manipulate one control and predict the visible change.

Commit to what Pre-Norm Transformer Block should make visible before reading the result.

After The First Pass

Turn the concept into an inspected object.

Once the invariant is visible in the intuition, math, code, and demo, use these panels to inspect the mechanism visually, check source support, practice the idea, and attach a grounded research question.

Mechanism Storyboard

See the idea move before the page explains it

Where normalization sits inside a Transformer block, why attention and MLP branches must preserve residual-stream shape, and how pre-norm changes the local gradient route.

Prediction open01 / Intuition
Prediction lens

Start with the picture, metaphor, or geometric mechanism.

Commit first

Before reading further, choose the kind of change Pre-Norm Transformer Block should make visible.

Visual Inquiry

Make the image answer a mathematical question

Where normalization sits inside a Transformer block, why attention and MLP branches must preserve residual-stream shape, and how pre-norm changes the local gradient route.

4/4 stages readyLive demo connected
Prediction

Which visible object should carry the first intuition?

Commit first

Pick the cue that should make Pre-Norm Transformer Block easier to reason about before the page gives the answer.

Source Grounding

Canonical references for the mechanism on this page.

paper · 2017Attention Is All You NeedVaswani et al.

Primary Transformer source for residual connections around each sublayer, post-norm LayerNorm(x + Sublayer(x)), and the same d_model output dimension needed for residual addition.

Open source
book · 2026Dive into Deep Learning: Transformer ArchitectureZhang, Lipton, Li, and Smola

Pedagogical source for AddNorm, positionwise feed-forward networks, residual shape requirements, and Transformer block assembly.

Open source
paper · 2020On Layer Normalization in the Transformer ArchitectureXiong et al.

Primary source comparing Post-LN and Pre-LN placement and arguing via mean-field analysis that Post-LN can have large initialization gradients near the output layer while Pre-LN gradients are better behaved in their setting.

Open source
documentation · 2026PyTorch torch.nn.LayerNormPyTorch

Framework source for LayerNorm over the last D dimensions and same input/output shape.

Open source
paper · 2019Root Mean Square Layer NormalizationZhang and Sennrich

Primary RMSNorm source for RMS-only rescaling. This page mentions RMSNorm as a common normalization object that can occupy the same pre-branch placement.

Open source
course-notes · 2026Stanford CS336: Language Modeling from ScratchStanford CS336 course staff

Course-level route source for implementing Transformer model architecture as part of a language-model-from-scratch path.

Open source

Claim Review

Where normalization sits inside a Transformer block, why attention and MLP branches must preserve residual-stream shape, and how pre-norm changes the local gradient route.

Status1 substantive review recorded

Claims without a substantive review badge still need exact source-support review.

Sources6 references

vaswani-2017-transformer-sublayers, d2l-2026-transformer-addnorm, xiong-2020-layernorm-transformer, pytorch-layernorm, zhang-2019-rmsnorm, cs336-language-modeling-from-scratch

Local checks4 local checks

Use equations, runnable code, and demos to check whether the source support is operational.

Substantively reviewedPre-norm and post-norm Transformer blocks preserve `(B,T,d)` residual shape, but place normalization on different sides of residual addition, changing the local gradient route.Claim metadata: source checked

Vaswani and D2L support residual-plus-normalization sublayer structure and same-shape residual requirements; PyTorch supports LayerNorm same-shape behavior; Xiong et al. support the distinction between Post-LN and Pre-LN placement and the claim that placement affects initialization gradients and warm-up behavior in their analysis.

Sources: Attention Is All You Need, Dive into Deep Learning: Transformer Architecture, On Layer Normalization in the Transformer Architecture, PyTorch torch.nn.LayerNormThis is a one-block shape and local-gradient-route witness. It does not prove universal superiority of pre-norm, does not cover every modern normalization variant, and does not replace full training-stability experiments.A bounded review summary is present; still check caveats and exact reference scope.

Checked Vaswani et al. sections 3.1/5.4, D2L Transformer AddNorm material, Xiong et al. 2020 abstract/source page, PyTorch LayerNorm docs, RMSNorm source context, and Stanford CS336 route context. Page remains review-status pending GPT Pro publication critique.

Reviewer: codex-local-source-audit; reviewed 2026-07-02
source-span-vaswani-2017-transformer-sublayerssource-span-d2l-2026-transformer-addnormsource-span-xiong-2020-layernorm-transformerEquation 1Equation 2Runnable code 1Interactive demo

Source support candidates

paper 2017Attention Is All You Need

Primary Transformer source for residual connections around each sublayer, post-norm LayerNorm(x + Sublayer(x)), and the same d_model output dimension needed for residual addition.

book 2026Dive into Deep Learning: Transformer Architecture

Pedagogical source for AddNorm, positionwise feed-forward networks, residual shape requirements, and Transformer block assembly.

paper 2020On Layer Normalization in the Transformer Architecture

Primary source comparing Post-LN and Pre-LN placement and arguing via mean-field analysis that Post-LN can have large initialization gradients near the output layer while Pre-LN gradients are better behaved in their setting.

documentation 2026PyTorch torch.nn.LayerNorm

Framework source for LayerNorm over the last D dimensions and same input/output shape.

Mechanism checks

Equation 1
A(⋅)∈RB×T×d,M(⋅)∈RB×T×d.A(\cdot) \in \mathbb R^{B \times T \times d}, \qquad M(\cdot) \in \mathbb R^{B \times T \times d}.
Equation 2
Ul=Xl+A(N(Xl)),U_l = X_l + A(N(X_l)),
Runnable code 1import numpy as np rng = np.random.default_rng(4) B, T, d = 1, 4, 6 x = rng.normal(size=(B, T...Demo stateLive mechanism probe

Practice Loop

Try the idea before it explains itself

Where normalization sits inside a Transformer block, why attention and MLP branches must preserve residual-stream shape, and how pre-norm changes the local gradient route.

Readiness0/3 checks ready
Predict

Before touching the demo, predict one visible change that should happen in Pre-Norm Transformer Block.

Hint 1

Reveal when your model needs a nudge.

Hint 2

Reveal when your model needs a nudge.

Hint 3

Reveal when your model needs a nudge.

Grounded research drawerClose
ConceptPre-Norm Transformer BlockAttention & Transformers
Runnable code comparisonPre-Norm Transformer Block runnable code 1rng = np.random.default_rng(4)Prediction before revealPre-Norm Transformer Block interactive demoManipulate one control and predict the visible change.
Grounded room questionWhat is the smallest example that makes Pre-Norm Transformer Block click without losing the math?Local snapshot ready

Research Room

Attach the question to a claim, equation, code, or demo

Pick the concept, equation, source, runnable code, claim, misconception, or demo state before asking for help. The handoff keeps that page item in context.
Next local actionNo local draft saved yet

Open the draft below to save one note and next action in this browser.

conceptAttention & Transformers

Pre-Norm Transformer Block

Attached question

What is the smallest example that makes Pre-Norm Transformer Block click without losing the math?

Local action draftNo local draft saved yetExpand only when ready to capture one local next action
Local action draft

This draft stays in this browser, attached to the selected learning item.

No local draft saved.
Evidence to inspect
  • References to inspect: attached references on this page.
  • Definition, prerequisite, and contrast concept links
  • The equation or runnable code that makes the concept operational
  • One demo state that shows the invariant instead of a slogan
What would resolve this
  • The learner can state the mechanism in their own words
  • The learner can name the prerequisite that would repair confusion
  • The learner can predict how the mechanism changes under one perturbation
Grounded AI handoff

I am working in Continuous Function's research reading room. Object: concept - Pre-Norm Transformer Block Selected item key: recorded for copy. Context: Attention & Transformers Page anchor: recorded for copy. Open question: What is the smallest example that makes Pre-Norm Transformer Block click without losing the math? Evidence to inspect: - References to inspect: attached references on this page. - Definition, prerequisite, and contrast concept links - The equation or runnable code that makes the concept operational - One demo state that shows the invariant instead of a slogan What would resolve this: - The learner can state the mechanism in their own words - The learner can name the prerequisite that would repair confusion - The learner can predict how the mechanism changes under one perturbation Answer as a careful research tutor: stay source-grounded, separate verified evidence from assumptions, name the relevant math objects, and end with one next action.

View it in context
concept/concept-notebook/attention-transformers/pre-norm-transformer-block concept:attention-transformers/pre-norm-transformer-block