Skip to main content

Second-order optimization

Curvature-aware optimizers (Newton, natural gradient, KFAC) need second-order information that autodiff makes expensive. For Riccati-class one-layer fields, omnibias gives it to you in closed form.

Prerequisites​

pip install omnibias-curvature

Step 1 — the closed-form Hessian​

omnibias-curvature computes the one-layer parameter Hessian directly from the derivative tower, validated against jax.hessian.

from omnibias.curvature import parameter_hessian

H = parameter_hessian(params, batch, activation="tanh")   # exact, closed form

Unlike a dense autodiff Hessian, this does not materialize and differentiate a graph — it contracts the σ'' tower against the layer structure.

Step 2 — Fisher and KFAC factors​

from omnibias.curvature import gauss_newton_fisher, kfac_factors

F = gauss_newton_fisher(params, batch, activation="tanh")
A, G = kfac_factors(params, batch, activation="tanh")   # Kronecker factors

Step 3 — a natural-gradient step​

import jax.numpy as jnp

grad = loss_grad(params, batch)
# Natural gradient: F^{-1} g  (here via the KFAC factorization)
nat_grad = kfac_solve(A, G, grad)
params = tree_sub(params, scale=lr, update=nat_grad)
Validate the curvature

Compare parameter_hessian against jax.hessian of your loss on a small problem. They should agree to float64 round-off — closed form is exact, not approximate.

Why this is practical now​

Second-order methods are often abandoned because the Hessian is too expensive to form. When it is a closed-form contraction, the calculus changes: natural gradient and KFAC become a reasonable default rather than a luxury.

Composes with regularization​

The curvature surface composes with the Sobolev / Jacobian regularization in the field substrate, so you can add a smoothness penalty without leaving the closed-form path.

Next​

Last updated on