ROOT: Robust Orthogonalized Optimizer for Neural Network Training

TL;DR

ROOT is a new optimizer designed to improve the robustness and stability of training large language models by addressing precision and noise vulnerabilities, leading to faster convergence and better performance.

Contribution

We introduce ROOT, a robust orthogonalized optimizer with adaptive Newton orthogonalization and proximal optimization, enhancing training stability for large-scale models.

Findings

ROOT outperforms Muon and Adam optimizers in robustness and convergence speed.

ROOT maintains stability in noisy, non-convex training scenarios.

Experimental results show improved final model performance.

Abstract

The optimization of large language models (LLMs) remains a critical challenge, particularly as model scaling exacerbates sensitivity to algorithmic imprecision and training instability. Recent advances in optimizers have improved convergence efficiency through momentum orthogonalization, but suffer from two key robustness limitations: dimensional fragility in orthogonalization precision and vulnerability to outlier-induced noise. To address these robustness challenges, we introduce ROOT, a Robust Orthogonalized Optimizer that enhances training stability through dual robustness mechanisms. First, we develop a dimension-robust orthogonalization scheme using adaptive Newton iterations with fine-grained coefficients tailored to specific matrix sizes, ensuring consistent precision across diverse architectural configurations. Second, we introduce an optimization-robust framework via proximal optimization that suppresses outlier noise while preserving meaningful gradient directions. Extensive experiments demonstrate that ROOT achieves significantly improved robustness, with faster convergence and superior final performance compared to both Muon and Adam-based optimizers, particularly in noisy and non-convex scenarios. Our work establishes a new paradigm for developing robust and precise optimizers capable of handling the complexities of modern large-scale model training. The code will be available at https://github.com/huawei-noah/noah-research/tree/master/ROOT.