API REFERENCE/omnibias-curvature

omnibias.curvature.one_layer

0.1.0a2AlphaPython >=3.10AGPL / commercial

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

import omnibias.curvature.one_layer

Module contract & mathematical conventions

Closed-form parameter Gradient / Hessian / Fisher for one-layer fields.

Symbol map

================== ============ ================================================= Tensor Shape Meaning ================== ============ ================================================= x (B, D) input batch W (H, D) hidden-layer weights beta (H,) hidden-layer bias c (H,) output-layer weights b () output-layer bias ================== ============ =================================================

Parameter vector layout (a flat (P,) array of length P = 1 + H + H + H*D):

• index [0] : b
• indices [1 .. 1 + H) : c
• indices [1 + H .. 1 + 2*H) : beta
• indices [1 + 2*H .. 1 + 2*H + H*D) : W flattened row-major

The flat layout matches jax.tree_util.tree_leaves order on a (b, c, beta, W) tuple of (scalar, (H,), (H,), (H, D)) arrays, so consumers can call pack_params / unpack_params to interconvert with their own pytree.

Per-sample gradient

For z_n = W \cdot x_n + \beta:

• \partial f / \partial b = 1
• \partial f / \partial c_h = \sigma(z_{n,h})
• \partial f / \partial \beta_h = c_h\,\sigma'(z_{n,h})
• \partial f / \partial W_{h,j} = c_h\,\sigma'(z_{n,h})\,x_{n,j}

Per-sample Hessian

Block-diagonal in h (parameters for different hidden units do not interact). Non-zero entries within a hidden-unit block:

• \partial^2 f / \partial c_h\,\partial \beta_h = \sigma'(z_{n,h})
• \partial^2 f / \partial c_h\,\partial W_{h,j} = \sigma'(z_{n,h})\,x_{n,j}
• \partial^2 f / \partial \beta_h^2 = c_h\,\sigma''(z_{n,h})
• \partial^2 f / \partial \beta_h\,\partial W_{h,j} = c_h\,\sigma''(z_{n,h})\,x_{n,j}
• \partial^2 f / \partial W_{h,j}\,\partial W_{h,j'} = c_h\,\sigma''(z_{n,h})\,x_{n,j}\,x_{n,j'}

All others (including all interactions with b) are zero.

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

kfac_kron_factors

omnibias.curvature.kfac_kron_factorsSource ↗

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

Closed-form KFAC (A, G) factors for the hidden W block.

For the linear layer z_n = W \cdot x_n + \beta with downstream output f_n = b + \sigma(z_n) \cdot c, the KFAC approximation to the parameter Fisher of W reads F_W \approx A \otimes G with

• A = \tfrac{1}{B}\sum_n x_n x_n^T \in \mathbb R^{D \times D} — the input second moment.
• G = \tfrac{1}{B}\sum_n (c \odot \sigma'(z_n))(c \odot \sigma'(z_n))^T \in \mathbb R^{H \times H} — the pre-activation gradient covariance.

Closed-form via \sigma'(z_n) from the omnibias activation fast-path: no autograd backward needed. This is the building block that lifts the omnibias Hessian primitive into a KFAC-style preconditioner usable with a FermiNet trainer.

Notes

The Kronecker assumption F_W \approx A \otimes G is an approximation; the closed-form bit here is the *exact* construction of G (via \sigma'(z)), not the Kronecker factorisation itself. The Kronecker assumption typically incurs a few-percent approximation error per layer (Martens & Grosse 2015) but is the standard KFAC trade-off.

Parameters

Name	Type & default	Contract
Xpositional-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]

Explicit validation & failure cases

• When spec.fastpath is None, raises ValueError(f'activation {activation!r} has no fast-path kernel.').

These are explicit source guards. Backend operations may also raise for invalid shapes, dtypes, or unsupported devices.

