메모리 사용량이 적으면서 learning rate도 알아서 찾아주는 Adafactor에 대해서 알아본다.
AdamW
AdamW는 Adam에 weight decay를 추가한 옵티마이저다. Adam은 RMSProp과 Adagrad(Momentum)를 합친 방식이다. RMSProp과 Momentum은 각각 gradient를 추가로 저장해야한다. 따라서 Adam을 쓸 경우 Optimizer는 모델 크기의 3배의 Gradient를 메모리에 점유하게 된다.
Adafactor
Adafactor는 Momentum을 사용하지 않고, RMSProp을 메모리를 더 적게 쓰며 계산하는 row-rank approximation으로 대신했다.
Momentum 제거
Adafactor가 Momentum을 사용하지 않음으로써, 불안전한 학습을 발생시킨다. 저자들은 이 이유를 RMSProp에서 out-of-date된 second-momentum 이라고 생각했다. 이것은 decay rate를 시간에 따라 증가시키는 방식과 Gradient Clipping으로 해결했다.
slow decay(beta2 = 0.999)를 하면 out-of-date로 인해 학습이 불안정해지고, fast decay(beta2 = 0.9)를 하면, 이 문제는 warm-up step을 추가해서 해결했다. warmup은 relative step과 함께 사용해야한다. 이 부분은 뒤에서 구체적으로 설명.
위의 식으로 gradient clipping을 수행하고
beta2는 시간이 지남에 따라 위 식으로 결정된다. 여기서 t는 optimizer step이다.
RMSProp 근사
위 식과 같이 RMSProp이 더 적은 메모리로 근사된다. 가중치 행렬이 RxC 크기일 때, 모멘텀은 동일한 RxC 복사본을 만든다. 하지만 Adafactor는 Row-wise sum과 column-wise sum을 저장한다. 그렇기 때문에 R+C크기의 복사본을 만든다.
Relative Step Size
Adafactor의 다른 장점 중 하나는 초기 learning rate가 주어지면, 스스로 learning rate 를 조절할 수 있다. 이것을 relative step size 라고 한다.
weight의 Root Mean Square와 epsilon(최소값) 중 큰 값으로 진행된다. X_t-1은 전체 파라미터가 아닌 single matrix, vector마다 다르다. weight마다 실제로 적용되는 learning rate은 a_t이다. 논문에서 최종적으로 제안된 방법과 하이퍼파라미터는 이렇다.
Adafactor 사용 주의
T5 Finetuning할 때: relative step, warmup_init, scale_parameter 사용하지 않고, LR 고정한 다음 별도 스케줄러로 warmup 시도. 이건 사실상 근사된 RMSProp만 사용하는 방식이다. 경험상 T5는 pretraining objective와 downstream task가 달라서 warmup이 중요한 것 같다. 만약 relative step을 쓴다면 warup_init과 scale_parameter를 같이 써야한다.
Huggingface Trainer를 쓸 때는, AdafactorScheduler를 써야하는 것 같다. 왜 스케듈러가 필요한지 이유를 살펴보니, trainer가 로깅할 때 learning rate를 알아내기 위해이며, 그렇지 않은 경우에는 불필요하다.
AdamW vs Adafactor
T5는 adafactor로 학습했다. 내가 파인튜닝할 때도 AdamW가 좋을 때도 있고, Adafactor가 좋을 때도 있다.
