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

Mikalai Korbit; Mario Zanon

cs.LGMar 26, 2026

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

Authors: Mikalai Korbit, Mario Zanon

Organizations: IMT School for Advanced Studies Lucca Lucca, Italy

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 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.

Figures & tables

Symbol	Meaning
$\mathcal{D},\,n$	Dataset and size
$\mathcal{B}_{t},\,b$	Mini-batch at step $t$ and size
$\bm{w}\in\mathbb{R}^{d}$	Model parameters (weights)
$\mathcal{L}_{\mathcal{B}_{t}}(\bm{w})$	Mini-batch loss at step $t$
$\bm{g}_{t}$	Mini-batch gradient $\nabla\mathcal{L}_{\mathcal{B}_{t}}(\bm{w}_{t})$
$\bm{H}_{t}$	Curvature operator (exact or approximate)

Table 1: Notation (excerpt).

Method	Curv.	Lane	Solver	Estimation / preconditioning
AdaHessian [ 40 ]	Hessian	diag.	explicit	Hutchinson diag. estimator + diag. EMA preconditioner
Sophia-G [ 15 ]	GGN (CE)	diag.	explicit	GNB diag. estimator + diag. EMA preconditioner + clipping
Sophia-H [ 15 ]	Hessian	diag.	explicit	Hutchinson diag. estimator + diag. EMA preconditioner + clipping
Newton-CG [ 19 ]	Hessian	param.	CG/PCG (HVP)	Optional diag. precond.; damping
SGN [ 8 ]	GGN (MSE/CE)	param.	CG/PCG (GGN matvec)	Optional diag. precond.; damping
EGN [ 13 ]	GGN (MSE/CE)	row	row-Cholesky / row-CG	Row-space solve with internal damping conversion

Table 2: Selected second-order methods for large-scale stochastic optimization.

Panel of Figure 1 — Figure 1: Somax system architecture. The assembler and planner resolve a user configuration into a static execution plan. This plan records one lane-specific execution path (diagonal, parameter-space, or row-space), the step-metric outputs, and cadence gates; the resulting step is then JIT-compiled. The executor manages the full update lifecycle – from linearization to Optax application and optional post-update control signals – ensuring that post-update metrics are consistent with the applied update.

Setting ( rho_every_k )	Median (ms)	p90 (ms)	Overhead vs. off
off ( -1 )	0.84	0.87	0.0%
every 10 steps ( 10 )	0.95	1.16	12.4%
every 5 steps ( 5 )	0.96	1.33	13.8%
every 2 steps ( 2 )	1.30	1.40	57.2%
every step ( 1 )	1.33	1.34	53.7%

Table 3: Microbench: steady-state overhead of cadence-gated rho probes. Overhead is reported relative to probes disabled.

Figure 2: Steady-state step time for matched parameter-space and row-space configurations under a shared training interface. Left: increasing model size at fixed batch size. Right: increasing batch size at fixed model size.

Solver	Precond.	Damping	CG budget	Time to 85% (s)	Final acc (%)
SGD	–	–	–	1.60	90.58
Adam	–	–	–	1.23	90.39
SGN	–	const	light	3.52	90.69
SGN	–	const	heavy	4.79	90.09
SGN	–	trust region	light	7.33	89.69
SGN	–	trust region	heavy	9.13	91.14

Table 4: Interactions between solver modules shape performance on Fashion-MNIST.

Figure 3: Wall-clock performance comparison of SGDM and Sophia-style optimizer regimes on CIFAR-10 with ResNet-20. Thick lines show mean test accuracy across 5 seeds; shaded bands indicate $\pm 1$ standard deviation.

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Spec Object	Role & Planner Effect
CurvatureSpec	Defines the operator family (e.g., hessian , ggn_ce ). Declares if row-primitives are available.
SolverSpec	Selects the algorithm (e.g., cg , row_cholesky ). Determines the execution lane (Param vs. Row).
DampingSpec	Defines the trust-region or regularization policy. Requests telemetry (e.g., rho-packs) that the planner must schedule.
PrecondSpec	(Optional) Defines diagonal scaling (e.g., diag_ema ) for the Diagonal lane or as a preconditioner for CG.
EstimatorSpec	(Optional) Layers randomized probing (e.g., Hutchinson) onto the step without extra linearizations.

