[논문 리뷰] Shampoo: Preconditioned Stochastic Tensor Optimization

Vineet Gupta, Tomer Koren외 1명|arXiv (Cornell University)|2018. 02. 26.

Stochastic Gradient Optimization Techniques참고 문헌 21인용 수 33

한 줄 요약

Shampoo는 딥러닝에서 텐서 구조를 가진 파라미터에 대해, 각 텐서 차원 별로 별도의 전처리 행렬을 유지함으로써 효율적인 2차 최적화를 가능하게 하는 새로운 전처리된 확률적 최적화 알고리즘입니다. 구조적 행렬 연산과 텐서플로우에서의 효율적 구현 덕분에, 단계당 런타임이 더 단순한 최적화기와 유사한 수준이지만, 실질적으로 SGD, Adam, AdaGrad보다 훨씬 더 빠른 수렴 속도를 달성합니다.

ABSTRACT

Preconditioned gradient methods are among the most general and powerful tools in optimization. However, preconditioning requires storing and manipulating prohibitively large matrices. We describe and analyze a new structure-aware preconditioning algorithm, called Shampoo, for stochastic optimization over tensor spaces. Shampoo maintains a set of preconditioning matrices, each of which operates on a single dimension, contracting over the remaining dimensions. We establish convergence guarantees in the stochastic convex setting, the proof of which builds upon matrix trace inequalities. Our experiments with state-of-the-art deep learning models show that Shampoo is capable of converging considerably faster than commonly used optimizers. Although it involves a more complex update rule, Shampoo's runtime per step is comparable to that of simple gradient methods such as SGD, AdaGrad, and Adam.

연구 동기 및 목표

• 고차원 기계학습 문제에서 전체 행렬 전처리의 확장성 한계를 해결하기 위해.
• 모델 파라미터의 텐서 구조(예: 행렬 및 4차원 컨볼루션 필터)를 활용하는 실용적인 전처리 방법을 개발하기 위해.
• 전체 헤시안 근사치 대신 각 차원별 전처리 행렬을 유지하여 딥러닝에서 효율적인 2차 최적화를 가능하게 하기 위해.
• 행렬 추적 부등식을 사용하여 스토하스틱 볼록 최적화 설정에서 이론적 수렴 보장을 확립하기 위해.
• 더 복잡한 업데이트 규칙에도 불구하고, Adam 및 SGD와 같은 표준 최적화기보다 수렴 속도에서 뛰어난 경험적 성능을 입증하기 위해.

제안 방법

• Shampoo는 각 텐서 차원 별로 전처리 행렬을 유지하며, 축적된 기울기의 2차 통계를 사용해 온라인으로 업데이트합니다.
• 각 전처리 행렬에 대해 대칭 행렬 제곱근 연산을 적용하고, 왼쪽 및 오른쪽 전처리 행렬의 기하 평균을 사용해 기울기를 스케일링합니다.
• 행렬 파라미터 $ W \in \mathbb{R}^{m \times n} $에 대해 업데이트 규칙은 $ W_{t+1} = W_t - \eta L_t^{-1/4} G_t R_t^{-1/4} $이며, 여기서 $ G_t $는 기울기이고 $ L_t, R_t $는 전처리 행렬입니다.
• 전처리 행렬 $ L_t $ 및 $ R_t $는 $ L_t = L_{t-1} + G_t G_t^\top $ 및 $ R_t = R_{t-1} + G_t^\top G_t $로 업데이트되어 기울기의 외적을 추적합니다.
• 이 원칙을 모든 텐서 모드에 적용함으로써 고차원 텐서로 일반화되며, 최적화 전반에 걸쳐 텐서 구조를 유지합니다.
• Shampoo는 표준 텐서 연산을 사용해 텐서플로우에서 효율적으로 구현되어 있어, 최소한의 코드 수정과 기존 딥러닝 프레임워크와의 호환성을 확보합니다.

실험 결과

연구 질문

• RQ1 딥러닝의 텐서와 같은 고차원 구조적 파라미터 공간에 대해 효율적으로 확장 가능한 전처리 최적화 방법을 설계할 수 있는가?
• RQ2 각 텐서 차원 별로 별도의 행렬을 유지하는 전처리 전략이, Adam 및 SGD와 같은 1차 최적화 방법보다 더 빠른 수렴을 달성할 수 있는가?
• RQ3 2차 정보를 사용하더라도, 이러한 방법이 낮은 단계당 계산 비용을 유지할 수 있는가?
• RQ4 스토하스틱 볼록 설정에서 이러한 구조적 전처리 접근법에 대해 어떤 이론적 수렴 보장을 확립할 수 있는가?
• RQ5 다양한 아키텍처와 데이터셋에서 대규모 딥러닝 모델에 대해 실질적으로 어떻게 성능을 발휘하는가?

주요 결과

• ResNet 및 Inception 아키텍처를 사용해 CIFAR-10 및 CIFAR-100에서 Shampoo는 Adam, SGD, AdaGrad보다 유의미하게 더 빠른 수렴 속도를 보이며, 더 적은 에포크 수로 낮은 훈련 손실을 달성합니다.
• LM1B 언어 모델링 벤치마크에서 Shampoo는 Adam 및 AdaGrad보다 낮은 테스트 퍼플렉서티를 달성했으며, 단계 수 기준으로 더 빠른 수렴 속도를 보였습니다.
• 행렬 역제곱근을 포함한 더 복잡한 업데이트 규칙에도 불구하고, Shampoo의 단계당 런타임은 SGD 및 Adam과 유사했으며, Tesla K40 GPU에서의 성능 저하는 작았습니다.
• 예를 들어 ResNet-55를 CIFAR-100에서 사용한 경우, Shampoo는 단계당 1.249단위로 Adam의 1.203보다 더 빠른 속도를 기록했습니다.
• 지연된 전처리 업데이트 및 모멘타임을 사용함으로써, 정확도를 희생시키지 않은 채 평균화된 런타임을 향상시켜 실용적 구현을 가능하게 했습니다.
• 모든 테스트된 모델에서 최신 기술 수준의 수렴 속도를 달성했으며, 딥러닝 최적화에서 구조 인식 전처리의 효과성을 입증했습니다.