API REFERENCE/omnibias-curvature

omnibias.curvature.sharpness

0.1.0a2AlphaPython >=3.10AGPL / commercial

Scoped parameter-curvature formulas, Fisher and Gauss–Newton constructions, natural steps, and sharpness utilities.

import omnibias.curvature.sharpness

Module contract & mathematical conventions

Exact-curvature sharpness-aware training for one-layer Riccati fields.

Motivation -- "SAM done right"

Sharpness-Aware Minimisation (Foret et al. 2021) replaces the training loss L(\theta) by its worst case inside an \ell_2-ball of radius \rho,

L^{\mathrm{SAM}}_\rho(\theta) \;=\; \max_{\lVert\varepsilon\rVert_2 \le \rho} L(\theta+\varepsilon),

and approximates the inner maximisation by a *single gradient-ascent step* \varepsilon^\star \approx \rho\,\nabla L / \lVert\nabla L\rVert. That approximation keeps only the **linear** term of the Taylor expansion and needs a second forward/backward pass; it never sees the curvature that actually makes a basin sharp or flat.

For the one-hidden-layer Riccati field f_\theta(x) = b + \sum_h c_h\,\sigma(W_h\!\cdot x + \beta_h) omnibias gives the **closed-form** \sigma', \sigma'' (and, through autodiff of the closed form, \sigma''' -- the third derivative the *gradient* of a curvature penalty needs). That lets us assemble the **exact** parameter Hessian of the loss and its curvature functionals analytically, so the SAM inner-max can be expanded to second order *exactly* instead of being estimated by an ascent step:

L^{\mathrm{SAM}}_\rho(\theta) - L(\theta) \;=\; \rho\,\lVert\nabla L\rVert \;+\; \tfrac12\,\rho^2\,\lambda_{\max}(\nabla^2 L) \;+\; O(\rho^3).

What this module ships

• mse_loss_hessian -- the **exact** full parameter Hessian of the MSE loss (Gauss-Newton term *plus* the residual \tfrac{2}{B}\sum_n r_n\,\nabla^2_\theta f(x_n) term). This is the "full Newton" curvature the one_layer roadmap flagged.
• hessian_trace, hessian_frobenius_sq, hessian_top_eigenvalue -- the three curvature functionals that the loss-landscape literature uses as sharpness proxies (\operatorname{tr}H=\sum_i\lambda_i, the Sobolev \lVert H\rVert_F^2=\sum_i\lambda_i^2, and the canonical \lambda_{\max}(H)).
• mse_curvature_sharpness -- assemble the exact loss Hessian and return one of those functionals.
• sharpness_aware_loss -- L + \lambda\,\mathcal S(\nabla^2 L), the differentiable regulariser used in a training loop.
• sam_sharpness_gap / sam_objective -- the ascent-free, exact second-order SAM surrogate.

Honesty caveats

• Everything here is **exact to floating point** *for the one-layer field*: the per-sample output Hessian and the \sigma^{(k)} are closed form, so mse_loss_hessian equals jax.hessian of the loss to ~1e-9 (see the tests). The *only* approximations are the ones you opt into: the SAM surrogate keeps the second-order Taylor term (error O(\rho^3)), and \lambda_{\max} is clamped at 0 so negative curvature is never rewarded.
• This is a **one-layer** primitive. The arbitrary-depth, torch port is omnibias.curvature.torch (matrix-free, via exact Hessian-vector products). The scientific claim is *not* "flat minima always generalise better" -- that is problem-dependent (No Free Lunch); the claim is that omnibias can measure and regularise curvature *exactly and cheaply*, where SAM only estimates its linear shadow.

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.

Module source ↗Source revision 9a2d73df0573Running reference examples →

FUNCTION

hessian_frobenius_sq

omnibias.curvature.hessian_frobenius_sqSource ↗

hessian_frobenius_sq(H: Array) -> Array

\lVert H\rVert_F^2 = \sum_i \lambda_i^2 -- the Sobolev sharpness.

Smooth and always differentiable (no eigen-decomposition), which makes it the most robust default for a training-time penalty.

Parameters

Name	Type & default	Contract
Hpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.

ReturnsArray

