Improving Neural Network Training by Decoupling the Magnitude and Direction of Weight Vectors

가중치 행렬의 크기(magnitude)와 방향(direction)을 분리해 업데이트하는 optimizer 개선. weight decay와 warmup 없이도 Adam/Muon baseline을 능가하며, large MoE에서 AdamW 대비 약 2배 효율적이다.

MD Decoupling은 각 가중치 행렬을 fixed-norm 방향과 학습 가능한 per-row/per-column magnitude gain으로 분해하고, 각각 다른 학습률로 업데이트하는 optimizer 수정이다. 방향은 hypersphere 상에 고정하고 magnitude gain을 별도로 학습하여, learning rate가 방향 변화를 직접 제어하게 한다. Adam과 Muon 모두에서 well-tuned baseline을 능가하며, 모델 폭이 변해도 최적 LR이 전이된다. large MoE 스케일에서 AdamW와 동일 loss에 약 2배 적은 compute로 도달한다.

핵심 요약

• 가중치 magnitude-direction 간섭을 수학적으로 분석: 방향 변화는 현재 magnitude에 반비례하고, magnitude는 방향 학습의 부산물로 증가한다.
• optimizer 내부에서 가중치를 fixed-norm direction + learnable gains로 분해하므로 모델 구조 변경 없이 적용 가능하다.
• weight decay와 warmup이 불필요해지며, 이는 large-scale 학습 레시피를 크게 단순화한다.
• 모델 폭이 달라도 최적 learning rate가 전이되어 LR 재튜닝이 필요 없다 (µP와 유사하나 직접 얻어짐).
• large MoE 모델에서 MuonMD는 AdamW와 동일 loss에 약 2배 적은 compute로 도달한다.
• 방향을 sphere에 고정하는 것이 대부분의 개선을 가져오고, learnable gains가 추가적인 정확도 향상을 제공한다.

논문 링크

외부 연구를 정리한 자료입니다. HDATF가 발표한 논문이나 제품 성능을 측정한 결과는 아닙니다.