논문 링크 : arxiv
저자 코드 : Author’s code

카카오 블라인드 공채에 지원해서 온라인, 오프라인 코딩테스트를 넘고 나니 회사별 과제가 주어졌습니다. 저는 카카오 브레인에 지원해서, RAdam 논문을 읽고 요약한 뒤 카카오에서 제공한 코드를 개선하라는 과제를 받았습니다. 비록 떨어지긴 했지만() 너무나도 재밌는 논문이었습니다. 논문 내용들을 요약해 보았습니다.

Abstract

• 이 논문에서는 기존 Adaptive learning rate 방법들이 bad local optima에 빠지던 문제가 학습초기에 adpative learning rate의 variance가 크기 때문임을 보입니다.
• 또한 warm-up learning rate 방법도 adpative learning rate의 variance를 줄여서 잘 작동하는 것임을 보입니다.
• 마지막으로 Adaptive learning rate의 variance를 일정하게 만들어주는 RAdam이라는 새로운 방법을 제시합니다.

Variance of adaptive rate in the early stage

저자는 학습 초기에 샘플이 부족하기 때문에 adaptive learning rate가 부적절하게 큰 variance를 갖고, 이것이 안좋은 local optima로 이끈다는 것을 수식적으로, 또 실험적으로 보여줍니다.

수식적 분석

저자는 먼저 time-step 마다 얻어지는 gradient들(\({g_1, \cdots, g_t}\))을 모두 정규분포 \( \mathcal{N}(0, \sigma^2)\)에서 i.i.d 하게 뽑힌 random variables 이라고 가정합니다. 가정에 의해 \( \frac{1}{g_1^2} \)은 the scaled inverse chi-squared distribution을 따르게 됩니다. 이때 sample이 1개인 첫번째 adaptive learning rate의 분산을 계산하면 다음의 식이 나옵니다. \[ \mathrm{Var}[\sqrt{\frac{1}{g_1^2}}] \propto \int_0^{\infty} x^{-1} e^{-x} dx \] 이 식을 통해 sample이 1개일 때 adaptive learning rate의 분산이 발산하는 것을 확인할 수 있습니다.

실험적 분석

저자는 Adam을 warm-up을 적용한 채로, 또 다른 한번은 warm-up을 적용 안한채로 각각 학습하며 iteration마다 gradient 값의 측정해보았습니다. x축은 gradient의 절댓값을 log scale로 나타내었고 그래프의 높이는 frequency를 나타냅니다. 그리고 왼쪽 그래프는 iteration 1번째부터 100번까지의 결과, 오른쪽 그래프는 7만번까지의 결과를 보여줍니다.
실험 결과를 보면 warm-up 없이 학습한 경우 iteration이 10번도 되기 전에 gradient 값들이 상대적으로 작은 값으로 변형되는 것을 볼 수 있습니다. 이러한 결과는 초기에 bad/suspicious local optima에 빠진다는 것을 의미합니다. warm-up은 초기의 작은 update로 초기에 bad/suspicious local optima에 빠지는 것을 방지합니다.

저자는 또한 Adam-2k, RAdam, Adam-eps로도 위 실험을 반복합니다.
Adam-2k는 처음 2000번 update동안 parameter를 update 시키지 않고 adaptive learing rate만을 계산하는 방법입니다. 따라서 다른 방법들과 비교를 위해 iteration이 -1999번부터 시작된다고 하고 실험을 진행하였습니다. 그 결과 bad local optima에 빠지지 않는 것을 확인 할 수 있습니다. 이는 저자의 가설에서 Adam이 초기에 sample이 부족한 것이 문제의 원인이라는 것을 입증합니다.
Adam-eps는 adaptive learning rate 식인 \(\hat{\psi}(g_1, \cdots, g_t) = \frac{\sqrt{1-\beta_2^t}}{\epsilon + \sqrt{(1 - \beta_2)\sum_{i = 1}^t \beta_2^{t-i}g_i^2}}\) 에서 엡실론 값을 키워 variance를 낮추는 방법입니다. 이 결과 역시 초기의 bad local optima에 빠지지 않는 것을 보입니다.

