将权重的大小与方向分开优化,让训练更稳定高效。
Improving Neural Network Training by Decoupling the Magnitude and Direction of Weight Vectors

- 把权重分解为方向和大小两部分,分别用不同学习率更新。
- 在多种模型上提升性能,且无需权重衰减和预热。
- 适合大规模模型训练,尤其对MoE架构有显著帮助。
现代神经网络训练依赖Adam、Muon等优化器,它们将每个权重矩阵视为整体进行更新。然而,每个权重矩阵包含两个独立成分:大小(magnitude)和方向(direction),现有优化器将两者耦合,导致方向变化受当前大小影响,而大小又随方向学习而漂移,二者均无法直接受学习率控制。因此,训练常依赖权重衰减和预热等技巧间接调节耦合。其他方法虽将权重约束在固定范数球面上,但无法学习可变大小,尺度控制全靠归一化层。本文提出大小-方向解耦(Magnitude–Direction, MD)优化策略,将权重分解为单位范数方向和可学习的行/列级大小增益,分别以独立学习率更新,同时模型仍感知单一融合权重张量。该方法不依赖基础优化器,无需权重衰减和预热。在Adam和Muon上,MD解耦均优于调优基线,在不同模型宽度间实现最优学习率迁移,且在大规模Mixture-of-Experts(MoE)模型上仍持续有效。将大小与方向作为独立可控变量,显著提升了训练动态的可预测性,并为现代优化器提供了一种简单、通用的改进方案。
原文摘要 · Abstract (English)
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. Then, neither is directly governed 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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。