FOAM Cuts SOAP Optimizer Memory and Decomposition Overhead with One-Sided Preconditioning
Summary:NAVER AI introduces FOAM, a lightweight second-order optimizer designed to retain the convergence benefits of SOAP without its full memory and eigendecomposition cost. FOAM keeps only the input-side Kronecker preconditioner and replaces SOAP's dense rotated second moment with a factored Adam-moment estimator based on row- and column-wise statistics. Across language-model pretraining runs ranging from 124M to 760M parameters, the paper reports validation-loss performance comparable to SOAP and M
Core Idea: Retain SOAP-Style Preconditioning Without the Full Overhead
Second-order optimizers such as SOAP and Shampoo approximate the Hessian through Kronecker factorization, correcting AdamW's uniform treatment of parameter directions. The cost is substantial: SOAP generally maintains dense second-moment state, performs periodic eigendecompositions, and incurs overhead when applying rotated parameters and gradients.
FOAM, or Factored One-sided Adam-Moment, removes two major cost centers:
- It keeps only the input-side Kronecker preconditioner, eliminating the output-side factor and one associated eigendecomposition.
- It replaces SOAP's dense rotated second moment with a factored Adam-moment estimator constructed from row- and column-wise statistics.
The theoretical claim is that input-only preconditioning introduces no additional approximation error when output-side curvature is diagonal, with limited loss under near-diagonal curvature. Under a K-FAC Fisher assumption with diagonal output curvature, the factored estimator can exactly recover SOAP's dense second moment. The paper measures a 15%–20% violation of the K-FAC assumption and reports little practical impact on the resulting updates.
Mechanism: Why State and Decomposition Costs Drop
A typical SOAP implementation stores a dense rotated second moment and its inverse or an equivalent representation. For a parameter matrix of shape m × n, the dense second-order state scales as O(m² + n²), compared with O(mn) for AdamW-style first moments. This difference can dominate training memory when either matrix dimension is large.
FOAM removes the output factor, leaving a state footprint closer to O(m²). Its Adam-type moment estimator maintains factorized lower-order statistics instead of a full dense rotated matrix. Reducing the number of eigendecompositions also lowers periodic computational overhead and makes the optimizer more practical for iterative training.
That said, O(m²) can still be substantially larger than O(mn). Actual memory savings therefore depend on weight shape, parameter sharing, optimizer sharding, and implementation details. A one-sided preconditioner does not guarantee the same reduction ratio for every layer.
Engineering Trade-offs: Faster, but Not Universally Equivalent
FOAM trades a full two-sided second-order structure for an approximation that assumes output-side curvature is diagonal or nearly diagonal. The resulting system offers:
- Lower optimizer-state memory;
- Fewer feature decompositions and lower periodic overhead;
- Wall-clock performance close to Muon-class optimizers in the reported training runs.
The cost is that the approximation is not exact for an arbitrary Hessian. Layers with strong output-side coupling, extreme matrix shapes, or rapidly changing curvature may require more iterations to match SOAP. Robustness over a threefold learning-rate range is encouraging, but it does not establish that learning rates can be scaled freely across different architectures or model sizes.
Empirical Results: Comparable Validation Loss, Limited Scale Coverage
The paper evaluates FOAM, SOAP, Muon, and other baselines in language-model pretraining from 124M to 760M parameters. FOAM matches the validation-loss performance of strong SOAP and Muon baselines while using approximately half of SOAP's optimizer state and running at roughly Muon-class wall-clock time.
The main engineering result is that SOAP-like gains do not necessarily require paying all of SOAP's memory and decomposition costs. Equal validation loss, however, is not identical to identical task quality. Multimodal data, long-context training, mixture-of-experts models, and tightly controlled fine-tuning may respond differently to second-order preconditioning and therefore require additional evaluation.
The reported model sizes are also relatively small to medium. Scaling to larger models may alter optimizer-state share, matrix shapes, communication behavior, and numerical conditioning, so the headline result should not be extrapolated directly to frontier-scale pretraining. Independent replication should preserve seeds, hardware configuration, throughput definitions, and data order before drawing strong conclusions.
Deployment Guidance
The implementation is available through naver-ai/foam. Engineering teams should benchmark FOAM against AdamW, Muon, and SOAP using identical training data and model configurations, recording:
- Per-step time and end-to-end training throughput;
- Peak optimizer-state memory;
- Eigendecomposition frequency and associated communication or recomputation overhead;
- Validation-loss stability across a learning-rate sweep;
- Tokens or training steps required to reach a target validation loss.
Start with a learning-rate sweep rather than assuming SOAP hyperparameters transfer unchanged. In distributed environments, inspect state sharding, additional all-gather traffic, and whether matrix and non-matrix parameters use separate kernels. FOAM is most compelling when optimizer memory and sustained training throughput dominate the cost model. Final selection should still be based on the target model, data mixture, and task-specific validation metrics.
Conclusion
FOAM offers a practical middle ground between AdamW-like efficiency and full SOAP-style second-order optimization. It replaces the output-side Kronecker factor with a diagonal-curvature approximation and reconstructs the update through factored moment statistics. The reported evidence supports comparable validation loss at roughly half of SOAP's optimizer-state footprint in small- to medium-scale language-model pretraining. The approach remains constrained by the K-FAC approximation, evaluation scale, and task coverage, making it best viewed as an optimizer infrastructure improvement worth testing rather than a universally proven SOAP replacement.
— END —Source: Hugging Face Daily Papers (2026-11-12)
Tags: #FOAM #SOAP #Second-Order Optimization #Kronecker Preconditioning #LLM Pretraining
Editorial & Fact-Checking Note: This article is compiled from primary research, official release documentation, and source papers by the ZICQ Newsroom pipeline with automated entity verification and human editorial review. If you notice any technical inaccuracy, please submit a correction via our corrections policy or email our editorial desk directly.
Community Comments