AMO: Adaptive Muon Orthogonalization

2026-05-17

AI Summary

This work addresses a key limitation of existing Muon optimizers, which apply a uniform orthogonalization schedule across all parameter matrices despite significant differences in orthogonalization difficulty among operators, leading to inconsistent orthogonality quality. To resolve this, the authors propose an adaptive orthogonalization scheduling method that, during early training, analyzes geometric properties—such as condition numbers and singular value distributions—of individual weight matrices. Leveraging dynamic heterogeneity arising from operator type, network depth, and training phase, the approach differentially allocates computational budgets for Newton–Schulz iterations. This strategy is the first to reveal structural disparities in matrix geometry and enables resource-efficient allocation, yielding average downstream performance gains of 0.76 and 0.51 on Llama3.1-1.4B and Qwen3-1.7B, respectively, substantially outperforming current state-of-the-art baselines.

Abstract

Muon has recently emerged as a competitive alternative to AdamW for large-scale pre-training, with orthogonalization via Newton-Schulz (NS) iterations as its core operation. Existing Muon variants apply a uniform NS schedule to all parameter matrices, overlooking possible differences in orthogonalization difficulty and its impact on performance. Through a systematic empirical study, we show that this per-matrix heterogeneity is pervasive and largely determined by matrix geometry, which evolves dynamically across operator types, training stages, and network depths. As a result, uniform NS schedules can lead to uneven orthogonalization quality across the model. Motivated by these findings, we propose Adaptive Muon Orthogonalization (AMO), an observe-then-commit method that measures weight geometry by operator type early in training and then uses these signals to allocate the NS budget for the remainder of training. AMO delivers consistent improvements over uniform-schedule Muon across standard, prolonged, and continual pre-training, surpassing the strongest baseline by +0.76 on Llama3.1-1.4B and +0.51 on Qwen3-1.7B in average downstream performance of 12 evaluation tasks.

Problem

• Muon
• orthogonalization
• Newton-Schulz iterations
• parameter matrices
• training optimization

Innovation

• Adaptive Muon Orthogonalization
• Newton-Schulz iteration
• matrix geometry
• heterogeneous optimization
• large language model pre-training