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

1. Problem Statement

Modern neural network training relies on optimizers like Adam and Muon that treat weight matrices as single, monolithic objects. However, every weight matrix $W$ inherently carries two distinct quantities: a magnitude ($\|W\|$) and a direction ($\hat{W} = W/\|W\|$). In standard optimization, these two quantities are coupled, leading to unstable and unpredictable training dynamics:

• Directional Change Dependency: The angular change induced by an update depends inversely on the current magnitude. Consequently, the learning rate (LR) does not directly control the rate of directional change, which can vary across layers and over time.
• Unintended Magnitude Drift: Updates tend to be perpendicular to current weights (due to scale-invariance or noise), causing the magnitude to grow even in the absence of a radial gradient. This "runaway growth" forces practitioners to rely on indirect mechanisms like weight decay to stabilize training and warmup to prevent early instability.
• Scaling Limitations: Current methods for scaling training (e.g., $\mu$P) often require re-tuning hyperparameters when changing model width or depth. Furthermore, recent approaches that constrain weights to a fixed-norm sphere often lack learnable magnitude parameters, leaving scale control solely to normalization layers.

2. Methodology: Magnitude–Direction (MD) Decoupling

The authors propose Magnitude–Direction (MD) Decoupling, an optimizer modification that factorizes each weight matrix into a fixed-norm direction and learnable magnitude gains, updating them at separate learning rates.

Core Factorization

Each weight matrix $W$ is reparameterized as:
$$W = \text{diag}(\gamma_{\text{row}}) \, \hat{W} \, \text{diag}(\gamma_{\text{col}})$$
Where:

• $\hat{W}$ is the direction, constrained to lie on a fixed hypersphere (e.g., Frobenius norm).
• $\gamma_{\text{row}}$ and $\gamma_{\text{col}}$ are learnable magnitude gains (vectors acting on rows and columns, respectively).

Optimization Procedure

The method operates entirely within the optimizer step, keeping the model's forward pass unchanged (the model sees a single fused tensor $W$):

1. Recovery: Recover $\hat{W}$ and gains $\gamma$ from the fused weight $W$.
2. Gradient Splitting: Split the gradient $G = \partial L / \partial W$ into a direction gradient ($\partial L / \partial \hat{W}$) and gain gradients (summed over the respective axes).
3. Direction Update: Update $\hat{W}$ using the base optimizer (e.g., Adam, Muon) with learning rate $\eta_W$, then project it back onto the sphere to maintain fixed norm.
4. Gain Update: Update the gains $\gamma$ (optionally via a smooth map like softplus to ensure positivity) using a separate learning rate $\eta_\gamma$.
5. Reassembly: Reconstruct $W$ for the next forward pass.

Key Properties

• Optimizer Agnostic: The method acts as a wrapper around any base optimizer (Adam, Muon, Shampoo, etc.).
• Direct LR Control: By fixing the direction on a sphere, the learning rate directly sets the relative angular update, removing the dependency on the current weight magnitude.
• Elimination of Heuristics: The method removes the need for weight decay (as weights are already norm-constrained) and warmup (as large, destabilizing early updates are regulated from step one).
• Fine-Grained Scale Control: Unlike methods relying solely on normalization layer gains, MD Decoupling provides explicit, learned control over the scale of individual rows and columns, allowing the model to amplify or dampen specific features.

3. Key Contributions and Findings

Empirical Results

The authors validate MD Decoupling across dense models (181M to 1.29B parameters) and large Mixture-of-Experts (MoE) models (up to 6.7B total parameters).

• Performance Gains: MD Decoupling improves upon well-tuned baselines (AdamW and Muon). On large MoE models, MuonMD (Muon with MD Decoupling) reaches the loss of AdamW with approximately 2$\times$ less compute.
• Learning Rate Transfer: The method enables zero-shot learning rate transfer across model width. The optimal matrix LR found on a small model remains optimal for significantly larger models without re-tuning. This holds for both width and depth scaling.
• Warmup-Free Training: Experiments show that dropping warmup not only maintains stability but improves final loss, as early training steps are no longer wasted on reduced learning rates.
• Ablation Insights:
  • Normalization Axis: Constraining the direction via Frobenius norm (whole matrix) is found to be the most flexible and performs comparably to row/column constraints.
  • Gain Parametrization: A combined per-row and per-column gain ($\gamma_{\text{row}} + \gamma_{\text{col}}$) outperforms scalar or single-axis gains. The softplus parameterization for gains provides a slight, consistent edge over direct or exponential updates.
  • Embeddings: Constraining embedding vectors to unit L2 norm performs slightly better than leaving them unconstrained.

Comparison to Related Work

The paper situates MD Decoupling within a growing body of work on spherical training and reparameterization:

• Unlike nGPT, which bundles spherical constraints with specific architectural changes (e.g., L2 normalization, interpolated residuals), MD Decoupling is an optimizer-side modification that can be applied to standard architectures. When applied to nGPT's architecture, MD Decoupling further improves performance.
• Unlike Muown or HyperP, which may fix norms without learnable magnitudes or tie the method to specific optimizers, MD Decoupling explicitly learns magnitudes and is optimizer-agnostic.
• It addresses the "magnitude-direction interference" identified in prior work (e.g., Kosson et al., 2024a) by separating the dynamics of the two quantities.

4. Significance and Claims

The paper claims that treating magnitude and direction as separately controlled quantities yields more predictable training dynamics and a simple, broadly applicable improvement to modern optimizers.

• Simplicity: The method requires no architectural changes and adds minimal computational overhead (element-wise operations during the optimizer step).
• Robustness: It simplifies the training recipe by removing the need for weight decay and warmup, reducing the hyperparameter search space.
• Scalability: It provides a principled mechanism for learning rate transfer across model width, addressing a key bottleneck in scaling large language models.

The authors conclude that while the additional logic in the optimizer step introduces a modest overhead, the benefits in training stability, compute efficiency, and hyperparameter transfer make MD Decoupling a significant advancement for training large-scale neural networks.