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

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.

• The paper presents MD Decoupling, which separates weight magnitude and direction to allow independent control via learning rates.
• It introduces a factorized weight representation that uses fixed-norm directions and distinct magnitude gains, eliminating the need for weight decay and warmup.
• Empirical studies demonstrate enhanced compute efficiency and zero-shot hyperparameter transfer across various architectures and model scales.

Decoupling Magnitude and Direction in Neural Network Training: A Detailed Analysis

Motivation and Theoretical Foundation

The paper "Improving Neural Network Training by Decoupling the Magnitude and Direction of Weight Vectors" (2606.25971) investigates the inherent coupling between weight magnitude and direction in gradient-based optimization of neural networks. Standard optimizers such as Adam, Muon, and derivatives operate on weight matrices as atomic objects, implicitly combining two distinct quantities: magnitude (norm) and direction. This coupling is empirically problematic; the angular update induced by a step is inversely proportional to the current magnitude, and magnitude increases even under scale-invariant losses due to stochastic perpendicular updates. This interferes with precise control via learning rate (LR), often necessitating auxiliary mechanisms like weight decay or warmup scheduling.

Figure 1: Illustration of weight norm-induced distortion in optimizer updates – identical steps rotate weights more at small magnitude and inflate the norm even when only direction matters.

The authors provide quantitative evidence for this interference, observing that the learning rate fails to independently regulate either directional or magnitude change. The equilibrium reached by standard methods is dictated by the interplay of LR and decay, not any explicit or optimal scale dictated by downstream normalization layers. This is particularly problematic for models featuring pervasive normalization (e.g., RMSNorm) where magnitude is often redundant yet still influences learning dynamics.

Magnitude–Direction Decoupling: Algorithmic Formulation

The central contribution is Magnitude–Direction (MD) Decoupling, an optimizer-side modification that factorizes each weight matrix into a fixed-norm direction (on a hypersphere) and separately learned per-row/per-column magnitude gains updated at distinct LRs:

$$W = \operatorname{diag}(\gamma_\text{row}) \cdot W_\text{dir} \cdot \operatorname{diag}(\gamma_\text{col}),$$

where $W_\text{dir}$ is strictly normalized (Frobenius or axiswise), and the $\gamma$ gains (initialized to unity) are updated independently (typically via Adam). The forward pass operates on the fused tensor, ensuring architectural compatibility and minimal overhead, while the optimizer internally manages the decoupled parameterization.

Figure 2: Two degrees of freedom for MD Decoupling: direction normalization axis (rows, columns, Frobenius) and gain axis (scalar, per row/column, row+column).

This design is agnostic to the underlying base optimizer, disables the need for weight decay and warmup, and allows for precise angular learning schedule control. Notably, fine-grained magnitude gains are not simply redundant with normalization; the model exploits them for activation scaling, as evidenced in gain dynamics monitored during training.

Empirical Ablations on Dense and Sparse Architectures

In-depth ablation studies are performed on GPT-style LMs (181M–1.29B parameters) to interrogate the sensitivity of normalization axis, embedding constraints, and gain parameterization:

• Normalization Axis: Row, column, and matrix (Frobenius) constraints yield nearly identical optimal losses; Frobenius is selected for maximal flexibility.

Figure 3: LR sweeps show normalization axis has negligible impact on final loss; Frobenius constraint is preferred for flexibility.

• Gain Axis: Among scalar, row-wise, column-wise, and combined row+column gains, the combined variant significantly outperforms alternatives. Higher-rank gains provide no additional benefit.
• Gain Parameterization: Direct, exponential, and smooth softplus maps perform similarly; softplus offers marginal stability and is adopted as default.
• Learning Rate Sensitivity: Matrix LR is dominant; embedding and gain group LRs show flat loss curves across orders of magnitude, further demonstrating the protocol's robustness.

Figure 4: LR sweeps confirm per-row+per-column gains provide best performance; gain parameterization is benign but softplus confers slight edge.

