뮤온부터 그래디언트 클리핑까지: QK 안정성에 대한 몇 가지 고찰

AI 요약

딥러닝 최적화 기법인 Muon 옵티마이저가 트랜스포머 모델의 Query(Q)와 Key(K) 가중치 행렬에 직접 적용될 때 훈련 불안정성과 붕괴 현상이 발생하는 문제를 분석한 기술 블로그 글이다. 저자는 Muon이 기존 Adam과 달리 함수 공간(function space)에서 업데이트 규칙을 구성하는 철학적 차이를 설명하며, QK 행렬에 특화된 변형 업데이트 방안을 이론적 유도와 엔지니어링 병목을 거쳐 제안하고 있다.

핵심 포인트

• Muon 옵티마이저는 Keller Jordan과 Jeremy Bernstein이 제안, 함수 공간 관점에서 업데이트 규칙 구성
• Kimi K2 모델에서 Muon이 성공적으로 적용되었으나 QK 가중치에는 직접 적용 시 훈련 붕괴 발생
• 표준 Muon은 Frobenius 노름 대신 함수 공간 제약 조건을 사용, 이로 인해 QK 행렬과 부적합
• 저자는 이론적 유도와 엔지니어링 병목을 거쳐 QK에 적합한 휴리스틱 방안을 모색

향후 전망

• Muon의 원리를 유지하면서 QK 안정성을 확보하는 변형 옵티마이저가 트랜스포머 훈련 효율성을 크게 개선할 가능성