Adam can converge without any modification on Update rules

 

모바일 앱 환경에서는 latex 수식이 깨져 나타나므로 가급적 웹 환경에서 봐주시길 바랍니다.

---

논문 제목 : Adam can converge without Any Modification On Update Rules

출판 연도 : 2022 Neurips

논문 저자 : Yushun Zhang et al.

---

오늘 리뷰할 논문은 Adam에 대한 convergence analysis를 수행한 논문으로 가장 최신 연도의 논문이다.

이 논문 이후 arxiv 기준으로 Adam에 대한 convergence analysis 논문이 몇 개 나왔으나 현재까지 conference에 accept된 논문으로는 위 논문이 가장 최신이다.

이 논문은 두 번에 걸쳐서 리뷰할 예정이다.

 

우선, 다음과 같은 (Non-convex) optimization problem에 Adam 알고리즘을 적용하면 다음과 같다.

 

$\min_x f(x) = \sum_{i=1}^n f_i(x)$

 

for $k = 1 \rightarrow \infty$ do

$m_k = \beta_1 m_{k-1} + (1 - \beta_1) \nabla f_k (x_k)$

$v_k = \beta_2 v_{k-1} + ( 1- \beta_2) \nabla f_k(x_k) \odot \nabla f_k(x_k)$

$x_{k+1} = x_k - \frac{\eta_k}{\sqrt{v_k} + \epsilon} \odot m_k$

end for

 

우선 $n$은 training data sample 개수 또는 mini-batch의 개수를 의미한다.

즉, $f_i(x)$는 $i$번쨰 data에 대한 loss, 또는 $i$번째 mini-batch에 대한 loss를 의미한다. 

즉, $\nabla f_i(x)$는 stochastic gradient로 생각할 수 있다.

---

Adam에 대한 기본적인 설명은 아래 글을 참고하기 바란다.

https://kyteris0624.tistory.com/30

Deep dive into optimization : Adaptive step-size (1)

https://kyteris0624.tistory.com/32

Deep dive into optimization : Adam

---

 

자, 우선 위 논문이 발표되기 이전 Adam 알고리즘이 소개된 이후 (Kingma, 2014), 2018년도 ICLR 논문 'On the

convergence of Adam and Beyond' (Reddi, 2018)에 의해 Adam의 문제점이 지적이 된다.

Kingma의 논문에서 저자들은 Adam 알고리즘의 convergence analysis를 유도하는데 (convex에 한해서)

이에 대한 비판이 이뤄지면서 Adam이 다음과 같은 상황에서는 train loss가 수렴하지 않을 수도 있고, 발산할 수도 있다는 주장이 제기된다.

 

$0 \le \beta_1 < \sqrt{\beta_2} < 1$, Adam이 발산할 수도 있는 문제들이 존재한다.

 

즉 first momentum의 하이퍼파라미터에 해당하는 $\beta_1$과 second momentum의 하이퍼파라미터에 해당하는 $\beta_2$가 위 부등식을 만족할 때, 단순한 convex optimization에서도 train loss가 수렴하지 않고 발산하는 경우들이 제기되었다. 특히 위 문제는 $\beta_2$가 작을 때 더 두드러진다. (그래서 Reddi는 large $\beta_2$를 사용할 것을 권장하였다.) 또한 추가적으로 Adam에서 $\beta_1$이 0일 때는 위 알고리즘이 RMSProp 알고리즘으로 바뀌게 되는데, Reddi는 이때 $\beta_2$의 값에 관계없이 단순한 convex optimization problem에서도 발산할 수도 있음을 지적한다.

즉, $\beta_1 = 0$은 사용하지 말고, $\beta_2$는 크게 설정하는 것을 실용적으로 조언해준다.

 

여기서 $0< \beta_1, \beta_2 <1$은 default이기 때문에 결국 핵심은 $\beta_1 < \sqrt{\beta_2}$이다.

그렇다면, 위 부등식을 만족하지 않게 $\beta_1$, $\beta_2$를 설정하면 되지 않는가?란 생각이 들 수 있다.

하지만 실제로 쓰이는 위 값들을 살펴보면, 위 부등식을 만족하는 경우가 상당히 많다는 것을 알 수 있다.

 

예를 들어, Kingma의 논문에서 저자들은 default setting으로 $(0.9, 0.999)$를 추천하였고, 이 값은 현재 파이토치나 텐서플로우의 default setting으로 이뤄져 있다.

