DASH: 배치 블록 조건화와 효율적인 역제곱근 솔버를 통한 더 빠른 Shampoo

DASH: Faster Shampoo via Batched Block Preconditioning and Efficient Inverse-Root Solvers

February 2, 2026

저자: Ionut-Vlad Modoranu, Philip Zmushko, Erik Schultheis, Mher Safaryan, Dan Alistarh

cs.AI

초록

샴푸(Shampoo)는 대표적인 준2차 최적화 도구 중 하나로, 그 변형 버전이 MLCommons AlgoPerf 대회에서 우승한 바 있으며, 압축이 용이한 낮은 활성화 이상치(activation outliers)를 가진 모델을 생성하는 것으로 입증되었습니다. 그러나 현재 샴푸를 적용하려면 복잡한 내부 연산으로 인해 상당한 계산 속도 저하를 감수해야 합니다. 본 논문에서는 이러한 단점을 해결하기 위한 중요한 단계로, 두 가지 새로운 핵심 기술을 기반으로 하는 분산 샴푸의 고속 구현체인 \method(분산 가속 샴푸)를 제안합니다. 첫째, 사전 조건자 블록(preconditioner blocks)을 3D 텐서로 쌓아 GPU 활용도를 크게 향상시킬 수 있음을 보여줍니다. 둘째, 샴푸에 필요한 역행렬 제곱근(inverse matrix roots) 계산을 위한 새로운 고속 접근법으로 Newton-DB 반복법과 체비쇼프 다항식 근사법을 소개합니다. 이러한 알고리즘적 기여와 함께, 행렬 스케일링(matrix scaling)이 샴푸의 수렴에 어떻게 결정적인 영향을 미치는지에 대한 첫 번째 심층 분석을 제공합니다. 실용적인 측면에서, 우리의 GPU 인식 구현은 최적화가 잘된 기존 분산 샴푸 대비 최대 4.83배 빠른 최적화 단계 속도를 달성했으며, Newton-DB는 검증된 모든 방법 중 반복당 가장 낮은 검증 퍼플렉서티(validation perplexity)를 기록했습니다. 우리의 코드는 https://github.com/IST-DASLab/DASH에서 확인할 수 있습니다.

English

Shampoo is one of the leading approximate second-order optimizers: a variant of it has won the MLCommons AlgoPerf competition, and it has been shown to produce models with lower activation outliers that are easier to compress. Yet, applying Shampoo currently comes at the cost of significant computational slowdown, due to its expensive internal operations. In this paper, we take a significant step to address this shortcoming by proposing \method (for Distributed Accelerated SHampoo), a faster implementation of Distributed Shampoo based on two main new techniques: First, we show that preconditioner blocks can be stacked into 3D tensors to significantly improve GPU utilization; second, we introduce the Newton-DB iteration and the Chebyshev polynomial approximations as novel and faster approaches for computing the inverse matrix roots required by Shampoo. Along with these algorithmic contributions, we provide a first in-depth analysis of how matrix scaling critically affects Shampoo convergence. On the practical side, our GPU-aware implementation achieves up to 4.83times faster optimizer steps compared to the well-optimized Distributed Shampoo, while Newton-DB attains the lowest validation perplexity per iteration among all tested methods. Our code is available at https://github.com/IST-DASLab/DASH.