[논문 리뷰] Scalable Second Order Optimization for Deep Learning

Rohan Anil, Vineet Gupta외 3명|arXiv (Cornell University)|2020. 02. 20.

Stochastic Gradient Optimization Techniques참고 문헌 49인용 수 29

한 줄 요약

이 논문은 인공지능 학습을 가속화하기 위해 인자 분해된 조정자(preconditioners)를 사용하는 두 번째 차수 적응형 방법(전체 행렬 Adagrad의 변종)의 확장 가능하고 하드웨어 최적화된 구현을 제시한다. CPU-가속기 파이프라인, 효율적인 행렬 루트 계산, 아키텍처 확장 기법을 활용하여, Transformer 및 BERT와 같은 대규모 모델에서 최대 47%의 벽시계 시간 절감을 달성하였으며, 수렴 속도와 학습 효율성 면에서 Adam과 같은 일阶 방법보다 뚜렷이 뛰어나다.

ABSTRACT

Optimization in machine learning, both theoretical and applied, is presently dominated by first-order gradient methods such as stochastic gradient descent. Second-order optimization methods, that involve second derivatives and/or second order statistics of the data, are far less prevalent despite strong theoretical properties, due to their prohibitive computation, memory and communication costs. In an attempt to bridge this gap between theoretical and practical optimization, we present a scalable implementation of a second-order preconditioned method (concretely, a variant of full-matrix Adagrad), that along with several critical algorithmic and numerical improvements, provides significant convergence and wall-clock time improvements compared to conventional first-order methods on state-of-the-art deep models. Our novel design effectively utilizes the prevalent heterogeneous hardware architecture for training deep models, consisting of a multicore CPU coupled with multiple accelerator units. We demonstrate superior performance compared to state-of-the-art on very large learning tasks such as machine translation with Transformers, language modeling with BERT, click-through rate prediction on Criteo, and image classification on ImageNet with ResNet-50.

연구 동기 및 목표

• 높은 메모리, 계산 및 통신 비용으로 인해 대규모 딥러닝에서 두 번째 차수 방법이 실용적이지 않다는 점을 해결한다.
• 이론적으로는 두 번째 차수 최적화의 이점과 실제 대규모 배포 간 격차를 메운다.
• 최신 이종 하드웨어(CPU + 가속기)에서 최첨단 모델을 위한 효율적이고 확장 가능한 두 번째 차수 학습을 가능하게 한다.
• 특히 임bedding과 같은 큰 레이어에 대해 전체 행렬 조정자 구현의 수치적 및 인프라 문제를 해결한다.
• 다양한 대규모 작업에서 Adam 및 LAMB와 같은 일阶 기준선 대비 뚜렷한 수렴 및 벽시계 시간 향상을 입증한다.

제안 방법

• 가속기에서 순방향/역방향 전파를 처리하는 동안 CPU에 조정자 계산을 오프로드하는 파이프라인 최적화 파이프라인을 설계한다.
• 비용이 많이 들고 수치적으로 불안정한 스펙트럼 분해(SVD 등)를 대체하기 위해, 행렬 루트(예: $L^{-1/4}$, $R^{-1/4}$) 계산을 위한 효율적이고 수치적으로 안정된 반복 방법을 도입한다.
• 예측 가능한 조정자 구조와 업데이트 규칙을 수정하여 Shampoo 알고리즘을 임베딩 행렬과 같은 큰 높은 질서의 레이어를 지원하도록 확장한다.
• AdaGrad에서 유도된 레이어별 학습률을 사용하여 학습을 안정화하고, AdaGrad보다 더 큰 학습률을 허용한다.
• TPU 코어 간 중복 계산을 줄이기 위해 조정자 통계를 크로스-레플리카 분할(sharding)한다.
• 비동기 CPU 오프로딩을 활용하여 최소한의 지연으로 학습 루프에 조정된 기울기 계산을 통합한다.

실험 결과

연구 질문

• RQ1 현대 딥러닝 워크로드에 대해 두 번째 차수 적응형 방법을 대규모에서 실용적으로 구현할 수 있는가?
• RQ2 전체 행렬 조정자 구현의 계산 및 메모리 오버헤드를 수렴 이점 손실 없이 줄일 수 있는가?
• RQ3 CPU-가속기 시스템에서 두 번째 차수 방법을 효율적으로 만들기 위해 필요한 아키텍처 및 알고리즘 최적화는 무엇인가?
• RQ4 Adam 및 LAMB와 같은 일阶 기준선 대비 두 번째 차수 방법이 학습 시간과 단계 수를 얼마나 줄일 수 있는가?
• RQ5 수치적 안정성과 확장성 간의 트레이드오프는 대규모 두 번째 차수 최적화기 설계에 어떻게 영향을 미치는가?

주요 결과

• WMT’14 영어-프랑스어 번역 작업에서 표준 Transformer는 12.0시간에서 6.7시간으로 학습 시간이 45% 감소했고, 더 큰 Transformer-Big 모델은 47.0시간에서 29.5시간으로 37% 감소하였다.
• 32K 배치 크기에서 BERT-Large 언어 모델링 작업에서 Shampoo는 하이퍼파ram터 조정 없이도 16% 적은 단계 수로 더 높은 마스크된-LM 정확도를 달성했고, 벽시계 시간은 3.8시간에서 3.65시간으로 4% 감소시켰다.
• ResNet-50을 사용한 CIFAR-10에서 Shampoo는 143 에포크 만에 93.45% 정확도를 달성했고, 기준선은 300 에포크를 소요하여 총 학습 시간을 42% 절감(1428초에서 827초로)하였다.
• Criteo 1TB 클릭-through 비율 예측 작업에서 Shampoo는 학습 시간을 13분에서 8.2분으로 줄여 대규모 희소 모델에서 뛰어난 효율성을 입증하였다.
• 조정자 계산이 아직 완전히 최적화되지 않은 모델에서도 단계 수뿐만 아니라 실제 벽시계 시간에서도 뚜렷한 속도 향상을 달성하였다.
• 현재 구현은 BERT-Large에서 단계 시간이 14% 증가하고 있으므로, 크로스-레플리카 분할 등을 통한 향후 최적화로 더 큰 성능 향상이 기대된다.