AI 4200 · Deep Learning

Residual connections and normalization

Keeping the forward and backward variance under control as networks get deep

MLPXₗ = F(Xₗ₋₁)
→ResidualXₗ₋₁ + F(Xₗ₋₁)
→ScaledXₗ₋₁ + αₗF(Xₗ₋₁)
→Pre-normXₗ₋₁ + F(Norm(Xₗ₋₁))
→Final normNorm(XL)

All slides, one per page

Quiz inbound

What to review from the last few lectures

Expect questions relating to
- Backpropagation
- Why deep linear models collapse to shallow
- Why we init with 1/k variance in a FFN
- Why sigmoid gradients vanish (both problems)
- Why tanh gradients vanish
- Why ReLU is still problematic even though gradients don’t vanish

Residual connections won’t be on the quiz.

X

· W₁ᵀ

ReLU

· W₂ᵀ

ReLU

· W₃ᵀ

ReLU

· W₄ᵀ

Output

· W₂ᵀ

ReLU

· W₁ᵀ

+

· W₄ᵀ

ReLU

· W₃ᵀ

+

· W₆ᵀ

ReLU

· W₅ᵀ

+

· W₈ᵀ

ReLU

· W₇ᵀ

+

X₀

X₄

Residual MLP vs MLP

A plain MLP replaces its input at every layer; a residual MLP adds to a running sum

MLP

Residual MLP

residual stream (running sum)

X₁

X₂

X₃

F₁

F₂

F₃

F₄

X₄ = X₀ + F₁(X₀) + F₂(X₁) + F₃(X₂) + F₄(X₃)

A residual MLP keeps a running sum: each block adds to the stream instead of replacing it.

· W₂ᵀ

ReLU

· W₁ᵀ

+

· W₄ᵀ

ReLU

· W₃ᵀ

+

· W₆ᵀ

ReLU

· W₅ᵀ

+

· W₈ᵀ

ReLU

· W₇ᵀ

+

X₀
e×f

X₄
e×f

Forward Pass

Residual stream (flows left to right)

n×f

f×n

e×f

P̂₁

e×n

Ŷ₁

e×n

X₀

e×f

n×f

f×n

e×f

P̂₂

e×n

Ŷ₂

e×n

X₁

e×f

n×f

f×n

e×f

P̂₃

e×n

Ŷ₃

e×n

X₂

e×f

n×f

f×n

e×f

P̂₄

e×n

Ŷ₄

e×n

X₃

e×f

X₁
e×f

X₂
e×f

X₃
e×f

Block by block

X₁ = X₀ + ReLU(X₀W₁ᵀ) W₂ᵀ
X₂ = X₁ + ReLU(X₁W₃ᵀ) W₄ᵀ
X₃ = X₂ + ReLU(X₂W₅ᵀ) W₆ᵀ
X₄ = X₃ + ReLU(X₃W₇ᵀ) W₈ᵀ

Xₗ = Xₗ₋₁ + ReLU(Xₗ₋₁W₂ₗ₋₁ᵀ) W₂ₗᵀ

Ŷₗ = Xₗ₋₁W₂ₗ₋₁ᵀ (ReLU input) · P̂ₗ = ReLU(Ŷₗ) (ReLU output)

X: e×f (examples × features) · W: neurons × features, so W₂ₗ₋₁ is n×f and W₂ₗ is f×n

Every block reads Xₗ₋₁ and adds an e×f update back onto the stream.

· W₂ᵀ

ReLU

· W₁ᵀ

+

· W₄ᵀ

ReLU

· W₃ᵀ

+

· W₆ᵀ

ReLU

· W₅ᵀ

+

· W₈ᵀ

ReLU

· W₇ᵀ

+

X₀
σ²

X₄
16σ²

Variance Forward Pass

Variance along the residual stream (σ² = Var(X₀))

σW²
= 1/n

E[P̂₁²] =
½·2σ²

σW²
= 2/f

σ²

P̂₁

