들어가며: 샴푸가 컨디셔너 전에 오듯
왜 이름이 Shampoo일까? 헤어케어에서 샴푸(shampoo)는 컨디셔너(conditioner) 전에 온다. 이 옵티마이저는 기울기를 사전조건화(preconditioning) 하는 것 — 즉, "프리컨디셔너" 다. 세련된 말장난이다.
핵심 아이디어는 단순하면서 강력하다. Adam은 가중치의 각 원소를 독립적으로 취급한다 — 행렬의 (i,j) 위치에 있는 숫자가 어떤 행, 어떤 열에 속하는지 무시한다. Shampoo는 가중치 행렬의 행과 열 구조를 보존하면서 사전조건화한다.
이것이 왜 중요한가? 행렬의 행은 출력 뉴런에 대응하고, 열은 입력 뉴런에 대응한다. 이 구조를 활용하면 기울기 공간의 곡률을 더 정확히 포착할 수 있다 — Adam의 대각 근사보다 훨씬 풍부하게.
1부: 행렬의 구조를 보는 눈
세 가지 시각
사전조건화 방법의 스펙트럼
Adam (대각선)
메모리: O(mn)
각 파라미터를 독립 처리. 행렬 구조 무시
Shampoo (크로네커)
메모리: O(m² + n²)
행과 열의 상관관계 포착. 최적의 균형
뉴턴법 (전체)
메모리: O(m²n²)
완전한 곡률. 현실적으로 불가능
m × n 가중치 행렬에 대해:
- Adam: mn개의 대각 원소만 추적 → 파라미터 간 상관관계 0
- Shampoo: m×m 행 공분산 + n×n 열 공분산 추적 → 구조적 상관관계 포착
- 뉴턴법: mn × mn 전체 헤시안 → 완벽하지만 저장 불가
크로네커 분해의 직관
가중치 행렬 W: m × n
↓
L: m × m (행 공분산)
⊗
R: n × n (열 공분산)
↓
L ⊗ R ≈ 전체 mn × mn 사전조건화
두 개의 작은 행렬의 크로네커 곱이 거대한 전체 사전조건화 행렬을 근사한다. 비유하자면: 도시 전체의 교통 흐름을 파악하기 위해 모든 교차로를 감시하는 대신, 각 도로의 동서 흐름과 남북 흐름만 따로 추적해도 상당히 정확한 그림을 얻는 것이다.
2부: Shampoo 알고리즘
"Shampoo: Preconditioned Stochastic Tensor Optimization"
Vineet Gupta, Tomer Koren, Yoram Singer
ICML 2018 / arXiv: 1802.09568 (2018년 2월 26일)
업데이트 규칙
매 스텝 t에서 기울기 G_t (m × n)가 주어질 때:
1. 행 공분산 갱신: L_t = L_{t-1} + G_t · G_tᵀ (m × m)
2. 열 공분산 갱신: R_t = R_{t-1} + G_tᵀ · G_t (n × n)
3. 업데이트: W_{t+1} = W_t - η · L_t-1/4 · G_t · R_t-1/4
핵심: L-1/4과 R-1/4이 행/열 방향의 곡률에 맞게 기울기를 스케일링
-1/4 지수가 핵심이다. 전체 사전조건화 행렬은 (L ⊗ R)^{-1/2}이어야 하는데, 크로네커 곱의 성질에 의해 각 인자에 -1/4를 적용하면 동일한 효과를 얻는다.
K-FAC과의 차이
| 속성 | Shampoo | K-FAC (Martens & Grosse, 2015) |
|---|
| 적용 범위 | 모든 텐서 최적화 | 신경망 레이어 구조에 특화 |
| 근사 대상 | 기울기 외적의 크로네커 분해 | 피셔 정보 행렬의 크로네커 분해 |
| 수렴 보장 | 볼록 경우 이론적 보장 있음 | 이론적 보장 없음 |
| 구현 | 범용적 | 활성화/기울기 통계에 의존 |
3부: Google의 대규모 확장 (2021)
Scalable Shampoo
"Scalable Second Order Optimization for Deep Learning"
Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, Yoram Singer
arXiv: 2002.09018 (2020년 2월 → 2021년 3월 개정)
대규모 모델에 Shampoo를 적용하려면 행렬 역 p-차 근(matrix inverse p-th root) 계산이 병목이다. Google 팀은 5가지 혁신으로 이를 해결했다:
Scalable Shampoo의 5가지 혁신
CPU-가속기 파이프라이닝
비동기 실행
행렬 연산을 CPU에 오프로드, TPU는 학습 계속
주기적 업데이트
매 k 스텝
수백~수천 스텝마다 갱신. 정확도 손실 미미
결합 뉴턴 반복
SVD 대체
수치적으로 안정적, 이전 결과로 워밍스타트
블록 분할
차원 128
큰 레이어를 블록으로 나눠 메모리 절감 (<3% 비용)
그래프팅 (Grafting)
레이어별 스케일 교정
참조 옵티마이저(Adam)의 업데이트 크기에 맞춰 스케일링
실험 결과
Scalable Shampoo 속도 향상
ResNet-50
1.35× 벽시계 향상 (31.7% 적은 스텝)
BERT-Large
16% 적은 스텝, 벽시계 단축
번역 (93.3M)
1.95× 적은 스텝, 40% 벽시계 단축
Google은 Shampoo를 검색 광고 CTR 모델에 프로덕션으로 배치했다. 2024년에는 Distributed Shampoo가 MLCommons AlgoPerf 대회에서 우승 — $25,000 상금과 함께, 비대각 사전조건화의 실용성을 공식적으로 입증했다.
4부: 저자들의 이야기
Shampoo 핵심 저자들
Vineet Gupta
Google Research
프로그래밍 언어 이론 + ML 최적화
Tomer Koren
텔아비브 대학 부교수
ERC Starting Grant(2022), Krill Prize(2023)
Yoram Singer
Google Brain + Princeton 교수
AAAI Fellow(2010), POEM 그룹 리드
Rohan Anil
前 Google DeepMind → Anthropic
Scalable Shampoo 주저자, Gemini 사전학습 리드
Rohan Anil의 궤적이 특히 주목할 만하다. Scalable Shampoo의 주저자로서 대규모 옵티마이저 연구를 이끈 뒤, Google DeepMind의 Distinguished Engineer로서 Gemini 사전학습을 총괄했다. 2025년 1월 Anthropic으로 이적했다.
5부: SOAP — Shampoo의 후계자 (ICLR 2025)
"Shampoo + Adam = SOAP"
"SOAP: Improving and Stabilizing Shampoo using Adam"
= ShampoO with Adam in the Preconditioner's eigenbasis
Nikhil Vyas, Depen Morwani 외 (Harvard University)
arXiv: 2409.11321 (2024년 9월) / ICLR 2025
SOAP의 핵심 통찰: Shampoo(1/2 지수)는 Shampoo 사전조건화의 고유기저(eigenbasis)에서 Adafactor를 실행하는 것과 동치다.
Shampoo 사전조건화의 고유벡터 계산 (L, R)
↓
기울기를 회전된 고유기저 공간으로 투영
↓
그 공간에서 AdamW 업데이트 수행
↓
원래 파라미터 공간으로 역투영
결과: 360M~660M 언어 모델에서 AdamW 대비 반복 40%+ 감소, 벽시계 35%+ 향상. Shampoo 대비로도 ~20% 향상. 사전조건화 갱신 빈도에 대한 강건성도 Shampoo보다 우수하다.
NVIDIA NeMo가 SOAP을 공식 통합했다.
6부: Muon과의 연결 — 스펙트럼 노름의 가장 가파른 하강
2024년, 번스타인과 뉴하우스(Bernstein & Newhouse) 의 "Old Optimizer, New Norm: An Anthology"가 결정적 연결을 밝혔다:
Shampoo에서 사전조건화 누적을 제거하면 = 스펙트럼 노름에서의 최급강하법
Muon (Keller Jordan, 2024년 10월)은 이 통찰에 기반한다:
Shampoo
→
누적 제거
→
모멘텀 추가
→
Muon
Shampoo가 행렬 역제곱근(O(m³) SVD)을 사용하는 대신, Muon은 뉴턴-슐츠 반복(5차 다항식)으로 직교화한다 — bfloat16에서 안정적으로 동작하며 비용이 훨씬 낮다. NanoGPT 스피드런에서 AdamW 대비 35% 속도 향상.
7부: "가장 빠른 옵티마이저는 모두 행렬 사전조건화를 쓴다"
Fantastic Pretraining Optimizers (Stanford, 2025년 9월)
"Fantastic Pretraining Optimizers and Where to Find Them"
Kaiyue Wen, David Hall, Tengyu Ma, Percy Liang (Stanford)
arXiv: 2509.02046 (2025년 9월)
11개 옵티마이저를 0.1B~1.2B 규모에서 1x~8x Chinchilla 계산으로 비교
핵심 발견: "가장 빠른 옵티마이저(Muon, SOAP)는 모두 행렬을 사전조건화로 사용한다 — 기울기에 행렬을 곱하지, 원소별 스칼라를 곱하지 않는다."
| 계산 영역 | 최고 옵티마이저 |
|---|
| 1x~4x Chinchilla | Muon |
| 8x+ Chinchilla (과훈련) | SOAP, Kron |
그러나 중요한 단서: 규모가 커지면 이점이 줄어든다 — 0.1B에서 1.4배이던 것이 1.2B에서는 ~1.1배. 잘 튜닝된 AdamW/NAdamW도 큰 규모에서는 행렬 방법에 근접한다.
8부: 한계와 도전
Shampoo 계열의 한계
계산 오버헤드
O(m³ + n³)
행렬 역제곱근/고유분해. CPU 오프로딩으로 완화
메모리
m² + n² per layer
Adam의 mn보다 큼 (m, n이 비슷할 때). 블록 분할로 완화
하이퍼파라미터 복잡도
갱신 빈도, 그래프팅, 블록 크기, 지수...
Adam보다 조절할 것이 많음
스케일링 수렴
1.2B+에서 ~1.1배
규모가 커지면 AdamW와 격차 줄어듦
맺으며: 대각선을 넘어서
Shampoo의 가장 중요한 공헌은 구체적인 알고리즘이 아니라 패러다임의 전환이다.
2018년 이전, 딥러닝 최적화는 대각선의 세계였다. Adam, RMSProp, AdaGrad — 모두 파라미터를 독립적인 숫자로 취급했다. 가중치가 행렬이라는 사실, 그 행렬에 행과 열이라는 구조가 있다는 사실을 무시했다.
Shampoo는 이 구조를 보는 법을 가르쳤다. 크로네커 분해로 행과 열의 상관관계를 포착하면, 같은 기울기 정보에서 더 풍부한 곡률 추정을 얻을 수 있다.
이 아이디어는 이후 세 가지 방향으로 발전했다:
- SOAP (2024): Shampoo의 고유기저에서 Adam을 실행 → 더 안정적이고 더 빠름
- Muon (2024): 사전조건화 누적을 제거하고 뉴턴-슐츠로 직교화 → 더 가볍고 빠름
- Distributed Shampoo (2023): Meta의 PyTorch 구현, AlgoPerf 대회 우승
그리고 2025년의 결론: "가장 빠른 옵티마이저는 모두 행렬 사전조건화를 사용한다."
Adam이 아직 왕좌에 앉아 있지만, Shampoo가 개척한 행렬 사전조건화의 세계가 그 왕좌를 흔들고 있다. 다음 세대의 LLM을 학습시키는 것은, 대각선이 아니라 행렬을 보는 눈을 가진 옵티마이저일 수 있다.
