Technical Summary: Exploiting Weight-Space Symmetries for Approximating Curvature

Problem Statement

Accurate estimation of a loss function's curvature (the Hessian matrix) is fundamental to many machine learning sub-fields, including second-order optimization, Bayesian deep learning, continual learning, and model compression. However, constructing, storing, and inverting the full Hessian is computationally and memory-prohibitive for modern deep networks. While existing approximation techniques like (E)KFAC, Shampoo, and Soap have achieved scalability by assuming specific block-diagonal or Kronecker-factored structures, the field lacks a principled theoretical framework explaining why these specific structures work or how to systematically derive them based on model architecture. Furthermore, previous work has largely overlooked the structural constraints on curvature imposed by the well-known weight-space symmetries inherent in neural networks (e.g., permutation invariance of hidden units).

Methodology

The authors propose a novel framework for approximating the Hessian by analytically averaging over the "group orbit" generated by weight-space symmetries. The core insight relies on the fact that many loss functions are invariant under specific parameter transformations (group actions), which implies that the loss gradient is equivariant under these same transformations.

1. Orbit Averaging and Gradient Equivariance

Given a current parameter set $w$ and gradient $g$, the authors consider the set of all pairs $(w^{(i)}, g^{(i)})$ reachable by applying symmetry transformations from a group $G$. Because the loss is invariant, the gradient at any point in the orbit is related to the original gradient by the group action. By averaging these gradients over the entire orbit, the method aggregates curvature information without requiring multiple gradient computations.

2. Structured Hessian Approximation

The method derives a structured approximation of the Hessian, $\hat{H}$, by solving a secant condition based on the orbit averages of the parameter differences and gradient differences.

• Commutant Algebra: The resulting Hessian approximation belongs to the commutant algebra of the symmetry group. This means the approximation can be represented as a linear superposition of a small number of basis elements (sparse, binary tensors formed by Kronecker products of identity and all-ones matrices).
• Factorization: Instead of storing the full matrix, the method stores only the scalar or matrix "factors" that weigh these basis elements. This allows for tractable estimation, storage, and inversion.
• Simplification: The authors propose a practical simplification where the orbit average of the parameter differences is assumed to be proportional to the identity matrix. This leads to a curvature estimate $\hat{H}_g = S_g^{1/2}$, where $S_g$ is the orbit average of the outer product of the gradient.

3. The "Symo" Optimizer

The framework is instantiated as a new optimizer called Symo (Symmetry-based optimizer). The update rule is $w_{t+1} = w_t - \eta (\hat{H}_g)^{-1} g_t$. The authors provide a just-in-time compiler that takes symbolic specifications of a network's symmetries and generates efficient PyTorch routines to compute the necessary factors.

Key Contributions

1. Novel Hessian Approximation: A method to derive a structured Hessian estimate from a single gradient computation by averaging over the symmetry group orbit. The approximation is provably highly structured and resides in the commutant algebra.
2. Theoretical Unification of Shampoo and Muon: The paper proves that specific choices of symmetry groups within this framework exactly recover the update rules of the Shampoo and Muon optimizers.
   • By restricting the symmetry group to specific signed permutation subgroups (e.g., applying identity symmetry to input/output layers and signed permutations to hidden layers), the derived curvature matrix becomes block-diagonal with Kronecker factors, mathematically equivalent to Shampoo/Muon.
   • This provides a principled explanation for the effectiveness of these optimizers: they implicitly exploit weight-space symmetries, striking a "sweet spot" between approximation accuracy and computational cost.
3. Complexity Control: The framework demonstrates that the complexity of the Hessian approximation is directly governed by the size of the chosen symmetry group. Larger groups yield lower-dimensional commutant algebras (fewer factors, cheaper computation) but larger orbits (potentially less local accuracy). Smaller groups yield more terms but better local fidelity.
4. Empirical Validation:
   • Hessian Quality: Experiments on a small MLP show that the proposed approximations (specifically the simplified $\hat{H}_g$) achieve high cosine similarity with the true Hessian along the negative gradient direction, outperforming orbit-averaged parameter estimates and matching Shampoo's performance in this metric.
   • Optimizer Equivalence: On MNIST autoencoders and nanoGPT (a small language model), the Symo optimizer produces training curves nearly identical to Muon, validating the theoretical equivalence under specific group constraints. Both significantly outperform Adam.

Results

• Accuracy vs. Cost Trade-off: The paper empirically explores different group choices. A "larger group" (allowing permutations of all layers including input/output) still outperforms Adam but is less effective than a "smaller group" (restricting permutations to internal layers and using block-diagonal approximations), which aligns with Shampoo's performance. This suggests that restricting the group to match the architecture's specific invariances is crucial for optimal performance.
• Rank Deficiency: The authors note that the full orbit average matrix $S_w$ can be rank-deficient, which destabilizes inversion. The proposed simplification ($S_w \approx I$) avoids this issue and empirically yields better curvature estimates than the full positive semi-definite solution.
• Scalability: The method successfully scales to Transformer architectures (nanoGPT), demonstrating that the symmetry-aware approach is applicable beyond simple MLPs.

Significance and Claims

The paper claims to provide a unifying theoretical lens for viewing existing second-order optimizers. By framing Shampoo and Muon as specific instances of symmetry-based Hessian approximation, the work moves these methods from heuristic engineering to principled mathematical derivation.

The authors state that their framework:

• Offers a systematic recipe for constructing curvature approximations tailored to specific computational budgets by selecting appropriate symmetry groups.
• Provides a principled route from model symmetries to structured curvature, potentially inspiring new optimizers for different model classes.
• Suggests applications beyond optimization, including uncertainty estimation, continual learning, and model compression, though these are presented as potential future directions rather than demonstrated results.

The paper remains modest regarding limitations, acknowledging that the accuracy of the Taylor expansion around the orbit average is not fully bounded for trained weights, and that the formal guarantee of eigenvector alignment between the preconditioner and the true Hessian remains an open problem. Additionally, while the framework automates the derivation of factors, selecting the optimal group for arbitrary new architectures remains a user-dependent task.