Gradient Memory Math

What the gradient actually is

The gradient of a scalar loss L with respect to the parameters θ is the vector of partial derivatives ∇L = (∂L/∂θ_1, …, ∂L/∂θ_n). For a model with n weights it is itself an n-vector living in the same space as θ. Each component answers one question: how much would L change per unit change in this one weight, holding all others fixed? Stacked together, those answers point in the direction of steepest ascent, and gradient descent steps the other way, θ ← θ − η∇L.

Crucially the gradient has the same shape as the thing it differentiates. A query projection W_Q of shape [d, d] has a gradient ∂L/∂W_Q that is also [d, d]. This shape-matching is why storing gradients costs exactly one parameter-set worth of memory — a fact we return to.

The gradient as a vector field

Fix nothing and let θ vary: then ∇L(θ) is a vector field over parameter space, assigning to every point a vector that says which way is uphill. Training is a trajectory that follows −∇L through this field, one step at a time. Thinking of the gradient as a field rather than a single arrow is what makes optimizer choices legible.

The local geometry of that field is the Hessian, the matrix of second derivatives; its curvature decides whether a fixed-size step glides downhill or overshoots. Where the field is smooth and successive vectors agree, large steps are safe; where it swirls or reverses — high curvature, saddle regions — the same step diverges. Momentum exploits this by averaging recent samples of the field, cancelling components that flip sign and accumulating those that agree.

The gradient norm

The single most useful scalar summary of a gradient is its norm. For one parameter tensor the L2 norm is ||g|| = √(Σ_i g_i^2). For the whole model you take the global norm: concatenate every parameter’s gradient into one long vector and measure that, which is the same as g_global = √(Σ_layers ||g_layer||^2).

This one number tracks how big a step the optimizer wants to take, before the learning rate scales it — the cheapest health check in all of training. A healthy transformer run shows a global norm that starts moderate, spikes early, then settles into a slowly decaying band. A norm drifting toward zero means learning has stalled or a path has been cut; a norm spiking into the hundreds means the next update is about to wreck the weights.

Per-layer gradient scale

The global norm usefully summarizes, but it also hides that different layers receive wildly different gradient magnitudes. Log ||∂L/∂W|| for each block of a deep transformer and you routinely see values that differ by an order of magnitude across depth.

The pattern has a cause. Early layers sit far from the loss, so their gradient is the product of many intermediate Jacobians and tends to be attenuated; the final layers, one hop from the loss, see larger gradients. So a single global learning rate is always a compromise — too large for the loud late layers or too small for the quiet early ones. It is why practitioners reach for careful initialization (scaling residual branches by roughly 1/√(2N) for N blocks), normalization layers, and sometimes layer-wise learning rates: all of them work to equalize the effective step each layer takes.

A worked numeric example

Take one weight matrix W of shape [4, 4], so 16 parameters. Suppose after backprop the gradient entries have a root-mean-square value of 0.05. Then ||g|| = √(16 × 0.05^2) = √0.04 = 0.2. With a learning rate η = 1e-3 the raw update magnitude is η × ||g|| = 2e-4 — tiny next to typical weight values of 0.02 to 0.1, a nudge of well under one percent: a healthy, stable step.

Now let one bad batch produce an outlier gradient with RMS 5.0. The norm jumps to ||g|| = √(16 × 25) = 20, and the update becomes η × 20 = 0.02 — a hundred times larger, easily enough to knock the matrix off its operating point in a single step. Same learning rate, same math, radically different outcome: this is the concrete reason the norm, not the learning rate alone, governs stability.

Vanishing gradients: the chained-Jacobian math

Backpropagation computes the early-layer gradient by multiplying per-layer Jacobians together:

∂L/∂h_0 = (∂L/∂h_L) · ∏_(k=1..L) ∂h_k/∂h_(k-1)

If each factor has a spectral norm (largest singular value) consistently below 1, the product shrinks geometrically. A modest factor of 0.9 over 48 layers gives 0.9^48 ≈ 0.006: the earliest layer’s gradient is roughly 170× smaller than the last layer’s, so early weights barely move and low-level features never form. Saturating nonlinearities make it worse — a sigmoid’s derivative maxes out at 0.25 — which is why deep RNNs were so hard to train. Transformers fight this structurally with residual connections: an identity skip makes each factor ∂h_k/∂h_(k-1) ≈ I + (something small), so the product stays near 1 instead of decaying to nothing.

