Matrix optimizers improve large AI model training performance
Normalize-Then-Precondition: A Hierarchical Approach to Marginal Scale and Interaction Geometry for LLM Training
Machine LearningArtificial Intelligence
Summary
Training large language models requires smart ways to update many parameters efficiently. The authors studied a type of advanced optimizer called matrix optimizers and found a new way to organize their information hierarchically. Their new approach, called Normalize-Then-Precondition, first normalizes simple scale information, then adjusts for complex interactions. This method beats common optimizers like AdamW and Muon in experiments with popular AI models. The authors also provide tools to help others use this new optimizer in large-scale model training.
What this means in practice
- •For machine learning engineers: Improve training speed and stability of large language models by employing the Normalize-Then-Precondition optimizer during model training.
- •For ai infrastructure teams: Enable efficient deployment of advanced matrix optimization in distributed training setups by integrating approximation methods like Newton-Schulz iterations and randomized sketching.
Authors
Zixuan Gong, Zeyu Gan, Jiaye Teng, Yong Liu
Abstract
Matrix optimizers have emerged as a promising direction, with Muon standing out as a prominent design. Revisiting Muon through its full-Gram representation, we observe that it jointly processes marginal-scale and interaction information. This opens an alternative way to organize geometric information hierarchically, motivating the Normalize-Then-Precondition framework. Specifically, it first uses diagonal-Gram information to construct a marginally normalized update, then applies spectral preconditioning to its directional interaction geometry. Building on this framework, we develop NormPre with NormPre-G and NormPre-L adopting global and localized spectral preconditioning, grounded in spectral-norm steepest descent and a regularized formulation followed by leading mode selection, respectively. To enable large-scale training, NormPre-G uses Newton-Schulz iterations and NormPre-L employs randomized sketching to approximate the leading interaction eigenspace. Theoretically, we establish $\mathcal{O}(T^{-1/2})$ convergence guarantees for simplified versions of NormPre. Across extensive pretraining experiments on GPT-2 Small, LLaMA and Qwen3, both variants consistently outperform AdamW, Muon and MANO under matched training budgets. Further efficiency and spectral analyses reveal the complementary strengths of two variants and characterize their performance-efficiency trade-off. We open-source our code through a GitHub repository at https://github.com/zx-gong/NormPre.