srinivas raghav blog's

Alice in Wonderland, All the Way to CS312 — Part 1

residual-stream-night

I have recently been working through Stanford's deep learning CS31 and plan to discuss my reasoning and rationale on the questions of the the quizzes and One of the Question of the quiz asked us to predict what happens to a transformer's validation loss when different architectural components are removed. I initially underestimated how important residual connections are relative to normalization. Working through the question changed my working mental model: some mechanisms determine what the model can compute, while others determine how reliably the model can learn that computation.

Qn1

exp1_2

The full model is the natural baseline. It contains attention to mix information across tokens, an MLP to transform token representations, positional information, normalization, and residual connections. The interesting question is not merely which modification hurts the model most, but why.

decoder_model

Consider removing PreNorm while keeping the residual connections. It is tempting to say that the hidden states will explode and training must fail. That conclusion is too strong. We can instead understand precisely what normalization changes, then distinguish the mathematical facts from an empirical prediction about validation loss.

A token's hidden representation is a vector of dimension d. Every nonzero vector can be separated into a length and a direction:

h=ru,r=‖h‖,u=h‖h‖.

For example, if h=[3,4], then its length is 5 and its unit direction is u=[0.6,0.8]. The vector [9,12] points in exactly the same direction but is three times longer. Should attention suddenly behave very differently just because the residual stream has become three times larger? Not necessarily: the relative pattern of features has not changed.

RMSNorm attempts to reduce this particular sensitivity. Its actual definition is

N(h)=γ⊙h1d‖h‖2+ϵ.

Here γ contains learned per-coordinate gains, ϵ avoids division by zero, and ⊙ denotes elementwise multiplication. To see its geometry, temporarily ignore γ and ϵ. Then

N(h)=h‖h‖/d=dh‖h‖=du.

01_magnitude_vs_direction

The simplified operation maps inputs of different positive scales to the same vector of length d. In particular, for positive c,

N(ch)=N(h).

02_rmsnorm_onto_sphere

This identity is exact in the simplified model and only approximate with nonzero ϵ; nonuniform learned gains can also alter the output direction.

Recollect: magnitude versus direction. Any nonzero vector can be written as h=ru. The scalar r is its length; the unit vector u is its direction. Idealized RMSNorm removes the overall positive input scale while retaining the direction, up to a fixed multiplier.

Why is scale important for attention? Recall the score between tokens i and j:

sij=qi⊤kjdk,qi=WQhi,kj=WKhj.

Assume linear query and key projections without biases or extra query-key normalization. If both relevant input vectors become ten times larger, their queries and keys also become ten times larger, so the attention score becomes one hundred times larger:

sij′=100sij.

Softmax is sensitive to this and if the Scores of [1,2] produce probabilities of about [0.269,0.731]. Scores of [10,20] produce probabilities of about [0.000045,0.999955]. Their ordering did not change, yet the computation went from combining two sources to selecting almost exclusively one. RMSNorm helps make the input scale seen by attention and MLP branches more predictable. It does not make them completely insensitive to all changes in their inputs.

03_attention_scale_sensitivity

The positioning of normalization matters as much as its presence. A simplified PreNorm residual branch is

hl+1=hl+Fl(N(hl)).

The learned operation Fl reads a normalized view of the current state. Separately, the unnormalized state hl continues directly through the residual path. This means the branch receives a scale-controlled input without forcing the entire residual stream to be normalized after each update. The stream remains free to evolve in both magnitude and direction.

05_prenorm_vs_postnorm_paths

Recollect: two different jobs. PreNorm controls the input read by the branch. The residual connection preserves the direct path carrying the state forward. Removing either one changes a different part of the computation.

The same separation matters for gradients. The Jacobian of the PreNorm block, using the chain rule, is

Jblock=I+JFJN.

The identity term I is contributed by the direct residual path. It provides a local route through which changes and gradients can propagate without traversing every branch transformation. This does not guarantee perfectly stable gradients across an arbitrarily deep model, but it changes the optimization problem in a useful way.

In conventional PostNorm, normalization is applied after the residual addition:

hl+1=N(hl+Fl(hl)),

and the Jacobian becomes

Jblock=JN(I+JF).

Now even the direct residual contribution passes through the normalization Jacobian. This helps explain why PreNorm is often easier to optimize at depth, though PostNorm can also work with suitable training choices.

If we remove PreNorm but retain residual connections, the branch becomes