Executed regression · kfac factors shape and psd

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_kfac_factors_shape_and_psd():
    """KFAC ``(A, G)`` factors are symmetric PSD; ``A ⊗ G`` matches the
    ``W``-block of the full Gauss-Newton Fisher to a Kronecker-factor
    approximation tolerance."""
    W, beta, c, b, _ = _rand_one_layer(D=6, H=4, seed=7)
    rng = np.random.default_rng(7)
    X = jnp.asarray(rng.normal(scale=0.4, size=(40, W.shape[1])))

    A, G = kfac_kron_factors(X, W, beta, c, b, activation="tanh")
    D, H = W.shape[1], W.shape[0]
    assert A.shape == (D, D)
    assert G.shape == (H, H)
    # Both symmetric & PSD.
    for M, name in [(A, "A"), (G, "G")]:
        np.testing.assert_allclose(
            np.asarray(M), np.asarray(M.T),
            rtol=1e-12, atol=1e-12,
            err_msg=f"{name} should be symmetric",
        )
        eigs = np.linalg.eigvalsh(np.asarray(M))
        assert eigs.min() >= -1e-12, (
            f"{name} not PSD: min eig = {eigs.min():.3e}"
        )

python -m pytest "packages/omnibias-curvature/tests/test_one_layer.py::test_kfac_factors_shape_and_psd" -q

Read the implementation

def kfac_kron_factors(
    X: Array,
    W: Array, beta: Array, c: Array, b: Array,
    activation: str = "tanh",
) -> tuple[Array, Array]:
    r"""Closed-form KFAC :math:`(A, G)` factors for the hidden ``W`` block.

    For the linear layer :math:`z_n = W \cdot x_n + \beta` with downstream
    output :math:`f_n = b + \sigma(z_n) \cdot c`, the KFAC approximation
    to the parameter Fisher of ``W`` reads
    :math:`F_W \approx A \otimes G` with

    * :math:`A = \tfrac{1}{B}\sum_n x_n x_n^T \in \mathbb R^{D \times D}`
      — the input second moment.
    * :math:`G = \tfrac{1}{B}\sum_n (c \odot \sigma'(z_n))(c \odot \sigma'(z_n))^T \in \mathbb R^{H \times H}`
      — the pre-activation gradient covariance.

    Closed-form via :math:`\sigma'(z_n)` from the omnibias activation
    fast-path: no autograd backward needed. This is the building block
    that lifts the omnibias Hessian primitive into a KFAC-style
    preconditioner usable with a FermiNet trainer.

    Notes
    -----
    The Kronecker assumption :math:`F_W \approx A \otimes G` is an
    approximation; the closed-form bit here is the *exact* construction
    of :math:`G` (via :math:`\sigma'(z)`), not the Kronecker
    factorisation itself. The Kronecker assumption typically incurs
    a few-percent approximation error per layer (Martens & Grosse 2015)
    but is the standard KFAC trade-off.
    """
    spec = get_activation(activation)
    if spec.fastpath is None:
        raise ValueError(
            f"activation {activation!r} has no fast-path kernel."
        )
    del b
    B = X.shape[0]
    Z = X @ W.T + beta[None, :]                # (B, H)
    sigma_p = spec.fastpath(Z, 1)              # (B, H)
    grad_pre = sigma_p * c[None, :]            # (B, H) (c ⊙ σ'(z_n))
    A = (X.T @ X) / B                          # (D, D)
    G = (grad_pre.T @ grad_pre) / B            # (H, H)
    return A, G

Related APIsomnibias.jax.get_activation↗

Other public import paths

• omnibias.curvature.kfac_kron_factors
• omnibias.curvature.one_layer.kfac_kron_factors

FUNCTION

mse_gauss_newton_fisher

omnibias.curvature.mse_gauss_newton_fisherSource ↗

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

Gauss-Newton Fisher and gradient for MSE loss L = \tfrac{1}{B}\sum_n (f(x_n) - y_n)^2.

Returns

F : (P, P) F = \tfrac{2}{B}\sum_n g_n g_n^T — the positive-definite Gauss-Newton approximation to the loss Hessian. Equivalent to the Fisher information of a Gaussian observation model with unit variance. g_loss : (P,) Gradient of the loss \nabla_\theta L = \tfrac{2}{B}\sum_n r_n g_n, where r_n = f_\theta(x_n) - y_n.

