On the Principles Behind Neural Network Optimizers

Adam의 수렴 조건과 트랜스포머에서 강한 이유를 이론과 실험으로 해석하고, 그 원리를 메모리 절약형 Adam-mini 설계로 연결한 박사학위 논문 축약본이다.

이 글은 새 단일 실험 논문이 아니라 Adam 연구를 네 단계로 묶은 박사학위 논문의 축약본이다. 먼저 Adam의 수렴과 발산 사이에 문제와 배치 크기에 따라 달라지는 경계가 있음을 이론적으로 보인다. 이어 헤시안이라는 손실 곡률 행렬이 트랜스포머 학습 중 거의 블록 대각 구조와 강한 블록별 차이를 갖게 되며, 이 구조가 Adam의 좌표별 보정을 유리하게 만든다고 설명한다. 마지막으로 블록당 하나의 학습률만 쓰는 Adam-mini로 옵티마이저 상태 메모리를 50퍼센트 줄이면서 AdamW 수준의 성능을 유지했다. 다만 정확한 수렴 경계와 현대 구조 전체의 엄밀한 분석은 남아 있다.

핵심 요약

• Adam은 항상 수렴하거나 항상 발산하는 것이 아니다. 문제와 미니배치 수에 따라 안전 구간이 달라지며, 배치가 작을수록 두 번째 모멘텀 계수인 베타2를 더 크게 잡아야 한다는 조건을 제시한다.
• 트랜스포머의 헤시안은 학습하면서 거의 블록 대각 구조로 바뀌고 블록마다 곡률 크기도 크게 달라졌다. 저자는 이 구조가 좌표별 학습률을 쓰는 Adam이 SGD보다 유리한 이유라고 설명한다.
• 연속된 큰 행렬 곱이 이런 헤시안 구조를 만든다는 점을 단순 신경망과 무작위 행렬 이론으로 분석했다. 이 관점은 어텐션 헤드마다 Muon 보정을 나누는 설계로도 이어진다.
• Adam-mini는 두 번째 모멘텀 값의 99.9퍼센트 이상을 블록별 스칼라로 바꿔 옵티마이저 상태 메모리를 50퍼센트 줄인다. 기존 Adam은 130억 매개변수 모델에서 두 모멘텀에만 약 104GB가 필요하다고 계산했다.
• C4에서 3,900만부터 10억 매개변수의 Llama 2 구조를 학습하자 Adam-mini가 AdamW의 검증 손실을 비슷하게 따라갔다. 다만 이 원고는 분량 제한이 있는 축약본이라 증명이 빠져 있고, 어텐션과 전문가 혼합 구조의 엄밀한 분석도 후속 과제로 남았다.

논문 링크

외부 연구를 정리한 자료입니다. HDATF가 발표한 논문이나 제품 성능을 측정한 결과는 아닙니다.