Executed regression · measures match numpy eigendecomposition

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_measures_match_numpy_eigendecomposition():
    W, beta, c, b, X, Y = _rand_problem(seed=3)
    H, _ = mse_loss_hessian(X, Y, W, beta, c, b, "tanh")
    ev = np.linalg.eigvalsh(np.asarray(0.5 * (H + H.T)))
    assert np.isclose(float(hessian_trace(H)), ev.sum(), rtol=1e-10)
    assert np.isclose(float(hessian_frobenius_sq(H)), (ev**2).sum(), rtol=1e-10)
    assert np.isclose(float(hessian_top_eigenvalue(H)), ev.max(), rtol=1e-10)

python -m pytest "packages/omnibias-curvature/tests/test_sharpness.py::test_measures_match_numpy_eigendecomposition" -q

Read the implementation

def hessian_frobenius_sq(H: Array) -> Array:
    r""":math:`\lVert H\rVert_F^2 = \sum_i \lambda_i^2` -- the Sobolev sharpness.

    Smooth and always differentiable (no eigen-decomposition), which
    makes it the most robust default for a training-time penalty.
    """
    return jnp.sum(H * H)

Other public import paths

• omnibias.curvature.hessian_frobenius_sq
• omnibias.curvature.sharpness.hessian_frobenius_sq

FUNCTION

hessian_top_eigenvalue

omnibias.curvature.hessian_top_eigenvalueSource ↗

hessian_top_eigenvalue(H: Array) -> Array

\lambda_{\max}(H) -- the canonical "sharpness".

Symmetrises H and returns the largest eigenvalue via jax.numpy.linalg.eigvalsh. Differentiable wherever the top eigenvalue is simple.

Parameters

Name	Type & default	Contract
Hpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.

ReturnsArray

Executed regression · measures match numpy eigendecomposition

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_measures_match_numpy_eigendecomposition():
    W, beta, c, b, X, Y = _rand_problem(seed=3)
    H, _ = mse_loss_hessian(X, Y, W, beta, c, b, "tanh")
    ev = np.linalg.eigvalsh(np.asarray(0.5 * (H + H.T)))
    assert np.isclose(float(hessian_trace(H)), ev.sum(), rtol=1e-10)
    assert np.isclose(float(hessian_frobenius_sq(H)), (ev**2).sum(), rtol=1e-10)
    assert np.isclose(float(hessian_top_eigenvalue(H)), ev.max(), rtol=1e-10)

python -m pytest "packages/omnibias-curvature/tests/test_sharpness.py::test_measures_match_numpy_eigendecomposition" -q

Regression example · sharpness aware training finds flatter minimum

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_sharpness_aware_training_finds_flatter_minimum():
    """From the *same* init and data, curvature-regularised training reaches a
    minimum with substantially smaller top-eigenvalue curvature than plain MSE
    -- while still fitting the (realisable) data and generalising at least as
    well. This is the SAM promise, delivered via exact curvature."""
    D, H_teacher, H_student, n_train = 3, 4, 8, 30
    Xtr, Ytr, Xte, Yte = _teacher_data(seed=17, D=D, H_teacher=H_teacher,
                                       n_train=n_train, noise=0.1)
    H = H_student
    theta0 = _init_student(seed=17, D=D, H=H)

    def plain_g(th):
        return jax.grad(
            lambda t: _loss_of_theta(t, mse_loss, Xtr, Ytr, H, D, "tanh")
        )(th)

    def flat_g(th):
        return jax.grad(
            lambda t: _loss_of_theta(t, sharpness_aware_loss, Xtr, Ytr, H, D,
                                     "tanh", lam=5e-3, measure="frobenius")
        )(th)

    theta_plain = _adam(jax.jit(plain_g), theta0, steps=500, lr=5e-3)
    theta_flat = _adam(jax.jit(flat_g), theta0, steps=500, lr=5e-3)

    def report(theta):
        bb, cc, be, WW = unpack_params(theta, H=H, D=D)
        Hh, _ = mse_loss_hessian(Xtr, Ytr, WW, be, cc, bb, "tanh")
        return (
            float(mse_loss(Xtr, Ytr, WW, be, cc, bb, "tanh")),
            float(mse_loss(Xte, Yte, WW, be, cc, bb, "tanh")),
            float(hessian_top_eigenvalue(Hh)),
        )

    tr_p, te_p, curv_p = report(theta_plain)
    tr_f, te_f, curv_f = report(theta_flat)

    # (1) both fit the training data
    assert tr_p < 0.05 and tr_f < 0.05, (tr_p, tr_f)
    # (2) the flat run lands in a clearly flatter basin
    assert curv_f < 0.7 * curv_p, f"curv flat={curv_f:.3f} vs plain={curv_p:.3f}"
    # (3) and generalises at least as well (here: better)
    assert te_f <= te_p * 1.05, f"test mse flat={te_f:.4f} vs plain={te_p:.4f}"

