MeqMuon이 Muon 업데이트의 행·열 불균형을 자동으로 잡는다

무엇인가

이 논문은 LLM 사전학습용 옵티마이저 Muon의 업데이트 불균형 문제를 다룬다. AdamW가 오랫동안 표준이었지만 Muon은 은닉층의 2D 가중치에 대해 모멘텀 행렬을 직교화해 더 나은 학습 효율과 정확도를 보여왔고, DeepSeek-V4나 Kimi K2/K3 같은 프런티어 모델 사전학습에도 쓰였다. 최근 NorMuon은 Muon의 직교화 업데이트에서 행 방향 크기 불균형이 크다는 점을 관찰하고 Adam 스타일의 행 단위 정규화를 추가했다. 그러나 저자들은 Muon의 직교화 업데이트가 행이 아니라 열 방향으로 불균형한 경우도 있다는 것을 확인한다. 즉 어떤 업데이트는 행 크기는 비교적 균형 잡혀 있지만 열 크기가 불균형하고, 다른 업데이트는 그 반대다. 따라서 행 단위 정규화만으로는 서로 다른 불균형 패턴을 모두 감당할 수 없다.

어떻게 동작하나

저자들은 이 불균형이 직교화의 구조적 제약에서 나온다고 설명한다. 완전 랭크 행렬 X의 정확한 직교화를 Y라 하면, m > n일 때 Y의 전치와 Y의 곱이 단위행렬이 되어 모든 열 크기가 같아지고 열 변동계수는 0이 되지만 행 크기는 제약을 받지 않는다. m < n이면 반대로 행만 균일해진다. 정방행렬은 행·열 변동계수가 모두 0이어야 하지만, 실제로는 뉴턴-슐츠 5회 반복의 근사 오차 때문에 잔여 불균형이 남고 행과 열의 변동계수가 서로 다른 속도로 줄어든다. Llama-350M 4번째 트랜스포머 레이어, 학습 스텝 1,144에서 저장된 업데이트를 보면 q_proj, k_proj, v_proj, gate_proj, up_proj는 열은 균형 잡혀 있고 행이 불균형한 반면, o_proj와 down_proj는 정반대 패턴을 보인다. v_proj의 경우 반복 0회에서 (행 0.673, 열 0.773)이던 변동계수가 5회에서 (0.195, 0.041)로 역전된다. 한편 직교화되지 않는 나머지 2D 파라미터인 embed_tokens와 lm_head의 모멘텀 행렬은 행과 열 양쪽 모두 불균형하다.

무엇과 다른가

MeqMuon은 파라미터마다 모멘텀 버퍼 하나를 유지하고 파라미터 종류에 따라 정규화를 다르게 적용한다. 은닉층 2D 가중치에서는 Muon처럼 모멘텀을 직교화한 뒤, 그 업데이트의 행 변동계수와 열 변동계수를 비교해 더 큰 쪽을 정규화한다. 행이 더 불균형하면 행 단위로, 열이 더 불균형하면 열 단위로 각 행렬마다 매 최적화 스텝 독립적으로 방향을 고른다. 토큰 임베딩과 LM 헤드 같은 나머지 2D 파라미터는 AdamW 업데이트를 버리고 직교화하지 않은 모멘텀에 행·열 양방향 정규화를 연속 적용한다. RMSNorm 가중치나 바이어스 같은 1D 파라미터는 모멘텀 벡터를 전역 RMS로 나눈다. 최종 업데이트는 가중치 감쇠 항과 정규화된 업데이트를 ρη로 스케일한 항으로 구성된다. 행 또는 열 정규화를 하면 업데이트 RMS가 행렬 모양과 무관해지므로 Muon이 쓰던 sqrt(max(m,n)) 보정이 필요 없어지고, 나머지 파라미터에서 AdamW를 정규화된 모멘텀으로 대체하면서 2차 모멘트 추정치 저장이 사라진다.

어떻게 쓰나

