A Convergence Analysis of Adaptive Optimizers under Floating-point Quantization

이 논문은 그라디언트, 가중치, 옵티마이저 상태의 부동소수점 양자화 하에서 Adam 과 Muon 같은 적응형 옵티마이저의 수렴성을 분석하는 최초의 이론적 프레임워크를 제시하며, 양자화 오차가 수렴률에 미치는 영향을 규명하고 Muon 이 Adam 보다 양자화에 더 강건할 수 있음을 보여줍니다.

핵심

이 논문은 **"왜 거대한 인공지능 (LLM) 을 만들 때, 컴퓨터가 계산하는 숫자의 정밀도를 낮춰도 (저정밀도) 여전히 잘 작동하는지?"**에 대한 수학적 해답을 제시합니다.

기존 이론은 "모든 계산이 완벽하게 정확해야 한다"고 가정했지만, 실제로는 메모리를 아끼기 위해 숫자를 잘게 자르거나 (양자화) 반올림하는 방식을 씁니다. 이 논문은 그 '잘못된' 계산이 왜 오히려 큰 문제를 일으키지 않는지, 그리고 어떤 최적화 알고리즘이 더 튼튼한지를 증명했습니다.

이 복잡한 내용을 일상적인 비유로 쉽게 설명해 드릴게요.

---

1. 배경: 거대한 도시 건설과 '빠른 공사' (저정밀도 학습)

거대한 언어 모델 (LLM) 을 훈련시키는 것은 거대한 도시를 건설하는 것과 같습니다.

• 기존 방식 (고정밀도): 모든 자재 (데이터) 를 미터 단위로 정확히 재서 쓰면 실수는 없지만, 시간이 너무 오래 걸리고 자재비 (메모리) 가 너무 비쌉니다.
• 새로운 방식 (저정밀도): "대충 1 미터면 돼, 1.0001 미터까지 재지 마!"라고 합니다. 자재는 줄고 공사 속도는 빨라집니다. 하지만 **정확하지 않게 재면 건물이 무너지지 않을까?**라는 걱정이 있었습니다.

실제로는 건물이 잘 지어지는데, 왜 그런지 이론적으로 설명하는 사람이 없었습니다. 이 논문이 바로 그 **'이유'**를 찾아낸 것입니다.

2. 두 명의 건축 감독 (Adam vs Muon)

도시 건설에는 두 가지 다른 감독 (최적화 알고리즘) 이 있습니다.

① 감독 '아담 (Adam)': 꼼꼼하지만 예민한 베테랑

• 특징: 과거의 실수 (기울기) 를 기억해서 다음 계획을 세웁니다. 특히 "과거의 실수가 얼마나 컸는지"를 아주 세심하게 기록합니다 ($\beta_2$ 파라미터).
• 문제점: 이 감독은 정밀도가 조금만 떨어져도 매우 예민하게 반응합니다. 마치 "과거의 실수 기록을 100% 정확히 기억해야만 다음 계획을 잘 세울 수 있다"고 믿는 것처럼, 숫자를 반올림하면 그 오차가 증폭되어 건물이 흔들릴 수 있습니다.
• 논문이 말한 것: 아담이 잘 작동하려면, 과거 기록 (2 차 모멘트) 과 자재 (가중치) 의 정밀도는 아주 높아야 합니다. 그래야만 "대충" 계산해도 건물이 무너지지 않습니다.

② 감독 '뮤온 (Muon)': 유연하고 튼튼한 천재

• 특징: 아담처럼 과거의 실수 크기를 세세하게 쫓기보다, 방향 (벡터) 을 잘 맞추는 데 집중합니다. 수학적으로 '특이값 분해 (SVD)'라는 강력한 도구를 써서 방향을 잡습니다.
• 장점: 이 감독은 정밀도가 조금 떨어져도 크게 신경 쓰지 않습니다. 숫자를 대충 반올림해도 방향만 맞으면 건물이 잘 지어집니다.
• 논문이 말한 것: 뮤온은 아담보다 오차에 훨씬 덜 민감합니다. 그래서 저정밀도 환경 (메모리가 부족한 상황) 에서 더 잘 견디고, 더 빠르게 건물을 지을 수 있습니다.

3. 핵심 발견: "숫자 자르는 정도"가 중요해

논문은 **"숫자를 얼마나 잘게 자르느냐 (부동소수점의 맨티사 길이)"**가 핵심이라고 말합니다.