python -m pytest packages/omnibias-curvature/tests/test_sharpness.py -k test_sharpness_aware_training_finds_flatter_minimum -q

Read the implementation

def hessian_top_eigenvalue(H: Array) -> Array:
    r""":math:`\lambda_{\max}(H)` -- the canonical "sharpness".

    Symmetrises ``H`` and returns the largest eigenvalue via
    :func:`jax.numpy.linalg.eigvalsh`. Differentiable wherever the top
    eigenvalue is simple.
    """
    Hs = 0.5 * (H + H.T)
    return jnp.linalg.eigvalsh(Hs)[-1]

Other public import paths

• omnibias.curvature.hessian_top_eigenvalue
• omnibias.curvature.sharpness.hessian_top_eigenvalue

FUNCTION

hessian_trace

omnibias.curvature.hessian_traceSource ↗

hessian_trace(H: Array) -> Array

\operatorname{tr}(H) = \sum_i \lambda_i -- total curvature.

Parameters

Name	Type & default	Contract
Hpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.

ReturnsArray

Executed regression · measures match numpy eigendecomposition

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_measures_match_numpy_eigendecomposition():
    W, beta, c, b, X, Y = _rand_problem(seed=3)
    H, _ = mse_loss_hessian(X, Y, W, beta, c, b, "tanh")
    ev = np.linalg.eigvalsh(np.asarray(0.5 * (H + H.T)))
    assert np.isclose(float(hessian_trace(H)), ev.sum(), rtol=1e-10)
    assert np.isclose(float(hessian_frobenius_sq(H)), (ev**2).sum(), rtol=1e-10)
    assert np.isclose(float(hessian_top_eigenvalue(H)), ev.max(), rtol=1e-10)

python -m pytest "packages/omnibias-curvature/tests/test_sharpness.py::test_measures_match_numpy_eigendecomposition" -q

Read the implementation

def hessian_trace(H: Array) -> Array:
    r""":math:`\operatorname{tr}(H) = \sum_i \lambda_i` -- total curvature."""
    return jnp.trace(H)

Other public import paths

• omnibias.curvature.hessian_trace
• omnibias.curvature.sharpness.hessian_trace

FUNCTION

mse_curvature_sharpness

omnibias.curvature.mse_curvature_sharpnessSource ↗

mse_curvature_sharpness(X: Array, Y: Array, W: Array, beta: Array, c: Array, b: Array, activation: str='tanh', measure: str='frobenius') -> Array

Exact curvature-based sharpness of the MSE loss at theta.

Assembles the exact loss Hessian (mse_loss_hessian) and returns the chosen functional -- "trace" (\operatorname{tr}H), "frobenius" (\lVert H\rVert_F^2, the default) or "top_eig" (\lambda_{\max}(H)).

Parameters

Name	Type & default	Contract
Xpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
Ypositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
Wpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
betapositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
cpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
bpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
activationpositional-or-keyword	strDefault: 'tanh'	See the operation contract and implementation below.
measurepositional-or-keyword	strDefault: 'frobenius'	See the operation contract and implementation below.

ReturnsArray

Executed regression · unknown measure raises

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_unknown_measure_raises():
    W, beta, c, b, X, Y = _rand_problem(seed=1)
    with pytest.raises(ValueError, match="unknown sharpness measure"):
        mse_curvature_sharpness(X, Y, W, beta, c, b, "tanh", measure="nope")

python -m pytest "packages/omnibias-curvature/tests/test_sharpness.py::test_unknown_measure_raises" -q