For large B, F is (P, P) so memory is O(P^2). Use kfac_kron_factors for the Kronecker-factored variant.

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 · gauss newton fisher is psd

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_gauss_newton_fisher_is_psd():
    """``F = (2/B) g g^T`` must be positive-semidefinite for every batch."""
    W, beta, c, b, _ = _rand_one_layer(seed=11)
    rng = np.random.default_rng(11)
    X = jnp.asarray(rng.normal(scale=0.5, size=(50, W.shape[1])))
    Y = jnp.asarray(rng.normal(scale=1.0, size=(50,)))
    F, _ = mse_gauss_newton_fisher(X, Y, W, beta, c, b, activation="tanh")
    eigs = np.linalg.eigvalsh(np.asarray(F))
    # F is sum of outer products → PSD up to round-off.
    assert eigs.min() >= -1e-10, f"min eig = {eigs.min():.3e} (should be ≥ 0)"

python -m pytest "packages/omnibias-curvature/tests/test_one_layer.py::test_gauss_newton_fisher_is_psd" -q

Regression example · gauss newton is the psd part of full 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_gauss_newton_is_the_psd_part_of_full_hessian():
    """At a residual-free (interpolating) fit the full Hessian == GN Fisher."""
    from omnibias.curvature.one_layer import mse_gauss_newton_fisher
    from omnibias.curvature.sharpness import _batch_forward
    W, beta, c, b, X, _ = _rand_problem(seed=7)
    # Targets == model outputs -> residuals are zero -> residual term vanishes.
    Y = _batch_forward(X, W, beta, c, b, "tanh")
    H_full, _ = mse_loss_hessian(X, Y, W, beta, c, b, "tanh")
    F_gn, _ = mse_gauss_newton_fisher(X, Y, W, beta, c, b, "tanh")
    np.testing.assert_allclose(np.asarray(H_full), np.asarray(F_gn), atol=1e-10)

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

Read the implementation

def mse_gauss_newton_fisher(
    X: Array, Y: Array,
    W: Array, beta: Array, c: Array, b: Array,
    activation: str = "tanh",
) -> tuple[Array, Array]:
    r"""Gauss-Newton Fisher and gradient for MSE loss
    :math:`L = \tfrac{1}{B}\sum_n (f(x_n) - y_n)^2`.

    Returns
    -------
    F : ``(P, P)``
        :math:`F = \tfrac{2}{B}\sum_n g_n g_n^T` — the positive-definite
        Gauss-Newton approximation to the loss Hessian. Equivalent to
        the Fisher information of a Gaussian observation model with
        unit variance.
    g_loss : ``(P,)``
        Gradient of the loss :math:`\nabla_\theta L = \tfrac{2}{B}\sum_n r_n g_n`,
        where :math:`r_n = f_\theta(x_n) - y_n`.

    For large ``B``, ``F`` is :math:`(P, P)` so memory is :math:`O(P^2)`.
    Use :func:`kfac_kron_factors` for the Kronecker-factored variant.
    """
    spec = get_activation(activation)
    B = X.shape[0]

    def per_sample(x_n):
        sigma_z = spec.forward(W @ x_n + beta)
        f = b + sigma_z @ c
        g = one_layer_param_grad(x_n, W, beta, c, b, activation)
        return f, g

    f_vals, gs = jax.vmap(per_sample)(X)            # (B,), (B, P)
    r = f_vals - Y                                  # (B,)
    F = (2.0 / B) * (gs.T @ gs)                     # (P, P)
    g_loss = (2.0 / B) * (gs.T @ r)                 # (P,)
    return F, g_loss

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

Other public import paths

• omnibias.curvature.mse_gauss_newton_fisher
• omnibias.curvature.one_layer.mse_gauss_newton_fisher

FUNCTION

mse_newton_step

omnibias.curvature.mse_newton_stepSource ↗

mse_newton_step(X: Array, Y: Array, W: Array, beta: Array, c: Array, b: Array, activation: str='tanh', learning_rate: float=1.0, damping: float=0.0001) -> tuple[Array, Array, Array, Array]

One Gauss-Newton step on the MSE loss.

Updates (b, c, beta, W) via

