Muon 옵티마이저가 첫 MoE 층의 전문가 균형을 두 배 어렵게 만든다
Muon Makes the First Mixture-of-Experts Layer Harder to Balance
무엇인가
2026년 주요 공개 MoE 언어모델 상당수는 가중치 행렬을 AdamW 대신 Muon으로 학습하고, expert별 라우팅 점수에 bias를 더한 뒤 매 스텝 갱신하는 auxiliary-loss-free 방식으로 로드를 맞춘다. 문제는 레시피마다 라우터까지 Muon으로 학습할지 의견이 갈린다는 점이다. DeepSeek-V4, Kimi K3, GLM-5.3은 라우터를 Muon에서 제외하지 않고, Qwen3.8-Flash-Next, MiMo-V2.6, IQuest-Q1은 제외한다. 이유를 밝힌 것은 Qwen3.8-Flash-Next뿐으로, Muon이 학습 초반 라우터를 불안정하게 만들고 후반에는 이득이 없었다고 적었다. 두 선택지를 데이터로 비교한 보고는 없었고, 기존 옵티마이저 비교 연구들은 bias를 쓰지 않거나 bias 자체를 측정하지 않았다. 이 논문은 그 빈칸을 겨냥한다.
어떻게 동작하나
저자가 택한 지표는 bias의 span, 즉 한 층에서 가장 큰 bias와 가장 작은 bias의 차이다. bias는 매 옵티마이저 스텝마다 공정 몫 1/E보다 많은 토큰을 받은 expert는 u만큼 내려가고 적게 받은 expert는 올라가며, 논문은 DeepSeek-V3와 같은 u=0.001을 쓴다. 로드 자체는 bias가 강제로 공정 쪽으로 밀기 때문에 라우터의 선호를 잘 드러내지 못한다. 반면 span은 모든 토큰이 같은 방향으로 특정 expert를 선호할수록 커지므로, 그 층이 얼마나 균형 잡기 어려운지를 보여준다. 점수가 (0,1) 구간에 갇히는 sigmoid 라우터에서는 bias 차이가 1을 넘으면 순위가 고정되어 span이 1+2u를 넘지 못한다. 함께 측정한 값으로는 라우터 입력들의 평균 코사인 유사도(공유 방향), 배치 평균 점수의 폭(score offset), 가장 붐비는 expert의 공정 몫 대비 배수(overload, E=32·k=4에서 상한 8), 스텝 간 로드 변동(load churn), 라우터 업데이트 중 공유 방향 성분의 비율(common share)이 있다.
무엇과 다른가
실험은 FineWeb-Edu sample-10BT로 학습한 decoder-only MoE 두 개로 진행했다. 작은 모델은 폭 512 블록 8개에 활성 파라미터 55M(전체 303M), 큰 모델은 폭 768 블록 12개에 활성 108M(전체 496M)이고, 블록 0은 dense FFN, 이후 블록은 라우팅 expert 32개 중 4개를 고르는 MoE 층이다. 파라미터를 라우터·routed expert·shared expert·attention·dense FFN·기타(embedding, output head, norm gain) 여섯 그룹으로 나눠 그룹별로 AdamW와 Muon을 따로 배정할 수 있게 했고, 학습률은 0.5B 토큰 짧은 런에서 옵티마이저별로 따로 튜닝해 AdamW 2e-3, Muon 4e-3, 대조군 5e-4를 골랐다. 같은 seed의 런은 같은 가중치와 같은 배치에서 출발하므로 seed 3개 안에서 비교한다.
어떻게 쓰나
핵심 결과는 Muon이 첫 MoE 층의 span을 모든 seed에서 끌어올렸다는 것이다. 평균 0.084에서 0.194로 약 두 배가 됐고, 이후 세 실험의 9개 seed에서도 반복됐다. 저자는 반대를 예상했다. Muon은 expert가 토큰을 얼마나 받든 업데이트 크기가 거의 같아서, 바쁜 expert가 더 바빠지는 피드백이 약해질 것으로 봤기 때문이다. 그러나 상승은 라우터가 아니라 나머지 네트워크(backbone)를 Muon으로 학습할 때 발생했고, 라우터 옵티마이저를 바꿔도 거의 그대로였다. backbone이 Muon이면 첫 라우터 입력들이 한 방향을 공유해 평균 코사인 유사도가 작은 모델에서 0.4대(AdamW는 0.07)까지 올라갔고, 이 평균이 모든 토큰에 동일한 expert 선호를 만들어 bias가 따라잡는 수백 스텝 동안 가장 붐비는 expert가 공정 몫의 4~7배를 받았다. 두 모델의 첫 3,000 스텝을 반복한 probe에서 이 라우터 입력에서 배치 평균을 빼면 overload가 사라지고 손실 변화는 0.003 미만이었다. 작은 모델에서는 Muon 학습률을 절반으로 줄이면 상승의 약 4분의 3이 사라지고 1B 토큰에서 손실이 약간 좋아졌다.
전제와 한계
큰 모델에서도 Muon은 첫 MoE 층 span을 모든 seed에서 평균 0.152(AdamW span의 112%)만큼 올렸고, 초반 overload는 6.3과 7.2로 상한 8에 근접했다. 학습이 끝난 시점에도 첫 MoE 층은 두 Muon 조합 모두에서 2.2~2.8배로 편중돼 있었고 AdamW는 1.07이었다. 반대로 backbone을 Muon에 둔 채 라우터만 AdamW로 옮기면 두 모델 크기, 모든 학습률에서 warmup 구간 load churn이 커졌고 튜닝된 학습률에서는 seed마다 손실이 0.001~0.004 올랐다. 큰 모델에서는 Muon+AdamW 라우터 조합에서만 손실 스파이크가 seed별로 2·6·2회 발생했고 AdamW 전체나 Muon 전체 런에는 없었다. 저자는 이를 common share로 설명한다. 라우터 입력이 한 방향을 공유하면 모든 행의 그래디언트가 그 방향 성분을 크게 갖는데, AdamW는 항목별 정규화로 이 성분을 상당히 남기는 반면 Muon은 특이값을 고르게 만들어 shared direction 비중이 0.160에 그쳤다(AdamW 라우터는 0.497).
실무 관점의 결론은 세 가지다. 첫째, Muon으로 MoE를 학습한다면 첫 MoE 층의 bias span과 최다 expert overload를 별도 지표로 감시해야 한다. 다른 층은 상대적으로 안정적이었고 문제가 첫 층에 집중됐다. 둘째, 이 규모에서는 backbone이 Muon일 때 라우터도 Muon에 두는 편이 손실과 warmup 로드 변동 모두에서 유리했다. Qwen3.8-Flash-Next가 라우터를 AdamW로 둔 근거는 데이터 없이 서술된 초반 불안정성이며, 이 논문의 모델 크기에서는 재현되지 않았다. 셋째, 첫 라우터 입력의 배치 평균을 빼는 centering이나 더 빠른 bias, 낮은 Muon 학습률이 편중을 줄이는 수단으로 확인됐다. 다만 저자는 배치 평균이 토큰 라우팅을 배치에 의존하게 만들므로 실용 버전은 running mean이어야 한다고 지적한다.
한계는 저자가 분명히 밝힌다. 모델이 활성 55M·108M으로, 문제를 제기한 레시피의 주인공들보다 50~2,000배 작고 학습 토큰도 최대 3B로 네 자릿수 배 적다. 아키텍처 하나, 데이터셋 하나, sigmoid 라우터 하나, balancer 하나(bias, 학습률 두 종), 배치 크기 하나(65,536 토큰), warmup 하나(500 스텝)만 시험했다. sigmoid가 점수를 묶어 span에 상한을 만드는데 첫 MoE 층이 실제로 그 상한에 근접했고, DeepSeek-V4처럼 상한이 없는 점수 함수에서는 다를 수 있다. 학습률은 seed 하나와 0.5B 토큰에서 튜닝해 1B에 그대로 썼고, 비교는 seed 3개에 다중비교 보정이 없으며 여러 probe는 3,000 스텝만 덮는다. 또한 왜 backbone 전체를 Muon으로 둘 때만 첫 라우터 입력의 공유 방향이 생기는지, router 행들이 왜 공유 방향을 따라 벌어지는지는 규명하지 못했다. attention과 dense FFN만 Muon으로 둔 경우 라우터 입력 코사인은 0.054에 그쳤는데도 초반 overload는 5.07에 이르렀고, routed expert만 Muon으로 둔 조합은 오히려 가장 큰 span을 냈다.