Adafactor (Shazeer & Stern, 2018) answers a blunt accounting problem: Adam keeps two extra floats per parameter, so the optimizer state of a large transformer costs more memory than the model itself. Adafactor’s move is surgical — keep Adam’s adaptive-scaling behavior, but stop storing the second moment v as a full matrix. For a weight matrix of shape [n, m] it stores only a row vector and a column vector (n + m numbers instead of n · m) and reconstructs the per-element scale as a rank-1 outer product. Add a decaying β_2 schedule, update clipping, relative step sizes, and (optionally) dropping momentum entirely, and the optimizer state for a billion-parameter model shrinks from gigabytes to megabytes. This article derives the factorization, walks the full update rule, works a numeric example, and does the memory math that made Adafactor the training optimizer behind T5 and PaLM.
Why Adam's memory bill comes due
Adam maintains two exponential moving averages per parameter: the first moment m_t = β_1 m_{t-1} + (1-β_1) g_t and the second moment v_t = β_2 v_{t-1} + (1-β_2) g_t^2. Both have exactly the shape of the parameter tensor, and both are conventionally kept in fp32. That is 8 extra bytes per parameter on top of the weights and gradients.
For a 1B-parameter model, Adam’s state alone is 8 GB. With fp32 master weights and gradients the full training footprint is roughly 16 bytes per parameter before you store a single activation. On a GPU that pressure forces sharding or offload; on a CPU box with 16–32 GB of RAM it simply forbids training anything beyond a small model. The target of Adafactor is precisely those two O(n · m) state tensors — and its claim is that one of them can be compressed by a factor of thousands with almost no loss in optimization quality.
The core idea: factor the second moment
The second moment V for a weight matrix W: [n, m] is itself an [n, m] matrix of nonnegative numbers — a running average of squared gradients. Adafactor’s observation: you never need the exact entries of V, only a good per-element scale for dividing the gradient. So approximate V by a rank-1 matrix:
V ≈ (R C) / (1_n^T R)
R = V 1_m # row sums, shape [n, 1]
C = 1_n^T V # column sums, shape [1, m]
1_n^T R = 1_n^T V 1_m = total sum of V (a scalar)Each reconstructed entry is v̂_ij = R_i C_j / total. Storage drops from n · m numbers to n + m. For a square 4096×4096 matrix that is 16.8M numbers down to 8,192 — a 2048× compression of the second-moment state.
Why row and column sums are the right factorization
Rank-1 approximation usually means SVD, which minimizes squared error — but squared error is the wrong loss here. V feeds a division, entries are nonnegative, and relative error matters more than absolute. Adafactor instead minimizes the generalized Kullback–Leibler divergence (I-divergence) between V and a nonnegative rank-1 candidate r c^T:
d(V, rc^T) = Σ_ij [ V_ij log(V_ij / (r_i c_j)) - V_ij + r_i c_j ]This objective has a remarkable property: unlike SVD, its minimizer has a closed form. Setting the partial derivatives to zero gives r_i ∝ Σ_j V_ij and c_j ∝ Σ_i V_ij, and the normalization works out to exactly V̂ = R C / (1^T R) — the row sums times the column sums over the grand total. No iteration, no eigen-decomposition: two reductions per step. That closed form is what makes the factored estimate cheap enough to recompute at every single training step.
The full Adafactor update, step by step
Putting the pieces together for a matrix parameter X with gradient G_t: [n, m]:
β̂_2t = 1 - t^(-0.8) # decaying decay rate
R_t = β̂_2t R_{t-1} + (1-β̂_2t) (G_t^2 + ε_1) 1_m # row EMA [n]
C_t = β̂_2t C_{t-1} + (1-β̂_2t) 1_n^T (G_t^2 + ε_1) # col EMA [m]
V̂_t = (R_t C_t) / (1_n^T R_t) # rank-1 reconstruction
U_t = G_t / sqrt(V̂_t) # scaled update
Û_t = U_t / max(1, RMS(U_t) / d) # update clipping, d = 1
α_t = max(ε_2, RMS(X_{t-1})) · ρ_t # relative step size
X_t = X_{t-1} - α_t Û_twith ρ_t = min(10^-2, 1/sqrt(t)), ε_1 = 10^-30, ε_2 = 10^-3. Every line replaces something Adam does with a cheaper or more robust variant; the next sections take them one at a time.
When is the reconstruction exact?
The factored estimate is not a blind guess — it is exact whenever V is genuinely rank-1, i.e. when V_ij = a_i b_j separates into a per-row scale times a per-column scale. Check: R_i = a_i Σ_j b_j, C_j = b_j Σ_i a_i, total = (Σ a)(Σ b), so R_i C_j / total = a_i b_j = V_ij exactly.
Why should squared-gradient statistics be near rank-1? In a linear layer Y = X W, the gradient is an outer-product sum ∇W = X^T ∇Y; its second-moment structure is dominated by per-input-feature scales times per-output-feature scales — how active input unit i is, times how large the error signal on output unit j is. Those multiplicative row/column effects are exactly what a rank-1 model captures. The approximation degrades only when specific (i, j) cells behave very differently from their row and column trends, which empirically matters little for preconditioning.
A worked numeric example
Take a tiny 2×3 second-moment matrix:
V = [ 4 8 4 ] R = row sums = [16, 4]
[ 1 2 1 ] C = col sums = [5, 10, 5], total = 20
V̂_11 = 16·5/20 = 4 V̂_12 = 16·10/20 = 8 V̂_13 = 4
V̂_21 = 4·5/20 = 1 V̂_22 = 4·10/20 = 2 V̂_23 = 1Reconstruction is perfect because this V is rank-1 (a = [4, 1], b = [1, 2, 1]). Now perturb one cell: set V_22 = 4. Then R = [16, 6], C = [5, 12, 5], total = 22, and V̂_22 = 6·12/22 ≈ 3.27 instead of 4 — an 18% underestimate on the perturbed cell, while its row and column neighbors shift by only a few percent. Errors are local, bounded, and smoothed across the row and column — benign for a quantity that only sets a division scale under a square root.
Memory accounting on a real transformer layer
Concrete numbers for one FFN up-projection in a d = 4096 model with the standard 4× expansion, W: [4096, 16384], 67.1M parameters:
| State | Adam | Adafactor (β_1 = 0) |
|---|---|---|
| Second moment | 67.1M floats = 268 MB (fp32) | 4096 + 16384 = 20,480 floats = 82 KB |
| First moment | 67.1M floats = 268 MB | none |
| Total per layer matrix | 537 MB | 0.08 MB |
That is a ~6500× reduction in optimizer state for this tensor. Across a whole model the ratio is diluted by embeddings and vector parameters, but the headline holds: Adam’s state scales as 2 · N_params, Adafactor’s factored state scales as roughly Σ (n_i + m_i), which for large square-ish matrices is O(sqrt(N_params)) per tensor — effectively negligible next to the weights themselves.
Decaying beta2 instead of bias correction
Adam fixes β_2 = 0.999 and patches the resulting cold-start bias with the correction factor 1/(1-β_2^t). Adafactor instead makes the decay rate itself a schedule: β̂_2t = 1 - t^(-c) with c = 0.8. At t = 1 this is 0, so the EMA starts as exactly the first observed squared gradient — no bias, no correction term needed. As t grows, β̂_2t → 1 and the average lengthens, giving stable long-run statistics.
The exponent matters. The paper shows c = 1 (equivalent to a plain running average of all history) adapts too slowly late in training, while a fixed small window forgets too fast; c = 0.8 sits between, keeping the estimator responsive without the instability that fast-decay Adam variants exhibit. It is a small change, but it removes a hyperparameter interaction (warmup vs. β_2) that plagues Adam at large batch sizes.
Update clipping: taming stale second moments
Any EMA of squared gradients is slightly out of date: if the gradient distribution shifts suddenly, V̂ underestimates the new scale and G/sqrt(V̂) can spike. Adam users paper over this with learning-rate warmup. Adafactor instead clips the update by its root-mean-square:
RMS(U) = sqrt( mean_ij (U_ij^2) )
Û = U / max(1, RMS(U) / d) # d = 1Interpretation: for a perfectly calibrated second moment, each entry of U = G/sqrt(V̂) should have magnitude around 1, so RMS(U) ≈ 1. An RMS well above the threshold d = 1 is direct evidence the denominator is stale; scaling the whole update down restores the intended step length without touching its direction. This is cheap (one reduction), self-tuning, and in the paper’s ablations it substitutes for warmup entirely.
Dropping momentum: the other half of the savings
Factoring only compresses v. The first moment m cannot be factored the same way — it carries signs, and signed matrices have no nonnegative rank-1 structure to exploit. So Adafactor’s default is more radical: set β_1 = 0 and keep no momentum at all, eliminating the second O(n · m) tensor outright.
Surprisingly, for transformer pre-training with large batches this costs little: the batch gradient is already a low-variance average, and update clipping supplies much of the stability momentum normally provides. When momentum does help (small batches, vision models, some fine-tuning), Adafactor can run with β_1 > 0, storing a full m — still saving the factored v, i.e. roughly half of Adam’s state. The two knobs are independent: factored second moment for memory, optional momentum for optimization quality.
Relative step sizes
The last departure from Adam is the learning rate itself. Adafactor scales the step for each tensor by that tensor’s own magnitude: α_t = max(ε_2, RMS(X_{t-1})) · ρ_t, where RMS(X) is the root-mean-square of the current weights and ρ_t = min(10^-2, 1/sqrt(t)) is a dimensionless relative rate.
The effect is scale invariance: a tensor initialized with standard deviation 0.02 takes proportionally smaller absolute steps than one initialized at 1.0, so every parameter moves a similar fraction of its own scale per step. This plays well with transformers, whose initialization deliberately varies per-layer scale (e.g. depth-scaled residual projections). The ε_2 = 10^-3 floor keeps parameters that start at zero — biases, some gates — from being frozen by a zero step size. One caveat: tied or specially-scaled embeddings sometimes need this feature disabled, which is why frameworks expose a scale_parameter flag.
What stays unfactored
Factorization applies only to tensors with two or more substantial dimensions. Vectors — biases, LayerNorm gains and offsets, scalar gates — keep an ordinary full second moment, because a length-d vector’s state is already only d numbers; there is nothing worth compressing. Embedding matrices and attention projections W_Q, W_K, W_V, W_O are factored like any other matrix.
Higher-rank tensors (e.g. a convolution kernel [k, k, c_in, c_out]) are handled by flattening to a matrix over the last two effective dimensions, factoring rows against columns. In a standard transformer essentially all parameters live in 2-D matrices, so in practice >99% of the second-moment state gets the factored treatment, and the unfactored remainder — a few d-sized vectors per layer — is measured in kilobytes. This is why the whole-model savings track the per-matrix analysis so closely.
Adafactor on a CPU SLM budget
Run the numbers for a 125M-parameter small language model trained or fine-tuned on a CPU box. Adam in fp32: weights 500 MB + gradients 500 MB + m 500 MB + v 500 MB = 2 GB before activations. Adafactor with β_1 = 0: weights 500 MB + gradients 500 MB + factored state on the order of 1–2 MB — a full gigabyte returned to the activation and batch-size budget on a 16 GB machine.
The compute overhead is trivial: per matrix, two reductions (row and column sums of G^2), one outer-product-style broadcast for V̂, and one RMS reduction for clipping — all O(n · m) memory-bandwidth-bound passes that vanish next to the matmuls of the forward and backward pass. Adafactor is one of the rare optimizations that is essentially free in time and enormous in space, which is exactly the trade a memory-starved CPU setup wants.
Pitfalls, and when AdamW still wins
Convergence per step: with β_1 = 0 and relative steps, Adafactor can trail AdamW early in training or on small-batch fine-tuning; if quality matters more than memory, enable momentum or revert to AdamW. Silent config traps: popular implementations (e.g. Hugging Face) expose relative_step, scale_parameter, and warmup_init; passing an external learning-rate schedule while relative steps are still on multiplies two schedules together and quietly mis-trains — the classic T5 fine-tuning bug.
Weight decay: apply it decoupled (as in AdamW), scaled consistently with the relative step, or regularization strength drifts with parameter norm. Rank-1 mismatch: tensors whose gradient statistics are genuinely high-rank — sparse embedding rows where a few tokens dominate — are approximated worst; monitoring update RMS per tensor catches this. None of these negate the design; they are the price of an optimizer that made billion-parameter training fit in memory at all.
[n, m] state collapses to n + m numbers — exact when gradient statistics factor into row and column scales, and benignly smoothed when they don’t. Around that core, a decaying β_2 = 1 - t^-0.8 replaces bias correction, RMS update clipping replaces warmup, relative step sizes make updates scale-invariant, and dropping momentum removes the other state tensor entirely. The result trains T5-class models with optimizer state measured in megabytes instead of gigabytes — and on a CPU SLM budget, that reclaimed gigabyte is the difference between fitting a training run and not. Watch the config flags: relative steps plus an external LR schedule is the one famous way to hold it wrong.