Appendix

Table 5: Primary configuration objects (Specs) in the Somax API.

Component	Consumes	Produces	Planner effect (why it matters)	Persistent state
Curvature operator	parameters, batch	linearization state; matrix-free primitives (and optionally row primitives); loss/grad consistent with snapshot	determines feasibility of the row lane; declares whether post-update probes can reuse the step-local snapshot	none (the linearization snapshot is step-local)
Estimator wrapper (optional)	curvature primitives and RNG (cadence-gated)	additional curvature summaries (e.g., diagonal, trace, spectrum)	adds optional probes without creating a second curvature snapshot; updates the fixed metric schema	optional estimator state (e.g., probe buffers), else none
Preconditioner (optional)	diagonal estimate and damping	diagonal scaling (used by diagonal lane or by PCG)	controls whether the param lane specializes to CG vs PCG; avoids runtime switching inside traced code	EMA buffers (optional)
Linear solver	lane-specific operator primitives, damping, right-hand side	scaled direction; solver diagnostics (if enabled)	fixes the lane (diag/param/row); enables cadence-gated solver telemetry keys	warm start (solution to the previous iteration, optional)
Damping policy	plan-selected telemetry (e.g., gain-ratio bundle) at declared cadence	updated damping state	turns on post-update probes only when required; ensures probes are computed against the applied update	damping state
Optax post-transform	scaled direction and Optax state	applied update and updated Optax state	forces step-level semantics: post-update probes must match the exact applied update, not the raw direction	Optax state

Appendix

Table 6: Planner-aware contracts (math

\rightarrow

components). Each component declares requirements that the planner resolves into a static plan (lane, metric schema, cadences), so the executor can fuse the step without redundant recomputation.

Curvature-Aware Optimization

See all 21 papers

Second-Order Optimization with Asynchronous Curvature
Oct 17, 2025El Mahdi Chayti, Martin JaggiSecond-Order Optimization Curvature-Aware Optimization
Machine Learning and Optimization Laboratory EPFL
Accelerated Gradient Descent for Faster Convergence with Minimal Overhead
May 15, 2026Manuel Graca, L. Miguel Silveira, Arlindo Oliveira +1Deep Learning Optimization Gradient Descent
IST, ULisboa Lisbon, Portugal · Old Dominion University Norfolk, VA USA
Why Muon Outperforms Adam: A Curvature Perspective
Jun 3, 2026Shuche Wang, Fengzhuo Zhang, Jiaxiang Li +2Loss Landscape Sharpness Neural Network Optimization
National University of Singapore · Yale University · University of Minnesota

Oct 17, 2025J/K move · Enter open · S save

Second-Order Optimization with Asynchronous Curvature

El Mahdi Chayti, Martin Jaggi

Machine Learning and Optimization Laboratory EPFL

Second-Order Optimization Curvature-Aware Optimization Optimization Convergence Analysis

Open paper AskPDF arXiv

Second-order methods make fast progress per iteration, but forming and factorizing the Hessian often dominates their wall-clock time. We study Split-Client, a task-parallel design in which a curvature client factorizes Hessians in the background while a gradient client takes cubic-regularized steps with the most recently published matrix, so that curvature cost appears as a delay rather than as blocking time. For delayed exact Hessians, we prove a first-order stationarity bound that depends on the average age of the curvature rather than its maximum, using an overlap count and a weighted Hardy inequality that allows any bounded-error initial matrix; the analysis extends to subsampled and uniformly inexact curvature. Under a fixed-cost model, the resulting wall-clock bound matches the scaling of optimally tuned Lazy Hessian without choosing a reuse period. In experiments on two LIBSVM datasets, an implementation decreases the objective earlier in wall-clock time than Lazy and Vanilla cubic Newton, with the largest gains at intermediate accuracy.

Abstract

Curvature-Aware Optimization

Explore similar work

Second-Order Optimization with Asynchronous Curvature

Accelerated Gradient Descent for Faster Convergence with Minimal Overhead

Why Muon Outperforms Adam: A Curvature Perspective