옵티마이저 SM3 (2019, 메모리 효율적)

관련글

---

옵티마이저 정리 및 논문 리뷰

 

Optimizer SM3 (2019)

---

논문 제목: Memory-Efficient Adaptive Optimization
연도: 2019
링크:  https://arxiv.org/pdf/1901.11150.pdf
https://arxiv.org/pdf/1901.11150v1.pdf
관련 개념:
1. 차원
2. Adafactor

서론

---

기존의 Adam 계열의 optimizers는 second momentum 항 때문에 memory overhead가 크다고 합니다. 저자는 효율적이고, 쉽게 적용이 가능하면서 memory overhead가 적은 새로운 adaptive optimization을 소개합니다. 저자가 고안한 방법은 adaptive의 성질을 만족한다고 합니다. 또한 수렴성을 보장하고, 매우 큰 translation and language model에서 효율적이라고 합니다.
SM3와 Adafactor 둘 모두 메모리를 효율적으로 다루는 optimizer라고 할 수 있지만, Adafactor의 경우 matrix 단위에서, SM3는 tensor 단위까지 지원한다는 점이 다른 점입니다. 즉, SM3는 1차원 ~ d차원의 모든 차원에 대해 적용할 수 있고, Adafactor는 1차원 ~ 2차원까지만 적용할 수 있다는 점이 차이입니다. (그 외에서 다른 차이점이 있습니다만, 제가 생각한 주된 차이는 "커버할 수 있는 차원이 SM3가 더 높다"였습니다.) Adafactor가 m x n 크기를 갖는 matrix를 길이가 m인 벡터 R과 길이가 n인 벡터 S로 factorizing했다면, SM3는 더 높은 차원, 예를 들면 n1 x n2 x n3 .... nN 차원의 tensor를 길이가 n1, n2 n3 ... nN인 차원의 벡터로 표현한다는 점이 다릅니다. 즉, 메모리의 사용량을 Adafactor는 m x n에서 m + n으로 줄였다면, SM3는 n1 x n2 x n3 ... nN에서 n1 + n2 + n3 ... + nN로 획기적으로 줄였습니다.

Adafactor는 factorizing을 통해 m x n크기의 matrix를 근사화 했다면, SM3에서는 그런 수식은 보이지 않고 있습니다. 제가 보았을 때는 수식적으로 도출된 optimizer는 아닌 것 같았습니다. (아니면 제가 못 찾았던 것일 수도 있습니다.)

SM3는 Google에서 만든 것으로 보인다는 것이 또 다른 특징인 것 같습니다. 소스 코드는 다음 아래에 있습니다.
Tensorflow version by google-research
PyTorch version by Enealor


SM3의 버전은 두 개가 있으며, 위에 올린 github links는 SM3-II에 해당하는 Optimizer입니다. (SM3-I을 구현한 코드는 아직 못 찾았습니다.)

본론

---

Adafactor와 달리 SM3의 알고리즘은 제 입장에서는 납득하기 힘든 알고리즘이었습니다. 수학적으로 도출된 식도 아니고, 직관적으로 원리를 알 수 있는 알고리즘도 아니었습니다. 하지만, 최대한 이해한 바로 설명해 보겠습니다.

https://arxiv.org/pdf/1901.11150.pdf

SM3-I의 알고리즘입니다. r, j, i, S_r라는 변수가 보이며, u라는 변수도 보입니다. 논문에서 이들에 관해 친절히 설명하지 않고 있기에 다음과 같이 제 생각을 곁들여 정리했습니다.

위의 설명에서 S_r과 v_t는 임시적인 변수이기 때문에 training이 끝나면 사라집니다.



SM3-I을 한 줄 한 줄 보면 다음과 같습니다.

각각의 축 방향을 제외한 방향에서 가장 큰 값 찾기

max 부분을 이해하기 어려울 수 있습니다. max 부분은 "S_r에 속한 j 번째 index의 gradient의 제곱에서 가장 큰 값을 얻는다."라고 해석할 수 있습니다. 예를 들어 2차원 matrix가 있다고 하고, r이 1이라고 할 때 max 부분의 의미는 다음과 같습니다.

max 항의 의미

즉, max 부분의 의미는 "r 방향 외에 다른 방향에서 가장 큰 값을 찾아라"입니다. 예시로 r=1로 놓았기에, max 부분의 의미는 "2방향에서 가장 큰 값들을 찾아라 또는 row 방향에서의 가장 큰 값들을 찾아라"입니다. 만약 2차원 matrix가 아니라 3차원 tensor이고 ,r이 1이면, "2 방향과 3방향에서 가장 큰 값들을 찾아라 또는 plane에서의 가장 큰 값들을 찾아라"일 것입니다. 2차원이든 3차원 tensor이든 결과는 r 방향의 vector이고요.