Regression example · penalty only descent monotonically reduces curvature

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_penalty_only_descent_monotonically_reduces_curvature():
    """Gradient descent on the pure curvature penalty must reduce it."""
    W, beta, c, b, X, Y = _rand_problem(seed=2)
    H, D = W.shape
    theta = pack_params(b, c, beta, W)

    def S(th):
        return _loss_of_theta(th, mse_curvature_sharpness, X, Y, H, D, "tanh",
                              measure="frobenius")

    vg = jax.value_and_grad(S)
    prev = float(S(theta))
    for _ in range(15):
        val, g = vg(theta)
        theta = theta - 1e-3 * g
        cur = float(S(theta))
        assert cur <= prev + 1e-9, f"curvature increased: {prev} -> {cur}"
        prev = cur
    assert prev < float(S(pack_params(b, c, beta, W)))

python -m pytest packages/omnibias-curvature/tests/test_sharpness.py -k test_penalty_only_descent_monotonically_reduces_curvature -q

Read the implementation

def mse_curvature_sharpness(
    X: Array, Y: Array,
    W: Array, beta: Array, c: Array, b: Array,
    activation: str = "tanh",
    measure: str = "frobenius",
) -> Array:
    r"""Exact curvature-based sharpness of the MSE loss at ``theta``.

    Assembles the exact loss Hessian (:func:`mse_loss_hessian`) and
    returns the chosen functional -- ``"trace"``
    (:math:`\operatorname{tr}H`), ``"frobenius"`` (:math:`\lVert H\rVert_F^2`,
    the default) or ``"top_eig"`` (:math:`\lambda_{\max}(H)`).
    """
    H, _ = mse_loss_hessian(X, Y, W, beta, c, b, activation)
    return _apply_measure(measure, H)

Related APIsomnibias.curvature.mse_loss_hessian↗

Other public import paths

• omnibias.curvature.mse_curvature_sharpness
• omnibias.curvature.sharpness.mse_curvature_sharpness

FUNCTION

mse_loss

omnibias.curvature.mse_lossSource ↗

mse_loss(X: Array, Y: Array, W: Array, beta: Array, c: Array, b: Array, activation: str='tanh') -> Array

Mean-squared-error loss L = \tfrac1B\sum_n (f(x_n)-y_n)^2.

Uses the same 1/B normalisation as omnibias.curvature.one_layer.mse_gauss_newton_fisher, so the Gauss-Newton Fisher there is exactly the positive-definite part of mse_loss_hessian.

Parameters

Name	Type & default	Contract
Xpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
Ypositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
Wpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
betapositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
cpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
bpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
activationpositional-or-keyword	strDefault: 'tanh'	See the operation contract and implementation below.

ReturnsArray

Executed regression · loss gradient matches jax grad

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_loss_gradient_matches_jax_grad(activation):
    """The returned loss gradient must equal jax.grad(L)."""
    W, beta, c, b, X, Y = _rand_problem(seed=(hash(activation) >> 3) & 0xFFFF)
    H, D = W.shape
    _, g_closed = mse_loss_hessian(X, Y, W, beta, c, b, activation)
    theta = pack_params(b, c, beta, W)
    g_jax = jax.grad(_loss_of_theta)(theta, mse_loss, X, Y, H, D, activation)
    np.testing.assert_allclose(
        np.asarray(g_closed), np.asarray(g_jax), rtol=1e-9, atol=1e-11,
    )

python -m pytest "packages/omnibias-curvature/tests/test_sharpness.py::test_loss_gradient_matches_jax_grad[exp]" -q

Regression example · full loss hessian matches jax hessian

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_full_loss_hessian_matches_jax_hessian(activation):
    """The full loss Hessian (GN + residual) must equal jax.hessian(L)."""
    W, beta, c, b, X, Y = _rand_problem(seed=hash(activation) & 0xFFFF)
    H, D = W.shape

    H_closed, _ = mse_loss_hessian(X, Y, W, beta, c, b, activation)
    theta = pack_params(b, c, beta, W)
    H_jax = jax.hessian(_loss_of_theta)(theta, mse_loss, X, Y, H, D, activation)

    assert H_closed.shape == H_jax.shape == (theta.shape[0], theta.shape[0])
    assert jnp.allclose(H_closed, H_closed.T, atol=1e-12), "Hessian not symmetric"
    np.testing.assert_allclose(
        np.asarray(H_closed), np.asarray(H_jax),
        rtol=1e-7, atol=1e-8,
        err_msg=f"loss Hessian mismatch for activation={activation!r}",
    )

python -m pytest packages/omnibias-curvature/tests/test_sharpness.py -k test_full_loss_hessian_matches_jax_hessian -q