hl+1=hl+Fl(hl),Jblock=I+JF.

Notice what remains: the existing representation still passes directly to the next layer, and the Jacobian still contains an identity contribution. What disappears is the branch's explicit protection from changes in residual-stream scale.

A toy model illustrates the distinction. Suppose every layer computes the update F(h)=0.1h. Without normalization,

h1=1.1h0h2=1.1h1=1.12h0h3=1.1h2=1.13h0⋮hL=1.1Lh0h16=1.116h0≈4.595h0

After sixteen steps, the state is about 4.595 times its starting value. Each new update is larger because it reads the accumulated state of the preceding layer.

For comparison, take a positive scalar input, idealized RMSNorm N(h)=1, and the same branch rule applied to the normalized input. Then

h1=h0+0.1h2=h1+0.1=h0+0.2h3=h2+0.1=h0+0.3⋮hL=h0+0.1Lh16=h0+1.6

If h0=1, the final value is 2.6. In this narrow example, normalization changes repeated proportional updates into repeated constant updates. But this is not a general law of transformers. Learned branches can amplify, shrink, rotate, or cancel parts of the state. Neither exponential growth nor instability follows automatically from removing PreNorm.

Recollect: two recurrences. If every update is proportional to the current state, hl+1=(1+α)hl and the result compounds as (1+α)Lh0. If every update is a constant c, then hL=h0+Lc. Normalization reduces one path by which residual magnitude directly controls branch input, but it does not promise constant updates or bounded states in real networks.

There is a deeper geometric way to understand the normalization operation: differentiate it. If h changes a little, how much does N(h) change? A derivative answers that question. For a scalar function f(x)=x2,

f(x+Δx)=x2+2xΔx+(Δx)2.

For small Δx, the squared term is comparatively tiny, so

f(x+Δx)≈f(x)+f′(x)Δx.

This is the first-order Taylor approximation. A smooth curve behaves locally like its tangent line. A vector-valued function behaves locally like a linear transformation, represented by the Jacobian matrix:

09_taylor_tangent_zoom

f(x+Δx)≈f(x)+f′(x)Δx

Now consider a vector-valued function:

h=[xy],F(h)=[f1(x,y)f2(x,y)]

A small input perturbation is

Δh=[ΔxΔy]

Unlike a scalar function, each output can depend on multiple input coordinates. We therefore use partial derivatives:

f1(x+Δx,y+Δy)≈f1(x,y)+∂f1∂xΔx+∂f1∂yΔyf2(x+Δx,y+Δy)≈f2(x,y)+∂f2∂xΔx+∂f2∂yΔy

We can collect the partial derivatives into a matrix called the Jacobian:

JF(h)=[∂f1∂x∂f1∂y∂f2∂x∂f2∂y]

Multiplying this matrix by the input change gives:

JF(h)Δh=[∂f1∂x∂f1∂y∂f2∂x∂f2∂y][ΔxΔy]=Δx[∂f1∂x∂f2∂x]+Δy[∂f1∂y∂f2∂y]=[∂f1∂xΔx+∂f1∂yΔy[6pt]∂f2∂xΔx+∂f2∂yΔy]

Notice that these are exactly the first-order changes we calculated for both output coordinates.

Therefore, we can write the entire approximation compactly as:

F(h+Δh)≈F(h)+JF(h)Δh

Recollect — The Jacobian as a local linear transformation

In one dimension, f′(x)Δx predicts the output change caused by a small input change.

In multiple dimensions, JF(h)Δh does the same thing, accounting for how every input coordinate influences every output coordinate.

Geometrically, the Jacobian transforms a small arrow in input space into the corresponding approximate arrow in output space.

The approximation becomes increasingly accurate as the perturbation approaches zero, provided the function is differentiable.

We can derive the Jacobian of idealized RMSNorm without memorizing a matrix-calculus formula. Write

r=‖h‖,u=hr,N(h)=du

First, square the length:

r2=h⊤h=[h1h2⋯hd][h1h2⋮hd]=h12+h22+⋯+hd2=h·h=‖h‖‖h‖cos(0∘)=‖h‖2.

Recollect: familiar calculus rules. The product rule is (ab)′=a′b+ab′. The quotient rule is (a/b)′=(ba′−ab′)/b2. Matrix calculus applies these same principles while respecting multiplication order and dimensions. The expression h⊤h is simply the sum of squared coordinates.

Differentiate both sides using the ordinary product rule:

