Source-linked AI summary
ROOT: Robust Orthogonalized Optimizer for Neural Network Training
Wei He, Kai Han, Hang Zhou, Hanting Chen, Zhicheng Liu, Xinghao Chen, Yunhe Wang
TL;DR
Orthogonalization-based optimizers face imprecise orthogonalization across matrix dimensions and sensitivity to gradient outliers. ROOT combines dimension-specific adaptive Newton iterations with proximal outlier suppression. Experiments report faster convergence and superior performance, particularly in noisy and non-convex scenarios.
Problem
Orthogonalization-based optimizers have limited robustness to imprecise orthogonalization across matrix dimensions and outlier-induced gradient noise.
Source: §1 · pp. 1–2Method
ROOT combines dimension-specific adaptive Newton iterations for orthogonalization with proximal optimization that suppresses outliers through soft-thresholding.
Source: §1 · p. 2Results
ROOT achieves superior performance and faster convergence than state-of-the-art optimizers in LLM pre-training and finetuning, particularly in noisy and non-convex scenarios.
Source: §1 · p. 2Takeaways & Limitations
In the evaluated benchmark tasks, ROOT’s enhancements are reported to yield competitive or superior zero-shot performance and higher-quality final language models.
Source: §4.4 · p. 7
Abstract
from arXiv · showhide
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.
1. Introduction
ROOT addresses two robustness limitations in orthogonalization-based optimizers: imprecise orthogonalization across matrix dimensions and sensitivity to gradient outliers. It combines dimension-specific adaptive Newton iterations with proximal outlier suppression, and experiments report faster convergence and superior performance in noisy and non-convex scenarios.
1. Introduction
- Orthogonalization-based optimizers face imprecision across matrix dimensions and vulnerability to gradient outliers.Outlier-induced noise can corrupt update directions and destabilize training.Source: §1 · pp. 1–2
- ROOT uses adaptive Newton-Schulz iterations with fine-grained, dimension-specific coefficients to improve orthogonalization accuracy across network architectures.The introduction presents this as addressing structural uncertainty from varying matrix dimensions.Source: §1 · p. 2
- ROOT adds proximal optimization with soft-thresholding to suppress outlier-induced gradient noise while supporting stable convergence.The introduction reports theoretical guarantees for outlier suppression and stable convergence.Source: §1 · p. 2
- ROOT’s experiments report superior performance and faster convergence than state-of-the-art optimizers on LLM pre-training and finetuning, particularly in noisy and non-convex scenarios.The framework is also reported to maintain computational efficiency.Source: §1 · p. 2
2. Related Work
Related work contrasts coordinate-wise optimizers with matrix-aware approaches and reviews Muon variants that target efficiency, adaptivity, and precision.
2.1. Matrix-Aware Optimizers
- Coordinate-wise methods overlook weight-matrix structure, while K-FAC and Shampoo capture parameter geometry at often prohibitive computational and memory cost; Muon instead orthogonalizes momentum with O(N) complexity.Muon’s Newton-Schulz iteration is described as regulating update geometry without explicit curvature approximation.Source: §2.1 · p. 2
2.2. Recent Advances in Muon Variants
- Muon variants target distributed efficiency and update cost through amortized power iteration, randomized SVD, and randomized layer subsampling.Dion, LiMuon, and Drop-Muon are the examples named for these approaches.Source: §2.2 · p. 2
- Other variants reintegrate adaptivity or improve orthogonalization precision, while ROOT investigates precision across varying matrix dimensions.Examples include AdaGO, AdaMuon, and CANS.Source: §2.2 · p. 2
3. Approach
The approach builds on momentum orthogonalization and addresses shape-sensitive Newton–Schulz precision and gradient outliers. It combines dimension-specific adaptive coefficients with proximal soft-thresholding.
3.1. Preliminaries
- Muon applies Newton–Schulz iteration to momentum to approximate an orthogonal update, but fixed coefficients yield inconsistent orthogonalization quality across matrix dimensions.The section frames this as a robustness limitation of the one-size-fits-all scheme.Source: §3.2.1 · p. 3, §3.1 · p. 3
3.2.1. ENHANCING ROBUSTNESS AGAINST MATRIX DIMENSION VARIATIONS
- At m = 2048, n = 2048, adaptive coefficients reduce orthogonal-approximation MSE from 0.0499 to 0.0352.The table caption identifies orthogonalization error as the comparison, while the prose supplies this square-matrix example.Source: §3.2.1 · p. 3, Table 1 · p. 3
- AdaNewton learns separate coefficients for each matrix size, targeting consistent orthogonalization across layers with different dimensions.The coefficients can be optimized jointly with model parameters to adapt to each layer type’s spectral properties.Source: §3.2.1 · pp. 3–4
3.2.2. THEORETICAL GUARANTEES FOR ROBUST CONVERGENCE
- Because each shape’s characteristic singular-value set is a subset of the global set, adaptive coefficients can achieve strictly better approximation for non-square matrices under the stated conditions.The smaller maximum singular-value error also yields a tighter bound on the Frobenius-norm objective.Source: §3.2.2 · pp. 4–5
3.3.1. ENHANCING ROBUSTNESS AGAINST OUTLIER-INDUCED GRADIENT NOISE
- Extreme gradient outliers can distort Muon’s normalization and be amplified by the polynomial iteration, preserving erratic update directions.The passage notes that this may destabilize training, including through exploding attention logits in Transformers.Source: §3.3.1 · p. 5
3.3.2. SOFT-THRESHOLDING FOR OUTLIER SUPPRESSION
- ROOT uses proximal optimization with soft-thresholding to separate outliers from the base component, then orthogonalizes only the base and discards the outliers.The section describes soft-thresholding as smooth shrinkage that preserves relative gradient-magnitude ordering while dampening extreme values.Source: §3.3.2 · pp. 5–6, §3.3.1 · p. 5
4. Experiments
ROOT is evaluated on LLM training, academic benchmarks, and vision classification, with ablations testing its orthogonalization and outlier-suppression components. The reported results show competitive or superior benchmark performance and improved training loss relative to Muon.
4.1. Implementation Details
- The 1B Transformer experiments use FineWeb-Edu, with 10B tokens for ablations, 100B for the main run, and zero-shot academic evaluation.Training uses Ascend NPUs; both runs train for one epoch.Source: §4.1 · p. 6
4.3. LLM Pre-training Analysis
- ROOT reaches a final training loss of 2.5407, 0.01 below the Muon baseline, with full ROOT achieving the lowest loss.Both ROOT-SoftThresh and full ROOT have consistently lower training loss than Muon over the 10B-token trajectory.Source: §4.3 · p. 6, Figure 3 · p. 7
4.2. Validation on Real Training Dynamics
- With five iterations, ROOT maintains lower orthogonalization relative error than Muon and Classic Newton-Schulz across optimized parameter shapes.The comparison measures error relative to ground-truth SVD.Source: Figure 2 · p. 7
4.4. Benchmark Evaluation Results
- Across the evaluated zero-shot academic benchmarks, ROOT outperforms AdamW and Muon after 100B-token training.Table 2 summarizes benchmark performance for the trained 1B models.Source: §4.4 · p. 7, Table 2 · p. 8
4.5. Ablation Studies
- For outlier-threshold ablations, p = 0.90 balances noise suppression and signal preservation, unlike the under-suppressing 0.99 and overly aggressive 0.85 settings.The threshold is percentile-based and adapts to the evolving gradient scale.Source: §4.5 · p. 7, Figure 4 · p. 7
- Mixed calibration at 1:3 is adopted for balancing convergence speed and robustness; the 1:1 real-sample ratio has lower loss here but can spike at larger scale or with SoftThresh.Random calibration offers limited convergence acceleration.Source: §4.5 · pp. 7–8, Figure 5 · p. 8
4.6. Generalization to Vision Tasks
- On CIFAR-10, ROOT consistently outperforms the baseline in a from-scratch 6.3M-parameter Vision Transformer experiment.The reported improvement is most significant at quantile p = 0.85, where the baseline achieves 84.67% accuracy.Source: §4.6 · p. 8, Table 3 · p. 8
5. Conclusion
ROOT addresses orthogonalization precision across matrix dimensions and optimization robustness to outliers through adaptive Newton iterations and proximal outlier suppression. The conclusion reports superior performance in noisy and non-convex scenarios.
5. Conclusion
- ROOT combines dimension-robust adaptive Newton iterations with proximal outlier suppression to address two limitations in large-scale language model training.The conclusion describes the method as a robust orthogonalized optimizer.Source: §5 · p. 8
- The authors report superior performance in challenging noisy and non-convex scenarios, alongside theoretical guarantees and practical benefits.Source: §5 · p. 8
- The conclusion points to robust optimization frameworks as a direction for handling the increasing complexity and scale of future language models.Source: §5 · p. 8