CacheMuon: Using Temporal Preconditioning To Approximate Polar Factor

2026-06-15

AI Summary

This work addresses the high computational cost of computing the polar factor of the momentum matrix in Muon optimization, which remains substantial even when using Newton-Schulz iterations to avoid singular value decomposition due to redundant orthogonalization operations. The paper introduces, for the first time, a temporal preconditioning perspective that exploits the temporal continuity of the momentum matrix and its polar factor during training. By caching historical polar factors and controlling their freshness, the proposed cache-driven approximation strategy enables a controllable trade-off between accuracy and efficiency. Supported by a rigorous error analysis framework, the method significantly reduces orthogonalization FLOPs in both language modeling and vision tasks—achieving nearly lossless performance under conservative settings and substantial computational savings with only minor validation loss degradation under aggressive configurations.

Abstract

Muon is an optimizer that computes updates using the polar factor of the momentum matrix and has shown strong empirical performance across a range of training settings. A key component of Muon is the Newton-Schulz iteration used to compute this polar factor. Although this avoids the cost of an exact singular value decomposition, it remains expensive in practice because it is applied at every optimization step. At the same time, the momentum matrix changes smoothly over training, suggesting strong temporal correlation in the corresponding polar factors. In this paper, we exploit this structure and propose CacheMuon, a temporal preconditioning method that reuses information from previous optimization steps to approximate the polar factor at the current step. This reduces redundant orthogonalization computation across iterations. We analyze CacheMuon as an inexact Muon update, with error controlled by fresh-solver error and cache staleness. Empirically, CacheMuon provides a controllable quality-efficiency frontier: conservative thresholds closely match fresh Muon on language-model and vision training while reducing orthogonalization FLOPs, whereas more aggressive thresholds yield larger arithmetic savings at the cost of modest validation-quality degradation.

Problem

• optimization
• polar factor
• temporal correlation
• orthogonalization
• momentum matrix

Innovation

• temporal preconditioning
• polar factor approximation
• orthogonalization reuse
• inexact optimization
• momentum matrix