MD解耦:通过解耦权重向量的幅度与方向改进神经网络训练

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

精选理由

这篇论文提出了一种简单通用的优化器改进方案,能解耦权重幅度和方向,消除权重衰减和warmup,在Adam和Muon上都有效,值得关注。

AI 摘要

论文提出MD Decoupling优化器修改方法,将每个权重分解为超球面上的固定范数方向与可学习的每行每列幅度增益,以解耦幅度和方向的更新。该方法与Adam和Muon等基础优化器兼容,消除了对权重衰减和warmup的需求。实验表明,MD Decoupling在宽模型和大型MoE模型上均优于精心调优的基线,并允许跨模型宽度直接迁移学习率而不需重新调参。

原文 · arXiv cs.LG

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

Modern neural network training relies on optimizers such as Adam and Muon which act on each weight matrix as a single object. Yet every weight matrix carries two distinct quantities -- a \emph{magnitude} and a \emph{direction} -- and all optimizers stepping in the matrix as a whole couple their dynamics: the directional change from an update depends on the current magnitude, while the magnitude drifts as a byproduct of learning the direction, so neither is governed directly by the learning rate. Typical training therefore leans on surrounding recipes such as weight decay and warmup to keep learning stable at scale, though these regulate the coupling only indirectly; other recent methods instead constrain the weight to a fixed-norm sphere, but add no learnable magnitude, leaving scale control to normalization layers alone. We propose \emph{Magnitude--Direction (MD) Decoupling}, an optimizer modification that factorizes each weight into a fixed-norm direction on a hypersphere and learnable per-row and per-column magnitude gains, updated at separate learning rates, all while the model still sees a single fused weight tensor. The method is agnostic to the base optimizer and removes the need for weight decay and warmup. Across both Adam and Muon, MD Decoupling improves on well-tuned baselines, transfers the optimal LR across model width without retuning, and continues to help at scale on large Mixture-of-Experts (MoE) models. Treating magnitude and direction as separately controlled quantities thus yields more predictable training dynamics and a simple, broadly applicable improvement to modern optimizers.