Analysis of adaptive learning rate variance

이번에는 adaptive learning rate의 variance를 구해보고 sample의 수가 늘어날수록 variance가 감소함을 보입니다. 이때 분석의 용이함을 위해 저자는 총 3번의 근사를 시킵니다.
첫번째 근사는 exponential moving average를 simple average로 근사시키는 것입니다. 저자는 exponential moving average가 simple average보다 더 큰 분산을 갖을 뿐더러 t가 작을때는 simple average와 차이도 상대적으로 작으므로 이 근사가 무리가 없다고 주장합니다. \[p(\psi(.)) = p(\sqrt{\frac{1-\beta_2^t}{(1 - \beta_2)\sum_{i = 1}^t \beta_2^{t-i}g_i^2}}) \approx p(\sqrt{\frac{t}{\sum_{i=1}^t g_i^2}}) \] \( g_i \sim \mathcal{N}(0, \sigma^2) \)을 가졍했기 때문에 \( \frac{t}{\sum_{i=1}^t g_i^2} \sim Scale–inv–\mathcal{X}^2(t,\frac{1}{\sigma^2}) \) 입니다. 근사 시킨 값이 scaled inverse chi-square 분포를 따르므로 \(\frac{1-\beta_2^t}{(1 - \beta_2)\sum_{i = 1}^t \beta_2^{t-i}g_i^2}\) 또한 자유도 \( \sigma \)를 갖는 scaled inverse chi-square 분포로 가정하겠습니다. 이 가정하에서 저자는 분산을 구합니다. 여기서 \( \mathcal{B}(.) \)은 Beta function을 의미합니다. 또한 \( \mathrm{Var}[\psi(.)] \)의 미분 값을 통해 자유도 \( \sigma \)가 증가함에 따라 분산이 단조 감소하는 것을 보입니다. 이로써 학습 초기에 sample이 부족한 것 때문에 나중보다 분산이 크다는 것을 보였습니다.

Rectified adaptive learning rate

저자는 앞에서 구한 계산값을 이용해 adaptive learning rate의 분산을 일정하게 만들어주는 방법을 제시합니다. 이를 위해서 먼저 자유도인 \( \sigma \)를 \( t \)를 이용해 구합니다.

Estimation of \( \sigma \)

자유도를 구하기 위해 두번째 근사를 합니다. 경제학에서 exponential moving average(EMA)는 simple moving average(SMA)로 자주 근사된다고 합니다. \[ p(\frac{(1 - \beta_2)\sum_{i = 1}^t \beta_2^{t-i}g_i^2}{1-\beta_2^t}) \approx p(\frac{\sum_{i = 1}^{f(t, \beta_2)} g_{t+1-i}^2}{f(t, \beta_2)}) \] 이때 \( f(t, \beta_2) \)는 SMA의 길이로 SMA가 EMA와 같은 “center of mass”를 갖도록 구해집니다. 따라서 \(f(t, \beta_2) = \frac{2}{1 - \beta_2} - 1 - \frac{2 t \beta_2^t}{1 - \beta_2^t}\) 입니다. \[ \frac{1-\beta_2^t}{(1 - \beta_2)\sum_{i = 1}^t \beta_2^{t-i}g_i^2} \approx \frac{f(t, \beta_2)}{\sum_{i=1}^{f(t, \beta_2)} g_i^2} \sim Scale–inv–\mathcal{X}^2(\rho,\frac{1}{\sigma^2}) \] \[ \rho \approx f(t, \beta_2) = \frac{2}{1 - \beta_2} - 1 - \frac{2 t \beta_2^t}{1 - \beta_2^t} \] 위와 같이 \( \rho \)의 근사값으로 \( f(t, \beta_2) \)을 이용해 \( t \)에 대한 함수로 표현할 수 있게 되었습니다. 추가로 \( \rho_t = f(t, \beta_2) \)라고 표기했을 때 \( \rho_\infty = \frac{2}{1 - \beta_2} - 1 \)가 됩니다.

