Improving Neural Network Training by Decoupling the Magnitude and Direction of Weight Vectors
Abstract Overview
This paper argues that standard optimizers couple two distinct aspects of each weight matrix: its magnitude and its direction, which makes directional updates depend on current norm and causes magnitude growth even in scale-invariant settings. The authors propose Magnitude–Direction (MD) Decoupling, which factorizes each weight into a fixed-norm directional component on a hypersphere plus learnable per-row and per-column magnitude gains updated with separate learning rates. The method is implemented inside the optimizer so the model still operates on a fused weight tensor, and it is presented as compatible with different base optimizers such as Adam and Muon. The empirical study combines dense-model ablations with larger Mixture-of-Experts experiments to examine normalization choices, gain parameterizations, learning-rate transfer, schedule behavior, and scaling.
Novelty
The distinctive contribution is an optimizer-level reparameterization that simultaneously enforces fixed-norm directional training and restores fine-grained learnable magnitudes through row and column gains. Unlike prior sphere-based approaches that fix norms without explicit magnitude learning, this method decouples direction and scale while remaining largely optimizer-agnostic.
Results
Across dense language models, MD Decoupling improves over well-tuned AdamW and Muon baselines, with the paper reporting that the best matrix learning rate transfers across model width without retuning. Ablations indicate that most benefit comes from spherical direction control, while adding learned row-and-column gains gives an additional improvement and performs best among gain variants. In large Mixture-of-Experts experiments, the method remains ahead across the tested compute range and is reported to reach AdamW-level loss with roughly 2× less compute.
Key Points
- MD Decoupling separates each weight matrix into a fixed-norm directional component and learnable magnitude gains, with separate update rules and learning rates.
- The authors find the Frobenius-sphere constraint is a flexible default, unit-norm embeddings behave at least as well as unconstrained ones, and combined row-and-column gains outperform scalar or one-sided gains.
- The method reduces reliance on weight decay and warmup, preserves learning-rate transfer across width, and continues to show benefits when scaling to large Mixture-of-Experts models.