σ²

Ŷ₁

2σ²

X₀

σ²

σW²
= 1/n

E[P̂₂²] =
½·4σ²

σW²
= 2/f

2σ²

P̂₂

2σ²

Ŷ₂

4σ²

X₁

2σ²

σW²
= 1/n

E[P̂₃²] =
½·8σ²

σW²
= 2/f

4σ²

P̂₃

4σ²

Ŷ₃

8σ²

X₂

4σ²

σW²
= 1/n

E[P̂₄²] =
½·16σ²

σW²
= 2/f

8σ²

P̂₄

8σ²

Ŷ₄

16σ²

X₃

8σ²

X₁
2σ²

X₂
4σ²

X₃
8σ²

Var(X₁) = σ² + σ² = 2σ²

Var(X₂) = 2σ² + 2σ² = 4σ²

Var(X₃) = 4σ² + 4σ² = 8σ²

Var(X₄) = 8σ² + 8σ² = 16σ²

Log-scale plot titled Part 3: the residual stream variance grows unbounded. Representation variance against depth from 1 to about 250. Plain (blue) and BatchNorm (orange) stay flat near 10 to the 0; Residual with no norm (green) climbs steadily, passing about 10 to the 21 near depth 30 and reaching about 10 to the 116 at depth 256.

Variance diverges to infinity

For a block with input variance σₓ²: Var(Ŷ) = f·(2/f)·σₓ² = 2σₓ² · E[P̂²] = ½·2σₓ² = σₓ²

Var(block out) = n·(1/n)·E[P̂²] = σₓ²

Var(Xₗ) = Var(Xₗ₋₁) + Var(block out) = 2·Var(Xₗ₋₁), so Var(Xₗ) = 2ˡσ²

With no control, the stream’s variance doubles every block and diverges.

Xₗ

· Winᵀ

ReLU

· Woutᵀ

+

ReLU

Residual provides gradient identity path

Residual connections

Xₗ₊₁ = Xₗ + F(Xₗ)

F(Xₗ) = A(Xₗ Winᵀ) Woutᵀ

Variance of Xₗ directly passes through to next layer (no attenuation)

∂Xₗ₊₁/∂Xₗ = 1 + ∂Xₗ₊₁/∂Fₗ · ∂Fₗ/∂Xₗ

Gradient now has an identity path

∂L/∂Xₗ = ∂L/∂Xₗ₊₁ · ∂Xₗ₊₁/∂Xₗ

= ∂L/∂Xₗ₊₁ + ∂L/∂Xₗ₊₁ · ∂Xₗ₊₁/∂Fₗ · ∂Fₗ/∂Xₗ

Often ignore W_out. Purpose is to rescale output before residual. Important but doesn’t change the gradient or statistics.

Xₗ₊₁

skip
connection
(identity)

F(Xₗ)

Xₗ + F(Xₗ)

The 1 in the Jacobian gives the gradient a path that skips the block.

· W₂

⊙ 𝟙(Ŷ₁>0)

· W₁

+

· W₄

⊙ 𝟙(Ŷ₂>0)

· W₃

+

· W₆

⊙ 𝟙(Ŷ₃>0)

· W₅

+

· W₈

⊙ 𝟙(Ŷ₄>0)

· W₇

+

∂L/∂X₀
e×f

∂L/∂X₄
e×f

Backward Pass

Gradient stream (flows right to left)

∂L/∂W₂
f×n

∂L/∂W₁
n×f

∂L/∂X₁

e×f

∂L/∂P̂₁

e×n

∂L/∂Ŷ₁

e×n

e×f

∂L/∂W₄
f×n

∂L/∂W₃
n×f

∂L/∂X₂

e×f

∂L/∂P̂₂

e×n

∂L/∂Ŷ₂

e×n

e×f

∂L/∂W₆
f×n

∂L/∂W₅
n×f

∂L/∂X₃

e×f

∂L/∂P̂₃

e×n

∂L/∂Ŷ₃

