DASH: 배치 블록 프리컨디셔닝과 고속 역루트 솔버로 가속화된 샴푸 옵티마이저

초록

DASH는 블록 프리컨디셔너를 3차원 텐서로 스택해 GPU 활용도를 높이고, Newton‑DB와 Chebyshev 다항식 근사를 도입해 역행렬 루트 계산을 가속화한다. 이를 통해 기존 Distributed Shampoo 대비 최적화 단계가 최대 4.83배 빨라졌으며, 검증 퍼플렉시티도 개선되었다.

상세 분석

본 논문은 Shampoo 최적화기의 가장 큰 병목인 역행렬 루트 연산을 두 가지 관점에서 혁신한다. 첫 번째는 블록 프리컨디셔너를 순차적으로 처리하던 기존 구현을 탈피해, 동일 크기의 B×B 블록들을 3‑D 텐서로 결합하고 GPU에서 병렬 BMM(배치 매트릭스‑매트릭스 곱)으로 일괄 처리한다. 이 설계는 메모리 접근 패턴을 연속적으로 만들고, Tensor Core의 고속 행렬 연산을 최대한 활용해 연산량을 거의 B배 정도 감소시킨다. 두 번째는 역루트 계산을 전통적인 고유값 분해(EVD) 대신, Newton‑Denman‑Beavers(NDB)와 Chebyshev 다항식(CBSHV) 기반 반복법으로 교체한다. NDB는 √A와 A⁻¹ᐟ²를 동시에 수렴시키며, 초기 1회 행렬곱을 폐기해 첫 반복부터 2개의 곱셈만 필요하도록 최적화했다. 또한 논문은 스케일링 문제를 깊이 파고들어, 기존 Distributed Shampoo가 사용하던 Frobenius norm 기반 스케일링이 λ_max보다 10‑100배 크게 과대평가되어 수렴 속도를 저해함을 실험적으로 입증한다. 이를 해결하기 위해 half‑precision Power‑Iteration 기반 multi‑Power‑Iteration을 제안, 스펙트럼 반경을 정확히 추정해 최적 스케일링 상수를 제공한다. 실험에서는 Newton‑DB가 Coupled‑Newton(CN)과 EVD보다 적은 반복 횟수로 동일 혹은 더 낮은 검증 퍼플렉시티를 달성했으며, Chebyshev 근사는 FP16 환경에서 메모리와 연산 효율을 동시에 끌어올렸다. 전체 시스템은 grafting, 로드밸런싱, FP16 지원 등 기존 Distributed Shampoo의 기능을 그대로 유지하면서도, 블록 스택과 고속 역루트 솔버 덕분에 전체 최적화 단계가 평균 3‑5배 가속화되었다. 이러한 설계는 대규모 언어 모델 훈련뿐 아니라, 메모리·연산 제약이 있는 모바일/엣지 환경에서도 근사 2차 최적화기의 실용성을 크게 확대한다.