Releases a composable, Optax-native stack that makes high-overhead second-order optimization methods (like K-FAC) practical and swappable.
March 30, 2026
Original Paper
Second-Order, First-Class: A Composable Stack for Curvature-Aware Training
arXiv · 2603.25976
The Takeaway
Second-order methods offer faster convergence but are notoriously difficult to implement and tune. Somax decouples the optimization plan from execution, providing a JIT-compiled framework that allows researchers to drop curvature-aware training into existing JAX pipelines with minimal overhead.
From the abstract
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