e×n

e×f

∂L/∂W₈
f×n

∂L/∂W₇
n×f

∂L/∂X₄

e×f

∂L/∂P̂₄

e×n

∂L/∂Ŷ₄

e×n

e×f

∂L/∂X₁
e×f

∂L/∂X₂
e×f

∂L/∂X₃
e×f

Block by block

∂L/∂X₃ = ∂L/∂X₄ + [(∂L/∂X₄ W₈) ⊙ 𝟙(Ŷ₄>0)] W₇
∂L/∂X₂ = ∂L/∂X₃ + [(∂L/∂X₃ W₆) ⊙ 𝟙(Ŷ₃>0)] W₅
∂L/∂X₁ = ∂L/∂X₂ + [(∂L/∂X₂ W₄) ⊙ 𝟙(Ŷ₂>0)] W₃
∂L/∂X₀ = ∂L/∂X₁ + [(∂L/∂X₁ W₂) ⊙ 𝟙(Ŷ₁>0)] W₁

From the forward pass: Ŷₗ = Xₗ₋₁W₂ₗ₋₁ᵀ (ReLU input) · P̂ₗ = ReLU(Ŷₗ) (ReLU output)

∂L/∂P̂ₗ = ∂L/∂Xₗ W₂ₗ · ∂L/∂Ŷₗ = ∂L/∂P̂ₗ ⊙ ∂P̂ₗ/∂Ŷₗ, with ∂P̂ₗ/∂Ŷₗ = 𝟙(Ŷₗ > 0) (both e×n)

∂L/∂W₂ₗ = (∂L/∂Xₗ)ᵀ P̂ₗ (f×n) · ∂L/∂W₂ₗ₋₁ = (∂L/∂Ŷₗ)ᵀ Xₗ₋₁ (n×f)

On the way back, the stream hands ∂L/∂Xₗ straight to the block before, and each block adds its own term.

· W₂

⊙ 𝟙(Ŷ₁>0)

· W₁

+

· W₄

⊙ 𝟙(Ŷ₂>0)

· W₃

+

· W₆

⊙ 𝟙(Ŷ₃>0)

· W₅

+

· W₈

⊙ 𝟙(Ŷ₄>0)

· W₇

+

∂L/∂X₀
16σ²

∂L/∂X₄
σ²

Variance Backward Pass

Gradient variance along the stream (σ² = Var(∂L/∂X₄))

σW²
= 1/n

½·8(f/n)σ²

σW²
= 2/f

∂L/∂X₁

8σ²

∂L/∂P̂₁

8(f/n)σ²

∂L/∂Ŷ₁

4(f/n)σ²

8σ²

σW²
= 1/n

½·4(f/n)σ²

σW²
= 2/f

∂L/∂X₂

4σ²

∂L/∂P̂₂

4(f/n)σ²

∂L/∂Ŷ₂

2(f/n)σ²

4σ²

σW²
= 1/n

½·2(f/n)σ²

σW²
= 2/f

∂L/∂X₃

2σ²

∂L/∂P̂₃

2(f/n)σ²

∂L/∂Ŷ₃

(f/n)σ²

2σ²

σW²
= 1/n

½·(f/n)σ²

σW²
= 2/f

∂L/∂X₄

σ²

∂L/∂P̂₄

(f/n)σ²

∂L/∂Ŷ₄

½(f/n)σ²

σ²

∂L/∂X₁
8σ²

∂L/∂X₂
4σ²

∂L/∂X₃
2σ²

Var(∂L/∂X₀) = 8σ² + 8σ² = 16σ²

Var(∂L/∂X₁) = 4σ² + 4σ² = 8σ²

Var(∂L/∂X₂) = 2σ² + 2σ² = 4σ²

Var(∂L/∂X₃) = σ² + σ² = 2σ²