Read the implementation

def mse_loss(
    X: Array, Y: Array,
    W: Array, beta: Array, c: Array, b: Array,
    activation: str = "tanh",
) -> Array:
    r"""Mean-squared-error loss :math:`L = \tfrac1B\sum_n (f(x_n)-y_n)^2`.

    Uses the same ``1/B`` normalisation as
    :func:`omnibias.curvature.one_layer.mse_gauss_newton_fisher`, so the
    Gauss-Newton Fisher there is exactly the positive-definite part of
    :func:`mse_loss_hessian`.
    """
    f = _batch_forward(X, W, beta, c, b, activation)
    r = f - Y
    return jnp.mean(r * r)

Other public import paths

• omnibias.curvature.mse_loss
• omnibias.curvature.sharpness.mse_loss

FUNCTION

mse_loss_hessian

omnibias.curvature.mse_loss_hessianSource ↗

mse_loss_hessian(X: Array, Y: Array, W: Array, beta: Array, c: Array, b: Array, activation: str='tanh') -> tuple[Array, Array]

Exact full parameter Hessian and gradient of the MSE loss.

For L = \tfrac1B\sum_n r_n^2 with residual r_n = f_\theta(x_n) - y_n,

\nabla^2_\theta L \;=\; \underbrace{\tfrac2B\sum_n g_n g_n^\top}_{\text{Gauss-Newton}} \;+\; \underbrace{\tfrac2B\sum_n r_n\,\nabla^2_\theta f(x_n)}_{\text{residual}},

where g_n = \nabla_\theta f(x_n). The Gauss-Newton term is ~omnibias.curvature.one_layer.mse_gauss_newton_fisher; the residual term uses the closed-form per-sample output Hessian ~omnibias.curvature.one_layer.one_layer_param_hessian. Unlike Gauss-Newton, this is the *true* loss Hessian (it can be indefinite away from a minimum), so it is the right object for a curvature / sharpness penalty.

Returns

H : (P, P) symmetric The exact loss Hessian. g_loss : (P,) The loss gradient \nabla_\theta L = \tfrac2B\sum_n r_n g_n.

Parameters

Name	Type & default	Contract
Xpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
Ypositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
Wpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
betapositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
cpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
bpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
activationpositional-or-keyword	strDefault: 'tanh'	See the operation contract and implementation below.

Returnstuple[Array, Array]

Executed regression · unknown measure raises

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_unknown_measure_raises():
    W, beta, c, b, X, Y = _rand_problem(seed=1)
    with pytest.raises(ValueError, match="unknown sharpness measure"):
        mse_curvature_sharpness(X, Y, W, beta, c, b, "tanh", measure="nope")

python -m pytest "packages/omnibias-curvature/tests/test_sharpness.py::test_unknown_measure_raises" -q

Regression example · measures match numpy eigendecomposition

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_measures_match_numpy_eigendecomposition():
    W, beta, c, b, X, Y = _rand_problem(seed=3)
    H, _ = mse_loss_hessian(X, Y, W, beta, c, b, "tanh")
    ev = np.linalg.eigvalsh(np.asarray(0.5 * (H + H.T)))
    assert np.isclose(float(hessian_trace(H)), ev.sum(), rtol=1e-10)
    assert np.isclose(float(hessian_frobenius_sq(H)), (ev**2).sum(), rtol=1e-10)
    assert np.isclose(float(hessian_top_eigenvalue(H)), ev.max(), rtol=1e-10)

python -m pytest packages/omnibias-curvature/tests/test_sharpness.py -k test_measures_match_numpy_eigendecomposition -q

Read the implementation