\theta \;\leftarrow\; \theta - \eta\,(F + \lambda I)^{-1}\,\nabla L

with the closed-form Gauss-Newton Fisher F from mse_gauss_newton_fisher.

Returns

(b_new, c_new, beta_new, W_new) : same shapes as the inputs.

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.
learning_ratepositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.
dampingpositional-or-keyword	floatDefault: 0.0001	See the operation contract and implementation below.

Returnstuple[Array, Array, Array, Array]

Executed regression · gaussian step equals mse newton at zero damping

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_gaussian_step_equals_mse_newton_at_zero_damping() -> None:
    X, W, beta, c, b, rng = _problem(D=3, H=3, B=60, seed=5)
    Y = jnp.asarray(rng.normal(size=X.shape[0]))
    g_b, g_c, g_be, g_W = glm_natural_gradient_step(
        X, Y, W, beta, c, b, family="gaussian", learning_rate=1.0, damping=0.0
    )
    n_b, n_c, n_be, n_W = mse_newton_step(X, Y, W, beta, c, b, learning_rate=1.0, damping=0.0)
    assert jnp.allclose(g_b, n_b, atol=1e-8)
    assert jnp.allclose(g_c, n_c, atol=1e-8)
    assert jnp.allclose(g_be, n_be, atol=1e-8)
    assert jnp.allclose(g_W, n_W, atol=1e-8)

python -m pytest "packages/omnibias-curvature/tests/test_natural_gradient.py::test_gaussian_step_equals_mse_newton_at_zero_damping" -q

Regression example · gauss newton reduces loss monotonically

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_gauss_newton_reduces_loss_monotonically():
    """Gauss-Newton with appropriate damping must reduce MSE loss
    monotonically over several iterations.

    The one-layer field is bilinear in ``(W, c)`` so Newton is *not*
    one-shot exact, but the loss trajectory must be monotonically
    non-increasing (any sign of divergence indicates a bug in the
    Fisher or the gradient).
    """
    rng = np.random.default_rng(0)
    D, H, B = 4, 3, 60
    activation = "tanh"

    # True parameters (small scale to stay in the well-behaved regime).
    W_star = jnp.asarray(rng.normal(scale=0.3, size=(H, D)))
    beta_star = jnp.asarray(rng.normal(scale=0.1, size=(H,)))
    c_star = jnp.asarray(rng.normal(scale=0.4, size=(H,)))
    b_star = jnp.asarray(0.2)

    X = jnp.asarray(rng.normal(scale=0.4, size=(B, D)))
    from omnibias.jax.activations import get_activation
    sigma_star = get_activation(activation).forward(
        X @ W_star.T + beta_star[None, :]
    )
    Y = b_star + sigma_star @ c_star  # zero observation noise

    # Init at the true params + perturbation (so Fisher is well-conditioned).
    perturb_scale = 0.2
    W = W_star + jnp.asarray(rng.normal(scale=perturb_scale, size=W_star.shape))
    beta = beta_star + jnp.asarray(
        rng.normal(scale=perturb_scale, size=beta_star.shape)
    )
    c = c_star + jnp.asarray(rng.normal(scale=perturb_scale, size=c_star.shape))
    b = b_star + jnp.asarray(rng.normal(scale=perturb_scale))

    def loss(W_, beta_, c_, b_):
        sigma_ = get_activation(activation).forward(X @ W_.T + beta_[None, :])
        return float(jnp.mean((b_ + sigma_ @ c_ - Y) ** 2))

    losses = [loss(W, beta, c, b)]
    for _step in range(8):
        b, c, beta, W = mse_newton_step(
            X, Y, W, beta, c, b, activation=activation,
            learning_rate=0.5, damping=1e-3,
        )
        losses.append(loss(W, beta, c, b))

    # Loss must decrease monotonically; a *jump up* is a sign of bug.
    for i, (l_prev, l_now) in enumerate(zip(losses, losses[1:], strict=False)):
        assert l_now <= l_prev * 1.01, (
            f"Gauss-Newton increased loss at step {i}: "
            f"{l_prev:.3e} → {l_now:.3e} (trajectory: {losses})"
        )

    # Final loss should be much lower than initial.
    assert losses[-1] < 0.1 * losses[0], (
        f"Gauss-Newton didn't converge: loss[0]={losses[0]:.3e}, "
        f"loss[-1]={losses[-1]:.3e}"
    )

