Exploiting weight-space symmetries for approximating curvature

Many machine learning techniques rely on approximating a loss function's curvature, but this is notoriously hard to do at the scale of modern deep networks. Surprisingly, no previous work has exploited the curvature constraints that arise from well known weight-space symmetries in loss landscapes. By analytically averaging over group actions that leave the loss invariant, we construct structured Hessian approximations from single gradients that can be tractably estimated, stored, and inverted. The choice of user-specified symmetry group directly governs the trade-off between approximation accuracy and computational cost. Moreover, our framework provides a unifying theoretical lens for viewing existing methods; in particular, a specific choice of symmetry group recovers Shampoo/Muon-like curvature estimates. We validate our method on a range of network architectures, and deploy it to second-order optimization benchmarks, including a small language model. Our curvature estimation framework might find applications in other machine learning problems such as uncertainty estimation, continual learning, compression/pruning, training data attribution, and more.

Training neural networks faster requires knowing not just the direction to update parameters (the gradient), but also the curvature of the loss landscape. Computing this curvature exactly is infeasible for large models. We observe that neural networks have built-in symmetries. For instance, swapping two neurons and their connections leaves the network unchanged. From a single gradient, these symmetries let us analytically infer gradients at many equivalent configurations, and we show how to distill this information into a curvature estimate. Larger symmetry groups yield cheaper but coarser estimates; smaller groups yield richer ones. We prove that a specific middle-ground choice produces updates mathematically equivalent to Shampoo and Muon, two popular optimizers, revealing a previously unknown connection between these methods and the network's architectural symmetries.