Core ML

Adaptive Optimizer Landscape

Adam's per-parameter scaling is one point on a family of optimizers, each trading memory or compute for a different way to tame wildly different gradient sizes.

Four adaptive optimizers, positioned by how much per-parameter state they keep against how they use it: RMSprop keeps one running average, Adam keeps two, Adafactor compresses Adam's matrix into row and column factors, and Lion keeps one sign-based average instead

Why Does This Exist?

Adam solved the "one learning rate for a thousand differently-scaled parameters" problem by giving each parameter its own effective rate, built from two running averages per parameter. That's the default almost everyone reaches for first, and for most training runs it's the right call. But "keep two full-size state vectors per parameter" is a specific design choice, not the only one — and at the scale a modern large language model trains at, that choice starts to cost real memory: a 7-billion-parameter model's Adam state alone is 56GB in float32, before the model's own weights or gradients take up a single byte.

Adam is one point in a family of optimizers that all attack the same problem — wildly different gradient magnitudes across parameters — with different tradeoffs between memory, compute, and robustness to a huge batch size. Knowing where Adam sits in that family, and why the alternatives exist, is what lets you actually pick when Adam stops being the obvious choice.

Think of It Like This

Four ways to remember a route

Memorizing directions for a familiar drive. One way is to memorize the exact GPS coordinates of every turn — precise, but a lot to hold in your head. Another is to remember just "which of two streets you're between" at each point — far less detail kept, almost as useful for actually driving.

A third way ditches distances entirely and remembers only which direction to turn at each landmark — left, right, straight — throwing away exactly how far each turn goes, keeping only the sign. Every one of these gets you there. They differ in how much detail about the route they insist on holding onto versus how much they're willing to compress or discard.

How It Actually Works

RMSprop — the second moment alone

RMSprop is Adam with the first moment removed: it keeps a running average of the squared gradient, $v_t = \beta v_{t-1} + (1-\beta) g_t^2$ , and divides the raw gradient by $\sqrt{v_t}$ . No momentum term, one state vector per parameter instead of two. It predates Adam and is still reasonable when memory is tight and the extra smoothing from momentum isn't worth its cost.

Adafactor — compressing the matrix, not dropping it

A weight matrix's Adam state is two matrices the same shape as the weights themselves — for a 4096×4096 layer, that's two 16-million-entry arrays. Adafactor's trick is to approximate that second-moment matrix as the outer product of a row vector and a column vector, storing only those two much smaller vectors instead of the full matrix. For a 4096×4096 layer this drops the optimizer state for that matrix from roughly 134MB to about 33KB — a factor of thousands — at the cost of a coarser, factored approximation rather than the exact per-entry history Adam keeps.

LAMB — Adam's step, rescaled per layer

Training with an enormous batch size lets you use far more parallel hardware, but past a certain batch size a fixed learning rate stops scaling with it and training degrades. LAMB takes Adam's usual update and additionally rescales it per layer, using the ratio of that layer's weight norm to its update norm, so every layer takes a proportionally sized step regardless of how large the global batch has pushed the raw gradients. It's Adam's mechanism plus a layer-wise trust-region correction, aimed specifically at very large batch training.

Lion — sign-based and cheap

Lion keeps a single momentum-like running average, like plain momentum, but takes the sign of that average as the update direction rather than its magnitude — every parameter moves by the same step size, only its direction differs. One state vector instead of two, and a cheaper per-step computation than Adam, at the cost of losing the fine-grained magnitude information Adam's second moment provides.

The actual decision

Adam or AdamW remains the default worth starting from for nearly everything. Reach for Adafactor when optimizer memory is the binding constraint on a very large model; reach for LAMB when batch size is pushed unusually high; reach for Lion when per-step compute and memory both matter more than squeezing out the last bit of quality. None of these displaces Adam by default — they each answer a specific constraint Adam runs into at a specific scale.

Show Me the Code

The memory difference Adafactor's factoring buys, on one realistic layer size.

1def optimizer_state_bytes(rows: int, cols: int, bytes_per_val: int = 4) -> tuple[float, float]:2    adam_mb = 2 * rows * cols * bytes_per_val / 1e6      # full m and v matrices3    adafactor_mb = (rows + cols) * bytes_per_val / 1e6   # row + column factors only4    return adam_mb, adafactor_mb5
6
7adam_mb, adafactor_mb = optimizer_state_bytes(4096, 4096)8print(f"Adam state:      {adam_mb:.2f} MB")9print(f"Adafactor state: {adafactor_mb:.4f} MB")10print(f"ratio:           {adam_mb / adafactor_mb:.0f}x smaller")11# -> Adam state:      134.22 MB12# -> Adafactor state: 0.0328 MB13# -> ratio:           4096x smaller

That ratio scales with the layer's width — the wider the matrix, the more Adafactor's factoring saves relative to storing it in full.

Watch Out For

Switching optimizers mid-project without re-tuning the learning rate

Adam, RMSprop, Lion, and LAMB do not share a learning rate scale — a rate tuned for Adam applied unchanged to Lion, in particular, tends to be too aggressive, because Lion's sign-based step behaves differently at the same nominal rate. Swapping the optimizer is a re-tuning event, not a drop-in replacement, and skipping that step is the most common way a switch looks like a regression.

Reaching for a memory-saving optimizer before checking where the memory is actually going

Adafactor addresses optimizer-state memory specifically. If a training run is actually memory-bound on activations, the KV cache, or the model weights themselves, switching optimizers changes none of that and the run still runs out of memory. Profile which part of the memory budget is binding before picking the fix aimed at it.

The Quick Version

Adam's two-state-per-parameter design is one point in a family of adaptive optimizers, not the only option.
RMSprop drops momentum and keeps only the second moment — one state vector instead of two.
Adafactor factors the second-moment matrix into row and column vectors, cutting memory by orders of magnitude on large layers.
LAMB rescales Adam's update per layer to stay stable at very large batch sizes.
Lion keeps one state vector and steps by sign rather than magnitude, cheaper per step than Adam.
Adam or AdamW is still the right default; the alternatives answer specific constraints Adam runs into at scale.