Adam's state, shape by shape, and what its division buys

STATEMENT

Give the shape of every tensor Adam holds for a concrete model, then establish the property those shapes exist to provide: that the step taken in each coordinate is unchanged when the loss is multiplied by a constant.

GIVEN

The network of I.5.B03 — widths $(784, 512, 256, 128, 10)$, four affine layers, $567\,434$ parameters. Adam as in (I.7.6), with $\beta_1 = 0.9$, $\beta_2 = 0.999$, $\varepsilon = 10^{-8}$.

FIND

The shape of $m$ and $v$ for each parameter tensor and the totals; then, for a loss rescaled to $c\mathcal{L}$, the effect on $g$, $m$, $v$ and the update.

STRATEGY

Every moment estimate is an average of things with the shape of the gradient, and a gradient has the shape of its parameter — so the shapes follow without any new argument. The scale question is answered by tracking a factor of $c$ through (I.7.6) and seeing where it cancels.

SOLUTION

Step 1 — the shapes. $m$ and $v$ are elementwise running averages of $g$ and $g^2$, so each has exactly the shape of the parameter it belongs to:

Two extra copies, $1\,134\,868$ values, $4.54$ MB in fp32 against the model’s own $2.27$ MB — the $2\times$ of I.7.B04, on a concrete model.

Nothing is per-layer or per-tensor: the state is per scalar parameter. That is what “adaptive” means here, and it is the whole reason the state is this large.

Step 2 — rescale the loss. Replace $\mathcal{L}$ by $c\mathcal{L}$ for a constant $c > 0$. Gradients are linear, so every $g_t$ becomes $cg_t$. Both moments are linear in their inputs:

The first moment picks up $c$ and the second picks up $c^2$, because one averages the gradient and the other its square.

Step 3 — the cancellation. Bias correction divides both by constants, so the factors survive it: $\hat{m}_t \to c\hat{m}_t$ and $\hat{v}_t \to c^2\hat{v}_t$. The step is

For $\varepsilon$ negligible against $c\sqrt{\hat{v}_t}$ the $c$ cancels and the update is unchanged. Adam’s step is invariant to rescaling the loss, and by the same argument to rescaling any individual coordinate’s gradient.

Step 4 — what plain descent does instead. Its update is $-\eta cg_t$, which scales with $c$ directly. Doubling the loss doubles every step, so a rate tuned for one loss scale is wrong for another — and since a sum over a batch and a mean over a batch differ by a factor of $B$, that is not a hypothetical.

Step 5 — where the invariance fails. $\varepsilon$ is an additive constant and does not scale, so the cancellation is approximate. When $c$ is small enough that $c\sqrt{\hat{v}} \lesssim \varepsilon$, the denominator is dominated by $\varepsilon$ and the step becomes $\approx c\hat{m}/\varepsilon$ — linear in $c$ again, like plain descent. So the invariance holds over a wide range and breaks at the bottom of it, which is one reason $\varepsilon$ is a number worth knowing rather than a formality: it sets where the adaptive behaviour stops.

Check — sanity

The state has no shape of its own. Every entry of $m$ and $v$ corresponds to exactly one scalar parameter, so a memory estimate that treats optimiser state as a per-layer overhead is wrong by orders of magnitude.

The exponents on $c$ are 1 and 2. They must be: $m$ averages $g$ and $v$ averages $g^2$. Getting $c$ and $c$ would mean the square had been dropped, and the cancellation would then leave a stray factor rather than none.

The step’s magnitude is about $\eta$. With $\hat{m} \approx \sqrt{\hat{v}}$ for a steady gradient, the ratio is near 1 and each coordinate moves by roughly $\eta$ per step regardless of its gradient’s size. That is the property that makes a single $\eta$ serviceable across parameters of very different scales, and it is also why Adam’s $\eta$ is typically $10^{-3}$ where SGD’s might be $10^{-1}$ — they are measured in different things.