AdamW and the spectral condition

1. Executive verdict

2. What the original paper actually establishes

An elementary re-derivation of AdamW feature motion, with two candidate arguments audited along the way: transferring the SGD learning rate directly, and deriving the Adam rate from its normalized update. The rejected arguments are explained here; no prior drafts are required.

Its proposed spectral condition assigns the same operator-scale budget to the current weight and to one optimizer update:

This norm condition is only an upper-scale budget until the matrix acts strongly on the relevant vector. Define action alignment

If κ = Θ(1), then a spectrally correct update causes a correctly scaled feature change:

The paper's explicit hyperparameter conversion in the implementation section is for ordinary gradient descent. There, ΔW = -η g, so knowing the gradient norm determines η. Adam breaks exactly that proportionality. The paper says its spectral condition is optimizer-independent; it does not say the raw optimizer learning-rate multiplier is identical for SGD and Adam.

3. Three questions the drafts merge together

These properties can coexist, but no one of them implies the other two. Stable rank controls the conversion from entry RMS to spectral norm. Action alignment controls feature motion. Gradient inner-product alignment controls the first-order loss response. The Adam learning rate follows from the first property, not the third.

4. Audit of candidate A: alignment argument

5. Audit of candidate B: parameterization argument

6. Independent optimizer-agnostic derivation

Let the layer matrix have shape d_out × d_in. Write an arbitrary optimizer step before its scalar learning rate as Pₜ:

For Adam, Pₜ = m̂ₜ / (√v̂ₜ + ε), with every operation in the quotient entrywise. Define two directly measurable quantities:

Feature motion is a second condition

Define κₜ = κ(Pₜ,hₜ). Once η satisfies the spectral budget,

Therefore the update learns features at maximal scale precisely when κₜ = Θ(1). Low stable rank is not enough: the dominant right singular space could be orthogonal to hₜ.

7. Exact first-step Adam calculation

For batch one at one layer, write the gradient as g = a hᵀ. After bias correction, m̂ = g and v̂ = g⊙g. If ε = 0 and no coordinate product vanishes,

It follows exactly that r=1, q=1, and ‖P‖_* = √(d_out d_in). Thus

The action alignment is also exact at finite width:

If the coordinates of h are iid with E|H| = μ₁ and E[H²] = μ₂ ∈ (0,∞), the law of large numbers gives

Gradient inner-product alignment, kept in its proper role

For dense iid Gaussian factors this tends to 2/π. This supports a first-order descent calculation. It does not alter the learning-rate derivation, which already followed from ‖P‖_*.

8. Multiple steps, multiple examples, and why full rank is compatible with low stable rank

First define the objects locally

Consider one batch-one training step s at a linear map zₛ = W hₛ. Here

Your interpretation is therefore right about aₛ and hₛ. The one correction is that Aᵢ is not a tuple of whole column vectors. It is the scalar time-history of coordinate i across the column vectors a₁,…,aₜ. Likewise, Hⱼ is the scalar time-history of coordinate j across h₁,…,hₜ:

Collecting all Aᵢ gives a d_out × t history table; collecting all Hⱼ gives a d_in × t history table.

The explicit formula for Qₜ

For ordinary Adam with bias correction, the scalar function applied to row history A=(A₁,…,Aₜ) and column history H=(H₁,…,Hₜ) is

Substituting row i and column j gives exactly

Omitting bias correction only removes the two width-independent prefactors. It does not change the structural argument.

Why separate a and h instead of keeping g?

For computing one Adam entry, you do not have to separate them. Since gₛ[i,j] = Aᵢ[s]Hⱼ[s], the scalar history (g₁[i,j],…,gₜ[i,j]) is sufficient to compute Pₜ[i,j]. Adam itself stores moment matrices, not Aᵢ and Hⱼ.

The separation is useful for understanding the whole matrix. It reveals that every entry is obtained by evaluating the same function Qₜ on one row history and one column history. This shared row/column structure is what can concentrate singular energy into a few modes. It also keeps the current activation coordinate hₜ[j]=Hⱼ[t] visible when we analyze Pₜhₜ. A history of product matrices gₛ hides both facts, and gₛ alone does not uniquely recover its two factors because (c aₛ)(hₛ/c)ᵀ = aₛhₛᵀ.

Because Qₜ contains a square root of a sum followed by an entrywise quotient, it generally does not factor as one row function times one column function. Thus Pₜ is a nonlinear kernel matrix, not a finite sum of the original rank-one gradients.

The elementary action-alignment condition

Fix output row i. Multiplying the Adam update by the current activation adds one contribution from every input coordinate:

Now make the elementary approximation that, after conditioning on row history Aᵢ, the column histories behave like comparable draws from some random history H. Define the mean and variance of one summand:

A sum of d_in such terms has mean d_in μₜ(Aᵢ) and standard deviation approximately √d_in σₜ(Aᵢ):

How this becomes an alignment statement