def mse_loss_hessian(
    X: Array, Y: Array,
    W: Array, beta: Array, c: Array, b: Array,
    activation: str = "tanh",
) -> tuple[Array, Array]:
    r"""Exact full parameter Hessian and gradient of the MSE loss.

    For :math:`L = \tfrac1B\sum_n r_n^2` with residual
    :math:`r_n = f_\theta(x_n) - y_n`,

    .. math::

        \nabla^2_\theta L
          \;=\; \underbrace{\tfrac2B\sum_n g_n g_n^\top}_{\text{Gauss-Newton}}
              \;+\; \underbrace{\tfrac2B\sum_n r_n\,\nabla^2_\theta f(x_n)}_{\text{residual}},

    where :math:`g_n = \nabla_\theta f(x_n)`. The Gauss-Newton term is
    :func:`~omnibias.curvature.one_layer.mse_gauss_newton_fisher`; the
    residual term uses the closed-form per-sample output Hessian
    :func:`~omnibias.curvature.one_layer.one_layer_param_hessian`. Unlike
    Gauss-Newton, this is the *true* loss Hessian (it can be indefinite
    away from a minimum), so it is the right object for a curvature /
    sharpness penalty.

    Returns
    -------
    H : ``(P, P)`` symmetric
        The exact loss Hessian.
    g_loss : ``(P,)``
        The loss gradient :math:`\nabla_\theta L = \tfrac2B\sum_n r_n g_n`.
    """
    spec = get_activation(activation)
    B = X.shape[0]

    def per_sample(x_n: Array, y_n: Array) -> tuple[Array, Array, Array]:
        z = W @ x_n + beta
        f = b + spec.forward(z) @ c
        r = f - y_n
        g = one_layer_param_grad(x_n, W, beta, c, b, activation)     # (P,)
        Hf = one_layer_param_hessian(x_n, W, beta, c, b, activation)  # (P, P)
        return r, g, Hf

    rs, gs, Hfs = jax.vmap(per_sample)(X, Y)         # (B,), (B,P), (B,P,P)
    gauss_newton = (2.0 / B) * (gs.T @ gs)           # (P, P)
    residual = (2.0 / B) * jnp.einsum("n,nij->ij", rs, Hfs)  # (P, P)
    H = gauss_newton + residual
    g_loss = (2.0 / B) * (gs.T @ rs)                 # (P,)
    return H, g_loss

Related APIsomnibias.curvature.one_layer_param_grad↗omnibias.curvature.one_layer_param_hessian↗omnibias.jax.get_activation↗

Other public import paths

• omnibias.curvature.mse_loss_hessian
• omnibias.curvature.sharpness.mse_loss_hessian

FUNCTION

sam_objective

omnibias.curvature.sam_objectiveSource ↗

sam_objective(X: Array, Y: Array, W: Array, beta: Array, c: Array, b: Array, activation: str='tanh', *, rho: float=0.05) -> Array

Ascent-free "SAM done right": L(\theta) + \text{gap}_\rho(\theta).

Equivalent to the SAM objective expanded to exact second order (see sam_sharpness_gap). Minimising it simultaneously lowers the loss and the worst-case curvature of the basin.

Parameters

Name	Type & default	Contract
Xpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
Ypositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
Wpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
betapositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
cpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
bpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
activationpositional-or-keyword	strDefault: 'tanh'	See the operation contract and implementation below.
rhokeyword-only	floatDefault: 0.05	See the operation contract and implementation below.

ReturnsArray

Read the implementation

def sam_objective(
    X: Array, Y: Array,
    W: Array, beta: Array, c: Array, b: Array,
    activation: str = "tanh",
    *,
    rho: float = 0.05,
) -> Array:
    r"""Ascent-free "SAM done right": :math:`L(\theta) + \text{gap}_\rho(\theta)`.

    Equivalent to the SAM objective expanded to exact second order (see
    :func:`sam_sharpness_gap`). Minimising it simultaneously lowers the
    loss and the worst-case curvature of the basin.
    """
    base = mse_loss(X, Y, W, beta, c, b, activation)
    return base + sam_sharpness_gap(
        X, Y, W, beta, c, b, activation, rho=rho
    )

Related APIsomnibias.curvature.mse_loss↗omnibias.curvature.sam_sharpness_gap↗

Other public import paths

• omnibias.curvature.sam_objective
• omnibias.curvature.sharpness.sam_objective

FUNCTION

sam_sharpness_gap

omnibias.curvature.sam_sharpness_gapSource ↗

sam_sharpness_gap(X: Array, Y: Array, W: Array, beta: Array, c: Array, b: Array, activation: str='tanh', *, rho: float=0.05) -> Array

Exact second-order surrogate of the SAM inner-max gap.

\max_{\lVert\varepsilon\rVert\le\rho} L(\theta+\varepsilon) - L(\theta) \;\approx\; \rho\,\lVert\nabla L\rVert \;+\; \tfrac12\,\rho^2\,\max\!\big(\lambda_{\max}(\nabla^2 L),\,0\big),

