Technical Summary: CacheMuon

Problem Statement

Muon is a geometry-aware optimizer designed for matrix-valued hidden-layer parameters in deep learning. Unlike coordinate-wise optimizers (e.g., Adam), Muon replaces standard momentum updates with an orthogonalized direction derived from the polar factor of the momentum matrix. While Muon has demonstrated strong empirical performance and scalability to large language models, its practical efficiency is bottlenecked by the orthogonalization step.

Current implementations approximate the polar factor using Newton–Schulz iterations (specifically Gram Newton–Schulz) to avoid the high cost of exact Singular Value Decomposition (SVD). However, this approximation still requires multiple matrix multiplications at every optimization step. Consequently, a significant portion of Muon's computational cost is spent repeatedly solving closely related matrix orthogonalization problems. While recent work has optimized the efficiency of a single orthogonalization call (e.g., PolarExpress, Gram Newton–Schulz), these methods still recompute the orthogonalizing transform from scratch at every step, ignoring the temporal correlation inherent in the training process.

Methodology: CacheMuon

The authors propose CacheMuon, a temporal preconditioning method that exploits the smooth evolution of the momentum matrix during training. The core hypothesis is that if the momentum matrix changes gradually, the orthogonalization transform computed at a previous step remains informative for the current step.

Core Mechanism

CacheMuon modifies the orthogonalization routine of Muon by introducing a caching mechanism:

1. Fresh Solve (Refresh Step): At specific "refresh" steps (determined by a trigger), the method runs a full Fresh Gram Newton–Schulz (FreshGNS) solver. This produces the approximate polar factor $\tilde{D}_t$ and, crucially, a cached left transform matrix $Q_t$.
2. Cached Reuse (Intermediate Steps): At non-refresh steps, instead of running a full solver, the method reuses the cached transform $Q_{a_t}$ (where $a_t$ is the index of the most recent refresh) to approximate the polar factor: $\hat{D}_t = Q_{a_t} X_t$, where $X_t$ is the normalized momentum matrix.
3. Residual-Based Trigger: To determine when the cached transform is no longer accurate, the authors introduce a lightweight, residual-based reuse trigger. Before reusing a cache, the algorithm computes the orthogonality residual of the candidate update:
   $$\text{err}_t = \frac{\| \hat{D}^{\text{cand}}_t (\hat{D}^{\text{cand}}_t)^\top - I \|_F}{\|I\|_F}$$
   If $\text{err}_t \leq \gamma$ (a prescribed threshold), the cached transform is accepted. Otherwise, a fresh solve is triggered to update the cache.

Theoretical Analysis

The paper analyzes CacheMuon as an inexact Muon update. The deviation between the cached direction $\hat{D}_t$ and the exact Muon direction $D_t$ is bounded by an additive error term $\delta_t$:
$$\| \hat{D}_t - D_t \|_{\text{sp}} \leq \delta_t = L_Q \tau_t + \varepsilon_{\text{ref}}$$
where:

• $\varepsilon_{\text{ref}}$ is the error of the fresh solver (due to finite Newton–Schulz iterations).
• $\tau_t$ is the "staleness" of the cached Gram matrix relative to the current one.
• $L_Q$ is a Lipschitz constant derived from the polynomial map induced by the solver.

Under standard assumptions (smoothness, bounded variance, bounded staleness), the authors prove a stochastic convergence guarantee for nonconvex optimization. The convergence rate depends on the aggregate inexactness parameter $\bar{\delta} = L_Q \bar{\tau} + \varepsilon_{\text{ref}}$. As long as $\bar{\delta} < 1$, the method converges, with the rate scaling as $O((1-\bar{\delta})^{-1} N^{-1/4})$.

Key Contributions

1. CacheMuon Algorithm: A variant of Muon that amortizes the cost of polar-factor computation by reusing cached orthogonalization transforms across optimization steps, leveraging temporal smoothness.
2. Convergence Theory: A theoretical framework treating CacheMuon as an inexact optimizer, providing stochastic convergence guarantees controlled by fresh-solver error and cache staleness.
3. Residual-Based Trigger: A practical, low-cost heuristic that dynamically decides when to refresh the cache based on the orthogonality of the candidate update, serving as a proxy for update fidelity.
4. Empirical Validation: Demonstration of a controllable quality–efficiency frontier across language model (GPT-2) and vision (ResNet-18) workloads.

Experimental Results

The authors evaluated CacheMuon on GPT-2 Large/Small (OpenWebText) and ResNet-18 (CIFAR-10).

• Quality–Efficiency Tradeoff: By varying the residual threshold $\gamma$, the authors demonstrate a smooth frontier.
  • Conservative Reuse ($\gamma=2$): Closely tracks the performance of fresh Muon (GramMuon and PolarExpMuon) while reducing orthogonalization FLOPs by 13.6% (vs. GramMuon) and 30.9% (vs. PolarExpMuon) on GPT-2 Large.
  • Aggressive Reuse ($\gamma=15$): Achieves up to 65.0% reduction in orthogonalization FLOPs (vs. GramMuon) with only a modest degradation in validation loss (e.g., 3.096 vs. 3.037 for GPT-2 Large).
• Diagnostic Insights: Experiments show that the residual trigger effectively correlates with the alignment between cached and fresh updates. Early in training, residuals are high, forcing frequent refreshes; as training stabilizes, residuals drop, allowing for higher cache hit rates and greater efficiency.
• Scalability: The method maintains competitive performance across different model scales (GPT-2 Small vs. Large) and modalities (Vision vs. Language).

Significance and Claims

The paper claims that CacheMuon offers a practical approach to reducing the computational overhead of Muon without altering its underlying geometric update rule. By treating orthogonalization as a temporal problem rather than a purely spatial one, the method reduces redundant computation.

The authors emphasize that while the method significantly reduces orthogonalization FLOPs, the realized wall-clock speedup in their current high-level implementation is modest due to system-level overheads (e.g., host-device synchronization for the residual check). They posit that a more optimized low-level implementation (e.g., fusing probes into device kernels) could translate a larger fraction of these arithmetic savings into actual runtime gains.

Ultimately, the work suggests that temporal reuse is a promising direction for optimizing geometry-aware optimizers, providing a tunable knob to balance computational efficiency with training quality.