python -m pytest packages/omnibias-curvature/tests/test_one_layer.py -k test_gauss_newton_reduces_loss_monotonically -q

Read the implementation

def mse_newton_step(
    X: Array, Y: Array,
    W: Array, beta: Array, c: Array, b: Array,
    activation: str = "tanh",
    learning_rate: float = 1.0,
    damping: float = 1e-4,
) -> tuple[Array, Array, Array, Array]:
    r"""One Gauss-Newton step on the MSE loss.

    Updates ``(b, c, beta, W)`` via

    .. math::

        \theta \;\leftarrow\; \theta - \eta\,(F + \lambda I)^{-1}\,\nabla L

    with the closed-form Gauss-Newton Fisher ``F`` from
    :func:`mse_gauss_newton_fisher`.

    Returns
    -------
    (b_new, c_new, beta_new, W_new) : same shapes as the inputs.
    """
    F, g_loss = mse_gauss_newton_fisher(X, Y, W, beta, c, b, activation)
    # Delegates to the shared eps -> 0 collapse kernel; bit-identical at this damping.
    delta = regularized_solve(F, g_loss, eps=damping)
    theta = pack_params(b, c, beta, W)
    theta_new = theta - learning_rate * delta
    return unpack_params(theta_new, H=W.shape[0], D=W.shape[1])

Related APIsomnibias.curvature.mse_gauss_newton_fisher↗omnibias.curvature.one_layer.pack_params↗omnibias.curvature.one_layer.unpack_params↗omnibias.curvature.regularized_solve↗

Other public import paths

• omnibias.curvature.mse_newton_step
• omnibias.curvature.one_layer.mse_newton_step

FUNCTION

one_layer_param_grad

omnibias.curvature.one_layer_param_gradSource ↗

one_layer_param_grad(x: Array, W: Array, beta: Array, c: Array, b: Array, activation: str='tanh') -> Array

Flat per-sample gradient \nabla_\theta f(x) of length P.

See module docstring for the parameter vector layout.

Parameters

Name	Type & default	Contract
xpositional-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 · 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 · one layer param grad matches jax grad

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_one_layer_param_grad_matches_jax_grad(activation):
    """``one_layer_param_grad`` must equal ``jax.grad`` to ~1e-12."""
    W, beta, c, b, x = _rand_one_layer(seed=hash(activation) & 0xFFFF)
    H, D = W.shape

    g_closed = one_layer_param_grad(x, W, beta, c, b, activation)
    theta = pack_params(b, c, beta, W)
    g_jax = jax.grad(_f_unpack)(theta, x, H, D, activation)

    assert g_closed.shape == g_jax.shape == (theta.shape[0],)
    np.testing.assert_allclose(
        np.asarray(g_closed), np.asarray(g_jax),
        rtol=1e-10, atol=1e-12,
        err_msg=f"param grad mismatch for activation={activation!r}",
    )

python -m pytest packages/omnibias-curvature/tests/test_one_layer.py -k test_one_layer_param_grad_matches_jax_grad -q

Read the implementation

def one_layer_param_grad(
    x: Array,
    W: Array, beta: Array, c: Array, b: Array,
    activation: str = "tanh",
) -> Array:
    r"""Flat per-sample gradient :math:`\nabla_\theta f(x)` of length P.

    See module docstring for the parameter vector layout.
    """
    del b  # not used in gradient
    sigma_z, sigma_p, _ = _per_sample_blocks(x, W, beta, c, activation)
    grad_b = jnp.asarray(1.0)
    grad_c = sigma_z                                # (H,)
    grad_beta = c * sigma_p                         # (H,)
    grad_W = jnp.outer(c * sigma_p, x)              # (H, D)
    return jnp.concatenate([
        grad_b.reshape(-1),
        grad_c.reshape(-1),
        grad_beta.reshape(-1),
        grad_W.reshape(-1),
    ])