2rdr=d(h⊤h)=(dh)⊤h+h⊤dh=[dh1dh2⋯dhd][h1h2⋮hd]+[h1h2⋯hd][dh1dh2⋮dhd]=(dh1h1+⋯+dhdhd)+(h1dh1+⋯+hddhd)=2(h1dh1+⋯+hddhd)=2h⊤dh.2rdr=(dh)⊤h+h⊤dh=2h⊤dh=2(dh)⊤h

The two dot products are the same scalar, so

dr=h⊤dhr=u⊤dh

Geometrically, the first-order change in a vector's length depends only on how much of the displacement points along the vector's current direction.

Imagine a vector as an arrow extending from the origin. There are two fundamentally different ways we can move its endpoint.

If we move along the arrow's direction, we are either stretching or shortening it. Its length changes immediately, in direct proportion to the displacement. This is the parallel component of the movement.

Suppose we move the endpoint of a vector perpendicular to its original direction, or sideways. Intuitively, we are not directly stretching the vector outward. Instead, we are primarily changing the direction in which it points.

But something interesting happens: the vector's length still increases, just by a much smaller amount than the sideways displacement itself. Why?

Imagine a vector of length r=5. We move its endpoint sideways by a small distance ϵ, perpendicular to the original vector.

By the triangle rule, the new vector is the sum of the original vector and this perpendicular displacement. Since the two vectors are perpendicular, they form a right triangle. The original vector is one side of length r, the sideways displacement is another side of length ϵ, and the new vector is the hypotenuse.

By the Pythagorean theorem,

rnew2=r2+ϵ2

Taking the square root,

rnew=r2+ϵ2

Now, the important question is: How much longer has the vector become?

The change in length is

Δr=rnew−r

Substituting,

Δr=r2+ϵ2−r

We can understand this without immediately using a Taylor expansion. Multiply and divide by the conjugate:

Δr=(r2+ϵ2−r)(r2+ϵ2+r)r2+ϵ2+r

Using the identity (a−b)(a+b)=a2−b2, the numerator simplifies:

Δr=r2+ϵ2−r2r2+ϵ2+r

Therefore,

Δr=ϵ2r2+ϵ2+r

This expression is exact. We haven't approximated anything yet.

Now suppose ϵ is very small compared with r. The new length will be very close to the original length:

r2+ϵ2≈r

So the denominator becomes approximately 2r, giving

Δr≈ϵ22r

This is precisely the result we would obtain from a Taylor expansion:

r2+ϵ2=r1+ϵ2r2

Recall the first-order approximation

1+x≈1+x2

when x is small. Substituting x=ϵ2/r2,

rnew≈r(1+ϵ22r2)rnew≈r+ϵ22r

Both approaches tell us the same thing.

Let's put some numbers into this.

Our original vector has length r=5. Suppose we move its endpoint sideways by ϵ=0.1.

The new length is

rnew=25+0.01≈5.001

We moved the endpoint by 0.1, but the vector's length increased by only about 0.001.

Now halve the sideways displacement to ϵ=0.05.

rnew=25+0.0025≈5.00025

The length now increases by only about 0.00025, roughly one-fourth of the previous increase.

Why one-fourth? Because the change in length is proportional to the square of the sideways displacement:

Δr∝ϵ2

If we halve ϵ, its square becomes one-fourth as large. If we make ϵ ten times smaller, the increase in length becomes approximately a hundred times smaller.

This is what mathematicians mean when they say that a perpendicular displacement changes the vector's length only at second order.

To understand the distinction, recall that a first-order change is proportional to ϵ, whereas a second-order change is proportional to ϵ2. As ϵ approaches zero, ϵ2 approaches zero much faster.

We can also see this directly from calculus. Define the length as a function of the sideways displacement:

r(ϵ)=r2+ϵ2

Differentiating with respect to ϵ,

drdϵ=ϵr2+ϵ2

At the instant we begin moving sideways, ϵ=0, so

drdϵ|ϵ=0=0

The first derivative is zero! This tells us that an infinitesimally small perpendicular movement produces no first-order change in length.

But the length is not perfectly constant. The second derivative at ϵ=0 is 1/r, which gives the leading quadratic term ϵ2/(2r).

10_perpendicular_second_order

Geometrically, this makes sense. Moving directly along a vector immediately stretches or shortens it. Moving sideways initially changes its direction, and only gradually increases its length as the new endpoint moves away from the original circle of radius r.