Exploding gradients: the same math the other way

Flip the inequality and the identical mechanism runs in reverse. If the per-layer Jacobian factors have spectral norm above 1, the product grows geometrically: 1.1^48 ≈ 114. The gradient norm blows up, one update takes an enormous step, weights overflow to NaN, and training dies. Transformers are most vulnerable early in training and at high learning rates.

The direct remedy is gradient clipping — rescale so ||g_global|| ≤ c, covered in a sibling article — but clipping treats the symptom. The causes are addressed by good initialization, by normalization (LayerNorm or RMSNorm bounding the scale of activations, and thus of the Jacobians), and by learning-rate warmup, which keeps the earliest updates small while the field is roughest. Vanishing and exploding are not two problems; they are one product being less than or greater than one.

Gradient noise: what you compute is an estimate

The true gradient is the average over the entire dataset, ∇L = (1/N) Σ_(i=1..N) ∇L_i. You never compute it. You compute a mini-batch estimate over B examples, g_B = (1/B) Σ ∇L_i. That estimate is unbiased — its expectation equals the true gradient — but it is noisy, and the noise has a precise shape.

The variance of the estimate scales as σ^2 / B, where σ^2 is the per-example gradient variance. So the noise standard deviation falls only as 1/√B: to halve the gradient noise you must quadruple the batch. That noise is not purely a nuisance — its jitter helps the optimizer escape sharp minima and saddle points — but when it dominates the true signal, steps point in nearly random directions and convergence stalls.

The gradient noise scale

There is a principled way to decide how large a batch is worth it. The gradient noise scale is roughly B_noise ≈ tr(Σ) / ||∇L||^2 — the ratio of total gradient variance to the squared magnitude of the true gradient. It is the batch size at which noise and signal are comparable.

When your batch B is well below B_noise, the gradient is signal-dominated and a bigger batch buys faster progress. When B is well above it, you are past diminishing returns: extra examples mostly average down noise that was already small. The noise scale tends to grow over training — as ||∇L|| shrinks near a minimum, the ratio climbs — which is the theory behind ramping batch size during a run, and why very large models tolerate very large batches.

Gradient memory: shapes and cost

Because ∂L/∂W has the same shape as W, the gradient buffer costs exactly one parameter-set of memory: 2 bytes per parameter in bf16, 4 in fp32. For a 1B-parameter model that is about 2 GB in bf16 of gradients alone, on top of the parameters and the optimizer state — Adam keeps two more moment buffers, another 8 bytes per parameter in fp32. This is why training memory dwarfs inference memory for the same model.

Two standard tricks bend the cost. Sharding schemes such as ZeRO-2 and FSDP split the gradient buffer across N devices so each holds only 1/N of it. Gradient accumulation trades time for memory the other way: sum several micro-batch gradients into one buffer before stepping, obtaining a large effective batch — and its lower noise — without holding a large batch’s activations at once.

What this means for CPUs and small models

On a CPU-bound small-model workflow, both of the gradient’s costs bite at once: the memory to hold it and the arithmetic to produce it. The gradient buffer doubling parameter memory (and the optimizer state adding more) is frequently the hard limit on how large a model you can fine-tune at all. This is where LoRA earns its place: by making only small adapter matrices require gradients, it shrinks the gradient and optimizer buffers by orders of magnitude, since the frozen base weights carry no gradient at all.

On the compute side, the backward pass costs roughly 2× the FLOPs of the forward pass, so a full training step is about 3× a forward pass — a useful rule when budgeting CPU time. And the cheap diagnostic never stops paying off: logging the global gradient norm every step, plus a per-layer norm now and then, costs almost nothing and catches a vanishing or exploding layer long before the loss curve shows it.

Common pitfalls

A handful of mistakes recur:

• Confusing the gradient norm with the loss. A low loss can coexist with a large, noisy gradient, and a small gradient does not prove the loss is low — it may just be a flat saddle.
• Reading only the global norm. It can look healthy while one layer quietly vanishes or explodes, and a bug that zeros a subset of gradients can leave it looking reasonable; log per-layer norms too.
• Over-clipping. A clip threshold set too low silently throttles every step, turning a stability tool into a hidden learning-rate cut.
• Overpaying for batch size. Because noise falls only as 1/√B, doubling the batch for a 1.4× noise reduction is often not worth the memory.