GPT-3나 Bert 같은 Transformer 계열의 모델들은 Adam을 상당히 많이 사용하는데, 이때 $(0.9, 0.95)$를 많이 사용한다.

또한 GAN 모델에서도 $(0.5, 0.999)$를 사용하는 경우가 많다.

 

위 값들은 모두 $\beta_1 < \sqrt{\beta_2}$를 만족한다. 하지만 실제로는 위 값들이 최적의 성능을 내는 하이퍼파라미터로 보고되었다.

여기서 Zhang (논문의 저자)은 Reddi의 논문에서 한 가지 문제점을 지적하는데 그것은 다음과 같다.

Reddi가 Adam이 발산하는 경우의 예시들을 설정하는 상황에서, 하이퍼파라미터 $(\beta_1, \beta_2)$를 먼저 선택하고나서 obejctive function과 $n$의 크기를 선택하였다는 것이다. 

즉, 둘의 순서가 뒤바뀌었음을 지적하였는데 발산하는 예시를 설정하기 위해 $(\beta_1, \beta_2)$ 값을 설정하고, 그 다음 함수와 $n$의 값을 설정하였다는 것이다.

하지만 실제로 우리가 알고리즘을 사용할 때는 먼저 모델, 데이터셋, objective function을 설계하고 mini-batch size를 설정한 다음, 하이퍼파라미터 튜닝에 들어간다. 

하이퍼파라미터 $(\beta_1, \beta_2)$ 를 먼저 선택하고 모델, 데이터셋, objective function을 설정한 이후, mini-batch size를 선택하지 않는다.

 

Zhang은 이를 지적한 이후, objective function과 $n$의 크기를 먼저 설정하고 이후 $(\beta_1, \beta_2)$ 값에 대한 연구를 진행한다. 이는 결국 하이퍼파라미터는 'problem-dependent'하게 설정된다는 의미이다.

즉, $\beta_1, \beta_2$의 조합은 문제 상황과 독립적으로 항상 발산하는 값이 존재하는 것이 아니고, 특정한 문제 상황에 대하여서 그때에 발산할 수도 있는 $\beta_1, \beta_2$ 조합이 있을 수 있다고 이야기한다.

이를 $\beta_2$는 problem-dependet hyperparameter라고 표현한다.

 

사실 'problem-dependent hyperparameter'는 Adam에만 국한된 이야기는 아니다.

만약 objective function이 $L$-smooth하다면, 우리가 Gradient descent로 학습할 때, step size를 $\frac{2}{L}$ 이하로 설정해야 하는데 이 순서를 바꾸어서 step size를 먼저 설정하고 objectie function을 설정한다면 당연히 우리는 GD가 발산하는 예시를 찾아낼 수 있을 것이다.

즉, step-size는 problem-dependent hyperparameter이다.

 

위와 같이 저자들은 Reddi의 논문을 반박한 이후 자신들의 Adam convergence analysis 결과를 유도해낸다.

오늘 글에서는 '가정'과 '결과' 그리고 그 의미를 살펴보고 다음 번 글에서 증명에 대한 간단한 소개를 하도록 하겠다.

 

<Assumptions>

$1$. $f_i(x)$의 gradient는 constant $L$에 대해서 $L$-Lipschitz continuous하다. 즉, $f_i$는 smooth function이다.

$2$. $f(x)$는 $f^{*}$ lower bound가 존재한다.

$3$. $f_i(x)$와 $f(x)$는 다음의 부등식을 만족한다.

$\sum_{i=1}^n \lVert \nabla f_i(x) \rVert_2^2 \le D_1 \lVert \nabla f(x) \rVert_2^2 + D_0$

 

여기서 $3$번 가정은 일반적인 variance 가정 (또는 second moment of stochastic gradient)을 확장시킨 것으로 볼 수 있는데, $D_1 = \frac{1}{n}$이면 우리가 아는 일반적인 variance 가정이 나오고 $D_0 =0$이면 이 가정을 'strong growth condition'이라고 부른다. 이때는 $\lVert \nabla f(x) \rVert = 0$이면 $\lVert \nabla f_i(x) \rVert = 0$이고 overparameterized neural network에서는 이 상황이 성립할 수 있음이 알려져 있다.

만약 $D_0 \ne 0$이면 우리는 stationary point 근방의 bounded region으로 수렴함을 보일 수 있고, $D_0 = 0$이면 stationary point로 수렴함을 보일 수 있다. 