This is also the intuition behind the derivative of vector normalization. If we normalize a vector, we deliberately remove its magnitude and retain its direction. A small displacement parallel to the vector changes only its magnitude, so normalization ignores it. A small perpendicular displacement changes its direction, so normalization responds to it at first order.

The deeper insight is that length and direction respond differently to movement. A radial displacement changes length at first order, while a perpendicular displacement changes direction at first order and length only at second order.

That distinction is exactly what appears in the Jacobian of RMSNorm.

It also explains the derivative we derived earlier:

dr=u⊤dh

Here, u is the unit vector pointing along the original direction of h, and dh is a small displacement. The dot product u⊤dh measures the signed component of that displacement along u.

If the displacement is parallel, the dot product captures its full signed length. If it is perpendicular, the dot product is zero, meaning there is no first-order change in magnitude.

When we say that a perpendicular displacement produces no first-order change in magnitude, we are not saying that the vector's length does not change at all. We are saying something more precise: the length changes so slowly near the starting point that its change has no term proportional to the displacement itself.

To understand this, recall what a derivative measures. If we have a function f(x) and make a small change Δx, its output can be approximated by

f(x+Δx)≈f(x)+f′(x)Δx

The term f′(x)Δx is called the first-order change because it is proportional to the displacement Δx.

For example, if f′(x)=3, then a small displacement of 0.01 produces an output change of approximately 0.03. Halving the displacement approximately halves the change. This is characteristic of first-order behavior.

But what happens if the derivative is zero?

f′(x)=0

Then the first-order approximation becomes

f(x+Δx)≈f(x)

This does not necessarily mean the function is perfectly constant. It means the function is locally flat to first order, and we must examine smaller, higher-order terms to understand how it changes.

Now consider our vector h, whose length is

r=‖h‖

Suppose we perturb it by a small vector dh. We previously derived that the first-order change in its length is

dr=h⊤dhr=u⊤dh

where u=h/r is the unit vector pointing in the direction of h.

Recall the geometric interpretation of the dot product:

u⊤dh=‖u‖‖dh‖cosθ

Since u has unit length,

dr=‖dh‖cosθ

If the displacement is parallel to h, then θ=0∘ and cosθ=1. Therefore,

dr=‖dh‖

The entire displacement contributes to the change in length. Moving outward by 0.1 increases the length by 0.1.

But if the displacement is perpendicular to h, then θ=90∘ and cosθ=0. Therefore,

dr=0

This means that the linear contribution to the change in length is zero. The sideways movement does not immediately stretch the vector in its existing direction.

However, for a finite perpendicular displacement of size ϵ, the Pythagorean theorem gives

rnew=r2+ϵ2

And, as we derived earlier,

Δr=rnew−r≈ϵ22r

Notice the difference: the length still increases, but its leading change is proportional to ϵ2, not ϵ.

For a concrete example, suppose r=5 and we move sideways by ϵ=0.1. The length increases by approximately 0.001. If we halve the sideways movement to 0.05, the increase becomes approximately 0.00025—one-fourth as large.

So when we say there is no first-order change, we mean that the coefficient of the linear term is zero. The first nonzero contribution appears at second order.

This is an important distinction: the derivative tells us what happens infinitesimally, not that every finite movement leaves the length unchanged.

Geometrically, a perpendicular displacement initially moves the endpoint along the tangent direction of the circle centered at the origin. The tangent is perpendicular to the radius, so the movement initially changes the vector's direction rather than its distance from the origin. Only as we move a finite distance along that tangent does the endpoint move farther away from the circle.

In short, a parallel displacement changes the vector's length immediately, at first order. A perpendicular displacement changes its direction immediately, while its length increases only at second order. This is why the dot product with the unit direction vector captures exactly the first-order change in magnitude.

Recollect — The geometric meaning of dr=u⊤dh

The derivative measures how quickly the vector's length changes under an infinitesimal displacement.

The fundamental insight is that the first-order change in magnitude is exactly the projection of the displacement onto the vector's current direction.

Next, differentiate the unit direction u=h/r by treating r as a scalar denominator:

du=rdh−hdrr2

Substitute dr=u⊤dh and h=ru:

du=1r(dh−u(u⊤dh))

Because u(u⊤dh)=(uu⊤)dh and dh=Idh,

du=1r(I−uu⊤)dh

Finally N(h)=du, giving

dN=dr(I−uu⊤)dh

Hence the Jacobian is

