옵티마이저 Adafactor (2018, 메모리 효율적)

관련글

---

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

 

Optimizer Adafactor (2018)

---

논문 제목: Adafactor: Adaptive Learning Rates with Sublinear Memory Cost

연도: 2018

링크: https://arxiv.org/pdf/1804.04235.pdf

관련 개념:

1. Matrix factorization

2. Adam

3. Rank

4. SVD

5. KL-divergence

 

서론

---

 확률적 optimizers such as RMSProp, Adam, Adadelta는 지수 이동 평균 항(v)을 갖고 있으며, 이 항은 훈련 시 계속해서 저장되어야 합니다. v의 개수는 weight나 bias와 같은 parameter의 수와 동일하며, 따라서 parameter가 N개가 있다면, Adam은 2N개의 추가적인 메모리를 사용해야합니다. (여기서 v는 matrix입니다.)

 저자는 지수 이동 평균 항인 v matrix를 이용하는 방법을 택하지 않고, 대신 v의 row와 column만을 누적하며, matrix에서 두 개의 vector로 factorizing 되었기에 메모리를 아낄 수 있다고 합니다. 이 두 vector는 원래의 v를 근사화할 수 있으며, 저자는 이러한 방식을 사용해도 기존의 방식과 비슷한 성능을 발휘한다고 합니다. 

  기존의 메모리가 n x m의 matrix를 요구했다면, Adafactor는 n, m 길이의 vector를 요구한다고 합니다. 따라서 메모리 복잡도는 O(nm)에서 O(n+m)으로 변하게 됩니다. 이는 matrix를 row, column vectors로 factorizing했기 때문입니다.

 

 

 

 저자는 second momentum 즉, v항의 decay가 너무 느릴 경우, 이상적인 update보다 더 빠른 update를 보인다고합니다. 따라서 이들은 update clipping을 도입하고, 제안된 방식에 따라 decay rate를 증가시킨다고 합니다.

 

 

본론

---

  Moving average matrix의 크기가 너무 커서 전체 matrix를 저장하지 못하는 경우(즉, 메모리가 부족한 경우), 우리는 해당 matrix를 분해하여 두 개의 low-rank matrices를 만들 수 있습니다. 즉, matrix가 n x m의 크기라면, 우리는 이 matrix를 분해하여 R(n x k), S(k x m) 크기의 matrices를 만들 수 있습니다. 이 두 matrices R과 S를 곱하면 원래의 matrix와 비슷한 어떠한 matrix가 나오겠죠. 물론 원래의 matrix가 나오지는 않습니다. 근사화시킬 뿐이죠. 저자는 이러한 방식으로 메모리를 아끼고자 합니다. (극단적으로 k를 1로 놓을 수도 있습니다.)

 이렇게 어떠한 matrix를 두 개의 matrices로 분할할 때, 우리는 SVD를 선택할 수도 있습니다. SVD를 통해 factorized된 signular matrix에서 k개의 값을 취득하면 되죠. 하지만, SVD는 Exponential moving average에는 적합하지 않으며 게다가 SVD로 얻은 matrices에 음수가 없다는 보장도 없습니다. (음수가 없어야하는 이유는 우리가 root를 씌워야하기 때문이죠.)

 저자는 다른 방법으로 KL divergence를 고려합니다. 즉, 원래의 matrix V와 factorized R, S matrices가 주어졌을 때, V과 RS의 KL-divergence를 최소화 하도록 R과 S를 선택하는 것입니다. k가 2 이상일 때의 R과 S에서 KL-divergence가 최소가 되는 R과 S를 구하는 것은 어렵지만, k가 1일 경우는 논문에 소개되어 있습니다. 따라서 저자는 R이 (n x 1), S가 (1 x m)인 vector로 원래의 matrix V를 factorizing합니다. 이 과정에서 KL-divergence를 미분합니다. 왜냐하면 최소값을 찾는 건 결국 기울기가 0이 되는 지점을 찾는 것이기 때문이죠.

수학적으로 표현하면 다음과 같습니다.

V와 RS에 관한 KL-divergence

KL-divergence

KL-divergence 전개

decomposing

Ri와 Sj에 대해 각각 미분 (R과 S는 벡터이기 때문에 i,j가 아닌 i 따로 j 따로이다.)

derivation

정리된 R과 S

 

따라서 RS는 다음과 같습니다.

단순히 R과 S를 곱한 형태이죠. 

 

 

위의 식에서 V1m은 V를 row 방향으로 더한 column vector, 1nTV는 V를 column 방향으로 더한 row vector가 되겠네요.

 

알고리즘에서는 V 항이 다음의 두 항 R과 C로 나눠집니다. (R은 row vector, C는 column vector이며, C와 S는 상관이 없습니다. 단순히 RS에서 나온 V1m을 R로, 1nTV를 C로 놓은 것 뿐입니다.)

Algorithm

 여기서 1m은 길이가 m이고 모든 원소가 1인 vector이며, 1Tn은 길이가 n이고, 모든 원소가 1인 row vector입니다. 저자들은 일단 1차 momentum은 고려하지 않았다 하네요.

 

 

 

 

 Adam에서 beta2가 작으면, second-momentum이 빠르게 작아져서 모델이 수렴이되지 않는다고 하며, 이는 Adafactor에서도 마찬가지라고 합니다. 하지만, Adafactor에서 beta2가 클 경우, 학습이 불안정하다고 합니다. 저자들은 Adam이 의도대로 동작한다면, 다음의 식이 1과 유사해야 한다고 합니다.

하지만, 저자들이 beta2를 0.9와 0.999로 놓았을 때 위의 값은 심히 달랐다고 합니다. 이는 논문의 Figure 1에서 확인할 수 있습니다. beta2가 0.9일 때 확실히 1과 가깝다는 것을 알 수 있지만, beta2가 0.9라는 것은 0.999일 때보다 decay가 빨리 된다는 의미입니다. 즉, second momentum이 빠르게 작아진다는 의미이죠. 그렇기 때문에 우리는 0.999를 어쩔 수 없이 선택할 수밖에 없습니다. 하지만, 0.999를 선택했을 때는 값이 크게 요동치는 것을 확인할 수 있습니다. 따라서 저자들은 Update clipping이라는 개념을 도입합니다.

 

Update clipping은 위에서 구한 U를 이용합니다. Update clipping은 threshold 값인 d를 hyperparameter로 받으며, U와 d를 통해 새로운 값인 hat{U}를 계산합니다.

 

Algorithm

 

 

결론

---

 Adafactor는 기존의 v항을 R과 S항으로 분류하여 메모리 사용량을 아꼈으며, 알고리즘 상으로는 R과 C항을 V항 대신 사용했습니다. 본문에는 안 썼지만, Adafactor 방식을 사용하려면 warmup이 필요합니다. 또한 논문에서는 beta2를 위한 특별한 scheduling 방법도 제시하고 있습니다.

 

장점

1. 메모리를 절약할 수 있다. -> 더 큰 모델을 학습시킬 수 있다. 2. warmup이 된다면, Adam과 비슷한 성능을 낸다.

단점

1. 추가적인 matrix 연산이 필요하다. 2. warmup이 없을 경우, 불안정하다. 3. 추가적인 hyperparameter인 d가 요구된다.

 

 

 

참고

---

Adafactor: Adaptive Learning Rates with Sublinear Memory Cost

Adafactor: Adaptive Learning Rates with Sublinear Memory Cost