# Factored One-sided Adam-Moment for Practical and Scalable SOAP

> FOAM matches SOAP's second-order optimizer quality at roughly half the memory and quadruple the speed by dropping the output-side Kronecker factor and using factored Adam moments.

- **Source:** [arXiv](https://arxiv.org/abs/2611.16905)
- **Published:** 2026-11-14
- **Permalink:** https://picx.dev/p/M9pRMn
- **Whiteboard:** https://picx.dev/p/M9pRMn/image

## Summary

## 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

| | AdamW | Muon | SOAP | FOAM |
|---|---|---|---|---|
| **Optimizer state** | $2mn$ | $mn$ | $2mn + 2m^2 + 2n^2$ | $mn + 2n^2 + (m+n)$ |
| **Per-step cost** | $\mathcal{O}(mn)$ | $\mathcal{O}(mn^2)$ | $\mathcal{O}(mn(m+n))$ | $\mathcal{O}(mn^2)$ |
| **Amortized (per pf)** | - | - | $\mathcal{O}(m^3 + n^3)$ | $\mathcal{O}(n^3)$ |

## Empirical Validation / Results

### Pre-training Results (Table 2)

| Optimizer | small (124M) val | medium (350M) val | large (760M) val | State (large) | $t_A$ (large) |
|---|---|---|---|---|---|
| AdamW | 3.370±.009 | 3.040±.012 | 2.847±.004 | 5.27 GiB | 10.55 h |
| Muon | 3.313±.008 | 2.998±.010 | 2.806±.005 | 2.64 GiB | 11.12 h |
| AdaMuon | 3.311±.009 | 2.992±.012 | 2.799±.004 | 5.27 GiB | 10.85 h |
| SOAP | **3.305±.009** | **2.983±.012** | **2.791±.005** | 23.73 GiB | 22.66 h |
| **FOAM** | 3.309±.008 | 2.991±.012 | 2.796±.005 | 11.87 GiB | **10.53 h** |

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)

| Input R | Output L | Factored v | Val loss | State (GiB) | Wall (h) |
|---|---|---|---|---|---|
| ✓ | ✓ | | 3.295±.011 | 2.92 | 0.48 |
| ✓ | ✓ | ✓ | 3.293±.008 | 2.59 | 0.49 |
| ✓ | | | **3.302±.008** | 1.78 | 0.44 |
| | ✓ | | 3.322±.010 | 1.78 | 0.46 |
| ✓ | | ✓ | 3.304±.007 | **1.46** | **0.46** |

### 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

---

_Markdown view of https://picx.dev/p/M9pRMn, served by PicX — AI-generated visual whiteboard summaries of research papers._