Log-scale plot titled Part 3: residual trades vanishing for exploding gradients. First-layer gradient RMS against depth from 1 to about 250. Plain (blue) stays near 10 to the 0 and slowly shrinks (vanishing); Residual with no norm (green) grows from about 10 to the 0 to about 10 to the 56 at depth 256 (exploding).

Variance diverges to infinity

For a block whose incoming gradient ∂L/∂Xₗ has variance v: Var(∂L/∂P̂) = f·(1/n)·v = (f/n)v · Var(∂L/∂Ŷ) = ½·(f/n)v

Var(block out) = n·(2/f)·½(f/n)v = v

Var(∂L/∂Xₗ₋₁) = Var(∂L/∂Xₗ) + Var(block out) = 2·Var(∂L/∂Xₗ), so Var(∂L/∂X₀) = 2⁴σ²

The residual stream trades vanishing gradients for exploding ones.

× α₁

F₁

+

× α₂

F₂

+

× αₗ

Fₗ

+

X₀

Xₗ

Introduce a scaling factor

Multiply each block’s output by αₗ before it joins the stream

X₁

X₂ ⋯ Xₗ₋₁

General residual MLP

Xₗ = Xₗ₋₁ + αₗ · Fₗ(Xₗ₋₁)

Fₗ = · Wᵀ → ReLU → · Wᵀ with He init, so each block returns the variance it reads.

αₗ = 1/√L keeps the output variance below e·σ²
αₗ = 1/√l makes Var(Xₗ) = (l + 1)σ²

All of these configs can work. (a) and (b) are preferred.

A per-block factor αₗ controls how much variance each block adds.

+

× α

Fₗ

+

Fₗ

× α

× α

+

Fₗ

Where the scaling factor goes

All of these configs can work. (a) and (b) are preferred.

(a) Scale the block input

Xₗ₋₁

Xₗ

Xₗ = Xₗ₋₁ + Fₗ(α·Xₗ₋₁)

Scale the variance going in:
block reads α²·Var(Xₗ₋₁)
and returns α²·Var(Xₗ₋₁)
stream adds α²·Var(Xₗ₋₁)
existing stream variance untouched

Same as (b)

(b) Scale the block output

Xₗ₋₁

Xₗ

Xₗ = Xₗ₋₁ + α·Fₗ(Xₗ₋₁)

Scale the variance coming out:
block reads Var(Xₗ₋₁)
and returns α²·Var(Xₗ₋₁)
stream adds α²·Var(Xₗ₋₁)
existing stream variance untouched

Same as (a)

Xₗ₋₁

Xₗ

Xₗ = α·Xₗ₋₁ + Fₗ(Xₗ₋₁)

Scale the stream’s variance:
stream keeps α²·Var(Xₗ₋₁)
block still adds Var(Xₗ₋₁)
Var(Xₗ) = (α² + 1)·Var(Xₗ₋₁)
X₀ and earlier blocks shrink by α²

Avoid

Scale the block, not the stream: (c) shrinks X₀ and every earlier block by α².

Scaling for convergence

Where should the shrinking factor go: on the whole stream, or only on each layer?

Two options

Attenuate the residual and layer

Xₗ₊₃ = αₗ₊₂(Xₗ₊₂ + F(Xₗ₊₂Wᵀ))

= αₗ₊₂(αₗ₊₁(αₗ(Xₗ + F(XₗWᵀ)) + F(Xₗ₊₁Wᵀ)) + F(Xₗ₊₂Wᵀ))

Scaling the residual stream shrinks the variance of X in the stream
Gradient identity path is now attenuated by scaling

Attenuate the layer

Xₗ₊₃ = Xₗ₊₂ + F(αₗ₊₂Xₗ₊₂Wᵀ)

= Xₗ + F(αₗXₗWᵀ) + F(αₗ₊₁Xₗ₊₁Wᵀ) + F(αₗ₊₂Xₗ₊₂Wᵀ)

αₗ₊₂ < αₗ₊₁ < αₗ

Variance of X is preserved
Gradient identity path is preserved

Attenuating the layer is preferred