두 번째 보이는 식은 다음과 같습니다. (알고리즘에서 d는 tensor의 크기입니다. 즉, m x n x l이면 d는 m곱하기 n곱하기 l입니다.)

앞서 논문에서는 (1, 2)같은 index를 단순히 (2)라고 표현했음을 언급했습니다. 위의 식도 (2)같은 표현으로 묘사된 것입니다. 여기서는 위의 알고리즘을 쉽게 알아보기 위해 다음과 같이 변형하고, gradient tensor는 3차원이라고 가정하겠습니다. 즉, r의 범위는 1 ~ 3이 되고, u_t 또한 3개가 생기겠죠.

변형된 식

위의 두 식은 의미적으로 동일합니다. 단지 (i) 하나로 표현된 index를 3차원에 맞게 (i, j, l) index로 표현했으며, u(1)(i)는 u(1)의 i번째 원소(entry)를 의미하고, u(2)(j)는 u(2)의 j번째 원소를 의미합니다. 즉, u(1)(i)는 1번 방향에서 i 번째 원소(값)을 의미하는 것입니다. 위의 식을 통해 얻은 v_r 값을 최종 우리가 알고 있는 update 식에 넣으면 업데이트가 완료됩니다.

---

SM3-II의 알고리즘은 다음과 같습니다.

여기서 v_t는 임시 변수이며, matrix형태입니다. u_t는 training 동안 계속 살아 있는 변수이며, vector이고 그 의미는 SM3-I와 유사합니다. 다른 부분은 SM3-I를 보면 이해하실 수 있을 것인데, 마지막 줄은 설명이 필요할 것 같네요.

이 뜻은 다음과 같습니다.

r 방향으로 v를 검사하여 가장 큰 값들의 집합(즉, 벡터)를 u_t로 한다.

tensorflow로 설명하면
u_t(r) = tf.math.reduce_max(v_t, axis=r)
위의 github의 내용으로 설명하면
u_t(r) = _max_reduce_except_dim(v_t, r)이 됩니다.

결론

---

SM3는 Adafactor와 다르게 tensor 단위까지 고려한 알고리즘입니다. 논문의 설명에 따르면 warm-up을 통해 성능을 향상시킬 수 있다고 하며, SM3-I와 SM3-II의 차이는 수렴의 범위라고 할 수 있습니다. 저자의 말을 차용하면 다음과 같습니다.

SM3--II provides a tighter upper bound on the cumulative gradient squares than SM3-I.

논문은 대체적으로 설명이 어렵게 되어 있다 느꼈습니다. 특히 논문에서의 cover의 의미(제 생각에는 차원을 의미하는 것 같습니다.)를 이해하지 못해 읽는데 힘이 들었으며, 알고리즘 또한 직관적이지 못하다는 느낌을 받았습니다. (물론 이해한다면, 쉽게 적용할 수 있는 알고리즘입니다.)


제가 언어 모델은 잘 모르지만, Transformer-Big 모델에 WMT'14 en->fr를 사용했을 때, BLEU 점수가 SM3를 사용했을 때 비교적 낮았다고 합니다. (BLEU 점수는 낮을 수록 좋습니다.) 이 결과는 논문의 8 page, Figure 2.에 나와 있습니다.



N차원의 tensor를 갖고 있는 gradient의 크기가 n1 x n2 x n3 ... nN일 때, 기존의 Adagrad 등은 n1 x n2 x n3 ... nN의 memory가 필요했지만, SM3의 경우, n1 + n2 + n3 ... nN의 메모리만 필요하다는 장점이 있습니다. 즉, 메모리 복잡도가 다음과 같습니다.

메모리 복잡도. 왼쪽: 기존의 Adagrad, 오른쪽: SM3

 

장점

1. Tensor의 범위에서도 적용할 수 있는 memory efficient algorithm이다. 2. Tensor의 범위에서도 적용할 수 있기에 flexible하다. 3. Adafactor처럼 쉽게 적용할 수 있는 algorithm이다. 4. 특수한 경우에서 Gradient의 Upper boundary가 존재한다. 5. 몇몇 실험 결과, Adam보다 성능이 좋게 나왔네요.

단점

1. 알고리즘이 직관적이지 않다. 즉, 왜 저런 알고리즘이 나왔는지 이해하기가 어렵다. 2. 알고리즘만 봤을 때, 해석하기가 어렵다. (해석만 되면 구현은 비교적 쉽다.) 3. warm-up이 필요한 algorithm이다. 4. Momentum 항이 없기 때문에, local minima에 빠질 수 있다.

참고

---

Memory-Efficient Adaptive Optimization
Tensorflow version by google-research

GitHub - google-research/google-research: Google Research

PyTorch version by Enealor

GitHub - Enealor/PyTorch-SM3: Implements the SM3-II adaptive optimization algorithm for PyTorch.