In the low-stable-rank, order-one-entry regime, ‖Pₜ‖_* = Θ(√(d_out d_in)) and ‖hₜ‖₂ = Θ(√d_in). Their product—the alignment denominator—is therefore Θ(d_in√d_out).

If the root-mean-square of μₜ(Aᵢ) across output rows stays nonzero, each output coordinate is order d_in and ‖Pₜhₜ‖₂ = Θ(d_in√d_out). Numerator and denominator match, so κ(Pₜ,hₜ)=Θ(1). The signs of μₜ(Aᵢ) may differ between rows; the Euclidean norm squares them, so cross-row sign cancellation is irrelevant.

If every conditional mean is zero, each output coordinate is only order √d_in, giving ‖Pₜhₜ‖₂ = Θ(√(d_out d_in)). Dividing by the same Θ(d_in√d_out) denominator gives κ(Pₜ,hₜ)=Θ(1/√d_in).

9. What the numerical reproductions actually show

The script's multiple-step check is a controlled synthetic experiment, not a trained neural network. At each of eight steps it samples independent Gaussian vectors aₛ and hₛ, constructs the exact rank-one gradient gₛ=aₛhₛᵀ, and updates Adam moments with β₁=0.9, β₂=0.999, and ε=0. It omits bias correction because the resulting scalar prefactors do not change rank or alignment. This deliberately isolates the matrix geometry from network dynamics.

Observed multiple-step width sweep

Reading those results

1. Adam's later matrix is full algebraic rank. Numerical rank equals width at every tested size. Entrywise division by √v destroys the literal rank-at-most-eight property of the first moment.
2. Its stable rank nevertheless stays near four. The small singular values make the matrix technically full rank, but a few leading modes still hold most of the squared singular-value energy. This is the precise sense in which “full rank but low-effective-rank” occurs here.
3. The update acts strongly on the current activation. The current hₜ helped form the current gradient and therefore the moment matrices. Its alignment remains order one (0.84–0.96) rather than decaying.
4. The same update acts weakly on a fresh independent activation. Its RMS alignment falls from 0.242 to 0.097. For an isotropic independent vector, the expected squared alignment is approximately rₜ/d_in, so the expected scale is √(rₜ/d_in). With stable rank near four, that decays like 1/√d_in, exactly matching the zero-mean explanation above.

Other checks in the script

10. Rectangular gradient bound and epsilon regimes

A shape-general upper-bound theorem

Consider one example at layer l, with arbitrary rectangular shape. Let

Because Gₗ is rank one, these are exact identities, not independence or typical-coordinate approximations:

For a chain network, write the activation Jacobian ∂hₗ/∂zₗ as Dₗ and define a dimensionless spectral-weight constant and realized backward-action fraction:

At fixed optimizer time, let the bias-corrected Adam second moment use nonnegative weights wₛ summing to one. The RMS statement propagates exactly through the elementwise squares:

Temporal sign cancellation is absent here because Adam squares before averaging. Minibatch examples can still cancel inside Gₛ before the square, which only makes the upper bound easier to satisfy.

Define the epsilon-insignificant state directly

For the standard denominator √v̂+ε, define each coordinate's exact fractional attenuation relative to epsilon-free Adam. On coordinates with √v̂>0:

Choose tolerances δ,ρ∈(0,1). Call a layer I(δ,ρ), an epsilon-insignificant state, if at least a fraction 1-ρ of the relevant coordinates have Eₜ[i,j]≤δ. Equivalently, that bulk has √v̂ₜ[i,j]≥((1-δ)/δ)ε. For sparse updates, “fraction” should be weighted by epsilon-free update energy so that inactive zero coordinates do not dominate the diagnostic.

In the stronger empirical transfer model where the bulk distribution of d_out√v̂ is width-invariant, a reference choice transfers transparently:

If backward faithfulness and forward feature scale strengthen Fₗ,Bₗ=O(1) to Θ(1), then the exact identities give ‖Gₗ‖_*=Θ(√(d_in/d_out)) and RMS(Gₗ)=Θ(1/d_out). Calling individual coordinates “typical” still needs a dense, nonspiky distribution; an RMS lower bound can be carried by a small fraction of coordinates.

An exact crossover with no probability or asymptotic shorthand

Let m=d_out, n=d_in, choose signs sᵢ,tⱼ∈{−1,+1}, and set aᵢ=sᵢ/m, hⱼ=tⱼ. Then all of the paper's scales hold exactly up to unit constants:

Equal magnitudes make this toy rank one even when epsilon is nonzero; that choice isolates the scale transition. With heterogeneous |aᵢ| and |hⱼ|, the additive denominator is not separable, and P is generally higher rank. The exact q,r learning-rate formula from Section 6 then handles the changed geometry.

This also explains why fixed epsilon can look harmless throughout a finite sweep. In this unit-constant toy, ε=10⁻⁸ crosses over only near d_out=10⁸. That number is not a prediction for a network: loss reduction, batch averaging, clipping, sparsity, parameter multipliers, and the coordinate distribution all change the constant. The empirical comparison is always the distribution of √v̂ᵢⱼ/ε in the optimizer's actual coordinates.

