Isotropic Curvature Model for Understanding Deep Learning Optimization: Is Gradient Orthogonalization Optimal? (2511.00674v1)

Abstract: In this paper, we introduce a model for analyzing deep learning optimization over a single iteration by leveraging the matrix structure of the weights. We derive the model by assuming isotropy of curvature, including the second-order Hessian and higher-order terms, of the loss function across all perturbation directions; hence, we call it the isotropic curvature model. This model is a convex optimization program amenable to analysis, which allows us to understand how an update on the weights in the form of a matrix relates to the change in the total loss function. As an application, we use the isotropic curvature model to analyze the recently introduced Muon optimizer and other matrix-gradient methods for training LLMs. First, we show that under a general growth condition on the curvature, the optimal update matrix is obtained by making the spectrum of the original gradient matrix more homogeneous -- that is, making its singular values closer in ratio -- which in particular improves the conditioning of the update matrix. Next, we show that the orthogonalized gradient becomes optimal for the isotropic curvature model when the curvature exhibits a phase transition in growth. Taken together, these results suggest that the gradient orthogonalization employed in Muon and other related methods is directionally correct but may not be strictly optimal. Finally, we discuss future research on how to leverage the isotropic curvature model for designing new optimization methods for training deep learning and LLMs.

• The paper proposes an isotropic curvature model that leverages matrix-form gradients to explore single-step optimization dynamics.
• It demonstrates that optimal updates require spectrum homogenization rather than full gradient orthogonalization, validated with empirical results on GPT-2 layers.
• The study suggests that auto-preconditioning based on curvature estimation can enhance optimizer design for deep neural networks.

Isotropic Curvature Model for Deep Learning Optimization: Analysis of Gradient Orthogonalization

Introduction and Motivation

The paper introduces the isotropic curvature model as a framework for analyzing single-step optimization in deep learning, with a focus on the matrix structure of weights. The motivation arises from the empirical success of the Muon optimizer, which updates weights using the orthogonalized (unitary) part of the gradient matrix, in contrast to Adam and other optimizers that treat weights as vectors. The central question addressed is whether gradient orthogonalization, as implemented in Muon and related methods, is theoretically optimal for deep learning optimization.

Isotropic Curvature Model: Formulation

The isotropic curvature model is derived by Taylor expanding the per-sample loss function with respect to a matrix perturbation and then imposing two key isotropy assumptions:

1. Curvature Isotropy: The higher-order curvature (Hessian and beyond) is assumed to be isotropic, i.e., its spectral properties depend only on the norm of the perturbation, not its direction.
2. Input Isotropy: The input vectors to the weight matrix are assumed to be uniformly distributed on the sphere, reflecting the lack of preferred directions in large-scale models and diverse data.

Under these assumptions, the one-step optimization problem for the update matrix $Q$ (with $G$ the gradient) becomes:

$$\min_Q -\mathrm{Tr}(Q G^\top) + \mathbb{E}_{\zeta \sim \text{sphere}} H(\|Q\zeta\|),$$

where $H$ is a univariate, increasing curvature function capturing higher-order loss terms.

Empirical Characterization of Curvature Growth

The model's validity hinges on the empirical growth behavior of $H(r)$. The paper provides numerical evidence from GPT-2 (small) that $H(r)/r^2$ is nearly constant for small $r$ (quadratic regime), but grows as a power law $r^\alpha$ with $\alpha \approx 0.2$ for larger $r$, indicating super-quadratic growth of higher-order terms.

Figure 1: Numerical approximation of $H(r)/r^2$ on GPT-2 (small), showing super-quadratic growth for both fully connected and attention layers.

Theoretical Analysis: Singular Spaces and Spectrum Homogenization

Alignment of Singular Spaces

The isotropic curvature model is rotationally invariant: the optimal update $Q^*$ can always be chosen to have its left and right singular spaces aligned with those of the gradient $G$. This justifies the design of Muon and similar optimizers, which preserve the singular spaces of the gradient and only modify its spectrum.

Spectrum Homogenization

Under the assumption that $H(\sqrt{x})$ is convex (i.e., $H$ grows super-quadratically), the optimal update $Q^*$ is shown to have singular values that are more homogeneous than those of $G$, while preserving their ordering. Formally, for singular values $\sigma_i$ of $G$ and $\sigma_i^*$ of $Q^*$,

$$\frac{\max\{\sigma_i^*, \sigma_j^*\}}{\min\{\sigma_i^*, \sigma_j^*\}} \leq \frac{\max\{\sigma_i, \sigma_j\}}{\min\{\sigma_i, \sigma_j\}},$$

for all $i, j$. This spectrum homogenization improves the conditioning of the update and is a key theoretical insight into why matrix-gradient methods outperform vectorized ones.

Orthogonalization as an Extreme Case

If $H$ exhibits a sharp "kink" (i.e., its derivative jumps abruptly at some radius), the optimal update is exactly the orthogonalized gradient: $Q^* = c \cdot \mathrm{msgn}(G)$, where $\mathrm{msgn}(G)$ is the unitary part of the polar decomposition of $G$. This provides a theoretical justification for Muon's update rule in the limiting case of extremely rapid curvature growth. However, the paper emphasizes that in practical settings, the curvature function does not exhibit such a discontinuity, and thus full orthogonalization is not strictly optimal.

Implications for Optimizer Design

The analysis leads to several important implications:

• Matrix Structure Matters: Treating weights and gradients as matrices, rather than vectors, is not only conceptually sound but also theoretically justified for large-scale models.
• Gradient Orthogonalization is Directionally Correct but Not Strictly Optimal: While Muon's orthogonalization is justified in the extreme case, the optimal update generally requires only partial homogenization of the spectrum, depending on the empirical curvature function.
• Auto-Preconditioning: The isotropic curvature model suggests a new class of optimizers that automatically precondition the gradient based on the estimated curvature function, without explicit Hessian computation.
• Practical Design: For practical optimizers, the model suggests first empirically estimating the curvature function $H$ for each layer, then solving the convex program to determine the optimal spectrum transformation, ideally using monotone transformations to preserve singular value ordering.

Limitations and Future Directions

The isotropic curvature model is currently phenomenological and relies on strong isotropy assumptions. Future work should aim to:

• Rigorously justify the isotropy assumptions in the context of deep learning architectures and data distributions.
• Extend the model to account for mini-batch stochasticity, momentum, and multi-step optimization dynamics.
• Develop efficient GPU-friendly algorithms for spectrum homogenization that preserve monotonicity and are scalable to large models.
• Investigate the connection between curvature growth, edge-of-stability phenomena, and the empirical performance of different optimizers.

Conclusion

The isotropic curvature model provides a tractable and theoretically grounded framework for understanding the optimization dynamics of matrix-gradient methods in deep learning. It clarifies the role of gradient orthogonalization, showing that while it is optimal in an extreme curvature regime, the general optimal update is a spectrum-homogenized version of the gradient. This insight opens new avenues for the principled design of optimizers that leverage the matrix structure of weights and the empirical properties of curvature in large-scale neural networks.