Other public import paths

• omnibias.curvature.one_layer_param_grad
• omnibias.curvature.one_layer.one_layer_param_grad

FUNCTION

one_layer_param_hessian

omnibias.curvature.one_layer_param_hessianSource ↗

one_layer_param_hessian(x: Array, W: Array, beta: Array, c: Array, b: Array, activation: str='tanh') -> Array

Full per-sample parameter Hessian \nabla^2_\theta f(x).

Returns a dense (P, P) matrix. Mostly zeros — the non-zero block structure is described in the module docstring. For large H and D, use the Gauss-Newton Fisher (mse_gauss_newton_fisher) or the Kronecker factors (kfac_kron_factors) instead of materialising the full Hessian.

The (residual-free) full Hessian is useful as a correctness oracle against jax.hessian, and as the building block for a second-order optimiser that does not approximate via Gauss-Newton.

Implementation note: H and D are static at JIT-trace time so the assembly uses a Python for h in range(H) loop that is unrolled into H scatter ops. This is fine for H ≤ 128; for larger one-layer fields prefer the vectorised mse_gauss_newton_fisher (no Hessian assembly) or kfac_kron_factors (closed-form Kronecker).

Parameters

Name	Type & default	Contract
xpositional-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 · 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 · one layer param 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_one_layer_param_hessian_matches_jax_hessian(activation):
    """``one_layer_param_hessian`` must equal ``jax.hessian`` to ~1e-10."""
    W, beta, c, b, x = _rand_one_layer(seed=hash(activation) & 0xFFFF)
    H, D = W.shape

    Hess_closed = one_layer_param_hessian(x, W, beta, c, b, activation)
    theta = pack_params(b, c, beta, W)
    Hess_jax = jax.hessian(_f_unpack)(theta, x, H, D, activation)

    assert Hess_closed.shape == Hess_jax.shape
    # The reference autodiff Hessian is exactly symmetric — verify the
    # closed-form one is too (a numerical asymmetry would mean we mis-mapped
    # an off-diagonal index above).
    np.testing.assert_allclose(
        np.asarray(Hess_closed), np.asarray(Hess_closed.T),
        rtol=1e-12, atol=1e-12,
        err_msg=f"closed-form Hessian not symmetric for {activation!r}",
    )
    np.testing.assert_allclose(
        np.asarray(Hess_closed), np.asarray(Hess_jax),
        rtol=1e-9, atol=1e-10,
        err_msg=f"param Hessian mismatch for activation={activation!r}",
    )

python -m pytest packages/omnibias-curvature/tests/test_one_layer.py -k test_one_layer_param_hessian_matches_jax_hessian -q

Read the implementation

