거대 모델 학습이 터지는 진짜 이유: QK 안정성과 MuonClip

수십억 파라미터 트랜스포머를 학습하다 보면 어텐션 로짓(QK 내적)이 폭발하며 loss가 튀는 불안정성이 자주 발생한다. 이 글은 Moonshot AI의 Kimi K2에서 쓰인 Muon 옵티마이저와 MuonClip 기법을 실마리로, 이 문제를 파고든다. 핵심 통찰은 이렇다. Muon은 그래디언트를 직교화(orthogonalize)해 업데이트 크기를 균일하게 만드는데, 이 특성이 특정 방향으로 가중치가 과도하게 커지는 걸 막아 학습을 안정시킨다. 저자는 이를 전통적인 그래디언트 클리핑과 연결해, 두 방법이 결국 '업데이트가 폭주하지 않도록 제약을 거는' 같은 목적을 공유한다고 본다. 실용적 처방이 QK-clip이다. 어텐션 로짓의 최댓값이 임계치를 넘으면 그 층의 query·key 가중치를 직접 스케일 다운해, 로짓 폭발을 근본에서 억제한다. 손실 함수나 아키텍처를 바꾸지 않고도 대규모 학습의 발산을 막는 값싸고 직접적인 개입이라는 점이 매력이다. LLM 사전학습을 다루는 엔지니어라면 옵티마이저 선택과 안정화 기법을 함께 설계해야 함을 시사한다.