exact to O(\rho^3). This is what standard SAM *estimates* with one ascent step (linear term only); here both terms are computed from the closed-form gradient and Hessian in a single pass. The curvature term is clamped at 0 so negative curvature is never rewarded (the gap is \ge 0).

Parameters

Name	Type & default	Contract
Xpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
Ypositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
Wpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
betapositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
cpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
bpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
activationpositional-or-keyword	strDefault: 'tanh'	See the operation contract and implementation below.
rhokeyword-only	floatDefault: 0.05	See the operation contract and implementation below.

ReturnsArray

Executed regression · sam gap upper bounds sampled worst case

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_sam_gap_upper_bounds_sampled_worst_case():
    """The exact 2nd-order SAM gap must upper-bound the worst loss increase in
    an l2-ball (sound model of the SAM inner-max), and stay close to it. Uses a
    low-dim field so uniform sphere sampling densely covers the ball."""
    W, beta, c, b, X, Y = _rand_problem(D=1, H=2, B=8, seed=1)
    H, D = W.shape
    theta = pack_params(b, c, beta, W)
    P = theta.shape[0]
    L0 = float(_loss_of_theta(theta, mse_loss, X, Y, H, D, "tanh"))

    rho = 0.1
    gap2 = float(sam_sharpness_gap(X, Y, W, beta, c, b, "tanh", rho=rho))
    assert gap2 > 0.0

    # Vectorised sampled worst case over the sphere of radius rho.
    rng = np.random.default_rng(0)
    dirs = rng.normal(size=(6000, P))
    dirs = rho * dirs / np.linalg.norm(dirs, axis=1, keepdims=True)

    def delta(d):
        return _loss_of_theta(theta + d, mse_loss, X, Y, H, D, "tanh") - L0

    worst = float(jnp.max(jax.vmap(delta)(jnp.asarray(dirs))))
    assert gap2 >= worst - 1e-3, f"gap2={gap2} under-shoots sampled worst={worst}"
    assert gap2 <= 1.6 * worst, f"gap2={gap2} far above sampled worst={worst}"

python -m pytest "packages/omnibias-curvature/tests/test_sharpness.py::test_sam_gap_upper_bounds_sampled_worst_case" -q

Read the implementation

def sam_sharpness_gap(
    X: Array, Y: Array,
    W: Array, beta: Array, c: Array, b: Array,
    activation: str = "tanh",
    *,
    rho: float = 0.05,
) -> Array:
    r"""Exact second-order surrogate of the SAM inner-max gap.

    .. math::

        \max_{\lVert\varepsilon\rVert\le\rho} L(\theta+\varepsilon) - L(\theta)
          \;\approx\; \rho\,\lVert\nabla L\rVert
                 \;+\; \tfrac12\,\rho^2\,\max\!\big(\lambda_{\max}(\nabla^2 L),\,0\big),

    exact to :math:`O(\rho^3)`. This is what standard SAM *estimates*
    with one ascent step (linear term only); here both terms are computed
    from the closed-form gradient and Hessian in a single pass. The
    curvature term is clamped at 0 so negative curvature is never
    rewarded (the gap is :math:`\ge 0`).
    """
    H, g = mse_loss_hessian(X, Y, W, beta, c, b, activation)
    gnorm = jnp.linalg.norm(g)
    lam_max = hessian_top_eigenvalue(H)
    return rho * gnorm + 0.5 * rho * rho * jnp.maximum(lam_max, 0.0)

Related APIsomnibias.curvature.hessian_top_eigenvalue↗omnibias.curvature.mse_loss_hessian↗

Other public import paths

• omnibias.curvature.sam_sharpness_gap
• omnibias.curvature.sharpness.sam_sharpness_gap

FUNCTION

sharpness_aware_loss

omnibias.curvature.sharpness_aware_lossSource ↗

sharpness_aware_loss(X: Array, Y: Array, W: Array, beta: Array, c: Array, b: Array, activation: str='tanh', *, lam: float=0.01, measure: str='frobenius') -> Array

Curvature-regularised training loss L + \lambda\,\mathcal S(\nabla^2 L).