Doesn’t attenuate the input Xₗ as a function of L
Preserves the identity path

Shrink each layer’s contribution, never the identity path.

× 1/√1

· W₂ᵀ

ReLU

· W₁ᵀ

+

× 1/√2

· W₄ᵀ

ReLU

· W₃ᵀ

+

× 1/√3

· W₆ᵀ

ReLU

· W₅ᵀ

+

× 1/√4

· W₈ᵀ

ReLU

· W₇ᵀ

+

X₀
σ²

X₄
5σ²

Converging variance

Variance along the residual stream (σ² = Var(X₀), multiplier 1/√l at block l)

σW²
= 1/n

E[P̂₁²] =
½·2σ²

σW²
= 2/f

σ²

σ²

P̂₁

σ²

Ŷ₁

2σ²

X₀

σ²

σW²
= 1/n

E[P̂₂²] =
½·4σ²

σW²
= 2/f

σ²

2σ²

P̂₂

2σ²

Ŷ₂

4σ²

X₁

2σ²

σW²
= 1/n

E[P̂₃²] =
½·6σ²

σW²
= 2/f

σ²

3σ²

P̂₃

3σ²

Ŷ₃

6σ²

X₂

3σ²

σW²
= 1/n

E[P̂₄²] =
½·8σ²

σW²
= 2/f

σ²

4σ²

P̂₄

4σ²

Ŷ₄

8σ²

X₃

4σ²

X₁
2σ²

X₂
3σ²

X₃
4σ²

Var(X₁) = σ² + 1·σ² = 2σ²

Var(X₂) = 2σ² + ½·2σ² = 3σ²

Var(X₃) = 3σ² + ⅓·3σ² = 4σ²

Var(X₄) = 4σ² + ¼·4σ² = 5σ²

The variance grows like σ² ∗ 2^L
- What if we introduce a scaling to each layer
  - All layers multiplied by 1/√l

σ² ∏l=1L (1 + 1/l) = (L + 1)σ²

Early layers contribute more variance

Block l’s output is multiplied by 1/√l before the add, so the variance it adds is scaled by (1/√l)² = 1/l

For block input variance σₓ²: Var(Ŷ) = f·(2/f)·σₓ² = 2σₓ² · E[P̂²] = ½·2σₓ² = σₓ² · Var(block out) = n·(1/n)·σₓ² = σₓ²

Block l reads Var(Xₗ₋₁) = l·σ², so it adds (1/l)·l·σ² = σ²

so Var(Xₗ) = Var(Xₗ₋₁) + σ² = (l + 1)σ², growing linearly with depth

With 1/√l, variance grows linearly with depth instead of exponentially.

× 1/√L

· W₂ᵀ

ReLU

· W₁ᵀ

+

× 1/√L

· W₄ᵀ

ReLU

· W₃ᵀ

+

× 1/√L

· W₆ᵀ

ReLU

· W₅ᵀ

+

× 1/√L

· W₈ᵀ

ReLU

· W₇ᵀ

+

X₀
σ²

X₄
2.44σ²

Converging variance

Variance along the residual stream (σ² = Var(X₀), L = 4 blocks)

σW²
= 1/n

E[P̂₁²] =
½·2σ²

σW²
= 2/f

0.25σ²

σ²

P̂₁

σ²

Ŷ₁

2σ²

X₀

σ²

σW²
= 1/n

E[P̂₂²] =
½·2.5σ²

σW²
= 2/f

0.31σ²

1.25σ²

P̂₂

1.25σ²

Ŷ₂

2.5σ²

X₁

1.25σ²

σW²
= 1/n

E[P̂₃²] =
½·3.13σ²

σW²
= 2/f

0.39σ²

1.56σ²

P̂₃

1.56σ²

Ŷ₃

3.13σ²

X₂

1.56σ²

σW²
= 1/n

E[P̂₄²] =
½·3.91σ²

σW²
= 2/f

0.49σ²

1.95σ²

P̂₄

