ROOT: Robust Orthogonalized Optimizer for Neural Network Training

TL;DR

ROOT tackles two core issues in large-scale neural network optimization: (i) dimensional fragility in orthogonalization for momentum-based methods, and (ii) sensitivity to gradient outliers during training. It combines a dimension-aware adaptive Newton-Schulz orthogonalization with per-dimension coefficients and a proximal outlier suppression mechanism via soft-thresholding to stabilize updates. Theoretical guarantees accompany empirical results showing faster convergence and improved final performance versus Muon and Adam-based optimizers, particularly under noisy or non-convex conditions. This framework offers a robust, scalable pathway for reliable optimization in modern large-scale language models and beyond.

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.

Figures (5)

• Figure 1: Analysis of gradient distribution revealing outlier characteristics. (Left) Histogram with Gaussian reference shows long-tailed distribution. (Right) Q-Q plot quantifies deviation from normality, where points deviating from the diagonal indicate outliers. These outliers can disproportionately influence the optimization process.
• Figure 2: Orthogonalization precision relative to ground-truth SVD. The plot tracks the Relative Error averaged over all optimized parameters (Attention QKV/O and MLP Up/Down projections). Under a fixed 5 iterations, ROOT maintains lower approximation error compared to the Muon baseline and Classic Newton-Schulz. This indicates that shape-specific coefficients provide superior fidelity across varying matrix dimensions.
• Figure 3: Training loss curves for 10B tokens. ROOT variants demonstrate faster convergence and lower final loss compared to Muon baseline, with full ROOT achieving the best performance.
• Figure 4: Ablation on the quantile hyperparameter $p$. The curve with $p=0.90$ demonstrates the optimal equilibrium between suppressing gradient noise and preserving informative gradient signals.
• Figure 5: Ablation on the data composition for coefficient calibration. While a high ratio of real samples (Ratio 1:1) achieves lower loss here, it induces instability (loss spikes) in larger-scale experiments or when combined with ROOT (SoftThresh). The Mixed (1:3) strategy provides the optimal balance between convergence speed and robustness.