omnibias.curvature
Scoped parameter-curvature formulas, Fisher and Gauss–Newton constructions, natural steps, and sharpness utilities.
import omnibias.curvatureChoose curvature when its architecture and loss assumptions match your model. Exact parameter Hessians, Gauss–Newton matrices, Fisher information, empirical Fisher matrices, and KFAC approximations are different objects. Do not interchange their names or guarantees.
Several closed-form constructions target one-hidden-layer Riccati fields or particular likelihood families. Read parameter packing order, residual convention, reduction, and damping before using a returned matrix in an optimizer. A positive-semidefinite approximation can still be singular; the step solver's regularization and rank assumptions matter.
The optimization tutorial demonstrates a supported Gauss–Newton closure with a complete solve. Use it to understand the closure contract, then validate the specific curvature module you need. Generic optimizer interfaces elsewhere in the Torch and JAX packages may use autodiff-based primitives and are not automatically closed-form curvature algorithms.
Module contract & mathematical conventions
Closed-form second-order optimisation for one-layer Riccati fields.
This package promotes the omnibias closed-form Hessian
(omnibias.jax.neural_field_value_grad_hessian, which is the
D \times D Hessian *with respect to the input*) into a usable
**parameter** Fisher / Hessian for the one-hidden-layer scalar field
f_\theta(x) \;=\; b \;+\; \sum_h c_h\, \sigma(W_h \cdot x + \beta_h),
where \sigma is one of the Riccati-class activations exposed by
omnibias.jax.activations.get_activation (tanh, sigmoid,
softplus, gaussian, exp). The closed forms for
\sigma'(z) and \sigma''(z) mean every block of the
parameter Hessian / Fisher can be assembled algebraically — no autograd
needed.
What this v0.1 ships
The omnibias.curvature.one_layer module provides:
one_layer_param_grad— per-sample\nabla_\theta f(x)as a flat vector of lengthP = 1 + H + H + H \cdot D.one_layer_param_hessian— per-sample full Hessian\nabla^2_\theta f(x); block-diagonal-in-hidden-unit with(D+2) \times (D+2)blocks.mse_gauss_newton_fisher— Gauss-Newton FisherF = \tfrac{2}{B}\sum_n g_n g_n^Tfor MSE loss; the standard positive-definite curvature surrogate used by natural-gradient and KFAC.mse_newton_step— one Newton-Gauss step\theta \leftarrow \theta - \eta\,(F + \lambda I)^{-1} \nabla L.kfac_kron_factors— closed-form KFAC(A, G)Kronecker factors for the hidden block, computed fromx_n x_n^Tand(c \odot \sigma'(z_n))(c \odot \sigma'(z_n))^Tstatistics.
Exact-curvature sharpness ("SAM done right")
The omnibias.curvature.sharpness module builds on the closed-form
Hessian to measure and regularise the *sharpness* of a basin exactly,
where Sharpness-Aware Minimisation only estimates its linear shadow:
mse_loss_hessian— the **exact full loss Hessian** (Gauss-Newton *plus* ther_n\,\nabla^2_\theta fresidual term).hessian_trace/hessian_frobenius_sq/hessian_top_eigenvalue— the sharpness functionals\sum_i\lambda_i,\sum_i\lambda_i^2,\lambda_{\max}.sharpness_aware_loss—L + \lambda\,\mathcal S(\nabla^2 L), a differentiable curvature regulariser (its gradient pulls in\sigma'''in closed form).sam_objective— the ascent-free, exact second-order SAM surrogateL + \rho\lVert\nabla L\rVert + \tfrac12\rho^2\lambda_{\max}.
These are the **one-layer, JAX** primitives. The arbitrary-depth,
**PyTorch** port lives in omnibias.curvature.torch (matrix-free via
exact Hessian-vector products) and requires the optional torch extra;
it is intentionally not imported here so this namespace stays JAX-only.
What's still on the roadmap
- Integration with
kfac_jaxas a custom layer registration so that a FermiNet trainer can drop in the closed-form Fisher block for the omnibias-envelope layer.
These follow-ups are tracked on the public roadmap in
docs/roadmap.md.
Shared formulas do not imply universal floating-point identity. Verify dtype, shape, derivative-order support, and the stated numerical assumptions for the path you use.
glm_fisherModule ↗informationModule ↗natural_gradientModule ↗one_layerModule ↗operatorsModule ↗regularizeModule ↗sharpnessModule ↗torchModule ↗Imports & aliases
These names share a definition with another public module. Follow a name for its full contract, parameters, and examples.
| Import from this module | Canonical reference |
|---|---|
CollapseResult | omnibias.curvature.CollapseResult |
damped_solve | omnibias.curvature.damped_solve |
fisher_information_metric | omnibias.curvature.fisher_information_metric |
glm_loss_gradient | omnibias.curvature.glm_loss_gradient |
glm_natural_gradient_step | omnibias.curvature.glm_natural_gradient_step |
hessian_frobenius_sq | omnibias.curvature.hessian_frobenius_sq |
hessian_top_eigenvalue | omnibias.curvature.hessian_top_eigenvalue |
hessian_trace | omnibias.curvature.hessian_trace |
kfac_kron_factors | omnibias.curvature.kfac_kron_factors |
min_norm_solve | omnibias.curvature.min_norm_solve |
mse_curvature_sharpness | omnibias.curvature.mse_curvature_sharpness |
mse_gauss_newton_fisher | omnibias.curvature.mse_gauss_newton_fisher |
mse_loss | omnibias.curvature.mse_loss |
mse_loss_hessian | omnibias.curvature.mse_loss_hessian |
mse_newton_step | omnibias.curvature.mse_newton_step |
natural_gradient_step | omnibias.curvature.natural_gradient_step |
numerical_rank | omnibias.curvature.numerical_rank |
one_layer_param_grad | omnibias.curvature.one_layer_param_grad |
one_layer_param_hessian | omnibias.curvature.one_layer_param_hessian |
rank_collapse | omnibias.curvature.rank_collapse |
regularization_path | omnibias.curvature.regularization_path |
regularized_solve | omnibias.curvature.regularized_solve |
sam_objective | omnibias.curvature.sam_objective |
sam_sharpness_gap | omnibias.curvature.sam_sharpness_gap |
sharpness_aware_loss | omnibias.curvature.sharpness_aware_loss |