ARO 옵티마이저 해설: 회전 좌표계로 LLM 학습을 AdamW보다 1.3배 빠르게 만든 원리

핵심 메시지

• AI 학습에서 옵티마이저는 투입한 연산을 얼마나 효율적으로 지능으로 바꾸는지를 결정하는 핵심 요소다.
• ARO는 기존 행렬 옵티마이저들이 사실상 회전된 좌표계에서 돌아가는 Adam이라는 관찰에서 출발한다.
• 회전을 해석의 도구가 아니라 옵티마이저 설계의 기본 원리로 승격시킨 것이 이 연구의 핵심 주장이다.
• 회전이 통하는 근거로는 트랜스포머 가중치에 내재한 회전 대칭성, 즉 모델 출력을 바꾸지 않는 방향이 지목된다.
• 0.3억에서 80억 파라미터까지 실험에서 AdamW 대비 약 1.3배, Muon 대비 약 1.1배 우위가 규모를 키워도 줄지 않았다.

쉽게 이해하기

발표는 문제 정의부터 분명하다. 대규모 언어모델 학습은 비싸고, 그 비용 대비 성과를 결정적으로 좌우하는 것이 옵티마이저다. 그런데 업계는 10년 가까이 AdamW에 머물러 있었다. 최근 Muon 같은 행렬 기반 옵티마이저가 등장해 표본 효율을 끌어올리고 프로덕션 규모 학습까지 감당하기 시작했고, 마이크로소프트 리서치 케임브리지의 차오 마와 웬보는 이 흐름을 한 단계 더 밀어붙이는 프레임워크로 ARO를 내놓았다.

출발점은 기존 방법들을 다시 읽는 일이다. 가장 단순한 경사하강법은 기울기를 그대로 써서 가중치를 갱신한다. Adam 같은 적응형 옵티마이저는 기울기를 따르되 비선형 사상으로 크기를 조정한 뒤 적용한다. 직교화 계열, 샴푸, 소프 같은 최신 행렬 옵티마이저들은 훨씬 복잡해 보이지만, 발표자들은 단순화를 거치면 이들이 근본적으로 Adam과 닮았다고 말한다. 기울기를 회전시키고, 비선형 사상을 적용하고, 회전된 기저에서 처리한 뒤 다시 되돌리는 구조라는 것이다.

이 재해석 자체는 완전히 새로운 것은 아니다. 문헌에서 비선형 사상은 대개 Adam의 변형이고 회전 행렬은 기울기의 특이 벡터로 정해진다. ARO의 주장은 여기서 갈린다. 회전을 사후 설명의 언어로 남겨 두지 말고 옵티마이저를 설계하는 기본 원리로 삼자는 것이다. 이 관점을 취하면 Adam식 크기 조정을 넘어선 더 강력한 비선형 사상을 쓸 수 있고, 무엇보다 그 비선형 사상의 선택에 맞춰 더 나은 회전을 고를 수 있게 된다. 왜 회전이 중요한가에 대한 가설은 손실 지형의 대칭성에 있다. 트랜스포머를 비롯한 현재의 언어모델 구조는 가중치에 풍부한 회전 대칭성을 갖고 있어서, 제대로 적용된 회전은 모델의 예측을 전혀 바꾸지 않는다. ARO는 이 성질을 이용해 출력을 바꾸지 않는 대칭 궤도를 따라 움직이면서 손실 지형을 더 잘 헤쳐 나갈 갱신 방향을 찾는다. 같은 가설에서 레이어 간 결합을 활용하는 추가 기능도 따라 나온다.

검증은 평가 편향을 줄이는 벤치마크 규약을 따로 만드는 데서 시작했다. 마이크로소프트 리서치 아시아와 협력해 새로 개발한 시그마 모델에서 성능을 재고, 지표로는 AdamW 대비 학습 스텝의 상대 속도 향상을 썼다. 예컨대 1.36배라는 수치는 같은 시점에 AdamW와 동일한 손실에 도달하는 데 데이터 또는 GPU 시간의 약 73%만 썼다는 뜻이다. 결과는 직교화 계열 방법들이 아래쪽에, ARO 변형들이 위쪽에 모이는 두 무리로 갈렸고, 데이터 규모를 바꿔도 그 간격이 유지됐다.

규모를 키워도 이점이 남는지도 따로 확인했다. 3억에서 80억 파라미터까지 서로 다른 구조와 학습 조건에서 ARO는 AdamW 대비 약 1.3배, Muon 대비 약 1.1배의 우위를 일관되게 보였고, 모델이 커진다고 그 우위가 줄어들지 않았다. 80억 파라미터 설정에서 분산 구현의 스텝당 실행 시간이 AdamW와 대등했다는 점도 함께 보고됐는데, 이 조건이 성립해야 스텝 수 절감이 실제 벽시계 시간 절감으로 옮겨간다.

주요 인사이트

• 서로 다른 동기에서 유도된 옵티마이저들이 하나의 형식으로 정리된다면, 그 형식 자체가 새 옵티마이저를 설계하는 도구가 된다.
• 모델의 예측을 바꾸지 않는 방향이 존재한다는 사실은 낭비처럼 보이지만, 오히려 그 위에서 자유롭게 움직일 여지를 준다.
• 스텝 수 기준 속도 향상은 스텝당 실행 시간이 기존과 비슷할 때에만 실제 비용 절감이 되며, 발표는 이 조건을 별도로 확인했다.
• 옵티마이저 비교는 평가 설계에 따라 결론이 흔들리기 쉬워, 벤치마크 규약을 먼저 세운 뒤 측정한 순서 자체가 결과의 일부다.
• 규모를 키울수록 이점이 사라지는 기법이 많은 분야에서, 우위가 줄지 않았다는 관찰은 그 자체로 중요한 신호다.

자주 묻는 질문

관련 AI 소식