mlsanityData to 7 Oct 2026

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

arXiv 2606.25971 ↗First posted 24 Jun 2026cs.LGcited by 3

On the map: Mixture-of-experts → MoE theory and gating functions

Abstract (arXiv)

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 magnitude and a 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. Then, neither is directly governed 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 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.

Key facts. Cited by 3 papers in the map (data through 7 Oct 2026). Placed in MoE theory and gating functions, one of 12 directions in Mixture-of-experts. Builds on 86 of its 98 references that are in the map. First posted on arXiv 24 Jun 2026.

Builds on

86 of 98 references.

References · 86

  1. 2021LoRA: Low-Rank Adaptation of Large Language Modelscited by 14,084
  2. 2020Measuring Massive Multitask Language Understandingcited by 7,768
  3. 2016Layer Normalizationcited by 5,579
  4. 2018Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challengecited by 4,710
  5. 2021RoFormer: Enhanced Transformer with Rotary Position Embeddingcited by 4,581
  6. 2019HellaSwag: Can a Machine Really Finish Your Sentence?cited by 4,018
  7. 2017Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layercited by 3,413
  8. 2022Training Compute-Optimal Large Language Modelscited by 2,927
  9. 2019Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelismcited by 1,829
  10. 2020GLU Variants Improve Transformercited by 1,724
  11. 2018Spectral Normalization for Generative Adversarial Networkscited by 1,618
  12. 2017Accurate, Large Minibatch SGD: Training ImageNet in 1 Hourcited by 1,479
38 more
  1. 2019Root Mean Square Layer Normalizationcited by 1,425
  2. 2023GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpointscited by 1,311
  3. 2019ZeRO: Memory Optimizations Toward Training Trillion Parameter Modelscited by 1,278
  4. 2024The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scalecited by 1,104
  5. 2021Scaling Vision Transformerscited by 993
  6. 2020On Layer Normalization in the Transformer Architecturecited by 906
  7. 2024DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Modelscited by 834
  8. 2019Big Transfer (BiT): General Visual Representation Learningcited by 766
  9. 2019On the Variance of the Adaptive Learning Rate and Beyondcited by 763
  10. 2024MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategiescited by 694
  11. 2021CogView: Mastering Text-to-Image Generation via Transformerscited by 691
  12. 2023Scaling Vision Transformers to 22 Billion Parameterscited by 690
  13. 2019Large Batch Optimization for Deep Learning: Training BERT in 76 minutescited by 671
  14. 2016Weight Normalization: A Simple Reparameterization to Accelerate Training of Deep Neural Networkscited by 541
  15. 2017Large Batch Training of Convolutional Networkscited by 529
  16. 2025Kimi K2: Open Agentic Intelligencecited by 523
  17. 2018Shampoo: Preconditioned Stochastic Tensor Optimizationcited by 413
  18. 2023Analyzing and Improving the Training Dynamics of Diffusion Modelscited by 396
  19. 2025Muon is Scalable for LLM Trainingcited by 381
  20. 2022Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfercited by 272
  21. 2024Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Expertscited by 191
  22. 2023Small-scale proxies for large-scale Transformer training instabilitiescited by 182
  23. 2024SOAP: Improving and Stabilizing Shampoo using Adamcited by 179
  24. 2025Training Deep Learning Models with Norm-Constrained LMOscited by 177
  25. 2024Old Optimizer, New Norm: An Anthologycited by 173
  26. 2024Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durationscited by 160
  27. 2017L2 Regularization versus Batch and Weight Normalizationcited by 121
  28. 2019An Exponential Learning Rate Schedule for Deep Learningcited by 121
  29. 2018A Closer Look at Deep Learning Heuristics: Learning rate restarts, Warmup and Distillationcited by 110
  30. 2024Language models scale reliably with over-training and on downstream taskscited by 105
  31. 2022On the SDEs and Scaling Rules for Adaptive Gradient Algorithmscited by 100
  32. 2023Tensor Programs VI: Feature Learning in Infinite-Depth Neural Networkscited by 94
  33. 2024Scaling Exponents Across Parameterizations and Optimizerscited by 71
  34. 2023Why Do We Need Weight Decay in Modern Deep Learning?cited by 70
  35. 2025Don't be lazy: CompleteP enables compute-efficient deep transformerscited by 69
  36. 2024nGPT: Normalized Transformer with Representation Learning on the Hyperspherecited by 68
  37. 2018Theoretical Analysis of Auto Rate-Tuning by Batch Normalizationcited by 65
  38. 2018Norm matters: efficient and accurate normalization schemes in deep networkscited by 62

50 most cited of 86 All 86 (JSON)

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2026: 3.

3 papers.

Similar, not cited

Close in wording, no citation either way

  1. 2026Scale Weight Decay and Train Better
  2. 2026Spectral Weight Decay: Inducing Low-Rank Structure in Neural Network Weights
  3. 2026The Ups and Downs of Backprop Weights
  4. 2026LayerRoPE: Dynamic Depth-wise Magnitude & Angular Superposition
  5. 2026Balancing Stability and Plasticity in Sequentially Trained Early-Exiting Neural Networks
  6. 2025Greedy Alignment Principle for Optimizer Selection
  7. 2026Neural Quadratic Forms: A Unified Minimal Model for Sudden Learning and Scaling Laws
  8. 2026Symmetry-Compatible Principle for Optimizer Design: Embeddings, LM Heads, SwiGLU MLPs, and MoE Routers