[논문 리뷰] Memory-Efficient Adaptive Optimization

Rohan Anil, Vineet Gupta외 2명|arXiv (Cornell University)|2019. 01. 30.

Advanced Neural Network Applications인용 수 19

한 줄 요약

이 논문은 Adam과 Adagrad와 같은 적응형 최적화 방법의 메모리 오버헤드를 줄이기 위해 두 번째 순서 통계의 압축된, 저랭크 근사값만 유지함으로써 메모리 효율적인 적응형 최적화 방법 SM3를 제안한다. 이 방법은 유사하거나 더 나은 수렴 성능을 유지하면서 훨씬 큰 배치 크기와 모델 크기를 허용하며, 대규모 언어 모델 및 이미지 분류 작업에서 훈련 시간을 최대 2배 빠르게 한다. 수렴 보장을 유지한다.

ABSTRACT

Adaptive gradient-based optimizers such as Adagrad and Adam are crucial for achieving state-of-the-art performance in machine translation and language modeling. However, these methods maintain second-order statistics for each parameter, thus introducing significant memory overheads that restrict the size of the model being used as well as the number of examples in a mini-batch. We describe an effective and flexible adaptive optimization method with greatly reduced memory overhead. Our method retains the benefits of per-parameter adaptivity while allowing significantly larger models and batch sizes. We give convergence guarantees for our method, and demonstrate its effectiveness in training very large translation and language models with up to 2-fold speedups compared to the state-of-the-art.

연구 동기 및 목표

• 대규모 훈련에서 모델 크기와 배치 크기를 제한하는 Adam과 Adagrad와 같은 적응형 최적화 방법의 높은 메모리 오버헤드 문제를 해결하기 위해.
• 각 파라미터에 대한 적응성의 이점을 유지하면서 메모리 소비를 극적으로 줄이는 방법을 개발하기 위해.
• 특히 NLP와 비전 분야에서 매우 큰 모델의 훈련을 가능하게 하기 위해 메모리를 확보하여 더 큰 배치 크기를 허용하기 위해.
• 볼록 온라인 최적화 설정에서 이론적 수렴 보장을 제공하기 위해.
• 메모리 소비를 줄이고 빠른 월 타임을 달성하면서도 유사하거나 더 나은 수렴 성능을 경험적으로 입증하기 위해.

제안 방법

• SM3는 각 파라미터의 전체 2차 모멘트를 저장하는 대신, 두 번째 순서 기울기 통계의 저랭크 근사를 유지한다.
• 이 방법은 메모리 효율적인 스토하스틱 업데이트 규칙을 사용하여, 랭크-1 근사를 통해 적응 통계를 압축한다.
• 저장소를 최소화하면서도 각 파라미터의 적응성을 유지하기 위해 대각 스케일링 요소를 적용한다.
• 표준 딥 러닝 프레임워크와의 호환성을 고려하여 설계되었으며, 최소한의 코드 변경만 필요하다.
• 피셔 정보 매트릭스의 구조적 근사를 통해 자연 기울기 유사 업데이트를 활용한다.
• 활성화 구조에 대한 사전 지식이 필요 없이 기울기 패턴에 동적으로 적응한다.

실험 결과

연구 질문

• RQ1 메모리 효율적인 적응형 최적화 방법인 SM3는 메모리 사용량을 줄이면서도 Adam과 Adagrad 수준의 수렴 성능을 유지할 수 있는가?
• RQ2 SM3에서 얻는 메모리 절감 효과는 분산 훈련에서 배치 크기를 얼마나 늘릴 수 있는가?
• RQ3 2차 통계의 저랭크 근사가 대규모 모델의 수렴 속도에 영향을 미치는가?
• RQ4 Adafactor와 Shampoo와 같은 기존의 메모리 효율적 방법과 비교했을 때, SM3는 메모리, 속도, 정확도 측면에서 어떻게 다른가?
• RQ5 고정된 계산 예산 하에서 표준 적응형 최적화 방법보다 SM3가 더 빠른 월 타임 수렴을 달성할 수 있는가?

주요 결과

• WMT’14 en→fr 번역 작업에서 SM3는 Adam보다 40% 적은 메모리로 40.50의 BLEU 점수를 기록했으며, 더 큰 배치 크기를 허용했다.
• BERT-Large 언어 모델링에서, SM3는 이중 배치 크기를 사용했을 때 Adam보다 35% 더 빠른 월 타임에 70%의 마스킹 LM 정확도를 달성했다.
• 배치 크기가 2048일 때, SM3는 코어당 6.02 GiB의 메모리만 사용했으며, 이는 Adam이 1024일 때의 수준과 유사했지만, 2^16까지 확장해도 메모리 한계에 도달하지 않았다.
• AmoebaNet-D를 사용한 ImageNet에서 SM3는 상위-1 정확도 78.71%와 상위-5 정확도 94.31%를 기록했으며, 최신 기술 수준의 성능을 달성했다.
• 대규모 모델에서 SM3는 Adam 대비 40% 메모리 사용량을 줄였고, 유사한 수렴 속도와 정확도를 유지했다.
• SM3의 스텝당 시간은 Adam보다 3% 빠르며, 배치 크기를 두 배로 늘였을 때 수렴에 소요되는 월 타임은 최대 50% 감소했다.