Improving Neural Network Training by Decoupling the Magnitude and Direction of Weight Vectors
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
- 2021LoRA: Low-Rank Adaptation of Large Language Modelscited by 14,084
- 2020Measuring Massive Multitask Language Understandingcited by 7,768
- 2016Layer Normalizationcited by 5,579
- 2018Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challengecited by 4,710
- 2021RoFormer: Enhanced Transformer with Rotary Position Embeddingcited by 4,581
- 2019HellaSwag: Can a Machine Really Finish Your Sentence?cited by 4,018
- 2017Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layercited by 3,413
- 2022Training Compute-Optimal Large Language Modelscited by 2,927
- 2019Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelismcited by 1,829
- 2020GLU Variants Improve Transformercited by 1,724
- 2018Spectral Normalization for Generative Adversarial Networkscited by 1,618
- 2017Accurate, Large Minibatch SGD: Training ImageNet in 1 Hourcited by 1,479
38 more
- 2019Root Mean Square Layer Normalizationcited by 1,425
- 2023GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpointscited by 1,311
- 2019ZeRO: Memory Optimizations Toward Training Trillion Parameter Modelscited by 1,278
- 2024The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scalecited by 1,104
- 2021Scaling Vision Transformerscited by 993
- 2020On Layer Normalization in the Transformer Architecturecited by 906
- 2024DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Modelscited by 834
- 2019Big Transfer (BiT): General Visual Representation Learningcited by 766
- 2019On the Variance of the Adaptive Learning Rate and Beyondcited by 763
- 2024MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategiescited by 694
- 2021CogView: Mastering Text-to-Image Generation via Transformerscited by 691
- 2023Scaling Vision Transformers to 22 Billion Parameterscited by 690
- 2019Large Batch Optimization for Deep Learning: Training BERT in 76 minutescited by 671
- 2016Weight Normalization: A Simple Reparameterization to Accelerate Training of Deep Neural Networkscited by 541
- 2017Large Batch Training of Convolutional Networkscited by 529
- 2025Kimi K2: Open Agentic Intelligencecited by 523
- 2018Shampoo: Preconditioned Stochastic Tensor Optimizationcited by 413
- 2023Analyzing and Improving the Training Dynamics of Diffusion Modelscited by 396
- 2025Muon is Scalable for LLM Trainingcited by 381
- 2022Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfercited by 272
- 2024Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Expertscited by 191
- 2023Small-scale proxies for large-scale Transformer training instabilitiescited by 182
- 2024SOAP: Improving and Stabilizing Shampoo using Adamcited by 179
- 2025Training Deep Learning Models with Norm-Constrained LMOscited by 177
- 2024Old Optimizer, New Norm: An Anthologycited by 173
- 2024Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durationscited by 160
- 2017L2 Regularization versus Batch and Weight Normalizationcited by 121
- 2019An Exponential Learning Rate Schedule for Deep Learningcited by 121
- 2018A Closer Look at Deep Learning Heuristics: Learning rate restarts, Warmup and Distillationcited by 110
- 2024Language models scale reliably with over-training and on downstream taskscited by 105
- 2022On the SDEs and Scaling Rules for Adaptive Gradient Algorithmscited by 100
- 2023Tensor Programs VI: Feature Learning in Infinite-Depth Neural Networkscited by 94
- 2024Scaling Exponents Across Parameterizations and Optimizerscited by 71
- 2023Why Do We Need Weight Decay in Modern Deep Learning?cited by 70
- 2025Don't be lazy: CompleteP enables compute-efficient deep transformerscited by 69
- 2024nGPT: Normalized Transformer with Representation Learning on the Hyperspherecited by 68
- 2018Theoretical Analysis of Auto Rate-Tuning by Batch Normalizationcited by 65
- 2018Norm matters: efficient and accurate normalization schemes in deep networkscited by 62
50 most cited of 86 All 86 (JSON)
Built on by
Citing papers per year: 2026: 3.
3 papers.
Similar, not cited
Close in wording, no citation either way
- 2026Scale Weight Decay and Train Better
- 2026Spectral Weight Decay: Inducing Low-Rank Structure in Neural Network Weights
- 2026The Ups and Downs of Backprop Weights
- 2026LayerRoPE: Dynamic Depth-wise Magnitude & Angular Superposition
- 2026Balancing Stability and Plasticity in Sequentially Trained Early-Exiting Neural Networks
- 2025Greedy Alignment Principle for Optimizer Selection
- 2026Neural Quadratic Forms: A Unified Minimal Model for Sudden Learning and Scaling Laws
- 2026Symmetry-Compatible Principle for Optimizer Design: Embeddings, LM Heads, SwiGLU MLPs, and MoE Routers