1.95σ²

Ŷ₄

3.91σ²

X₃

1.95σ²

X₁
1.25σ²

X₂
1.56σ²

X₃
1.95σ²

Var(X₁) = σ² + ¼·σ² = 1.25σ²

Var(X₂) = 1.25σ² + ¼·1.25σ² = 1.56σ²

Var(X₃) = 1.56σ² + ¼·1.56σ² = 1.95σ²

Var(X₄) = 1.95σ² + ¼·1.95σ² = 2.44σ²

The variance grows like σ² ∗ 2^L
- What if we introduce a scaling to each layer
  - All layers multiplied by 1/√L

σ² (1 + 1/L)L < eσ²

All layers contribute equal variance

Each block’s output is multiplied by 1/√L before the add, so the variance it adds is scaled by (1/√L)² = 1/L = ¼

For block input variance σₓ²: Var(Ŷ) = f·(2/f)·σₓ² = 2σₓ² · E[P̂²] = ½·2σₓ² = σₓ² · Var(block out) = n·(1/n)·σₓ² = σₓ²

Var(Xₗ) = Var(Xₗ₋₁) + (1/L)·Var(Xₗ₋₁) = (1 + 1/L)·Var(Xₗ₋₁)

so Var(X_L) = (1 + 1/L)^L σ² = 2.44σ² for L = 4, and stays below eσ² ≈ 2.72σ² at any depth

With 1/√L, variance stays below eσ² at any depth.

· W₁ᵀ → ReLU → · W₂ᵀ

BatchNorm

+

− μⱼ

× 1/√(σⱼ² + ε)

× γⱼ + βⱼ

Batch Norm for automatic variance control

Batch norm on the branch into a residual block

Xₗ₋₁

Xₗ

v + 1

Var(Xₗ₋₁) = v

1

1

Inside BatchNorm

Xⱼ

1

2

3

mean 0

var 1

Yⱼ

f features

e examples (batch)

column j
per-column stats over the batch

For each feature j:

μⱼ = mean of column j

σⱼ² = variance of column j

X̂ᵢⱼ = (Xᵢⱼ − μⱼ) / √(σⱼ² + ε)

Yᵢⱼ = γⱼX̂ᵢⱼ + βⱼ

γⱼ = 1, βⱼ = 0 at init (learned)

at inference: running μⱼ, σⱼ²

ε prevents amplifying noise too much.

In model.train() mode, stats come from batch.

In model.eval() mode, stats come from the most recent window of inputs.

One feature’s distribution through the steps

Input column Xⱼ

mean μⱼ, var σⱼ²

1 · subtract μⱼ

mean 0, var σⱼ²

2 · divide by √(σⱼ² + ε)

mean 0, var 1

3 · scale γⱼ, shift βⱼ

mean βⱼ, var γⱼ²

BatchNorm standardizes each feature across the batch, then γ and β rescale it.

· W₁ᵀ → ReLU → · W₂ᵀ

LayerNorm

+

− μᵢ

× 1/√(σᵢ² + ε)

× γⱼ + βⱼ

Layer Norm for automatic variance control

Layer norm on the branch into a residual block

Xₗ₋₁

Xₗ

v + 1

Var(Xₗ₋₁) = v

1

1

Inside LayerNorm

Xᵢ

1

2

3

mean 0

var 1

Yᵢ

f features

e examples (batch)

row i = one example
per-row stats over the features

For each example i:

μᵢ = mean of row i

σᵢ² = variance of row i

X̂ᵢⱼ = (Xᵢⱼ − μᵢ) / √(σᵢ² + ε)

Yᵢⱼ = γⱼX̂ᵢⱼ + βⱼ

γⱼ = 1, βⱼ = 0 at init (per feature)

same at training and inference

ε prevents amplifying noise too much.

Stats are always based on the current vector across features.

One example’s features through the steps

Input row Xᵢ

mean μᵢ, var σᵢ²

1 · subtract μᵢ