Scaling Behavior and Hyperparameter Transfer

One of the strongest empirical claims is that the optimal matrix LR under MD Decoupling is invariant to model width (and also depth given proper scaling), enabling zero-shot LR transfer across sizes. This is validated across width, depth, and joint scaling sweeps.

Figure 5: Optimal matrix LR stays fixed as model is scaled; transferability without retuning is demonstrated.

Relative weight updates and activation statistics confirm constant dynamics when scaling, consistent with theory. In addition, standard warmup, which stabilizes large early updates, is rendered obsolete and even detrimental; dropping it yields improved losses both at initialization and after checkpoint re-warming.

Learning Rate Schedules: Direct Control Enabled by MD Decoupling

Training on the sphere makes the LR schedule directly control the weight update magnitude; the schedule's shape is thus more critical than under weight decay. The paper demonstrates that linear decay and Warmup-Stable-Decay (WSD) both have distinct profiles; annealing schedules need to be rethought for on-sphere optimization.

Figure 6: Relative weight update follows the LR schedule directly under MD Decoupling; schedule selection is critically important.

Large-Scale Validation on Mixture-of-Experts Models

MD Decoupling is benchmarked on DeepSeekMoE-style transformers (up to 6.7B total parameters, 810M active), showing persistent gains: MuonMD achieves the same loss as AdamW at roughly $2\times$ less compute. Scaling law fits demonstrate the improvement is a downward level shift (reduced coefficient $A$) without affecting scaling exponent $\alpha$.

Figure 7: Decoupling magnitude/direction yields improved compute-efficiency for MoEs; MuonMD reaches AdamW's loss at half the compute required.

Batch-size scaling and checkpoint re-warming experiments corroborate practical feasibility at scale.

Optimizer Configuration and Computational Efficiency

The implementation is compatible with memory-efficient sharding and distributed setups. Throughput measurements verify negligible runtime overhead, especially for MuonMD ($<2\%$ slowdown); AdamMD's overhead reduces rapidly with increased global batch size.

Practical and Theoretical Implications

• No Weight Decay/No Warmup: MD Decoupling obviates auxiliary hyperparameters, simplifying tuning and enabling direct schedule control.
• Hyperparameter Transfer: LR invariance across width/depth provides a practical recipe for scaling; tuning on small models generalizes to much larger systems.
• Architectural Compatibility: The fused-parameter approach ensures MD Decoupling integrates with standard architectures and optimizer stacks.
• Research Directions: The precise role of learned magnitudes, optimal LR schedules on the sphere, loss-landscape sharpness, behavior in RL regimes, and efficacy in low-precision training are posed as open questions.

Relation to Prior Work

The method situates itself among numerous recent advances: Spherical optimization (nGPT (Loshchilov et al., 2024), HyperP (Ren et al., 30 Mar 2026)), weight normalization (Salimans et al., 2016), scale-vector learning (Wang et al., 26 May 2026), and operator norm constraints (Xie et al., 13 Jan 2026), with the distinction of decoupled learning rates and explicit learnable gains inside the optimizer. Zero-shot LR transfer aligns with $\mu$P (Yang et al., 2022), but is achieved via decoupled direction/magnitude updates rather than tensor programs. Recent variants of Muon (Muown (Lion et al., 11 May 2026), AngularMuown (Hübler et al., 22 Jun 2026)) further corroborate the theoretical and empirical premises.

Conclusion

Magnitude–Direction Decoupling constitutes an optimizer-centric approach to disentangling weight magnitude and direction, conferring precise control over learning dynamics and enabling a robust, scalable recipe for large-scale pretraining. Empirical results across dense and sparse architectures validate significant improvements in both stability and compute efficiency, with ultralight hyperparameter tuning and direct schedule control. The method aligns with theoretical advances in scale-invariant dynamics and opens avenues for further research into optimizer parametrization, schedule design, and interpretability of learned magnitudes.