# Second-Order, First-Class: A Composable Stack for Curvature-Aware Training

arXiv: https://arxiv.org/abs/2603.25976 · HTML: https://mlsanity.com/paper/2603.25976/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 26 of its 38 references that are in the map.
- First posted on arXiv 26 Mar 2026.

## Builds on (most cited first, 26 of 26)

- [Optimizing Neural Networks with Kronecker-factored Approximate Curvature](https://mlsanity.com/paper/1503.05671/) (2015)
- [Shampoo: Preconditioned Stochastic Tensor Optimization](https://mlsanity.com/paper/1802.09568/) (2018)
- [PyHessian: Neural Networks Through the Lens of the Hessian](https://mlsanity.com/paper/1912.07145/) (2019)
- [Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training](https://mlsanity.com/paper/2305.14342/) (2023)
- [Efficient and Modular Implicit Differentiation](https://mlsanity.com/paper/2105.15183/) (2021)
- [ADAHESSIAN: An Adaptive Second Order Optimizer for Machine Learning](https://mlsanity.com/paper/2006.00719/) (2020)
- [SOAP: Improving and Stabilizing Shampoo using Adam](https://mlsanity.com/paper/2409.11321/) (2024)
- [Descending through a Crowded Valley - Benchmarking Deep Learning Optimizers](https://mlsanity.com/paper/2007.01547/) (2020)
- [BackPACK: Packing more into backprop](https://mlsanity.com/paper/1912.10985/) (2019)
- [Newton Sketch: A Linear-time Optimization Algorithm with Linear-Quadratic Convergence](https://mlsanity.com/paper/1505.02250/) (2015)
- [A Stochastic Quasi-Newton Method for Large-Scale Optimization](https://mlsanity.com/paper/1401.7020/) (2014)
- [Hutch++: Optimal Stochastic Trace Estimation](https://mlsanity.com/paper/2010.09649/) (2020)
- [Large-Scale Distributed Second-Order Optimization Using Kronecker-Factored Approximate Curvature for Deep Convolutional Neural Networks](https://mlsanity.com/paper/1811.12019/) (2018)
- [DeepOBS: A Deep Learning Optimizer Benchmark Suite](https://mlsanity.com/paper/1903.05499/) (2019)
- [Efficient Subsampled Gauss-Newton and Natural Gradient Methods for Training Neural Networks](https://mlsanity.com/paper/1906.02353/) (2019)
- [A Multi-Batch L-BFGS Method for Machine Learning](https://mlsanity.com/paper/1605.06049/) (2016)
- [Benchopt: Reproducible, efficient and collaborative optimization benchmarks](https://mlsanity.com/paper/2206.13424/) (2022)
- [On the Promise of the Stochastic Generalized Gauss-Newton Method for Training DNNs](https://mlsanity.com/paper/2006.02409/) (2020)
- [Optimistix: modular optimisation in JAX and Equinox](https://mlsanity.com/paper/2402.09983/) (2024)
- [Scalable and Practical Natural Gradient for Large-Scale Deep Learning](https://mlsanity.com/paper/2002.06015/) (2020)
- [The Full Spectrum of Deepnet Hessians at Scale: Dynamics with SGD Training and Sample Size](https://mlsanity.com/paper/1811.07062/) (2018)
- [Exact Gauss-Newton Optimization for Training Deep Neural Networks](https://mlsanity.com/paper/2405.14402/) (2024)
- [Training Neural Networks with Stochastic Hessian-Free Optimization](https://mlsanity.com/paper/1301.3641/) (2013)
- [Cockpit: A Practical Debugging Tool for the Training of Deep Neural Networks](https://mlsanity.com/paper/2102.06604/) (2021)
- [Stochastic quasi-Newton with line-search regularization](https://mlsanity.com/paper/1909.01238/) (2019)
- [Incremental Gauss-Newton Descent for Machine Learning](https://mlsanity.com/paper/2408.05560/) (2024)

## Abstract (arXiv)

Second-order methods promise improved stability and faster convergence, yet they remain underused due to implementation overhead, tuning brittleness, and the lack of composable APIs. We introduce Somax, a composable Optax-native stack that treats curvature-aware training as a single JIT-compiled step governed by a static plan. Somax exposes first-class modules -- curvature operators, estimators, linear solvers, preconditioners, and damping policies -- behind a single step interface and composes with Optax by applying standard gradient transformations (e.g., momentum, weight decay, schedules) to the computed direction. This design makes typically hidden choices explicit and swappable. Somax separates planning from execution: it derives a static plan (including cadences) from module requirements, then runs the step through a specialized execution path that reuses intermediate results across modules. We report system-oriented ablations showing that (i) composition choices materially affect scaling behavior and time-to-accuracy, and (ii) planning reduces per-step overhead relative to unplanned composition with redundant recomputation.