mean 0, var σᵢ²

2 · divide by √(σᵢ² + ε)

mean 0, var 1

3 · scale γⱼ, shift βⱼ

mean 0, var 1 at init

LayerNorm standardizes each example across its features, the same at training and inference.

Norm Gradients

What happens to the backward pass when a norm sits on the branch

We will ignore the gradient through the norms (for now)
PyTorch will take care of it for us

For the variance story, treat each norm as a rescaling to variance 1.

· W₂ᵀ

ReLU

· W₁ᵀ

Norm

+

· W₄ᵀ

ReLU

· W₃ᵀ

Norm

+

· W₆ᵀ

ReLU

· W₅ᵀ

Norm

+

· W₈ᵀ

ReLU

· W₇ᵀ

Norm

+

X₀
σ²

X₄
σ²+4

Forward variance with added normalization

Variance along the residual stream (σ² = Var(X₀), norm at the start of each block)

σW²
= 1/n

E[P̂₁²] =
½·2

σW²
= 2/f

σ² → 1

1

P̂₁

1

Ŷ₁

2

1

X₀

σ²

σW²
= 1/n

E[P̂₂²] =
½·2

σW²
= 2/f

σ²+1 → 1

1

P̂₂

1

Ŷ₂

2

1

X₁

σ²+1

σW²
= 1/n

E[P̂₃²] =
½·2

σW²
= 2/f

σ²+2 → 1

1

P̂₃

1

Ŷ₃

2

1

X₂

σ²+2

σW²
= 1/n

E[P̂₄²] =
½·2

σW²
= 2/f

σ²+3 → 1

1

P̂₄

1

Ŷ₄

2

1

X₃

σ²+3

X₁
σ²+1

X₂
σ²+2

X₃
σ²+3

Var(X₁) = σ² + 1

Var(X₂) = σ²+1 + 1 = σ²+2

Var(X₃) = σ²+2 + 1 = σ²+3

Var(X₄) = σ²+3 + 1 = σ²+4

Each norm scales the variance to 1.

For early layers the scaling is near 1.
For later layers the scaling is near 1/√l.

Similar to scaling by 1/√l.
Early layers contribute more!

Norm (BatchNorm or LayerNorm, γ = 1 and β = 0 at init) rescales whatever variance the stream has to 1

Var(Ŷ) = f·(2/f)·1 = 2 · E[P̂²] = ½·2 = 1 · Var(block out) = n·(1/n)·1 = 1

every block adds exactly 1: Var(Xₗ) = Var(Xₗ₋₁) + 1

so Var(Xₗ) = σ² + l, linear growth with no hand-picked multiplier

With a norm in front of every block, Var(Xₗ) = σ² + l, no multiplier to pick.

· W₂ᵀ

ReLU

· W₁ᵀ

Norm

+

· W₄ᵀ

ReLU

· W₃ᵀ

Norm

+

· W₆ᵀ

ReLU

· W₅ᵀ

Norm

+

· W₈ᵀ

ReLU

· W₇ᵀ

Norm

+

X₀
σ²

X₄
σ²+4

Norm

Output

Normalizing the output

Variance along the residual stream (σ² = Var(X₀), norm at the start of each block)

σW²
= 1/n

E[P̂₁²] =
½·2

σW²
= 2/f

σ² → 1

1

P̂₁

1

Ŷ₁

2

1

X₀

σ²

σW²
= 1/n

E[P̂₂²] =
½·2

σW²
= 2/f

σ²+1 → 1

1

P̂₂

1

Ŷ₂

2

1

X₁

σ²+1

σW²
= 1/n

E[P̂₃²] =
½·2

σW²
= 2/f

σ²+2 → 1

1

P̂₃

1

Ŷ₃

2

1

X₂

σ²+2

σW²
= 1/n

E[P̂₄²] =
½·2

σW²
= 2/f

σ²+3 → 1

1

P̂₄

1

Ŷ₄

2

1

X₃

σ²+3

