Summary (Overview)

• FOAM (Factored One-sided Adam-Moment) is a novel second-order optimizer that achieves SOAP-level quality at near-AdamW cost for LLM pre-training
• Employs two structural simplifications: input-only preconditioning (dropping the output-side Kronecker factor) and a factored Adam-moment estimator (replacing dense rotated second moments with row/column vectors)
• Reduces SOAP's optimizer state from $2mn + 2m^2 + 2n^2$ to $mn + 2n^2 + (m+n)$ per parameter and halves the number of eigendecompositions
• Matches strong Muon and SOAP baselines in validation loss across 124M–760M parameter scales while using half of SOAP's optimizer state and achieving Muon-class wall-clock time
• Retains robustness across a three-fold learning-rate range and handles large batch sizes better than AdamW

Introduction and Theoretical Foundation

Background and Motivation

The paper addresses the practical limitations of second-order optimizers like Shampoo and SOAP. While SOAP achieves 35–40% improvements over AdamW in wall-clock time and step count, it requires approximately 4.5× the optimizer state of AdamW and significantly slower per-step wall-clock time. This is due to:

1. Dense rotated second moment $v_{rot} \in \mathbb{R}^{m \times n}$
2. Two Kronecker factor matrices $L = GG^\top \in \mathbb{R}^{m \times m}$ and $R = G^\top G \in \mathbb{R}^{n \times n}$
3. Periodic eigendecomposition overhead

Theoretical Foundation

The work builds on the K-FAC (Kronecker-Factored Approximate Curvature) framework, where the empirical Fisher information is approximated as:

$$F \approx \mathbb{E}[aa^\top] \otimes \mathbb{E}[ss^\top]$$

For an affine layer with cross-entropy loss, the Hessian decomposes exactly as:

$$H = xx^\top \otimes (\text{diag}(p) - pp^\top)$$

where $x$ is the layer input and $p$ is the softmax output. The output-side factor has diagonal entries $p_i(1-p_i)$ and off-diagonal entries $-p_i p_k$, which are smaller by a factor of $p_k/(1-p_i)$ — making the output-side curvature approximately diagonal for large vocabularies.

Key Insight

The Kronecker structure of the Hessian shows that all output rows of an affine layer share the same input-side outer product $xx^\top$, while transformer layers exhibit near-block-diagonal Hessians along output neurons. This justifies keeping only the input-side factor $R$.

Methodology

FOAM Algorithm

The core update maintains an input-side factor $R_t$ with its eigenbasis $Q_R$, applies rotation only to the input axis ($G_{rot} = GQ_R$), and replaces the dense rotated second moment with factored estimators:

$$\hat{v}_{rot}[i,j] = \tau \cdot \hat{r}[i] \cdot \hat{c}[j], \quad \tau = \frac{\text{tr}(R)}{\|\hat{r}\|_1 \|\hat{c}\|_1}$$

where $\hat{r} \in \mathbb{R}^m$ and $\hat{c} \in \mathbb{R}^n$ are bias-corrected EMAs of row and column means of the squared rotated gradient.

Lemma 1 (Key Theoretical Result)

Under the K-FAC factorization $\mathbb{E}[gg^\top] = R \otimes L$, the covariance of the input-rotated gradient is:

$$\mathbb{E}[g'g'^\top] = \Lambda_R \otimes L$$

where $\Lambda_R = Q_R^\top R Q_R$ is the diagonal eigenvalue matrix of $R$. This implies that when $L$ is diagonal, the rotated covariance is diagonal and factorizes as:

$$\mathbb{E}\left[(GQ_R)[i,j]^2\right] = \lambda_R[j] \cdot L[i,i]$$

making the entire $m \times n$ second-moment surface capturable by $m + n$ numbers.

Algorithm Structure

Algorithm 1: FOAM update at step t for W ∈ R^{m×n}
1: R ← β₂R + (1-β₂)GᵀG                    (input-side factor)
2: if t mod pf = 0 and off(Q_RᵀRQ_R) > tol then
3:   Q_R ← eigh(R)                         (refresh eigenbasis)
4: end if
5: M ← β₁M + (1-β₁)G                       (first moment)
6: G_rot ← GQ_R, M_rot ← MQ_R             (rotate)
7: r ← β₂r + (1-β₂)G²_rot 1ₙ/n
   c ← β₂c + (1-β₂)1ₘᵀG²_rot/m           (factored second moment)
