정규화된 적응 모멘텀과 직교 모멘텀 통합 최적화

초록

NAMO와 그 확장인 NAMO‑D는 직교화된 모멘텀에 Adam식 노름 기반 적응 스텝 사이즈를 결합한 새로운 옵티마이저이다. NAMO는 단일 스칼라 적응계수를 사용해 직교화된 업데이트 방향을 보존하면서 노이즈에 따라 스텝을 조절하고, NAMO‑D는 열별 대각 행렬을 곱해 뉴런 단위의 세밀한 노이즈 적응을 제공한다. 두 알고리즘은 결정론적 상황에서 최적 수렴률을, 확률적 상황에서는 그래디언트 노이즈 수준에 맞춰 수렴 속도가 자동 조정됨을 이론적으로 증명했으며, GPT‑2 사전학습 실험에서 AdamW와 Muon을 능가하는 성능을 보였다.

상세 분석

본 논문은 대규모 언어 모델 학습에서 널리 사용되는 Adam 계열 옵티마이저와 최근 주목받는 Muon 옵티마이저의 장점을 구조적으로 결합한다는 점에서 의미가 크다. Muon은 가중치 행렬에 대해 직교화(Orth)된 모멘텀을 사용함으로써 업데이트 방향을 스펙트럴 노름 기준 최 steep descent 로 변환한다. 이는 행렬 구조를 활용해 수렴 속도를 높이고, 파라미터 스케일에 덜 민감한 특성을 제공한다. 그러나 직교화 연산은 무한대 노름을 갖는 경우가 있어, 노이즈가 큰 상황에서는 업데이트가 과도하게 증폭될 위험이 있다. 반면 Adam은 각 파라미터에 대해 1차 모멘텀(평균)과 2차 모멘텀(분산) 추정을 통해 노이즈에 적응하는 스텝 사이즈를 제공한다. 하지만 Adam의 업데이트는 단순히 부호와 스케일을 결합한 형태라, Muon이 제공하는 구조적 이점을 활용하지 못한다는 한계가 있다.

NAMO는 이러한 두 접근법을 통합한다. 구체적으로, 첫 번째 모멘텀 M_t 를 직교화하여 O_t = Orth(M_t) 를 얻고, 두 번째 모멘텀 v_t 로 전체 그래디언트의 Frobenius 노름 제곱을 추정한다. 이후 α_t = √((1-μ₂^t)/(1-μ₁^t))·‖M_t‖_F / (√v_t + ε) 라는 스칼라 적응계수를 도입해 O_t 를 스케일링한다. 이때 α_t 는 1차 모멘텀의 크기와 2차 모멘텀의 노이즈 수준을 동시에 반영하므로, 노이즈가 클수록 α_t 가 작아져 안정적인 수렴을 보장한다. 중요한 점은 α_t 가 스칼라이기 때문에 O_t 의 직교성은 그대로 유지된다는 것이다. 따라서 Muon이 제공하는 구조적 장점은 보존되면서, Adam식 노이즈 적응이 추가된다.

NAMO‑D는 보다 미세한 적응을 위해 열별(뉴런별) 스케일링을 도입한다. 각 열 j 에 대해 v_t^j 를 해당 열의 그래디언트 노름 제곱으로 추정하고, d_t^j = √((1-μ₂^t)/(1-μ₁^t))·‖M_t^{:j}‖ / (√v_t^j + ε) 로 개별 스케일을 계산한다. 이후 평균값 \bar d_t 와 클램핑 파라미터 c∈(0,1] 를 이용해 d_t 를