X₁
σ²+1

X₂
σ²+2

X₃
σ²+3

Var(X₁) = σ² + 1

Var(X₂) = σ²+1 + 1 = σ²+2

Var(X₃) = σ²+2 + 1 = σ²+3

Var(X₄) = σ²+3 + 1 = σ²+4

1

1

Norm (BatchNorm or LayerNorm, γ = 1 and β = 0 at init) rescales whatever variance the stream has to 1

Var(Ŷ) = f·(2/f)·1 = 2 · E[P̂²] = ½·2 = 1 · Var(block out) = n·(1/n)·1 = 1

every block adds exactly 1: Var(Xₗ) = Var(Xₗ₋₁) + 1, so Var(Xₗ) = σ² + l

The stream is never normalized (X₄ still has σ²+4); the final norm branches off it, just like each block’s norm, so the output leaving the network has variance 1

A final norm hands the next layer an input with variance 1.

Final norm for scaling

We end as if the network’s output is the input to another layer

Two plots against depth from 1 to about 250. Left, log scale, titled Bonus: the raw residual stream keeps growing internally, hidden by the terminal LN: raw stream variance before the final LayerNorm (red) rises steadily from about 1 at depth 1 to about 160 at depth 256, while a dashed line at 1 marks the unit variance the final LN reports. Right, titled Bonus: raw stream vs. final representation, does the terminal LN actually undo the collapse?: correlation between examples in the raw stream (red) rises from 0.59 at depth 1 to 0.99 at depth 256, while the final, terminal-LN correlation (purple) falls from 0.21 to 0.08.

The final norm rescales the output without touching the stream.

Stable at depth

Pre-norm blocks plus a final LayerNorm, against a plain MLP and an unnormalized residual MLP

Two plots against depth from 1 to about 250. Left, log scale, titled Part 4: Pre-LN keeps the gradient alive far deeper than Plain, without exploding like unnormalized Residual: first-layer gradient RMS for Plain (blue) stays near 10 to the 0 and dips slightly, Residual with no norm (green) explodes to about 10 to the 56, and Pre-LN plus terminal LN (red) stays flat near 10 to the 0. Right, titled Part 4 (bonus): Pre-LN fixes the gradient, not necessarily the correlation: correlation between different examples for Plain (blue) rises from 0.48 to 1.0, plus BatchNorm (orange) stays at 0, and Pre-LN plus terminal LN (red) falls from 0.21 to 0.08.

QR code linking to tinyurl.com/TnTech-AI4200-Normalization

See it in action:
https://tinyurl.com/TnTech-AI4200-Normalization

Pre-norm plus a final norm keeps the gradient stable out to 256 layers.

What’s inside

Two well-known deep networks, built from the pieces in this lecture

Resnet-V2

GPT-2

Clean identity path to preserve gradient

Layer input normalization

Normalize before the output

Both use the same three ideas: identity path, pre-norm, final norm.

Quiz details

Given a neural net, write the equation to compute the updated weights (don’t do the arithmetic)
Why do we need non-linearity in a deep net?
Ideally the gradient norm through a neural net will…
How does variance relate to norm and why does the variance of the gradient matter?
- Draw variance if gradient is exploding
- Draw variance if gradient is vanishing
- Draw variance if gradient is stable
Why is tanh a better activation function for hidden layers than sigmoid?

What is the variance of XWᵀ?
- X is e × f with variance σ_X²
- W is n × f with variance σ_W²
- What should σ_W² be to ensure the output variance is σ_X²
Same question but for the backward pass
Going beyond: Why is tanh modeled as y = x and when is this true?
Why does tanh fail in deep feed forward nets?
What is the mean and variance of P̂ = ReLU(XWᵀ)?
- X is e × f with variance σ_X²
- W is n × f with variance σ_W²
- What should σ_W² be to ensure the E[P̂²] is σ_X²
How does ReLU fail in deep feed forward nets?

Formats: short answer, multiple choice, draw the graph.