• 비유: 자를 때 1mm 단위까지 재는지 (정밀함), 1cm 단위만 재는지 (대충) 의 차이입니다.
• 결과:
  • 아담: 1cm 단위로 재면 건물이 흔들립니다. 최소 1mm 단위는 지켜줘야 합니다.
  • 뮤온: 1cm 단위로 재도 건물이 튼튼합니다.
  • 놀라운 사실: 만약 우리가 반복 횟수가 늘어날수록 자르는 정밀도를 아주 조금씩만 높여준다면 (로그 스케일), 아담도 뮤온도 원래의 완벽한 정밀도와 똑같은 속도로 건물을 지을 수 있다는 것을 수학적으로 증명했습니다.

4. 실험: 실제 시뮬레이션

저자들은 컴퓨터로 여러 실험을 해보았습니다.

• 인공 산 (로젠브록 함수): 아담은 정밀도가 낮아지면 산 정상에 도달하는 데 걸리는 시간이 길어지거나 아예 못 갔지만, 뮤온은 정밀도가 낮아도 정상에 잘 올라갔습니다.
• 실제 이미지 (CIFAR-10) 와 언어 모델 (nanoGPT): 실제 데이터를 이용해 훈련시켰을 때도 같은 결과가 나왔습니다. 뮤온이 아담보다 저정밀도 환경에서 더 강인하게 작동했습니다.

5. 결론: 왜 이 논문이 중요한가?

이 논문은 **"왜 우리가 메모리를 아끼기 위해 숫자를 대충 계산해도 AI 가 잘 작동하는지"**에 대한 첫 번째 확실한 이론적 근거를 제시했습니다.

• 기존의 오해: "정확하지 않으면 안 될 거야."
• 이 논문의 진실: "적당한 오차는 허용되지만, 어떤 알고리즘을 쓰느냐에 따라 그 오차를 견디는 능력이 다릅니다."

한 줄 요약:

"거대한 AI 를 만들 때, **아담 (Adam)**은 정밀한 자를 필요하지만, **뮤온 (Muon)**은 대충 자르는 자로도 잘 작동합니다. 그래서 앞으로는 메모리를 아끼고 빠르게 AI 를 훈련시킬 때 뮤온이 더 좋은 선택지가 될 수 있습니다."

이 연구는 앞으로 더 크고 빠른 AI 를 개발할 때, 하드웨어의 한계를 넘어서는 새로운 길을 열어줍니다.

기술 요약

1. 문제 제기 (Problem)

• 배경: LLM 의 규모가 급격히 커짐에 따라 메모리 사용량 감소와 계산 효율성 향상을 위해 저정밀도 (Low-precision) 훈련 (예: BF16, FP8) 이 필수적이 되었습니다.
• 현황: 실제 훈련에서는 그래디언트, 가중치, 옵티마이저 상태 (모멘텀, 2 차 모멘트 등) 를 모두 양자화하여 사용합니다.
• 연구 격차: 기존 이론적 연구들은 대부분 다음과 같은 한계가 있었습니다.
  • 단순화된 가정: 양자화를 '편향 없는 (unbiased)' 노이즈로 가정하거나, 오차 피드백 (error-feedback) 메커니즘이 존재한다고 가정했습니다. 이는 실제 하드웨어의 부동소수점 양자화 동작과 다릅니다.
  • 불완전한 모델링: 옵티마이저 상태 (모멘텀, 2 차 모멘트) 의 양자화를 고려하지 않거나, 그래디언트 양자화만 분석했습니다.
  • 이론적 부재: 실제 LLM 훈련에서 모든 구성 요소를 양자화했을 때 적응형 옵티마이저가 왜 여전히 효과적으로 수렴하는지에 대한 이론적 근거가 부족했습니다.

2. 방법론 (Methodology)

이 논문은 실제 하드웨어의 부동소수점 동작을 반영한 새로운 이론적 프레임워크를 제시합니다.

• 상대 오차 모델 (Relative Error Model):
  • 기존 연구의 '편향 없는 양자화' 가정을 버리고, 부동소수점의 특성을 반영한 상대 오차 (Relative Error) 모델을 도입했습니다.
  • 가정 3.1: 양자화 오차는 $|x_Q - x| \le q|x|$ 형태이며, 여기서 $q = \Theta(2^{-M})$ ($M$은 가수부 (mantissa) 의 비트 길이) 입니다. 이는 오버플로우/언더플로우가 발생하지 않는 범위 내에서 유효합니다.
• 포괄적인 양자화 프레임워크:
  • 훈련 루프 내의 모든 핵심 구성 요소 (가중치 $W$, 그래디언트 $G$, 1 차 모멘트 $M$, 2 차 모멘트 $V$) 에 대한 양자화를 명시적으로 모델링했습니다.
  • Adam 과 Muon 옵티마이저의 업데이트 규칙을 이 프레임워크에 적용하여 각 구성 요소의 양자화 오차가 수렴에 미치는 영향을 분석했습니다.