예를 들어, SGD가 constant step size일 때는 그 값이 작다면, bounded region으로 수렴함을 보일 수 있고, 정확히 stationary point로 수렴함을 보일 수는 없다.

하지만 $D_0 = 0$이라면 SGD가 constant step-size여도 그 값이 작다면 stationary point로 수렴함을 보일 수 있다.

 

추가적으로 위 가정에서는 'bounded gradient assumption'이 없는데 (즉, 어떤 상수 $G$에 대해서 $\lVert \nabla f(x) \rVert < G$라는 가정이 없음)$ 위 가정은 처음부터 optimization algorithm이 수렴하지 않을, 발산할 수 있는 가능성을 사전에 차단시키는 가정이다. 하지만 저자들은 Adam으로 학습을 할 때, train loss (또는 gradient norm)가 발산하는 상황들이 존재하기 때문에 (그리고 그것을 이론적, 실험적으로 보여줌) 위 가정 없이 결과를 유도하는 것이 더 현실적인 상황들과 들어맞는다고 판단하였고, bounded gradient 가정 없이 convergence analysis를 유도하였다.

이에 대해서는 다음 글에서 살펴보겠다.

 

마지막으로 convex 또는 strong convex 가정이 없으므로 non-convex에서의 convergence analysis이다.

 

<Theorem>

자, 그렇다면 이제 어떤 결과가 나왔을까?

 

 

$\beta_1 < \sqrt{\beta_2} < 1$이고 $\beta_2$가 어떤 threshold $\gamma_1(n)$보다 크고, step size $\eta_k = \frac{\eta_1}{\sqrt{nk}}$이고, $\beta_1^{k_m - 1}n \le \frac{\beta_1^n}{\sqrt{k_m -1}}$일 때 $k_m$보다 큰 모든 $T$에 대해서 다음이 성립한다.

 

결국 쉽게 말하면 convergence rate이 $O(\frac{\log T}{\sqrt{T}})$라는 것이다.

추가적으로, $D_0 = 0$이면, 즉 strong growth condition 하에서는 stationary point로 수렴함을 보일 수 있다.

 

자, 그러면 위 이론을 통해서 어떤 것을 알 수 있을까? (증명 및 실험적 결과들은 다음 글에서 다루겠다.)

 

1. $\beta_2$는 충분히 커야 한다.

위 이론에서는 $\gamma_1(n)$이라는 어떤 threshold를 제시하였고 여기서 $n$은 mini-batch 개수이다.

즉, mini-batch 개수가 결정되어야 $\beta_2$값을 결정할 수 있으므로, 앞선 Reddi 논문과는 다르게 조금 더 실용적인 "problem-dependent hyperparameter"에 대한 결과이다.

구체적으로 $\gamma_1(n)$은 $1 - O(\frac{1 - \beta_1^n}{n^2 \rho})$이고, $\rho \approx O(n)$는 어떤 constant이다.

$\gamma_1(n) \approx 1 - O(n^{-3})$이므로 예를 들어 CIFAR-10을 사용할 때, Mini-batch size를 128로 설정하면 $n$의 값은 대략 390이고 결국 $\beta_2$는 충분히 커야 한다.

 

2. $\beta_1 < \sqrt{\beta_2}$

1번의 해석에 의거하여, $\beta_2$는 충분히 커야 하기 때문에 결국 $\beta_1$은 0과 1사이의 대부분의 값이 위 범위를 만족한다. 

 

3. Bounded region으로의 convergence

$D_0 > 0$일 때 위 이론은 stationary point가 아닌, 그 근방 (bounded region)으로 수렴함을 보였다.

저자들은 실험적으로 심지어 convex qudratic function일 때 $D_0 >0$이면 Adam은 zero gradient로 수렴하지 않음을 보이는데, (diminishing step size) 이에 대해 다음과 같이 해석한다.

step-size $\eta_k$가 감소하게 설정하였다 하더라도, Adam과 같은 'Adaptive step-size' method에서 실제 step-size의 역할을 하는 것은 $\frac{\eta_k}{\sqrt{v_k}}$이고 이것은 학습이 진행할 때 감소하지 않을 수도 있다는 것이다.

하지만 $\beta_1$이 1에 가까울수록 $O(\sqrt{D_0})$가 감소해나간다는 것을 저자들은 보여주는데 이는 다음 글에서 자세하게 이야기하겠다.