A drop-in replacement for the plain MSE loss in a jax.grad / optimiser loop. Because \mathcal S is built from the exact loss Hessian, differentiating this objective pulls in \sigma''' **in closed form** (autodiff of the closed-form \sigma''), so the penalty gradient is exact -- no finite-difference sharpness estimate, no extra forward/backward pass.

Parameters

Name	Type & default	Contract
Xpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
Ypositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
Wpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
betapositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
cpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
bpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
activationpositional-or-keyword	strDefault: 'tanh'	See the operation contract and implementation below.
lamkeyword-only	floatDefault: 0.01	See the operation contract and implementation below.
measurekeyword-only	strDefault: 'frobenius'	See the operation contract and implementation below.

ReturnsArray

Executed regression · sharpness aware training finds flatter minimum

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_sharpness_aware_training_finds_flatter_minimum():
    """From the *same* init and data, curvature-regularised training reaches a
    minimum with substantially smaller top-eigenvalue curvature than plain MSE
    -- while still fitting the (realisable) data and generalising at least as
    well. This is the SAM promise, delivered via exact curvature."""
    D, H_teacher, H_student, n_train = 3, 4, 8, 30
    Xtr, Ytr, Xte, Yte = _teacher_data(seed=17, D=D, H_teacher=H_teacher,
                                       n_train=n_train, noise=0.1)
    H = H_student
    theta0 = _init_student(seed=17, D=D, H=H)

    def plain_g(th):
        return jax.grad(
            lambda t: _loss_of_theta(t, mse_loss, Xtr, Ytr, H, D, "tanh")
        )(th)

    def flat_g(th):
        return jax.grad(
            lambda t: _loss_of_theta(t, sharpness_aware_loss, Xtr, Ytr, H, D,
                                     "tanh", lam=5e-3, measure="frobenius")
        )(th)

    theta_plain = _adam(jax.jit(plain_g), theta0, steps=500, lr=5e-3)
    theta_flat = _adam(jax.jit(flat_g), theta0, steps=500, lr=5e-3)

    def report(theta):
        bb, cc, be, WW = unpack_params(theta, H=H, D=D)
        Hh, _ = mse_loss_hessian(Xtr, Ytr, WW, be, cc, bb, "tanh")
        return (
            float(mse_loss(Xtr, Ytr, WW, be, cc, bb, "tanh")),
            float(mse_loss(Xte, Yte, WW, be, cc, bb, "tanh")),
            float(hessian_top_eigenvalue(Hh)),
        )

    tr_p, te_p, curv_p = report(theta_plain)
    tr_f, te_f, curv_f = report(theta_flat)

    # (1) both fit the training data
    assert tr_p < 0.05 and tr_f < 0.05, (tr_p, tr_f)
    # (2) the flat run lands in a clearly flatter basin
    assert curv_f < 0.7 * curv_p, f"curv flat={curv_f:.3f} vs plain={curv_p:.3f}"
    # (3) and generalises at least as well (here: better)
    assert te_f <= te_p * 1.05, f"test mse flat={te_f:.4f} vs plain={te_p:.4f}"

python -m pytest "packages/omnibias-curvature/tests/test_sharpness.py::test_sharpness_aware_training_finds_flatter_minimum" -q

Read the implementation

def sharpness_aware_loss(
    X: Array, Y: Array,
    W: Array, beta: Array, c: Array, b: Array,
    activation: str = "tanh",
    *,
    lam: float = 1e-2,
    measure: str = "frobenius",
) -> Array:
    r"""Curvature-regularised training loss :math:`L + \lambda\,\mathcal S(\nabla^2 L)`.

    A drop-in replacement for the plain MSE loss in a ``jax.grad`` /
    optimiser loop. Because :math:`\mathcal S` is built from the exact
    loss Hessian, differentiating this objective pulls in
    :math:`\sigma'''` **in closed form** (autodiff of the closed-form
    :math:`\sigma''`), so the penalty gradient is exact -- no
    finite-difference sharpness estimate, no extra forward/backward pass.
    """
    H, _ = mse_loss_hessian(X, Y, W, beta, c, b, activation)
    base = mse_loss(X, Y, W, beta, c, b, activation)
    return base + lam * _apply_measure(measure, H)

Related APIsomnibias.curvature.mse_loss↗omnibias.curvature.mse_loss_hessian↗

Other public import paths

• omnibias.curvature.sharpness_aware_loss
• omnibias.curvature.sharpness.sharpness_aware_loss

← Package overviewRunnable learning examples →