AdaMeZO: GPU 메모리에 모멘트를 저장하지 않고 Adam 방식으로 LLM 파인튜닝하는 새 최적화기

연구자 Zhijie Cai, Haolong Chen, Guangxu Zhu는 AdaMeZO를 발표했습니다. 이는 GPU 메모리에 모멘트를 저장하지 않으면서도 인기 있는 Adam 알고리즘의 장점을 대형 언어 모델(LLM) 파인튜닝에 도입하는 제로 차수(zeroth-order) 최적화기입니다.

LLM 파인튜닝에서 GPU 메모리가 병목이 되는 이유

신경망 학습에 일반적으로 사용되는 표준 Adam 최적화기는 모델의 각 파라미터에 대해 두 가지 통계량을 추적합니다. 1차 모멘트(그래디언트의 이동 평균)와 2차 모멘트(그래디언트 제곱의 이동 평균)입니다. 수십억 개의 파라미터를 가진 모델에서는 이로 인해 메모리 요구량이 세 배로 증가합니다. 실제 그래디언트를 계산하지 않고 순전파만 사용하는 기존 접근 방식인 MeZO는 메모리 문제를 해결하지만, 파라미터별 학습률 적응이 없어 수렴 속도가 느립니다.

AdaMeZO는 두 접근 방식을 어떻게 결합하는가?

AdaMeZO는 Adam의 모멘트를 영구 저장 없이 추정합니다. 가중치에 무작위 섭동을 가하고 손실(loss) 변화를 측정하여 최적화 단계마다 즉석에서 적응형 모멘트 거동을 재구성합니다. 그 결과 손실 곡면의 기하학에 따라 학습률을 조정하는 Adam과 유사하게 동작하면서도 메모리 사용량은 MeZO와 동일하게 유지됩니다.

최적화 경로 시각화가 보여주는 것

저자들은 다양한 손실 곡면에서의 최적화 경로 시각화를 제시합니다. 이는 AdaMeZO가 균일하게 동작하는 MeZO와 달리, 파라미터 공간의 평탄한 부분과 곡률이 있는 부분을 적응적으로 탐색하는 방식을 보여줍니다. 정량적으로 AdaMeZO는 원래 MeZO 대비 향상된 수렴과 함께 동일한 성능 수준을 달성하는 데 최대 **70% 적은 순전파(forward pass)**를 기록합니다.