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

2026-06-24

AI Summary

This work addresses the limitations of conventional optimizers, which entangle the magnitude and direction of weight updates, leading to unstable training dynamics that necessitate indirect stabilization techniques such as weight decay and learning rate warmup. To overcome this, the authors propose a Magnitude-Direction (MD) decoupling mechanism that explicitly decomposes each weight matrix—without altering model architecture—into a unit-norm directional component and learnable row- and column-wise magnitude gains. These components are optimized independently using separate learning rates, enabling precise control over magnitude and direction dynamics. The MD framework is compatible with any base optimizer (e.g., Adam, Muon) and eliminates reliance on weight decay or warmup schedules. Experiments demonstrate consistent improvements over carefully tuned baselines across diverse model scales, support learning rate transfer across model widths, and remain effective in large-scale Mixture-of-Experts (MoE) architectures.

Abstract

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.

Problem

• neural network training
• weight magnitude
• weight direction
• optimizer coupling
• training stability

Innovation

• Magnitude-Direction Decoupling
• optimizer
• weight normalization
• learning rate
• neural network training