Dion2 옵티마이저 공개 - 무온(Muon)의 직교화 비용을 줄여 대규모 학습을 확장하는 방법

핵심 메시지

• 현재 AI 모델 학습의 사실상 표준은 AdamW이지만, 행렬 파라미터를 직교화하는 무온(Muon)이 새로운 대안으로 떠올랐고 키미와 GLM 같은 프런티어 모델에 채택됐다.
• 무온은 행렬 단위 연산에 의존해 가중치가 샤딩되는 분산 학습과 충돌하고, 직교화 자체가 초선형 복잡도라 규모를 키우기 어렵다.
• Dion2는 업데이트 행렬 중 얼마만큼을 직교화할지 정하는 조절 손잡이 베타를 도입해, 베타가 1이면 원래 무온과 같고 1보다 작으면 부분 직교화로 연산과 통신을 줄인다.
• 노름이 큰 상위 베타 비율의 행을 골라 직교화한 뒤 그 행을 감쇠시켜, 학습 단계마다 서로 다른 뉴런이 선택되도록 한다.
• 행렬의 25%만 직교화해도 무온과 대등한 업데이트 품질이 나왔고, 코드베이스에 무온이 이미 들어 있다면 Dion2는 그대로 갈아 끼울 수 있다.

쉽게 이해하기

마이크로소프트 리서치 포럼에서 뉴욕 랩의 AI 프런티어스 팀이 Dion2를 공개했다. 문제의식은 단순하다. AI 모델을 학습시키는 데 막대한 연산이 들어가는데, 지금보다 더 잘할 수 있느냐는 것이다. 현재 사실상의 표준은 AdamW이고 대단히 널리 쓰이지만, 발표는 이 기본 알고리즘을 넘어설 방법을 묻는다. 그 대안으로 등장한 것이 직교 옵티마이저 계열의 무온(Muon)이다. 무온은 행렬 파라미터를 위한 옵티마이저인데, 현대 신경망의 파라미터는 사실상 거의 전부가 행렬이다. 핵심 아이디어는 각 업데이트 행렬을 직교정규화하는 것으로, 이론적 동기가 깔끔하고 실증 성능도 강하다. 실제로 키미와 GLM 같은 프런티어 모델에 채택됐다.

그렇다면 모두 무온으로 갈아타면 되느냐 하면 그렇지 않다. 발표는 무온을 큰 규모로 확장하기가 까다롭다고 지적한다. 무온은 행렬 단위 연산에 의존하는데, 가중치가 여러 장치에 쪼개져 있는 분산 학습과 이 방식이 충돌한다. 게다가 직교정규화 연산 자체가 본질적으로 초선형 복잡도다. 원소 단위 연산만 쓰기 때문에 확장이 쉽고 모든 분산 학습 프레임워크와 잘 맞는 AdamW와는 대조적이다.

Dion2의 발상은 무온에 확장성 조절 손잡이를 다는 것이다. 업데이트 행렬 중 얼마만큼을 직교화할지 정하는 파라미터 베타를 0과 1 사이에서 고른다. 베타가 1.0이면 행렬 전체를 직교화해 원래 무온으로 되돌아가고, 1.0보다 엄격히 작으면 부분 직교화가 되어 연산과 통신 비용이 줄어든다. 목표는 행렬의 일부만 직교화하면서도 무온이 주는 업데이트 품질의 이점을 지키는 것이다.

동작 방식은 놀랄 만큼 간단하고, 발표자는 그 단순함이야말로 Dion2를 실용적으로 만드는 지점이라고 강조한다. 먼저 노름을 기준으로 상위 베타 비율의 행 또는 뉴런을 고르고, 선택된 행만 직교화한 뒤 그 행들을 감쇠시킨다. 감쇠 덕분에 학습 단계가 바뀌면 다른 다양한 뉴런이 선택된다. 이는 효율을 위해 다소 복잡한 저랭크 근사에 기댔던 Dion 1세대와 대비된다. Dion 1은 선형대수 함수에 의존해 작은 규모에서 오버헤드가 생겼고 그만큼 실용성이 떨어졌다.

결과는 두 축으로 제시됐다. 옵티마이저 스텝 벤치마크에서는 베타 비율이 작아질수록 Dion2가 무온 대비 속도 이득을 얻었고, 실제 학습 곡선 비교에서는 직교화 비율이 25%인 설정에서도 무온과 대등하게 경쟁력 있는 업데이트 품질이 나왔다. 발표자는 자신들의 경험상 Dion2가 Dion 1보다 업데이트 품질도 더 낫다고 덧붙였다. 도입 방법은 간단하다. 코드베이스에 이미 무온이 통합돼 있다면 같은 설정을 그대로 쓰면 된다. 경험상 베타를 0.5나 0.25로 두면 완전한 무온과 견줄 만한 알고리즘이 된다는 것이 권장값이다. 다시 말해 무온이 들어 있는 코드베이스에서 Dion2는 갈아 끼우기만 하면 되는 플러그 앤 플레이에 가깝다. 구현은 마이크로소프트의 dion 저장소에 공개돼 있다.

주요 인사이트

• 옵티마이저 경쟁의 승부처가 수렴 속도에서 분산 학습 적합성으로 옮겨갔다. 아무리 이론이 깔끔해도 가중치가 샤딩된 환경에서 행렬 단위 연산을 요구하면 대규모에서는 쓰기 어렵다.
• 직교화 비율을 0과 1 사이 연속값으로 노출한 설계는 비용과 품질 사이를 상황에 맞게 선택할 수 있게 만든다. 극단값에서 기존 방법으로 정확히 되돌아가는 구조라 도입 위험도 낮다.
• 선택된 행을 감쇠시켜 매 단계 다른 뉴런이 뽑히도록 한 장치가 핵심이다. 노름 상위 행만 계속 고르면 같은 방향만 반복해 다듬는 문제가 생기기 때문이다.
• 저랭크 근사를 쓴 Dion 1세대보다 단순한 행 선택 방식이 더 나은 결과를 냈다는 점은, 대규모 학습에서 이론적 정교함보다 오버헤드 없는 단순함이 이기는 경우가 많다는 사례로 읽힌다.
• 25% 직교화로 대등한 품질이 나온다는 결과는 무온의 이득 상당 부분이 행렬 전체의 완전한 직교화가 아니라 큰 노름을 가진 소수 방향의 정규화에서 온다는 것을 시사한다.

자주 묻는 질문

관련 AI 소식