실험은 Llama, SmolLM2, Qwen2 모델을 무작위 초기화에서 영어 C4로 사전학습하며 수행했다. Chinchilla 규칙에 따라 파라미터 수의 20배 토큰을 예산으로 잡았고, 시퀀스 길이 1,024, 전역 배치 512, 가중치 감쇠 0.1, 5% 선형 워밍업 후 코사인 감쇠를 사용했다. 비교 대상은 AdamW, SCALE, Muon, NorMuon이다. 저자들은 Muon, NorMuon, MeqMuon이 AdamW와 SCALE보다 최종 검증 퍼플렉시티가 크게 낮고, MeqMuon이 보고된 모든 모델 스케일에서 최고의 최종 퍼플렉시티를 얻었다고 서술한다. 다만 제공된 본문에는 Table 4의 개별 퍼플렉시티 수치가 실려 있지 않아 구체적인 값은 확인할 수 없다. 학습 중 검증 퍼플렉시티 곡선도 대체로 Muon과 NorMuon 아래에 머문다고 보고한다.

전제와 한계

메모리 측면에서는 Qwen2-0.5B 기준으로 MeqMuon이 Muon과 NorMuon 대비 옵티마이저 상태 메모리를 21.6% 줄인다. Muon이 은닉층 2D 가중치의 모멘텀과 나머지 파라미터의 AdamW 1·2차 모멘트를 저장하고 NorMuon이 여기에 행 단위 2차 모멘트를 추가로 저장하는 반면, MeqMuon은 은닉층과 나머지 파라미터 모두 모멘텀 값만 저장하기 때문이다. 정규화 방향 선택 빈도를 보면 Llama-350M과 SmolLM2-360M에서 세로로 긴 gate_proj와 up_proj는 항상 행 정규화를, 가로로 넓은 down_proj는 항상 열 정규화를 선택했고, SmolLM2-360M의 k_proj와 v_proj도 주로 열 정규화를 골랐다. 이는 직교화가 직사각 행렬에서 한쪽만 제약한다는 분석과 일치한다. 정방행렬인 q_proj는 주로 행, o_proj는 주로 열 정규화를 선택해 모듈과 모델에 따라 선호가 갈렸다. 구성요소별 절제 실험에서는 은닉층 2D 가중치만 바꾼 변형과 나머지 2D 파라미터만 바꾼 변형 모두 Muon보다 낮은 검증 퍼플렉시티를 냈고, 세 가지를 모두 적용한 MeqMuon이 가장 낮았다.

실무 관점에서 MeqMuon은 기존 Muon 학습 파이프라인에 옵티마이저만 교체하는 형태로 적용할 수 있고, AdamW 2차 모멘트 상태를 없애 옵티마이저 상태 메모리를 줄인다는 점이 가장 직접적인 이득이다. NorMuon 대비 추가 메모리도 늘지 않는다. 도입 전에 확인할 것은 세 가지다. 첫째, 매 스텝 행·열 RMS와 변동계수를 계산하는 비용이 학습 처리량에 얼마나 영향을 주는지다. 둘째, RMS 정렬 계수 ρ와 모멘텀 계수 μ 같은 하이퍼파라미터가 기존 Muon 설정과 어떻게 맞물리는지다. 셋째, 토큰 임베딩과 LM 헤드를 공유하는 모델 구조에서 나머지 2D 파라미터를 양방향 정규화로 처리하는 방식이 그대로 통하는지다.

저자들이 밝힌 전제와 한계도 분명하다. 실험은 NVIDIA RTX A6000에서 DDP로 수행됐고 평가 모델은 Llama-350M, SmolLM2-360M, Qwen2-0.5B 수준으로, 논문이 동기로 든 1조~2조 파라미터급 프런티어 규모에서의 검증은 제시되지 않았다. SCALE은 토큰 임베딩과 LM 헤드를 서로 다르게 업데이트해서 가중치를 공유하는 SmolLM2와 Qwen2에는 적용할 수 없어 Llama 모델에서만 평가했다. 실험 설정에서 그래디언트 클리핑과 활성화 체크포인팅은 비활성화했다. 또한 변동계수 비교로 정규화 방향을 고르는 절차는 경험적 관찰에 근거하며, 수렴에 대한 이론적 분석은 제시되지 않는다.