def one_layer_param_hessian(
    x: Array,
    W: Array, beta: Array, c: Array, b: Array,
    activation: str = "tanh",
) -> Array:
    r"""Full per-sample parameter Hessian :math:`\nabla^2_\theta f(x)`.

    Returns a dense :math:`(P, P)` matrix. Mostly zeros — the non-zero
    block structure is described in the module docstring. For large
    :math:`H` and :math:`D`, use the Gauss-Newton Fisher
    (:func:`mse_gauss_newton_fisher`) or the Kronecker factors
    (:func:`kfac_kron_factors`) instead of materialising the full
    Hessian.

    The (residual-free) full Hessian is useful as a correctness oracle
    against :func:`jax.hessian`, and as the building block for a
    second-order optimiser that does not approximate via Gauss-Newton.

    Implementation note: ``H`` and ``D`` are static at JIT-trace time so
    the assembly uses a Python ``for h in range(H)`` loop that is
    unrolled into ``H`` scatter ops. This is fine for ``H ≤ 128``; for
    larger one-layer fields prefer the vectorised
    :func:`mse_gauss_newton_fisher` (no Hessian assembly) or
    :func:`kfac_kron_factors` (closed-form Kronecker).
    """
    _, sigma_p, sigma_pp = _per_sample_blocks(x, W, beta, c, activation)
    Hh = int(W.shape[0])
    Dd = int(W.shape[1])
    P = 1 + 2 * Hh + Hh * Dd

    Hess = jnp.zeros((P, P))

    # Index map (all concrete):
    #   b @ 0
    #   c @ [1, 1+H)
    #   β @ [1+H, 1+2H)
    #   W flattened row-major @ [1+2H, P)

    idx_c = jnp.arange(1, 1 + Hh)
    idx_beta = jnp.arange(1 + Hh, 1 + 2 * Hh)

    # c[h] ↔ β[h]: σ'(z_h)  (diagonal in h).
    Hess = Hess.at[idx_c, idx_beta].set(sigma_p)
    Hess = Hess.at[idx_beta, idx_c].set(sigma_p)

    # β[h] ↔ β[h]: c_h σ''(z_h).
    Hess = Hess.at[idx_beta, idx_beta].set(sigma_pp * c)

    # c[h] ↔ W[h, j] and β[h] ↔ W[h, j]: vectorised scatter.
    # W[h, j] flat index = 1 + 2H + h*D + j.
    idx_W_flat = jnp.arange(1 + 2 * Hh, P)                # (H*D,)
    idx_c_rep = jnp.repeat(idx_c, Dd)                     # (H*D,) — c[h] repeated D times
    idx_beta_rep = jnp.repeat(idx_beta, Dd)               # (H*D,) — β[h] repeated D times
    val_c_W = (sigma_p[:, None] * x[None, :]).reshape(-1)  # (H*D,)
    val_beta_W = ((sigma_pp * c)[:, None] * x[None, :]).reshape(-1)  # (H*D,)

    Hess = Hess.at[idx_c_rep, idx_W_flat].set(val_c_W)
    Hess = Hess.at[idx_W_flat, idx_c_rep].set(val_c_W)
    Hess = Hess.at[idx_beta_rep, idx_W_flat].set(val_beta_W)
    Hess = Hess.at[idx_W_flat, idx_beta_rep].set(val_beta_W)

    # W[h, :] ↔ W[h, :]: c_h σ''(z_h) x x^T, block-diagonal in h.
    # Loop over the static dimension H (Python-time unroll).
    sppc = sigma_pp * c                                   # (H,)
    xxT = jnp.outer(x, x)                                 # (D, D)
    for h in range(Hh):
        W_start = 1 + 2 * Hh + h * Dd
        Hess = jax.lax.dynamic_update_slice(
            Hess, sppc[h] * xxT,
            (W_start, W_start),
        )
    return Hess

Other public import paths

• omnibias.curvature.one_layer_param_hessian
• omnibias.curvature.one_layer.one_layer_param_hessian

FUNCTION

pack_params

omnibias.curvature.one_layer.pack_paramsSource ↗

pack_params(b: Array, c: Array, beta: Array, W: Array) -> Array

Pack (b, c, beta, W) into a flat parameter vector.

Order matches unpack_params and the gradient / Hessian routines in this module.

Parameters

Name	Type & default	Contract
bpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
cpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
betapositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
Wpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.

ReturnsArray

Executed regression · pack unpack roundtrip

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_pack_unpack_roundtrip():
    W, beta, c, b, x = _rand_one_layer(D=8, H=4, seed=42)
    theta = pack_params(b, c, beta, W)
    b2, c2, beta2, W2 = unpack_params(theta, H=4, D=8)
    np.testing.assert_array_equal(np.asarray(b), np.asarray(b2))
    np.testing.assert_array_equal(np.asarray(c), np.asarray(c2))
    np.testing.assert_array_equal(np.asarray(beta), np.asarray(beta2))
    np.testing.assert_array_equal(np.asarray(W), np.asarray(W2))

python -m pytest "packages/omnibias-curvature/tests/test_one_layer.py::test_pack_unpack_roundtrip" -q

Regression example · natural gradient step formula

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_natural_gradient_step_formula() -> None:
    rng = np.random.default_rng(1)
    p = 5
    a = rng.normal(size=(p, p))
    fisher = jnp.asarray(a @ a.T)
    grad = jnp.asarray(rng.normal(size=p))
    theta = jnp.asarray(rng.normal(size=p))
    new = natural_gradient_step(theta, grad, fisher, learning_rate=0.5, damping=1e-3)
    expected = theta - 0.5 * damped_solve(fisher, grad, damping=1e-3)
    assert jnp.allclose(new, expected, rtol=1e-12, atol=1e-12)

