[논문 리뷰] Adafactor: Adaptive Learning Rates with Sublinear Memory Cost

Noam Shazeer, Mitchell Stern|arXiv (Cornell University)|2018. 04. 11.

Stochastic Gradient Optimization Techniques참고 문헌 10인용 수 163

한 줄 요약

Adafactor는 행렬 매개변수의 두 번째 모멘트 추정치를 행별 및 열별 합으로 분解(factoring)하여 메모리 효율적인 적응 최적화기를 도입하고, Transformer 학습에서 Adam과 견줄 만한 성능을 유지하며 서브선형 메모리 사용을 가능하게 한다; 또한 훈련 안정화를 위한 업데이트 클리핑과 감소율 스케줄을 제안하고 규모에 맞춘 업데이트를 위한 상대 스텝 크기(Adafactor)를 도입한다.

ABSTRACT

In several recently proposed stochastic optimization methods (e.g. RMSProp, Adam, Adadelta), parameter updates are scaled by the inverse square roots of exponential moving averages of squared past gradients. Maintaining these per-parameter second-moment estimators requires memory equal to the number of parameters. For the case of neural network weight matrices, we propose maintaining only the per-row and per-column sums of these moving averages, and estimating the per-parameter second moments based on these sums. We demonstrate empirically that this method produces similar results to the baseline. Secondly, we show that adaptive methods can produce larger-than-desired updates when the decay rate of the second moment accumulator is too slow. We propose update clipping and a gradually increasing decay rate scheme as remedies. Combining these methods and dropping momentum, we achieve comparable results to the published Adam regime in training the Transformer model on the WMT 2014 English-German machine translation task, while using very little auxiliary storage in the optimizer. Finally, we propose scaling the parameter updates based on the scale of the parameters themselves.

연구 동기 및 목표

• 모델 크기가 커짐에 따라 적응형 그래디언트 방법의 메모리 제약을 동기화한다.
• 행렬 매개변수에 대해 per-parameter 메모리를 O(nm)에서 O(n+m)로 줄이는 factored second-moment 추정치를 제안한다.
• 두 번째 모멘트의 감소 속도가 너무 느려질 때의 불안정성 문제를 식별하고 업데이트 클리핑 및 적응적 감소를 통해 해결책을 제시한다.
• Transformer 학습에서 Adafactor가 Adam과 유사한 성능을 달성하면서도 훨씬 더 적은 메모리를 필요로 함을 입증한다.
• 매개변수 규모에 따라 업데이트를 조정하는 상대 스텝 사이즈를 도입하여 최적화의 확장성을 높인다.

제안 방법

• 행렬 값 매개변수에 대해 근사식 V ≈ RS를 도입하고 R ∈ R^{n×k}, S ∈ R^{k×m}이며 k ≪ n,m 인 계층 분해를 제시한다.
• (rank-1) 케이스(k=1)에 대한 해를 해석적으로 도출하여 행/열 합의 지수적 평활화와의 호환성(V1_m1_n^T V / 1_n^T V 1_m)을 보장한다.
• per-row 및 per-column 누적값(R_t와 C_t)을 사용하여 Factored Second Moments를 적용한 Adam을 구현하고, 1_n^T R_t로 정규화하여 ㅤfㅤaㅤcㅤtㅤoㅤrㅤeㅤdㅤ V_t를 형성한다.
• RMS(U_t)가 임계값 d를 초과할 때 큰 미스케일 업데이트를 상한하는 업데이트 클리핑을 제안한다.
• 두 번째 순간의 감소 스케줄을 증가시키는 방법(또는 Reddi 등에 따른 ㅆ aㅜ)과 훈련을 안정화시키기 위한 대안적 스케줄을 제안한다.
• Adafactor를 상대 스텝 최적화기로 정의하며 실제 스텝 alpha_t는 매개변수 규모의 RMS와 상대 스텝 rho_t에서 계산되고, per-parameter 업데이트 U_t = G_t / sqrt(V_hat_t)이며 필요 시 클리핑을 허용한다.

실험 결과

연구 질문

• RQ1 사실화된(행별/열별) 두 번째 모멘트 추정이 완전한 두 번째 모멘트 누적기와 성능이 일치할 수 있는가?
• RQ2 행렬 매개변수의 팩터화를 통한 메모리 감소가 Transformer 학습과 같은 대규모 작업에서 수렴 및 모델 품질에 영향을 미치는가?
• RQ3 모멘텀을 분리하고 적응 학습률을 사용할 때 어떤 안정성 문제가 발생하며 업데이트 클리핑 및 감소율 스케줄이 이를 완화할 수 있는가?
• RQ4 매개변수 규모에 따라 크기에 맞춘 상대 스텝 크기가 서로 다르게 스케일되는 초기화에서 견고성을 개선하는가?
• RQ5 제안된 해결책(업데이트 클리핑, 감소 증가, 상대 스텝 크기)이 현대 신경 기계 번역 작업에서 실전에서 어떻게 상호작용하는가?

주요 결과

• 팩터링된 두 번째 모멘트 추정은 매트릭스 매개변수에 대해 O(nm)에서 O(n+m)로 메모리를 줄이면서도 Transformer BLEU 점수를 전체 누적기 Adam과 유사하게 제공합니다.
• 모멘텀을 제거하면 훈련이 불안정해질 수 있지만 업데이트 클리핑과 적절한 감소 스케줄이 안정성을 회복시킵니다.
• 임계값 d로의 업데이트 클리핑은 워밍업 비적용 설정에서 안정성을 향상시켰으며, d=1일 때 불안정을 크게 완화했고(d=2는 개선에 큰 효과 없음).
• 상대 스텝 크기는 Adafactor와 결합되었을 때 경쟁력 있는 성능을 유지했고 서로 다른 임베딩 매개변수 규모에서도 견고했습니다.
• 두 번째 모멘트를 증가시키는 스케줄(예: 1 - t^{-c})은 특정 c 값에서 안정적이고 수렴하는 결과를 보였으며, 특히 c=0.5(및 관련 변형)와 클리핑을 조합했을 때 그렇다.
• 랭크-1 또는 팩터화된 표현과 제안된 스케일링을 사용하면 BLEU 점수가 Adam 기반 기준에 근접하면서도 보조 저장소가 서브선형으로 줄어드는 Transformer 모델의 학습이 가능해졌다.