8: bias correction
9: v̂_rot ← τr̂ĉᵀ with τ = tr(R)/(‖r̂‖₁‖ĉ‖₁)
10: U_rot ← M̂_rot ⊘ (√v̂_rot + ε)          (Adam step)
11: W ← W - ηU_rotQ_Rᵀ                     (rotate back)

Complexity Comparison

Empirical Validation / Results

Pre-training Results (Table 2)

Key findings:

• FOAM trails SOAP by only 0.004–0.008 nats (within one seed std)
• FOAM uses half of SOAP's optimizer state at every scale
• FOAM reaches AdamW's final loss at equal wall-clock time at large scale (10.53 vs 10.55 h), while SOAP needs 2.15× longer
• Measured peak memory: +7.4% over AdamW for FOAM vs +12.9% for SOAP

Construction Ablation (Table 3)

Key Diagnostic Results

1. Hyperparameter robustness: FOAM tracks AdamW and SOAP within 0.05 nats across LR, WD, β₁, and β₂ sweeps. Notably, AdamW collapses to 3.86 at β₁=0.99 while FOAM degrades only 0.02 nats.

2. Gauss-Newton alignment: Input-only preconditioning keeps 75% of SOAP's curvature alignment (0.38 vs 0.10 for output-only), confirming the input-feature axis carries the dominant curvature structure.

3. Rank-1 separability: In SOAP's eigenbasis, the rotated second-moment matrix stays near 0.90 rank-1 separable (vs only 0.54–0.63 in Adam's raw basis), with Frobenius cosine ≈ 0.96–0.97, validating the factored update.

4. K-FAC violation: The K-FAC assumption is violated by only 15–20% relative error, costing ~0.002 nats in final validation loss.

Extended Results

• Scale: At GPT-2 XL (1.56B params), FOAM and SOAP within 0.001 nats
• Modality: On ViT-Tiny/Tiny-ImageNet, FOAM beats AdamW by ~7 points top-1 accuracy
• Large batches: At 2M tokens/step, FOAM reaches AdamW's 0.5M-batch quality in ~1/3 the wall-clock time

Theoretical and Practical Implications

Theoretical Contributions

1. Structural justification for input-only preconditioning: The paper provides a rigorous argument based on the Kronecker structure of the Hessian that the output-side curvature of transformer layers is near-diagonal, making the input side the natural one-sided approximation.

2. Exactness under K-FAC: The factored Adam-moment estimator exactly recovers SOAP's dense second moment under K-FAC with diagonal output curvature (Lemma 1), providing theoretical grounding for a purely empirical observation.

3. Implicit Fisher structure: FOAM's implicit Fisher is $R_t \otimes \text{diag}(\hat{r}_t)$, which bridges the gap between Muon's structure $((M_t^\top M_t)^{1/2} \otimes I_m)$ and SOAP's full Kronecker fit.

Practical Implications

• Memory efficiency: FOAM's state reduction from $2mn + 2m^2 + 2n^2$ to $mn + 2n^2 + (m+n)$ makes second-order optimization viable with only +7.4% memory over AdamW
• Wall-clock speed: The amortized eigendecomposition cost drops from $\mathcal{O}(m^3 + n^3)$ to $\mathcal{O}(n^3)$, achieving Muon-class speed
• Robustness: FOAM inherits SOAP's wider safe range on the momentum axis and maintains quality across 3× LR ranges
• Batch scaling: FOAM enables larger effective batch sizes, which is crucial for distributed training efficiency

Conclusion

FOAM successfully addresses the complementary problem of bringing SOAP-level quality close to Adam-level cost. The two structural simplifications—input-only preconditioning and factored Adam-moment—are both theoretically justified (via Kronecker Hessian structure and K-FAC assumptions) and empirically validated (via Gauss-Newton alignment and rank-1 separability diagnostics).

Main takeaways:

• SOAP's gains need not require SOAP's cost
• The dominant curvature structure in transformers lives on the input-feature axis
• Factored updates are nearly lossless when applied in the correct rotated basis

Future directions:

• Scaling to frontier-scale pre-training (≥7B parameters)
• Sharding $R$ and $Q_R$ across devices
• Extending FOAM to embeddings/lm-head layers
• Attention-aware factorizations for the key projection (where the factored update is least accurate)
• The authors note that Muon/AdaMuon remain preferable when optimizer state is the hardest constraint, while FOAM excels when SOAP-class quality is needed with bounded memory/wall-clock