Improving Neural Network Training by Decoupling the Magnitude and Direction of Weight Vectors
Alexander Hägele, Alejandro Hernández-Cano, Atli Kosson, Martin Jaggi
Read on arXiv →
Key claim
Decoupling weight magnitude and direction improves training stability.
In plain English
Imagine you're training a neural network, and you want it to learn effectively without running into issues that slow down or destabilize the process. Typically, optimizers like Adam and Muon treat the entire weight matrix as a single entity, which can lead to problems. For instance, when you update weights, the change in direction is influenced by the current magnitude, and vice versa. This coupling can create instability, especially as models grow larger and more complex. You might try using techniques like weight decay or warmup to manage this, but they only address the issue indirectly, which can still lead to unpredictable training behavior. This is what's called coupling dynamics failure.
What this paper introduces is a way to tackle these issues head-on by separating the control of weight magnitude and direction. Instead of treating the weight matrix as a whole, the authors propose a method that factors each weight into a fixed-norm direction and learnable magnitude gains, which can be updated independently. This means you can adjust the learning rates for magnitude and direction separately, leading to more predictable and stable training dynamics.
In practice, this approach has shown to improve performance over well-tuned baselines across different optimizers, and it simplifies the training process by removing the need for weight decay and warmup. For anyone building large models, especially those using Mixture-of-Experts architectures, this method could lead to more efficient training and better overall performance.
The proposed method introduces a new way to decouple weight dynamics in optimizers, which is a significant advancement.
The claims are supported by improvements over well-tuned baselines and applicability across different optimizers.
Deep reliability assessment
The methodology supports the claim that MD Decoupling improves pretraining loss and learning-rate transfer in the authors' tested dense LLM and MoE settings, especially around Adam/Muon optimizer variants. Broader claims like removing the need for warmup and weight decay, or improving downstream task quality, are less directly supported because downstream evaluations, broader modalities, and more exhaustive hyperparameter comparisons are not shown.
Reproducibility
No open-source code repository is mentioned in the provided paper text; a shorter explanatory project/blog page is linked at https://haeggee.github.io/posts/magnitude-direction-decoupling. The experiments appear to require large-scale LLM/MoE pretraining infrastructure, and the excerpts do not provide enough implementation and data details for straightforward reproduction.
Key figure
Figure 1 shows that MD Decoupling improves LR-sweep loss for both Adam and Muon, preserves gains across MoE scaling runs, and makes the optimal MuonMD learning rate transfer across model widths.