mlsanityData to 7 Oct 2026

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

arXiv 2603.25976 ↗First posted 26 Mar 2026cs.LGcited by 0

Not on the map. Why?

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.

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

Builds on

26 of 38 references.

References · 26

  1. 2015Optimizing Neural Networks with Kronecker-factored Approximate Curvaturecited by 776
  2. 2018Shampoo: Preconditioned Stochastic Tensor Optimizationcited by 413
  3. 2019PyHessian: Neural Networks Through the Lens of the Hessiancited by 291
  4. 2023Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-trainingcited by 260
  5. 2020ADAHESSIAN: An Adaptive Second Order Optimizer for Machine Learningcited by 201
  6. 2021Efficient and Modular Implicit Differentiationcited by 199
  7. 2024SOAP: Improving and Stabilizing Shampoo using Adamcited by 179
  8. 2020Descending through a Crowded Valley - Benchmarking Deep Learning Optimizerscited by 118
  9. 2019BackPACK: Packing more into backpropcited by 77
  10. 2015Newton Sketch: A Linear-time Optimization Algorithm with Linear-Quadratic Convergencecited by 72
  11. 2014A Stochastic Quasi-Newton Method for Large-Scale Optimizationcited by 68
  12. 2020Hutch++: Optimal Stochastic Trace Estimationcited by 62
14 more
  1. 2018Large-Scale Distributed Second-Order Optimization Using Kronecker-Factored Approximate Curvature for Deep Convolutional Neural Networkscited by 37
  2. 2019DeepOBS: A Deep Learning Optimizer Benchmark Suitecited by 37
  3. 2019Efficient Subsampled Gauss-Newton and Natural Gradient Methods for Training Neural Networkscited by 32
  4. 2016A Multi-Batch L-BFGS Method for Machine Learningcited by 30
  5. 2022Benchopt: Reproducible, efficient and collaborative optimization benchmarkscited by 27
  6. 2020On the Promise of the Stochastic Generalized Gauss-Newton Method for Training DNNscited by 18
  7. 2024Optimistix: modular optimisation in JAX and Equinoxcited by 17
  8. 2020Scalable and Practical Natural Gradient for Large-Scale Deep Learningcited by 15
  9. 2018The Full Spectrum of Deepnet Hessians at Scale: Dynamics with SGD Training and Sample Sizecited by 12
  10. 2024Exact Gauss-Newton Optimization for Training Deep Neural Networkscited by 10
  11. 2013Training Neural Networks with Stochastic Hessian-Free Optimizationcited by 9
  12. 2021Cockpit: A Practical Debugging Tool for the Training of Deep Neural Networkscited by 7
  13. 2019Stochastic quasi-Newton with line-search regularizationcited by 6
  14. 2024Incremental Gauss-Newton Descent for Machine Learningcited by 1

Built on by

Not cited yet.