Muon 최적화가 Transformer QK 가중치에서 붕괴하는 이유
From Muon to Gradient Clipping: Some Thoughts on QK Stability
Muon 최적화기는 함수 공간에서의 제약을 통해 설계된 혁신적인 접근법으로, Kimi K2에서 성공을 거두었지만 Transformer의 Query(Q)와 Key(K) 가중치에 직접 적용하면 훈련이 극도로 불안정해지고 붕괴할 수 있습니다. 이 글은 표준 Muon이 QK에 맞지 않는 기하학적 원인을 분석하고, Muon의 원리를 따르면서 QK에 적합한 수정된 업데이트 규칙을 도출하기 위한 이론적 시도와 엔지니어링적 한계, 그리고 실용적인 휴리스틱 방안까지의 사고 과정을 기록합니다.
Muon의 업데이트는 msign으로 처리되어 모든 특이값이 동일하므로 사실상 full-rank이며, 이는 저랭크인 Adam 업데이트보다 W의 특이 벡터와 충돌할 확률이 훨씬 높아 특이값을 증가시키기 쉽다.