AMUSE 논문 리뷰: 샤프니스와 리버 밸리 관점으로 본 Muon·스케줄 프리 옵티마이저

핵심 메시지

• 샤프니스(손실 함수 헤시안의 최대 고윳값)가 학습률의 2배 역수 경계를 넘지 않아야 경사하강법이 손실을 안정적으로 줄인다는 고전적 조건이 출발점이다.
• 실제 학습에서는 샤프니스가 그 경계까지 올라간 뒤 그 근처에서 계속 요동치는 '엣지 오브 스태빌리티' 현상이 아키텍처와 옵티마이저를 가리지 않고 나타난다.
• 헤시안 고윳값 중 크게 떨어져 나온 상위 몇 개의 개수는 분류 문제의 클래스 수와 정확히 일치하며, 이를 기준으로 파라미터 공간을 도미넌트와 벌크 부분공간으로 나눌 수 있다.
• 직관과 달리 학습을 빠르게 만드는 쪽은 도미넌트 부분공간이 아니라 벌크 부분공간이며, 이를 '강과 계곡' 지형으로 그리면 여러 현상이 한꺼번에 설명된다.
• AMUSE는 Muon의 업데이트 방향 설계와 스케줄 프리의 안정적인 기울기 평가 지점을 결합하고, 계수 베타를 학습 초반엔 작게 후반엔 1에 가깝게 스케줄링한 알고리즘이다.

쉽게 이해하기

이번 발표는 KAIST 김재철AI대학원 윤철희 교수 연구실에서 나온 AMUSE 논문을 다루지만, 본론에 도달하기 전에 딥러닝 최적화에서 관찰된 여러 현상을 먼저 훑는다. 발표자는 논문 결과 자체보다 그 배경이 되는 현상들을 함께 엮어 보는 편이 더 재미있다고 판단해 구성을 그렇게 잡았다고 밝혔다.

먼저 샤프니스 개념이 등장한다. 손실 함수의 헤시안에서 가장 큰 고윳값을 샤프니스라 부르는데, 이는 현재 위치에서 손실 지형이 가장 가파르게 휘어진 방향의 휘어짐 정도를 뜻한다. 1차원 이차 모델에서 경사하강법이 수렴하려면 학습률과 곡률의 곱이 일정 조건을 만족해야 하고, 고차원에서도 샤프니스가 학습률의 2배 역수보다 작아야 손실이 단조 감소한다는 사실이 잘 알려져 있다. 같은 학습률이라도 완만한 지형에서는 잘 내려가지만 가파른 지형에서는 진동한다는 그림으로 이해하면 쉽다.

그런데 ICLR 2021에서 보고된 '엣지 오브 스태빌리티'는 이 안전 경계가 지켜지지 않는 방식으로 학습이 진행됨을 보여준다. 어떤 학습률을 쓰든 샤프니스가 학습 과정에서 점점 증가해 그 학습률에 해당하는 경계 근처까지 올라간 뒤, 그 부근에서 계속 요동치며 학습이 이어진다. 2023년 연구는 이를 '자기 안정화' 메커니즘으로 규명했는데, 샤프니스 증가 → 파라미터 진동 확대 → 샤프니스 감소 → 진폭 축소 후 재상승이라는 네 단계가 반복된다는 설명이다.

두 번째 축은 부분공간 이야기다. 클래스가 K개인 분류 문제에서 헤시안의 고윳값을 보면 소수의 큰 값들이 나머지 뭉쳐 있는 값들과 확연히 떨어져 있는데, 이 상위 고윳값의 개수가 정확히 클래스 수 K와 같다. 이들이 펼치는 공간을 도미넌트 부분공간, 나머지를 벌크 부분공간이라 부른다. 2018년 연구는 SGD의 기울기 방향이 도미넌트 부분공간에 크게 정렬돼 있음을 관찰했고, 그래서 학습의 핵심이 도미넌트 쪽이라는 추측이 있었다.

그러나 2024년 NeurIPS 논문은 반대 결과를 보였다. 기울기를 도미넌트 쪽에 투영해 업데이트한 경우와 벌크 쪽 성분을 살려 업데이트한 경우를 비교하자, 일관되게 벌크 성분을 살린 쪽이 손실을 더 빠르게 줄이고 더 빠르게 안정화됐다. 이 두 갈래를 한꺼번에 설명하는 그림이 '리버 밸리' 손실 지형이다. 곡률이 큰 도미넌트 방향은 좁은 계곡(밸리), 곡률이 작은 벌크 방향은 완만하게 흘러 내려가는 강(리버)에 해당한다. SGD가 노이즈 때문에 계곡 방향으로 크게 진동하며 학습하는 모습도, 강 방향으로 곧장 나아가는 쪽이 더 빠른 이유도 이 그림에서 자연스럽게 읽힌다.

이 관점은 최근 주목받은 두 옵티마이저도 설명한다. 스케줄 프리 방식은 실제 업데이트 지점들의 이동평균을 만들어 계곡 방향 진폭을 걷어내고, 그 이동평균과 현재 지점 사이의 가중 평균에서 기울기를 계산한다. 이 방법을 AdamW에 적용한 스케줄 프리 AdamW는 2024년 AlgoPerf 대회에서 1위를 하며 주목받았는데, 학습률 스케줄링 없이도 좋은 성능을 냈다는 점이 핵심이었다. Muon은 모멘텀 행렬을 특이값 분해한 뒤 특잇값 행렬을 항등 행렬로 바꿔 업데이트하는 방식으로, 비싼 SVD 대신 뉴턴–슐츠 근사를 쓴다. 실제로 측정해 보니 Muon은 도미넌트 성분 비율이 유독 작아 강 방향으로 안정적으로 학습하고 있었고, 이것이 성능의 이유로 해석됐다.

AMUSE는 이 둘을 결합한다. 다만 단순 결합은 계수 튜닝에 지나치게 민감해 잘 동작하지 않았고, 저자들은 이동평균 쪽 가중치인 베타를 학습 초반에는 작게 두어 빠르게 학습하다가 특정 시점 이후 1에 가깝게 키워 후반에는 강 방향으로만 학습이 일어나도록 스케줄링했다. FineWeb 데이터로 언어 모델 최적화를 돌린 실험에서 AMUSE는 AdamW, Muon, 스케줄 프리 AdamW 등을 모두 앞서며 훨씬 빠르게 퍼플렉시티를 낮췄다.

주요 인사이트

• 최적화 알고리즘을 비교할 때 '업데이트 방향'과 '기울기를 계산하는 지점'을 별개의 설계 축으로 분리해 보면, Muon과 스케줄 프리가 서로 다른 문제를 풀고 있었음이 드러난다.
• 기울기가 특정 부분공간에 정렬돼 있다는 관찰이 곧 그 방향이 학습에 유용하다는 뜻은 아니다. 상관과 기여를 혼동하면 정반대 결론에 도달한다.
• 리버 밸리 지형이라는 하나의 시각화가 엣지 오브 스태빌리티, 벌크 우위, 스케줄 프리와 Muon의 강점을 모두 설명한다는 점에서, 좋은 비유는 그 자체로 연구 도구가 된다.
• Muon이 모든 특잇값을 1로 맞추는 방식에는 대가가 따른다. 학습에 도움이 되지 않는 방향까지 무조건 키워 오히려 방해가 될 수 있다는 점이 발표에서 지적됐다.
• 안정성과 수렴 속도의 트레이드오프를 고정된 하이퍼파라미터로 정하는 대신 학습 단계에 따라 옮겨가게 하는 것이 AMUSE의 실질적 기여다.

자주 묻는 질문

관련 AI 소식