새로운 세대의 AI 모델과 가장 강력한 연구 논문 중 하나

오늘 저는 Token AI에서 발표한 새로운 연구 논문에 대해 이야기하려고 합니다.
"Stable Training with Adaptive Momentum"

이 논문은 이론적으로나 결과적으로나 가장 강력한 최적화 도구 중 하나가 될 수 있는 알고리즘을 소개합니다.

수년간 우리는 Adam, AdamW, LAMB 등 잘 알려진 최적화 도구에 의존해 왔습니다. 의심할 여지 없이, 이들은 AI 모델을 훈련할 때 가장 먼저 선택하는 도구였습니다.

최적화 도구가 무엇인지 익숙하지 않은 분들을 위해 간단히 설명하자면, 이는 모든 AI 모델 훈련의 핵심 부분입니다. 훈련 중 손실을 줄이기 위해 모델의 가중치를 업데이트하는 알고리즘을 말합니다.

하지만 이러한 최적화 도구들은 훈련에 영향을 미치는 한계점을 가지고 있습니다.

예를 들어, Adam은 훈련 내내 고정된 beta1을 사용하는데, 이는 오래된 모멘텀을 유지하여 모델을 잘못된 방향으로 계속 밀어붙일 수 있습니다.

STAM은 현재 그래디언트와 이전 모멘텀의 차이(g - m)를 측정하여 이 문제를 해결합니다. 차이가 클 때 beta1을 줄여 노이즈가 많은 단계에서도 더 안정적인 훈련을 가능하게 합니다.

훈련 중에 변화나 노이즈가 발생할 때 또 다른 문제가 나타납니다. 오래된 모멘텀은 오히려 해로울 수 있습니다. STAM은 잔차 분산(residual variance)을 기반으로 적응형 beta1을 사용하여 이를 처리합니다.

SGD의 주요 문제 중 하나는 방향이 잘못되어도 고정된 모멘텀 때문에 계속 진행된다는 점입니다. STAM은 첫 번째 모멘텀이 스스로 교정되도록 하여 이 문제를 해결합니다.

이제 더 가벼운 버전인 STAMLite에 대해 이야기해 보겠습니다.

이것은 많은 경우에 AdamW를 대체할 기본 선택지로 설계되었습니다. 핵심 차이점은 beta1이 고정된 것이 아니라 동적이라는 점입니다:

• 그래디언트가 노이즈가 많으면 모멘텀을 줄입니다.
• 그래디언트가 안정적이면 모멘텀을 높게 유지합니다.

또한 최적화 도구 상태 메모리 측면에서도 효율성을 개선했습니다:

• AdamW는 파라미터 크기의 약 2배가 필요합니다.
• STAM Full은 AdamW와 비슷합니다.
• STAMLite는 파라미터 크기의 약 1배가 필요합니다.

실제로 STAMLite는 AdamW 및 STAM에 비해 리소스를 약 50% 절약하며, 이는 훈련 중 GPU 사용량이 현저히 줄어든다는 것을 의미합니다.

벤치마크를 살펴보면 결과가 스스로를 증명합니다.

하이퍼파라미터 스윕(Hyperparameter Sweep)에서 STAMLite는 다음을 달성했습니다:
정확도: 0.61
손실: 0.91

Long-Horizon Non-Stationary MLP에서 STAM은 NAdam과 함께 거의 동일한 결과로 1위를 차지했습니다:
정확도: 0.97
손실: 0.09

더 많은 벤치마크는 웹사이트와 연구 논문에서 확인할 수 있습니다.

이는 TokenAI가 알려진 문제를 가진 제한된 최적화 도구 세트에 대한 오랜 의존을 깨뜨리는 중요한 단계입니다.

초기 릴리스임에도 불구하고 강력하고 유망하다는 것을 증명합니다. 개인적으로 저는 이미 STAM으로 전환했으며, 현재 이를 사용하여 처음부터 끝까지 전체 LLM을 훈련하고 있습니다. 곧 결과를 공유하겠습니다.