[논문리뷰] Factored One-sided Adam-Moment for Practical and Scalable SOAP

링크: 논문 PDF로 바로 열기

저자: Jaemyung Yu, Byeongho Heo, Sangdoo Yun, Dongyoon Han (NAVER AI)

1. Key Terms & Definitions (핵심 용어 및 정의)

• SOAP: Shampoo의 eigenbasis 안에서 Adam을 실행하는 Kronecker-factored second-order optimizer입니다. 입력 쪽 factor R = GᵀG와 출력 쪽 factor L = GGᵀ를 모두 유지하고, 두 factor 각각에 periodic eigendecomposition을 수행합니다. Optimizer state는 2mn + 2m² + 2n²로, AdamW의 약 **4.5×**입니다.
• FOAM (Factored One-sided Adam-Moment): 본 논문이 제안하는 경량 SOAP variant입니다. Input-only preconditioning과 factored Adam-moment estimator를 결합했고, 암묵적인 Fisher는 R ⊗ diag(r̂) 형태입니다.
• Input-only Preconditioning: 출력 쪽 Kronecker factor L과 그 rotation을 버리고 입력 쪽 factor R의 eigenbasis Q_R로만 gradient를 회전시키는 방식입니다(G_rot = G·Q_R). m² 크기의 memory와 eigendecomposition 1회를 없앱니다.
• Factored Adam-Moment: dense rotated second moment v_rot(크기 mn)를 대체하는 estimator입니다. 회전된 gradient 제곱의 row mean r ∈ ℝᵐ과 column mean c ∈ ℝⁿ을 EMA로 추적하고, trace normalization τ = tr(R)/(‖r̂‖₁‖ĉ‖₁)를 적용한 outer product τ·r̂ĉᵀ로 second moment를 재구성합니다. Buffer 크기는 mn에서 m+n으로 줄어듭니다.
• K-FAC Fisher Assumption: empirical Fisher를 E[aaᵀ] ⊗ E[ssᵀ]라는 Kronecker product로 근사하는 가정입니다. 이 가정에 diagonal output curvature가 더해지면 FOAM의 factored estimator는 SOAP의 dense second moment를 정확히 복원합니다.

2. Motivation & Problem Statement (연구 배경 및 문제 정의)

본 논문의 핵심 문제는 SOAP의 optimization 품질을 유지하면서 비용을 AdamW 수준으로 낮출 수 있는지입니다. SOAP는 AdamW보다 per-step convergence가 빠르지만, 대가가 큽니다. Optimizer state가 AdamW의 약 **4.5×**이고, L과 R 두 factor의 periodic eigendecomposition이 width에 대해 cubic으로 증가해 wall-clock이 크게 늘어납니다. 그래서 실무에서는 여전히 AdamW나 Muon을 기본으로 쓰는 경우가 많습니다.

기존 연구에도 비슷한 시도는 있었습니다. 원 SOAP 논문은 one-sided variant와 factored variant를 단순한 cost 절감 ablation으로만 다뤘고, 어떤 조건에서 품질이 보존되는지에 대한 기준은 제시하지 않았습니다. Frans et al.은 각 기법이 품질을 대부분 유지한다는 것을 경험적으로만 보였습니다. 또한 SOAP의 one-sided variant는 "차원이 작은 쪽"을 남기는데, 이 규칙이 구조적으로 타당한지는 검증되지 않았습니다. 따라서 저자들은 세 가지 질문에 이론적·실증적으로 답하는 것을 목표로 합니다. 어느 쪽 factor를 남겨야 하는가, factorization은 언제 정당화되는가, 두 기법의 결합은 왜 효과적인가입니다. 이 목표가 실제로 어느 정도 달성되었는지는 Pareto 비교 결과에서 확인할 수 있습니다 [Figure 1].

3. Method & Key Results (제안 방법론 및 핵심 결과)

FOAM은 SOAP에 두 가지 구조적 단순화를 적용합니다. Input-only preconditioning과 factored Adam-moment를 통해 optimizer state를 mn + 2n² + m + n으로 줄이고, eigendecomposition 횟수를 2회에서 1회로 줄입니다 [Figure 2].

