LLM 학습이 갑자기 무너지는 순간 — Muon 옵티마이저와 QK-Clip 이야기

수천억 원짜리 GPU 클러스터에서 몇 주씩 돌아가는 대형 언어 모델 학습이 어느 날 갑자기 loss가 치솟으며 망가진다면 어떨까요? 실제로 대규모 학습 현장에서 심심찮게 벌어지는 일인데요. 그 주범 중 하나로 지목되는 게 '어텐션 로짓 폭주'라는 현상이에요. 최근 요즘 주목받는 옵티마이저 Muon과 이 안정성 문제의 관계를 파고든 연구 블로그 글이 올라와서, 그 내용을 바탕으로 이 주제를 차근차근 풀어볼게요.

어텐션 로짓이 뭐냐면요

트랜스포머의 핵심인 어텐션은 각 토큰이 다른 토큰들을 얼마나 참고할지 점수를 매기는 장치예요. 이때 쿼리(Q) 벡터와 키(K) 벡터를 내적해서 점수를 구하는데, softmax에 들어가기 전의 이 원시 점수를 로짓이라고 불러요. 문제는 학습이 진행되면서 Q와 K를 만드는 가중치 행렬이 점점 커지면, 이 로짓 값이 수백, 수천까지 치솟을 수 있다는 거예요. softmax는 입력이 너무 크면 출력이 한쪽으로 쏠려버리는 함수라서, 로짓이 폭주하면 어텐션이 사실상 한 토큰만 바라보는 상태로 굳어버려요. 이러면 그래디언트가 제대로 흐르지 않거나 수치가 불안정해져서 loss spike로, 심하면 학습 전체가 발산하는 사태로 이어지거든요.

Muon이 뭐냐면요

Muon은 AdamW의 아성에 도전하는 신흥 옵티마이저예요. 핵심 아이디어는 가중치 업데이트 행렬을 Newton-Schulz 반복이라는 계산으로 직교화해서 적용하는 건데요. 쉽게 말하면 업데이트가 특정 방향으로만 쏠리지 않게, 모든 방향으로 고르게 힘이 실리도록 다듬어주는 거예요. 덕분에 같은 데이터로 더 빠르게 학습된다는 보고가 이어지면서 주목을 받았고, 실제로 Moonshot AI가 1조 파라미터급 모델 Kimi K2를 학습할 때 채택하면서 대규모 검증까지 이뤄졌어요. 그런데 바로 그 대규모 학습에서 QK 로짓 폭주 문제가 불거졌고, Moonshot 팀은 QK-Clip이라는 장치를 더한 MuonClip을 만들어 대응했죠. 이번 블로그 글은 이 지점을 파고들어요. 왜 이런 문제가 생기고, QK-Clip은 우리가 익히 아는 그래디언트 클리핑과 무엇이 다른가 하는 질문이에요.

증상을 다스리는 약과 원인을 다스리는 약

전통적인 안정화 기법인 그래디언트 클리핑은 그래디언트가 일정 크기를 넘으면 잘라내는 방식이에요. 말하자면 '한 걸음의 보폭'을 제한하는 건데요. 문제는 보폭을 아무리 제한해도, 작은 걸음이 한 방향으로 꾸준히 쌓이면 가중치 자체는 계속 커질 수 있다는 거예요. 로짓 폭주의 원인은 걸음이 아니라 도착한 위치, 즉 가중치의 크기거든요. QK-Clip은 그래서 접근이 달라요. 학습 중에 어텐션 로짓의 최대값을 감시하다가 임계치를 넘으면, 쿼리와 키를 만드는 가중치 행렬 자체를 비율에 맞게 줄여버려요. 증상(그래디언트)이 아니라 원인(가중치 크기)에 직접 개입하는 거죠. 비슷한 고민에서 나온 기법의 계보도 있어요. Gemma 2가 쓴 logit soft-capping은 로짓을 tanh 함수로 부드럽게 눌러주는 방식이고, 여러 모델이 채택한 QK-norm은 내적하기 전에 Q와 K를 정규화해서 애초에 폭주할 수 없게 만들어요. 각각 개입하는 지점이 다른데, 이 차이를 비교해 보는 것만으로도 학습 안정성이라는 주제의 지형이 그려져요.

한국 개발자에게 주는 시사점

대규모 사전학습을 직접 하는 분은 많지 않겠지만, 파인튜닝에서도 loss spike는 흔한 골칫거리예요. 학습이 불안정할 때 학습률만 만지작거리지 말고, 어텐션 로짓의 최대값 같은 내부 지표를 모니터링해 보면 원인을 훨씬 빨리 찾을 수 있어요. 작은 모델을 직접 학습해 보는 분이라면 QK-norm 같은 기법은 지금 바로 몇 줄로 적용해 볼 수 있고요. Muon 자체도 오픈소스 구현이 나와 있어서, 작은 규모의 실험에서 AdamW와 비교해 보는 것도 좋은 공부가 돼요. 옵티마이저와 안정화 기법은 거대 연구소만의 이야기처럼 보이지만, 원리는 작은 실험에서도 그대로 재현되거든요.

정리하며

더 빠른 옵티마이저를 얻으면 새로운 불안정성이 따라오고, 그걸 다시 구조적으로 해결하는 것. 이게 요즘 LLM 학습 연구의 최전선 풍경이에요. 여러분은 모델 학습 중에 겪은 인상적인 불안정 사례가 있나요? 원인을 어떻게 추적했는지 경험을 나눠주시면 재미있을 것 같아요.