• 수렴성 증명:
  • 매끄러운 비볼록 (smooth non-convex) 목적 함수 하에서, 양자화된 Adam 과 Muon 의 수렴 속도를 유도했습니다.
  • 각 옵티마이저의 수렴 속도가 양자화 오차 ($q_W, q_G, q_M, q_V$) 에 어떻게 의존하는지 정량화했습니다.

3. 주요 기여 (Key Contributions)

1. 최초의 이론적 프레임워크: 부동소수점 양자화 하에서 적응형 옵티마이저 (Adam, Muon) 의 수렴성을 분석하는 최초의 이론적 체계를 구축했습니다. 기존 연구와 달리 편향 없는 양자화나 오차 피드백 가정을 사용하지 않습니다.
2. Adam 과 Muon 의 수렴 보장:
   • Adam: 가중치와 2 차 모멘트의 양자화 오차에 매우 민감함을 증명했습니다. 특히 $\beta_2 \to 1$로 설정될 때, 2 차 모멘트의 오차가 업데이트 단계의 역제곱근을 통해 증폭되어 수렴에 치명적인 영향을 줍니다.
   • Muon: Adam 에 비해 양자화 오차에 훨씬 덜 민감하며, 더 약한 오차 제어 조건 ($O(T^{-1/2})$) 으로도 수렴이 보장됨을 보였습니다. 이는 Muon 이 SVD 기반의 부호 (sign) 연산자를 사용하여 2 차 모멘트의 역제곱근에 의한 오차 증폭을 피하기 때문입니다.
3. 정밀도 요구사항의 정량화:
   • Adam 이 전정밀도 (full-precision) 와 유사한 수렴 속도 ($\tilde{O}(T^{-1/4})$) 를 달성하려면, 가수부 비트 길이가 반복 횟수 $T$에 대해 로그 스케일 ($M = \Omega(\log T)$) 로 증가해야 함을 보였습니다.
   • 구체적으로 Adam 은 2 차 모멘트와 가중치에 대해 더 엄격한 오차 조건 ($O(T^{-2})$) 을 요구하는 반면, Muon 은 상대적으로 완화된 조건 ($O(T^{-1/2})$) 으로 충분함을 규명했습니다.

4. 실험 결과 (Results)

• 합성 데이터 (Rosenbrock 함수):
  • Adam 은 $\beta_2$가 1 에 가까워질수록 양자화 오차에 극도로 민감해져 수렴된 그래디언트 노름이 커지는 것을 관찰했습니다.
  • Muon 은 동일한 조건에서도 Adam 보다 훨씬 안정적인 수렴 성능을 보였습니다.
• 실제 데이터 (CIFAR-10, nanoGPT):
  • CIFAR-10: 다양한 가수부 비트 길이 ($M$) 에서 Adam 과 Muon 을 훈련시켰으며, 이론적 예측과 일치하게 비트 수가 증가함에 따라 양자화 오차가 감소하고 수렴 성능이 향상됨을 확인했습니다.
  • nanoGPT (LLM): OpenWebText 데이터셋에서 AdamW 와 Muon 을 비교했습니다.
    • 낮은 정밀도 ($M=2$) 에서 AdamW 는 수렴이 느리고 최종 손실이 높았으나, Muon 은 상대적으로 견고한 성능을 보였습니다.
    • 정밀도가 높아짐에 따라 두 옵티마이저 모두 전정밀도 수준의 성능에 도달했습니다.
  • 이는 실제 대규모 LLM 훈련 환경에서도 Muon 이 저정밀도 훈련에 더 강건하다는 이론적 결론을 뒷받침합니다.

5. 의의 및 결론 (Significance)

• 이론과 실전의 간극 해소: 저정밀도 훈련의 성공적인 경험적 결과를 뒷받침하는 엄밀한 이론적 근거를 제공했습니다.
• 옵티마이저 선택에 대한 통찰:
  • Adam: 2 차 모멘트와 가중치의 정밀도 유지가 매우 중요하며, $\beta_2$ 설정에 따라 양자화 오차 증폭이 발생할 수 있음을 경고합니다.
  • Muon: SVD 기반의 구조적 특성으로 인해 저정밀도 환경에서 더 우수한 견고성 (Robustness) 을 가지므로, 차세대 LLM 훈련을 위한 유망한 옵티마이저임을 이론적으로 입증했습니다.
• 미래 방향: 이 연구는 저정밀도 최적화 알고리즘의 설계와 분석을 위한 기초를 마련하며, 하드웨어 친화적인 양자화 전략 개발에 중요한 지침을 제공합니다.

요약하자면, 이 논문은 **"왜 저정밀도 훈련이 작동하는가?"**에 대한 답을 제시하며, 특히 Muon 이 Adam 보다 저정밀도 환경에서 이론적으로 더 우월한 견고성을 가진다는 점을 수학적으로 증명하고 실험적으로 검증했습니다.