python -m pytest packages/omnibias-curvature/tests/test_natural_gradient.py -k test_natural_gradient_step_formula -q

Read the implementation

def pack_params(b: Array, c: Array, beta: Array, W: Array) -> Array:
    """Pack ``(b, c, beta, W)`` into a flat parameter vector.

    Order matches :func:`unpack_params` and the gradient / Hessian
    routines in this module.
    """
    return jnp.concatenate([
        jnp.asarray(b).reshape(-1),
        c.reshape(-1),
        beta.reshape(-1),
        W.reshape(-1),
    ])

FUNCTION

unpack_params

omnibias.curvature.one_layer.unpack_paramsSource ↗

unpack_params(theta: Array, H: int, D: int) -> tuple[Array, Array, Array, Array]

Inverse of pack_params for the given hidden / input sizes.

Parameters

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

Returnstuple[Array, Array, Array, Array]

Explicit validation & failure cases

• When theta.shape[-1] != 1 + 2 * H + H * D, raises ValueError(f'unpack_params: expected last dim = {1 + 2 * H + H * D} (= 1 + 2*H + H*D for H={H}, D={D}), got {theta.shape[-1]}').

These are explicit source guards. Backend operations may also raise for invalid shapes, dtypes, or unsupported devices.

Executed regression · pack unpack roundtrip

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_pack_unpack_roundtrip():
    W, beta, c, b, x = _rand_one_layer(D=8, H=4, seed=42)
    theta = pack_params(b, c, beta, W)
    b2, c2, beta2, W2 = unpack_params(theta, H=4, D=8)
    np.testing.assert_array_equal(np.asarray(b), np.asarray(b2))
    np.testing.assert_array_equal(np.asarray(c), np.asarray(c2))
    np.testing.assert_array_equal(np.asarray(beta), np.asarray(beta2))
    np.testing.assert_array_equal(np.asarray(W), np.asarray(W2))

python -m pytest "packages/omnibias-curvature/tests/test_one_layer.py::test_pack_unpack_roundtrip" -q

Regression example · glm loss gradient matches autodiff

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_glm_loss_gradient_matches_autodiff(family: str) -> None:
    X, W, beta, c, b, rng = _problem(seed=3)
    spec = get_activation("tanh")
    H, D = W.shape
    Y = jnp.asarray(rng.normal(size=X.shape[0]) ** 2 if family == "poisson" else rng.normal(size=X.shape[0]))
    A = _A_OF[family]

    def nll(theta):
        bb, cc, bb2, WW = unpack_params(theta, H=H, D=D)
        eta = bb + spec.forward(X @ WW.T + bb2[None, :]) @ cc
        return jnp.mean(A(eta) - Y * eta)

    theta = pack_params(b, c, beta, W)
    auto = jax.grad(nll)(theta)
    closed = glm_loss_gradient(X, Y, W, beta, c, b, activation="tanh", family=family)
    assert jnp.allclose(auto, closed, rtol=1e-9, atol=1e-9)

python -m pytest packages/omnibias-curvature/tests/test_natural_gradient.py -k test_glm_loss_gradient_matches_autodiff -q

Read the implementation

def unpack_params(
    theta: Array, H: int, D: int,
) -> tuple[Array, Array, Array, Array]:
    """Inverse of :func:`pack_params` for the given hidden / input sizes."""
    if theta.shape[-1] != 1 + 2 * H + H * D:
        raise ValueError(
            f"unpack_params: expected last dim = {1 + 2*H + H*D} (= "
            f"1 + 2*H + H*D for H={H}, D={D}), got {theta.shape[-1]}"
        )
    b = theta[..., 0]
    c = theta[..., 1:1 + H]
    beta = theta[..., 1 + H:1 + 2 * H]
    W = theta[..., 1 + 2 * H:].reshape(theta.shape[:-1] + (H, D))
    return b, c, beta, W

← Package overviewRunnable learning examples →