JN(h)=dr(I−uu⊤)

Why does the matrix I−uu⊤ appear? The answer comes from projection, one of the most useful geometric ideas in linear algebra.

Take a unit vector u=[0.6,0.8] and another vector v=[5,0]. The dot product is

u⊤v=u·v=0.6(5)+0.8(0)=3

For a unit vector, u⊤v=‖v‖cosθ is the signed length of v's shadow along u. The dot product gives a number, not a vector. Multiply by the unit direction to reconstruct the parallel vector:

v∥=(u⊤v)u=3[0.6,0.8]=[1.8,2.4]

The perpendicular remainder is

v⟂=v−v∥=[3.2,−2.4]

We can check that it is perpendicular:

u⊤v⟂=0.6(3.2)+0.8(−2.4)=0

06_projection_shadow

This is the triangle law: v=v∥+v⟂. Place the parallel and perpendicular displacements head-to-tail; together they reproduce the original vector. The dot product tells us how much of the vector lies along u, and subtraction gives the remainder.

Recollect: the dot product. The notations u⊤v and u·v mean the same operation. In general, u⊤v=‖u‖‖v‖cosθ. For unit u, the dot product u⊤u=‖u‖2=1. A zero dot product means perpendicular vectors.

We can rewrite projection as a matrix operation. First measure the scalar u⊤v, then reconstruct the vector by multiplying by u:

v∥=u(u⊤v)=(uu⊤)v

Thus uu⊤ is a projection matrix. Importantly, u⊤u is a scalar, while uu⊤ is a matrix. Their multiplication order makes them fundamentally different.

Subtracting the projection from the original vector gives

v⟂=v−(uu⊤)v=(I−uu⊤)v

07_projection_matrices_transform_space

In other words, uu⊤ keeps the component along u, while I−uu⊤ keeps the component perpendicular to u. This explains every term in the normalization Jacobian:

dN=drdh⟂

08_rmsnorm_jacobian_tangent

Idealized normalization discards the first-order radial change and reacts to changes of direction. If we double h=[3,4] to [6,8], its normalized direction stays exactly the same. In Jacobian language, JN(h)u=0. But a small sideways displacement changes its direction, so the Jacobian preserves and rescales that perpendicular component.

There is a shorter proof of the same geometry. Since the normalized direction always has unit length,

u⊤u=1

Differentiate:

(du)⊤u+u⊤du=0

Both terms are equal scalars, so

u⊤du=0

An infinitesimal movement of a unit vector must be perpendicular to the unit vector itself. In two dimensions the normalized output travels along a circle, and its instantaneous movement lies on the circle's tangent. In higher dimensions it lies in the tangent plane of a sphere. A finite chord between two points is not necessarily perpendicular to the original radius; the statement applies to the infinitesimal derivative.

Recollect: the geometry of the Jacobian. The factor I−uu⊤ removes a perturbation's component along the current direction. The factor d/r rescales the remainder. Therefore idealized RMSNorm's derivative responds to changes in direction, not pure positive rescaling. Actual RMSNorm has ϵ and learned gains, so the exact Jacobian differs.

This brings us back to PreNorm. The normalized branch receives a standardized view of the state, but the residual path does not pass through that normalization. In the simplified model, a radial perturbation u satisfies JNu=0, while the whole PreNorm block has

Jblocku=(I+JFJN)u=u

The branch is locally insensitive to that pure scaling direction; the direct residual path still carries it onward. This illustrates a useful division of responsibilities, not a guarantee that every piece of information survives through the full trained network.

With this mental model, removing PreNorm is not automatically catastrophic. It makes the learned branches more sensitive to variations in the residual stream's scale but leaves their computational mechanisms and direct state path intact. In a relatively shallow transformer, suitable initialization, parameters, and training settings may be enough for the model to remain competitive. How much performance changes is an empirical question, not something the algebra can settle by itself.

The practical test of the normalization explanation is to inspect residual-stream RMS, attention-logit distributions, activations, and gradient norms across layers. If a model without PreNorm remains well behaved on these quantities, there is no mathematical reason it must fail. If they drift or become unstable, they provide clues about the mechanism of failure.

The deeper lesson is about building a predictive mental model rather than memorizing which component is more important. Magnitude and direction explain normalization. Dot products and projections explain its geometry. Jacobians explain local sensitivity. Residual connections explain the direct state path. These ideas fit together into one design intuition:

Normalize what the branch reads, while preserving the residual path.