LayerNorm, Differentiated

ROB-GY 6013 · Supplement to Lecture 6

LayerNorm, Differentiated

One feature vector, three routes, two correction terms

Where $\nabla_hL$ comes from. One feature reaches the output three ways, and two of them leave a correction term.

Assumes: Lecture 6 on LayerNorm forward, through $z=\gamma\odot u+\beta$.

Pages 63–72 of the lecture, unchanged.

Mathematics for Robotics · ROB-GY 6013 · NYU Tandon · Jing Zhang · Fall 2026

Layer Normalization

Shared statistics couple the features

$$u_i=\frac{h_i-\mu(h)}{\sqrt{v(h)+\varepsilon}},\qquad \mu(h),\ v(h)\ \text{read all of }h$$

$h_i$: the feature in the numerator of $u_i$
$\mu(h),\ v(h)$: scalars built from every feature, so $u_i$ depends on all of $h$
three routes: direct, through $\mu$, through $v$: three terms in every derivative

Layer Normalization

Differentiate centering

$$d\mu=\frac1d\mathbf1\T dh,\qquad dc=dh-\mathbf1\,d\mu$$

$dh$: the change we make to the input, $d\times1$
$d\mu$: one scalar: the average of that change
$dc$: $d\times1$, and generally nonzero in every coordinate

Centering is linear in $h$, so its differential is the same subtraction applied to $dh$.

Vote

Layer Normalization

Centering removes a common shift

Add the same $\alpha$ to every feature: $h\mapsto h+\alpha\mathbf1$. Then $u$

Achanges
Bis unchanged✓
Cbecomes $0$

$dh=\alpha\mathbf1$ gives $d\mu=\alpha$, so look at $dc=dh-\mathbf1\,d\mu$.

$dc=0$, so $c$, $v$ and $u$ are untouched: $\alpha=5$ still gives $u=(-1,0,1)$. C would need $c=0$; a nonuniform change does move $u$.

Layer Normalization

Differentiate variance

$$dv=\frac2d\,c\T dc=\frac2d\,c\T dh$$

$dv$: one scalar: how the variance moves
$c$: the centred column, the same one the forward pass built
$\mathbf1\T c$: exactly $0$, which is why $dc$ may be replaced by $dh$

$$c\T dc=c\T dh-(\mathbf1\T c)\,d\mu,\qquad \mathbf1\T c=\sum_{j}(h_j-\mu)=0$$

$d\mu$: the same scalar as before

Layer Normalization

Differentiate the denominator

$$ds=\frac{dv}{2s},\qquad du=\frac{dc}{s}-c\,\frac{ds}{s^{2}}$$

$ds$: one scalar; it comes only from $dv$
$du$: $d\times1$: the numerator route plus the denominator route
$c,\ s$: saved by the forward pass, and reused here

Layer Normalization

Collect the input coefficient

$$dL=a\T du,\quad a=g_z\odot\gamma$$

$a$: arrives at $u$, $d\times1$

$$=\frac{a\T dc}{s}-\frac{(a\T c)\,ds}{s^2}$$

$a\T c$: one scalar; $-5$ here

$$=\frac{(a-\operatorname{mean}(a)\mathbf1)\T dh}{s}-\frac{(a\T c)(c\T dh)}{d\,s^3}$$

$dh$: the change we are collecting

$$g_h=\frac{a-\operatorname{mean}(a)\mathbf1}{s}-\frac{(a\T c)\,c}{d\,s^3}$$

$g_h$: $d\times1$, the column we return

Vote

Layer Normalization

Reconstruct the two correction terms

Which is the complete $g_h$?

A$g_h=a/s$
B$g_h=\big(a-\operatorname{mean}(a)\mathbf1\big)/s$
CB minus $u\operatorname{mean}(a\odot u)/s$✓

$dL=a\T du$, and $\mu$ and $v$ both moved.

$\mathbf1\T g_h=0$ rules out A, but not B: B already sums to zero.

The probe gives $(\tfrac43,-1,-\tfrac13)$; A gives $(2,-2,-3)$ and B gives $(3,-1,-2)$.

Layer Normalization

Gradients of learned scale and shift

$$g_\gamma=g_z\odot u,\qquad g_\beta=g_z$$

$g_\gamma,\ g_\beta$: $d\times1$ each, shaped like the parameters
$u$: saved by the forward pass; $g_\beta$ does not need it

$$g_\gamma=\sum_{n=1}^{B}g_z^{(n)}\odot u^{(n)},\qquad g_\beta=\sum_{n=1}^{B}g_z^{(n)}$$

$B$: samples sharing $\gamma$; a $1/B$ belongs in $g_z$

One $\gamma$ serves every sample, so the contributions add.

Layer Normalization

A check that can reveal a wrong formula

$$\frac{L(h+\tau e_j)-L(h-\tau e_j)}{2\tau}\ \approx\ (g_h)_j,\qquad L=g_z\T z$$

$\tau$: the probe step, not $\varepsilon$
$(g_h)_j$: what the formula claims

nonconstant $h$ — else $u=0$ hides the $u$-term
nonuniform $\gamma$ — else $\gamma$ factors out
$\varepsilon>0$ — else $\sqrt v$ and $s$ agree

Agreement is evidence, not a theorem; a gap says where to look.

Vote

Layer Normalization

Which invariance survives epsilon?

With $\varepsilon>0$, which always holds?

A$\mathbf1\T g_h=0$✓
B$u\T g_h=0$
C$u$ is unchanged by $h\mapsto\lambda h$, $\lambda>0$

Sum the formula: $\mathbf1\T u=0$, and $a-\operatorname{mean}(a)\mathbf1$ sums to $0$ too.

Here $u\T g_h=-5/3$, and $\varepsilon$ does not scale with $h$, so B and C both fail.