• Input-only의 이론적 근거: affine layer와 cross-entropy loss에서 Hessian은 H = xxᵀ ⊗ (diag(p) − ppᵀ)로 정확히 분해됩니다. Vocabulary가 크면 출력 쪽 factor는 diagonal이 지배적이므로, off-diagonal 구조는 입력 쪽에 남아 있습니다. 따라서 차원 크기와 무관하게 입력 쪽 factor를 남기는 것이 자연스럽습니다.
• Factorization의 근거: Q_R로 회전하면 gradient covariance가 Λ_R ⊗ L이 됩니다. 이 행렬의 diagonal을 m×n으로 reshape하면 rank-one 행렬 diag(L)·λ_Rᵀ가 되므로, r과 c만 저장해도 dense buffer 대비 손실이 없습니다. 실제 학습에서 측정한 K-FAC violation은 15–20% 수준이었고, update에 미치는 영향은 무시할 만했습니다.
• 알고리즘 구성: R의 EMA를 갱신하고, Purifying Shampoo의 adaptive tolerance criterion(pf=10, tol=0.1)을 만족할 때만 eigenbasis를 refresh합니다. 이후 회전된 공간에서 bias-corrected Adam step을 수행합니다.

주요 실험 결과 (GPT-2 계열, FineWeb, 1× Chinchilla budget, 124M/350M/760M):

• Validation loss
  • SOAP이 모든 scale에서 가장 낮았지만, FOAM과의 차이는 0.004/0.008/0.005 nats로 약 1 seed std 이내였습니다.
  • FOAM은 모든 scale에서 Muon, AdaMuon, AdamW를 앞섰습니다. 760M 기준 FOAM 2.796, Muon 2.806, AdamW 2.847, SOAP 2.791입니다.
• Memory
  • FOAM의 matrix-parameter state는 SOAP의 약 절반입니다(760M 기준 11.87 GiB vs. 23.73 GiB).
  • 124M 기준 peak memory는 AdamW 대비 +7.4%로, SOAP의 +12.9%보다 작았습니다.
• Wall-clock
  • AdamW의 최종 loss에 도달하는 시간(t_A)은 모든 scale에서 AdamW 자체의 7% 이내였습니다. 760M에서는 10.53h vs. 10.55h로 사실상 동일했습니다.
  • 반면 SOAP는 124M에서 1.35×, 760M에서 **2.15×**의 시간이 필요했습니다 [Figure 1].
  • FOAM은 SOAP 수준의 품질 대역에 SOAP wall-clock의 51%만으로 도달했습니다.
• Ablation 및 분석
  • Input-only는 full SOAP 대비 0.01 nats 이내의 손실만 보인 반면, output-only는 약 0.02 nats 손실이 발생했습니다.
  • Gauss-Newton alignment에서도 비대칭이 뚜렷했습니다. Input-only는 SOAP-full alignment의 75%를 유지했지만, output-only는 0.10 수준으로 떨어졌습니다.
  • Rotated second moment의 rank-1 share는 FOAM의 input-only basis에서 약 0.84, Adam의 raw basis에서는 0.54–0.63이었습니다. 이는 회전 후에 factorize하는 설계가 타당함을 뒷받침합니다.
• 확장 실험
  • GPT-2 XL(1.56B)에서 FOAM과 SOAP의 차이는 0.001 nats 이내였습니다.
  • Batch size를 2M tokens로 키웠을 때 AdamW의 loss는 3.370에서 3.956으로 올랐지만, FOAM은 3.309에서 3.364로 소폭 증가하는 데 그쳤습니다.
  • Learning rate를 약 3배 범위에서 바꿔도 AdamW 수준의 robustness를 유지했습니다.

4. Conclusion & Impact (결론 및 시사점)

FOAM은 SOAP 수준의 품질을 SOAP 절반의 optimizer state와 Muon급 wall-clock으로 달성합니다. 즉 SOAP의 이득을 얻는 데 반드시 SOAP만큼의 비용이 필요하지는 않습니다. 이 결과는 Hessian의 Kronecker 구조와 K-FAC 분석으로 "입력 쪽 factor를 남기고, 회전 후에 factorize한다"는 설계 원칙을 정당화한 데서 나옵니다. 따라서 one-sided·factored variant를 단순한 heuristic이 아닌 원리에 기반한 설계로 끌어올렸다는 점에서 학술적 의미가 있습니다. 산업적으로는 second-order optimizer가 AdamW와 Muon의 실용적 대안이 될 수 있음을 보여줍니다. 특히 큰 batch에서 높은 throughput을 얻는 학습에 유리합니다.

다만 한계도 명확합니다. 실험이 1.56B 이하 규모에 그쳐 frontier-scale에서의 효과는 아직 검증되지 않았습니다. 7B 이상에서 필요한 R과 Q_R의 device sharding이 구현되지 않았고, RoPE를 거치는 key projection에서 factored update의 정확도가 가장 낮습니다. Embedding과 lm-head는 여전히 AdamW로 학습합니다.

⚠️ 알림: 이 리뷰는 AI로 작성되었습니다.