A leaner optimizer keeps SOAP’s quality at lower cost
Factored One-sided Adam-Moment for Practical and Scalable SOAP · arXiv:2611.16905
In plain words
An optimizer adjusts a model’s settings as it learns. SOAP can reach lower validation loss than AdamW, but it keeps large extra matrices and periodically decomposes them, which costs memory and time.
FOAM is an optimizer for transformer matrix settings that aims to keep SOAP’s quality with less of that work.
FOAM keeps one matrix rotation and summarizes a large matrix with row and column statistics. In GPT-2-style pretraining from 124M to 760M parameters, its average validation loss stayed close to SOAP’s, though SOAP remained slightly lower.
The research
The trade-off matters because better validation loss can bring extra optimizer memory and time. In these pretraining runs, FOAM stayed close to SOAP’s validation loss while reaching AdamW’s final loss at the largest main size in nearly the same time as AdamW. That makes the quality and cost trade-off less stark in the tested range.
FOAM keeps SOAP’s input-side curvature estimate and replaces its dense moment with row and column statistics to cut optimizer memory and update overhead.
The researchers pretrained GPT-2-style models at 124M, 350M, and 760M parameters on FineWeb, selecting each optimizer’s learning rate and weight decay separately.
SOAP had the lowest mean validation loss at each size; FOAM’s mean was higher by just 0.004, 0.008, and 0.005 nats, each within about one seed standard deviation.
At 760M, FOAM reached AdamW’s final validation loss in 10.53 hours, compared with 10.55 for AdamW and 22.66 for SOAP. This supports near-SOAP quality and near-AdamW time to that particular threshold.
The main pretraining comparison used separately selected learning rates and weight decay for each optimizer, and the source does not specify whether runs shared starting states. The largest experiment was 1.56B parameters; frontier-scale performance remains open.
What the evidence shows
- A leaner SOAP-style optimizer
Factored One-sided Adam-Moment (FOAM) is an optimizer for transformer matrix parameters that aims to retain SOAP’s optimization quality with less memory and computation. It keeps the input-side factor and drops the output-side factor, then replaces a dense second-moment matrix with row and column statistics.
The paper argues that the input side carries more of the curvature structure in the affine layers it studies.
- The input side fits better
The mean favors input-side rotation: at step 2,000, researchers compared each update’s implicit preconditioner, the matrix that scales its parameter updates, with the layer’s Gauss-Newton matrix, a measure of loss curvature.
Alignment was measured as a Frobenius cosine between those matrices; averaged over 18 GPT-2 small matrices, input-only rotation scored 0.347 and output-only rotation 0.132. This supports FOAM’s choice for these layers, although some individual layers reverse the ordering.
- Testing the factored estimate
The row-column estimate tracks the dense rotated second-moment estimate reasonably closely on average across these sampled layers, with a similar fit at both model sizes. But the averages mask differences between layers, so they do not show that the estimate fits every matrix equally well.
- Quality close, time near AdamW
Across the tested sizes, FOAM’s validation loss remained close to SOAP’s, while at 760M it reached AdamW’s final-loss threshold in nearly the same elapsed time as AdamW. This supports a trade-off close to SOAP on loss and AdamW on time in this setup; it does not mean their final losses matched.
Deep
- Original paper authors
- Jaemyung Yu, Byeongho Heo, Sangdoo Yun, Dongyoon Han
- Submitted to arXiv
- License
- CC BY 4.0