[논문 리뷰] AdaX: Adaptive Gradient Descent with Exponential Long Term Memory

Wenjie Li, Zhaoyang Zhang외 2명|arXiv (Cornell University)|2020. 04. 21.

Advanced Vision and Imaging참고 문헌 27인용 수 24

한 줄 요약

AdaX는 기존 Adam의 지수 이동 평균을 지수 장기 기억 메커니즘으로 대체함으로써 훈련의 안정성과 수렴 성능을 향상시키는 새로운 적응 최적화 알고리즘을 제안한다. 볼록 및 비볼록 설정 모두에서 수렴성이 이론적으로 입증되었으며, 시각 및 NLP 벤치마크에서 Adam을 능가하고, 모멘텀이 있는 SGD와 경쟁 수준의 성능을 보인다. 특히 일반화 성능과 하이퍼파라미터 선택에 대한 강건성에서 뛰어나다.

ABSTRACT

Although adaptive optimization algorithms such as Adam show fast convergence in many machine learning tasks, this paper identifies a problem of Adam by analyzing its performance in a simple non-convex synthetic problem, showing that Adam's fast convergence would possibly lead the algorithm to local minimums. To address this problem, we improve Adam by proposing a novel adaptive gradient descent algorithm named AdaX. Unlike Adam that ignores the past gradients, AdaX exponentially accumulates the long-term gradient information in the past during training, to adaptively tune the learning rate. We thoroughly prove the convergence of AdaX in both the convex and non-convex settings. Extensive experiments show that AdaX outperforms Adam in various tasks of computer vision and natural language processing and can catch up with Stochastic Gradient Descent.

연구 동기 및 목표

• 비볼록 최적화에서 Adam의 불안정성과 수렴 불가 문제를 해결하기 위해, 특히 국소 최적해에 수렴하는 경향을 완화하고자 한다.
• 2차 모멘트 계산에서 Adam의 지수 이동 평균의 한계를 극복하여 불안정한 학습률과 열악한 일반화 성능을 개선하고자 한다.
• 빠른 수렴을 유지하면서 일반화 성능을 향상시키는 이론적으로 타당한 적응 최적화기를 설계하고자 한다.
• 다양한 딥러닝 작업에서 AdaX가 Adam, AdamW, AMSGrad, 모멘텀이 있는 SGD에 비해 우월함을 경험적으로 검증하고자 한다.
• 실제 훈련 환경에서 하이퍼파라미터, 특히 $eta_2$와 학습률 선택에 대해 AdaX가 강건한지 입증하고자 한다.

제안 방법

• AdaX는 기존 Adam의 이전 기울기 제곱의 지수 이동 평균을 장기적 기울기 정보 누적 메커니즘으로 대체한다.
• AdaX의 2차 모멘트는 모든 이전 제곱 기울기의 지수 가중 합으로 계산되며, 이는 과거 기울기 행동의 지속적인 기억을 보장한다.
• 학습률은 장기 기억 행렬의 역제곱근을 사용해 적응적으로 스케일링되며, Adam과 유사하지만 더 높은 안정성을 확보한다.
• AdaX는 1차 모멘텀을 유지하면서 적응적 분모를 더 안정적이고 누적된 기울기 기억으로 대체한 수정된 업데이트 규칙을 사용한다.
• 알고리즘은 볼록 및 비볼록 설정 모두에서 수렴성이 이론적으로 입증되었으며, AMSGrad와 유사한 수렴 속도를 보인다.
• AdaX는 기본 $eta_2 = 10^{-4}$로 구현되었으며, 이는 강건성과 계산 효율성이 뛰어나다는 게 입증되었다.

실험 결과

연구 질문

• RQ1 비볼록 문제에서 Adam의 빠른 수렴이 불안정한 2차 모멘트 추정으로 인해 국소 최적해에 수렴하는 하위최적해로 이어지는가?
• RQ2 Adam의 지수 이동 평균을 지수 장기 기억 메커니즘으로 대체하면 수렴 불가 문제를 제거하고 일반화 성능을 향상시킬 수 있는가?
• RQ3 시각 및 NLP 작업 전반에서 AdaX는 Adam, AdamW, AMSGrad, 모멘텀이 있는 SGD에 비해 수렴 속도와 최종 성능 측면에서 어떻게 비교되는가?
• RQ4 실제 훈련 환경에서 AdaX는 하이퍼파라미터, 특히 $eta_2$와 초기 학습률 선택에 대해 강건한가?
• RQ5 AdaX는 적응 방법의 빠른 수렴 성능을 유지하면서도 모멘텀이 있는 SGD 수준의 성능을 달성할 수 있는가?

주요 결과

• CIFAR-10에서 AdaX는 AdamW보다 뛰어난 성능을 보였으며, 테스트 Top-1 정확도 94.5%를 기록하여 AdamW의 92.1%를 상회했다.
• ResNet-18을 사용한 ImageNet에서 AdaX-W는 Top-1 정확도 75.58%를 기록하여 AdamW(68.27%)와 AMSGrad(W)(68.76%)를 크게 앞섰다.
• ResNet-50을 사용한 ImageNet에서 AdaX-W는 Top-1 정확도 75.58%를 기록했으며, SGDM의 77.12%에 가까워졌고, 더 빠른 수렴 속도와 높은 훈련 정확도를 보였다.
• AdaX는 학습률과 $eta_2$ 하이퍼파라미터에 대해 강건했으며, 다양한 설정에서도 성능 저하가 최소화되었고, $eta_2 = 10^{-4}$일 경우 尤히 두드러졌다.
• One Billion Word 데이터셋에서의 언어 모델링 실험에서 AdaX는 모멘텀이 있는 SGD와 동등한 성능을 보였으며, 뛰어난 일반화 성능과 빠른 수렴을 입증했다.
• 실행 시간 벤치마크 결과, AdaX-W는 AdamW와 거의 유사한 효율성을 보였으며, ImageNet에서 훈련 시간은 50.74시간으로 AdamW의 50.10시간 대비 약 1.8% 증가에 그쳐, 계산 오버헤드가 최소한이었다.