Variance estimation and rectification

위에서 분산은 자유도 \( \rho \)가 증가함에 따라 단조감소함을 보였으므로 \( \min_{\rho_t} \mathrm{Var}[\psi(.)] = \mathrm{Var}[\psi(.)]|_ {\rho_t = \rho_\infty} \)이고 이때의 분산을 \( C_{\mbox{var}} \)라고 표기하겠습니다. 저자는 학습 초기의 분산이 커서 bad local optima에 빠지는 문제를 방지하기 위해 매 timestep마다 분산이 \( C_{\mbox{var}} \)가 되도록 아래와 같이 rectification 해줍니다. \[ \mathrm{Var}[r_t \,\psi(g_1, \cdots, g_t)] = C_{\mbox{var}} \quad \mbox{where} \quad r_t = \sqrt{\frac{C_{\mbox{var}}}{\mathrm{Var}[\psi(g_1, \cdots, g_t)]}} \] 위에서 \( \mathrm{Var}[\psi(.)] \) 값을 구했지만 수치적으로 불안정하므로 \( \sqrt{\psi^2(.)} \) 에 대한 1차 근사(세번째 근사)를 이용해서 다시 구하면 다음과 같습니다. \[ \mathrm{Var}[\psi(.)] \approx \frac{\rho_t}{2(\rho_t - 2)(\rho_t - 4)\sigma^2} \] 분산을 \( C_{\mbox{var}} \)로 일정하게 유지시키는 \( r_t \) 값은 다음과 같습니다. \[ r_t = \sqrt{\frac{(\rho_t - 4)(\rho_t - 2)\rho_\infty}{(\rho_\infty - 4)(\rho_\infty - 2)\rho_t}} \] 단 \(\rho_t\)가 4보다 작을 경우 adaptive learning rate는 무시합니다. 이를 정리해 적용한 알고리즘은 다음과 같습니다. 위와 같이 Adam의 분산을 일정하게 만들어주는 방법을 통해 RAdam이라는 새로운 방법을 제안합니다.

In comparison with warmup

저자는 \( r_t \)가 결국 gradient에 \( \frac{min(t, T_w)}{T_w} \)을 곱하는 heuristic linear warmup과 비슷한 역할을 한다고 말합니다. 하지만 RAdam은 \( T_w \)와 같은 hyperparameter를 추가로 필요하지 않는 장점이 있습니다. 또한 뒤에서 살펴볼 실험 결과에 따르면 Adam-warmup은 learning rate를 증가시킬 step 수에 따라 learning rate에 민감함을 보이고, RAdam은 robust함을 보입니다.

Results

저자는 Language Modeling을 평가하기 위해 One Billion Word 데이터셋을, Image Classification을 평가하기 위해 Cifar10 와 ImageNet 데이터셋을 사옹했습니다.

RAdam vs Vanila Adam vs SGD

저자는 먼저 vanila Adam과 비교해 RAdam이 더 좋은 성능을 내는 것을 보입니다. 또한 learning rate의 변화에 robust함을 보입니다.

RAdam vs Adam-warmup

RAdam은 위에서 말했듯이 Adam-warmup과 달리 처음에 learning rate를 증가시킬 step수를 지정하지 않아도 되므로 hyper-parameter가 하나 더 적습니다. 또한 Adam-warmup은 learning rate를 증가시킬 step 수에 따라 learning rate에 민감함을 보이고, RAdam은 robust함을 보입니다.

느낀점

optimization에 관련된 논문은 처음으로 읽어봤습니다. Adam이 나온 이후로 모든 task에 general하게 잘 작동하는 optimizer가 나오지 않은것으로 알고 있습니다. 새로운 시각으로 접근해서 재밌었고 논문에서 이론적 증명을 통해 RAdam의 정당성을 보여준 점이 좋았습니다.

Reference

New State of the Art AI Optimizer: Rectified Adam (RAdam). - Less Wright