II · THE IDEA · ARTIFICIAL INTELLIGENCE
AdaFactor: Adaptive Learning Rates with Sublinear Memory
▶ Listen · narrated
When every parameter carries its own adaptive statistics, the optimizer's memory can match the model. AdaFactor keeps those rates while refusing a full moving average of squared gradients.
At a glance
- What it is
- An adaptive-rate optimizer built for massive models
- Memory aim
- Sublinear cost relative to a full second-moment table
- Core move
- Approximate second-order moments instead of storing them fully
- What it drops
- Need for a full moving average of squared gradients
Think of each weight in a network as a tap you must turn a little each training step. A fully adaptive optimizer keeps a private notebook for every tap, writing down how hard that tap has been pushed lately, so it can turn gentle taps more boldly and fierce ones more carefully. The notebooks work well, but for a huge model there are as many notebooks as taps, and the desk runs out of space.
AdaFactor keeps the idea of adjusting each tap with some memory of past force, yet it refuses to keep a full notebook per tap. It maintains a compressed sketch of those second-order statistics—an approximation of the second moments—rather than a complete moving average of squared gradients. The sketch still guides the step sizes, so training remains adaptive, but the pile of notes no longer has to match the model weight for weight. That is what “sublinear memory” is aiming at in plain terms: the optimizer’s paperwork grows more gently than the list of parameters.
Standard adaptive methods such as those that maintain a per-parameter moving average of squared gradients allocate second-moment state with the same cardinality as the parameter tensor. For large models that state is a first-class memory consumer alongside weights and activations.
AdaFactor targets that cost. It retains adaptive learning rates but replaces the full moving average of squared gradients with an approximation of the second-order moments. Optimizer memory is thereby intended to grow sublinearly relative to a dense second-moment table, which matters most for large parameter matrices where a one-to-one average is most expensive.
Operationally, the training loop is unchanged at the surface: compute gradients, form an adaptive update, apply it. The difference is internal representation. Second-order information used to scale steps is reconstructed from the approximated statistics rather than read from a complete per-weight accumulator of squared gradients. First-moment (momentum) handling and other stabilising choices are separate configuration details; the defining commitment is the refusal to store that full second-moment map.
Limitations follow from the same commitment. The approximation is not guaranteed to reproduce the preconditioning of a dense moving-average method on every tensor and every step. Runs may need retuning of learning rate and related hyperparameters. Memory savings also depend on actual parameter shapes: the relative win is largest where a full second-moment tensor would have been huge, and smaller for parameter groups that were cheap anyway. For practitioners, AdaFactor is best read as a memory–fidelity trade on second-order state, not as a drop-in promise of identical loss curves at lower RAM.
Look closer
The cost that scales with every weight
Adaptive methods that keep a moving average of squared gradients assign statistics to each parameter. For a model whose trainable weights already fill device memory, that second table is not a minor overhead: it is another structure the same size as the parameters themselves. AdaFactor is aimed at that pressure. Its claim is not that second-order information is useless, but that a full per-parameter moving average is more than the trainer can always afford.
Approximation in place of a full average
Rather than retain a complete moving average of squared gradients, AdaFactor works with an approximation of the second-order moments. The optimizer still forms adaptive step sizes from curvature-like statistics, yet those statistics are no longer stored as a dense companion to every weight. Memory therefore grows more gently than the parameter count, which is the sublinear behaviour the method is built around.
What remains visible in training
From the outside, the loop still looks familiar: gradients arrive, rates adapt, weights move. What changes is the bookkeeping underneath. The run no longer depends on holding a full second-moment map in step with the model. For large matrices of weights, that bookkeeping difference is the practical distinction between fitting an adaptive optimizer beside the model and having to shrink the model to make room for the optimizer's state.
The story
Adaptive optimisers earned their place by giving each parameter a sense of scale. Instead of one global step size fighting every layer at once, a moving record of squared gradients lets the update shrink where gradients are routinely large and grow where they are small. That record is useful. It is also expensive. Stored in full, it is essentially another copy of the model’s shape, kept only so the next step can be better scaled than the last.
For modest networks the extra table is tolerable. For massive models it becomes a first-order constraint. Device memory is already claimed by parameters and by activations needed for the backward pass. An optimizer that insists on a dense moving average of squared gradients competes directly with those needs. Reducing batch size or model width to free space is a real trade, not an accounting detail.
AdaFactor answers that pressure with a deliberate refusal. It still wants second-order information for adaptive rates, but it will not keep a full moving average of squared gradients. Second-order moments are approximated, and the memory required for those statistics is arranged to stay sublinear in the sense that matters for large parameter tensors: the optimizer’s state need not grow lockstep with every individual weight. The method is therefore not a rejection of adaptivity. It is a redesign of how much state adaptivity is allowed to demand.
The practical picture is straightforward. Gradients are computed as usual. From them, AdaFactor maintains a compressed account of second-order structure rather than a per-parameter ledger of squared values. Step sizes are still derived from that account, so different weights can still move at different effective rates. What disappears is the assumption that every weight must own a dedicated running average slot for the square of its gradient.
That design choice sits inside a broader family of decisions trainers make when models outgrow comfortable hardware. One can switch to a non-adaptive method and reclaim memory at the cost of harder tuning. One can shard optimizer state across devices and pay in communication and complexity. Or one can keep adaptivity and change the representation of the statistics themselves. AdaFactor is an instance of the third path: keep the spirit of per-parameter scaling, change the storage model underneath it.
Nothing in that move removes the need for care. An approximation of second-order moments is not the same object as the full moving average, and runs can differ when the omitted detail would have mattered. The point of the method is narrower and more engineering-minded. If the full table is what prevents the model from fitting, an optimizer that no longer requires that table reopens a range of model sizes that adaptive training would otherwise price out of reach.
Why it mattered then
As models grew, optimizer state stopped being an invisible implementation detail and became part of the capacity budget. Methods that kept full moving averages of squared gradients offered stable, adaptive steps, yet they did so by mirroring the parameter tensor in memory. For groups already pushing device limits, that mirror was increasingly hard to justify. AdaFactor arrived as a direct response: preserve adaptive learning rates, but rebuild the second-order bookkeeping so massive models were not forced to carry a complete second-moment table beside every weight. In its own moment, that was less a theoretical flourish than a way to keep adaptive training inside the memory envelope of the hardware then available.
Why it matters now
The same tension still governs large-model training. Parameters, activations, and optimizer state share a finite pool of memory, and adaptive methods remain attractive because they reduce the brittleness of a single global step size. Whenever a full moving average of squared gradients would dominate or equal the footprint of the weights, an approach that approximates second-order moments and aims for sublinear optimizer memory stays relevant. It is one of the levers—alongside sharding, lower precision, and simpler optimizers—that determines whether adaptivity is affordable at the scale actually being trained.
The surprising detail
The striking move is not a new kind of gradient, but a change in what the optimizer is allowed to remember. Second-order information is treated as something that can be approximated aggressively enough to drop a full moving average of squared gradients, without giving up adaptive rates altogether. Memory, not the update rule’s surface form, is the primary design target.
What is disputed
Public descriptions agree on the aim—adaptive rates with sublinear memory by approximating second-order moments—but the quality of that approximation depends on tensor shapes, hyperparameters, and the training run. Treat memory relief as the firm claim; treat identical optimization behaviour to a full moving-average method as something to verify, not assume.
Remember this
AdaFactor keeps adaptive rates for massive models by approximating second-order moments, so training need not store a full moving average of squared gradients.
Test yourself
A full adaptive optimizer stores a moving average of squared gradients beside every parameter. AdaFactor instead approximates second-order moments. If two training setups use the same model and batch size, what distinct pressure does AdaFactor ease, and what must you still treat as unsettled when you switch?
It eases memory pressure from optimizer state: you no longer need a dense second-moment table that scales with every weight, which is often what blocks a larger model or batch from fitting. What remains unsettled is behavioural equivalence. An approximation of second-order moments is not the same object as the full moving average, so step-size dynamics can differ; you still need to validate learning speed, stability, and final quality rather than assume the only change is footprint.
Go deeper
- [1804.04235] Adafactor: Adaptive Learning Rates with Sublinear Memory Cost · arxiv.org
- [1904.09779] Reduced Chandrasekhar mass limit due to the fine-structure constant · arxiv.org
Image: Original diagram, The Daily Triptych. Licence: Original work. Source.