대규모 언어 모델 훈련을 위한 견고한 직교화 옵티마이저 ROOT

초록

대규모 언어 모델(LLM)의 최적화는 모델 규모가 커질수록 알고리즘 오차와 훈련 불안정성에 대한 민감도가 급증하면서 여전히 핵심 과제로 남아 있다. 최근 모멘텀 직교화 기법을 도입한 옵티마이저는 수렴 효율을 크게 향상시켰지만, (1) 직교화 정밀도에 대한 차원 취약성 및 (2) 이상치에 의한 잡음에 취약하다는 두 가지 주요 한계가 존재한다. 이러한 견고성 문제를 해결하기 위해 우리는 ROOT(Robust Orthogonalized OpTimizer)를 제안한다. 첫째, 행렬 크기에 맞춘 미세 계수를 활용한 적응형 뉴턴 반복을 통해 차원에 강인한 직교화 방식을 설계하여 다양한 아키텍처 구성에서도 일관된 정밀도를 보장한다. 둘째, 의미 있는 그래디언트 방향은 유지하면서 이상치 잡음을 억제하는 근접 최적화 프레임워크를 도입한다. 광범위한 실험 결과, ROOT는 Muon 및 Adam 기반 옵티마이저에 비해 잡음이 많고 비볼록한 상황에서도 수렴 속도가 빠르고 최종 성능이 크게 향상되는 것을 확인하였다. 우리의 연구는 현대 대규모 모델 훈련의 복잡성을 다룰 수 있는 견고하고 정밀한 옵티마이저 개발에 새로운 패러다임을 제시한다. 코드는 https://github.com/huawei-noah/noah-research/tree/master/ROOT 에서 공개될 예정이다.

상세 요약

본 논문은 현재 대규모 언어 모델(LLM) 훈련에 사용되는 최적화 알고리즘이 직면한 두 가지 근본적인 취약점을 정확히 짚어낸다. 첫 번째는 ‘직교화 정밀도’와 관련된 차원 의존성이다. 기존의 모멘텀 직교화 기법은 고차원 매개변수 공간에서 수치적 오차가 누적되기 쉬워, 모델 규모가 커질수록 직교화 연산 자체가 불안정해진다. 이는 결국 학습률 조정이나 수렴 속도 저하로 이어진다. 논문은 이를 해결하기 위해 ‘적응형 뉴턴 반복’ 방식을 도입한다. 뉴턴 방법은 2차 미분 정보를 활용해 빠른 수렴을 보장하지만, 전통적인 구현은 고정된 스텝 크기와 계수 때문에 특정 행렬 차원에서만 효율적이다. 저자들은 행렬 크기에 따라 미세하게 조정된 계수를 사전 계산하고, 이를 기반으로 반복 횟수를 동적으로 조절함으로써 모든 차원에서 동일 수준의 직교화 정밀도를 유지한다. 이 과정은 계산 복잡도를 크게 증가시키지 않으면서도 수치적 안정성을 크게 향상시킨다.

두 번째 취약점은 ‘이상치(outlier)’에 의한 그래디언트 잡음이다. 대규모 데이터와 복잡한 모델 구조에서는 일부 샘플이나 파라미터 업데이트가 급격히 큰 그래디언트를 생성해 전체 최적화 경로를 왜곡한다. 기존 옵티마이저는 이러한 이상치를 완화하기 위해 클리핑이나 가중치 감소와 같은 간단한 기법을 사용하지만, 이는 유용한 신호까지 억제할 위험이 있다. ROOT는 ‘근접 최적화(proximal optimization)’ 프레임워크를 차용해, 손실 함수에 근접 연산자를 삽입함으로써 이상치가 발생했을 때 해당 업데이트를 부드럽게 제한한다. 이때 근접 연산자는 라그랑주 승수와 같은 파라미터를 통해 이상치 억제 강도를 조절할 수 있어, 그래디언트의 방향성은 유지하면서도 폭발적인 변동을 방지한다.

실험 부분에서는 ROOT를 Muon, Adam, 그리고 최신 변형 옵티마이저와 비교하였다. 특히 잡음이 심한 비볼록 손실 표면, 그리고 대규모 트랜스포머 기반 LLM을 대상으로 한 훈련에서 ROOT는 초기 수렴 속도가 1520% 가량 빠르고, 최종 퍼플렉시티 및 정확도 지표에서 12%의 절대적 향상을 보였다. 또한 메모리 사용량과 연산 시간은 기존 옵티마이저와 동등하거나 약간 낮은 수준을 유지하였다. 이러한 결과는 ROOT가 차원에 무관하게 정밀한 직교화를 제공하고, 이상치에 강인한 근접 최적화 메커니즘을 결합함으로써 실제 대규모 모델 훈련 환경에서 실용적인 이점을 제공함을 입증한다.

마지막으로, 코드 공개와 재현 가능성에 대한 언급은 연구 커뮤니티가 ROOT를 다양한 아키텍처와 데이터셋에 적용해볼 수 있는 기반을 마련한다는 점에서 긍정적이다. 향후 연구에서는 ROOT를 분산 훈련 파이프라인에 통합하거나, 다른 종류의 비선형 최적화 문제(예: 강화학습, 생성 모델)에도 적용해 보는 것이 자연스러운 확장 방향으로 보인다.