Scaling epsilon and scaling gradients are the same convention change

Suppose every gradient in one layer is multiplied by a positive width-dependent constant c before Adam updates its moments. Bias-corrected moments then scale as m̂→c m̂ and √v̂→c√v̂, so exactly

Thus raw gradients with ε∝1/d_out are equivalent, for the normalized quotient, to multiplying gradients by c∝d_out and keeping epsilon fixed. This is why an adaptive-optimizer theorem can use a width-independent entrywise function while an ordinary raw-gradient Adam implementation cannot: the paper's recalled abcd convention explicitly permits a width-dependent gradient multiplier before the entrywise optimizer.

A brief μP cross-check

This elementary result matches the role of scale in μP, but does not replace its theory. Tensor Programs IVb formalizes a broader notion of faithfulness: after the prescribed parameter and gradient scalings, optimizer inputs remain order one at each fixed training time under its representability and regularity assumptions. Its Remark 2.2.6 also notes the exact convention change above: multiplying gradients by a positive width power is equivalent to dividing epsilon by the same power.

Tensor Programs V, Appendix B.3 says its Adam derivation assumes epsilon is negligible and, when it is not, prescribes an outside-the-root epsilon scaling of 1/fan_in (or 1/fan_in² inside the root) in its μP optimizer coordinates. That is not a competing raw-weight d_out law: μP includes parameter and gradient multipliers. Once optimizer coordinates are translated with Q(cg;ε)=Q(g;ε/c), the statements are compatible. Neither μP citation by itself proves that the raw-weight bulk d_out√v̂ stays nonzero in a particular trained network. The same appendix recommends width-independent AdamW weight decay; with η=Θ(1/d_in) and growing d_in, that is the spectrally subleading ηλ=o(1) row of Section 11, not a claim that decay itself saturates the per-step spectral budget.

Parameter multipliers require the same bookkeeping. If the model uses W=αw and Adam optimizes w, then

The effective epsilon in actual-weight gradient units is ε_w/α, while the effective adaptive learning rate on W is αη. Therefore “scale epsilon by 1/d_out” is correct only after naming the parameter seen by Adam. For the convention used throughout this audit—Adam directly optimizes W—the rule is literal.

The AdamW consequence is immediate. The decay strength is ηλ, so a fixed-epsilon crossover also changes how weight decay transfers. In equal-width sign-normalized layers η=Θ(1/d); in the fixed-epsilon linear limit the exact toy instead has η→ε. Any width rule for λ must name which of these optimizer regimes it is intended to accompany.

11. The AdamW extension

For decoupled AdamW applied to the actual layer matrix, write

Let sₗ = √(d_out/d_in). If the current weights already obey ‖Wₜ‖_* = Θ(sₗ) and the adaptive learning rate is chosen so ‖ηₗPₜ‖_* = Θ(sₗ), then

Define the dimensionless per-step decay strength δₗ = ηₗλₗ. The spectral budget is transparent:

Per-step versus whole-run regularization

Ignoring the adaptive update, repeated AdamW decay multiplies the weight by

Consequently, fixed λ plus η=Θ(1/d_in) weakens decay over a fixed number of steps. There are three distinct transfer targets:

• Keep decay spectrally harmless: require only ηₗλₗ = O(1).
• Keep per-step fractional shrinkage constant: hold ηₗλₗ constant, so λₗ ∝ 1/ηₗ.
• Keep whole-run shrinkage constant: hold λₗ Σₜηₗ,t constant, which also depends on the number of steps and schedule.

12. Corrected claims and limitations

Architectural and asymptotic limits

• The derivation is for dense matrix parameters. Bias vectors, embeddings with sparse row access, convolution kernels, tied weights, attention tensors, and parameterized multipliers need their own shape map.
• The clean backward theorem is for a fixed-depth chain. Residual or attention branches introduce path sums; growing depth requires control of products and sums of layer constants rather than a fixed-number iteration.
• All Θ statements refer to a specified family of models as widths grow. They say nothing by themselves about constants, finite-width optima, depth scaling, batch scaling, or training-time scaling.
• Epsilon scaling is coordinate-dependent. Raw weights, normalized trainable parameters, and explicit gradient multipliers describe different optimizer inputs even when they implement the same forward matrix.
• Fixed finite time is materially easier than a number of optimizer steps that grows with width. History complexity and stable rank can grow with the horizon.
• Algebraic low rank, stable rank, activation alignment, and gradient-descent alignment are different measurements and should be logged separately.
• Weight decay mostly rescales existing feature maps. A leading decay update can satisfy an operator-norm budget while failing to represent task-directed feature learning.

13. Minimal empirical audit for a real training run

For every matrix layer and selected steps, log the following against width. These measurements directly test every assumption used above: