API REFERENCE/omnibias-curvature

omnibias.curvature.natural_gradient

0.1.0a2AlphaPython >=3.10AGPL / commercial

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

import omnibias.curvature.natural_gradient

Module contract & mathematical conventions

Natural-gradient (Fisher-scoring) optimisation on the closed-form GLM Fisher.

The natural gradient preconditions the loss gradient by the inverse Fisher information, \tilde\nabla = F^{-1}\nabla L, taking steps that are invariant to reparametrisation and -- for an exponential family with the canonical link -- coincide with Fisher scoring / IRLS and Newton's method.

Two generic helpers do the linear algebra on any (P, P) Fisher and (P,) gradient:

• damped_solve -- the regularised solve (F + lambda I)^{-1} grad.
• natural_gradient_step -- the parameter update theta - lr (F+lambda I)^{-1} grad.

and one driver wires them to the **closed-form one-layer GLM Fisher** (omnibias.curvature.glm_fisher.glm_fisher):

• glm_natural_gradient_step -- one Fisher-scoring step on a one-layer field eta = b + sigma(W x + beta) . c under a Bernoulli / Poisson / Gaussian observation model. The Fisher reuses the exact per-sample gradient and the GLM variance A''(eta); the loss gradient reuses the GLM mean A'(eta), so no autodiff backward pass is needed. For the Gaussian family this reduces (at zero damping) to the Gauss-Newton omnibias.curvature.one_layer.mse_newton_step.

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

damped_solve

omnibias.curvature.damped_solveSource ↗

damped_solve(fisher: Array, grad: Array, *, damping: float=0.001) -> Array

Natural-gradient direction delta = (F + damping I)^{-1} grad.

Solves the (Tikhonov-)regularised linear system for the preconditioned gradient. fisher is a (P, P) positive-semidefinite matrix and grad a (P,) vector; damping >= 0 keeps the system well-posed when F is singular or ill-conditioned.

Parameters

Name	Type & default	Contract
fisherpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
gradpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
dampingkeyword-only	floatDefault: 0.001	See the operation contract and implementation below.

ReturnsArray

Explicit validation & failure cases

• When fisher.ndim != 2 or fisher.shape[0] != fisher.shape[1], raises ValueError(f'fisher must be a square (P, P) matrix, got {tuple(fisher.shape)}').
• When grad.ndim != 1 or grad.shape[0] != fisher.shape[0], raises ValueError(f'grad must be (P,) with P = {fisher.shape[0]}, got {tuple(grad.shape)}').
• When damping < 0.0, raises ValueError(f'damping must be >= 0, got {damping}').

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

Executed regression · damped solve guards

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_damped_solve_guards() -> None:
    with pytest.raises(ValueError, match="square"):
        damped_solve(jnp.zeros((2, 3)), jnp.zeros(2))
    with pytest.raises(ValueError, match="grad must be"):
        damped_solve(jnp.eye(3), jnp.zeros(2))
    with pytest.raises(ValueError, match="damping must be"):
        damped_solve(jnp.eye(3), jnp.zeros(3), damping=-1.0)

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

Regression example · damped solve satisfies normal equations

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_damped_solve_satisfies_normal_equations() -> None:
    rng = np.random.default_rng(0)
    p = 7
    a = rng.normal(size=(p, p))
    fisher = jnp.asarray(a @ a.T)  # SPD
    grad = jnp.asarray(rng.normal(size=p))
    lam = 1e-2
    delta = damped_solve(fisher, grad, damping=lam)
    resid = (fisher + lam * jnp.eye(p)) @ delta - grad
    assert float(jnp.max(jnp.abs(resid))) < 1e-9

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

Read the implementation

def damped_solve(fisher: Array, grad: Array, *, damping: float = 1e-3) -> Array:
    r"""Natural-gradient direction ``delta = (F + damping I)^{-1} grad``.

    Solves the (Tikhonov-)regularised linear system for the preconditioned
    gradient. ``fisher`` is a ``(P, P)`` positive-semidefinite matrix and ``grad``
    a ``(P,)`` vector; ``damping >= 0`` keeps the system well-posed when ``F`` is
    singular or ill-conditioned.
    """
    if fisher.ndim != 2 or fisher.shape[0] != fisher.shape[1]:
        raise ValueError(
            f"fisher must be a square (P, P) matrix, got {tuple(fisher.shape)}"
        )
    if grad.ndim != 1 or grad.shape[0] != fisher.shape[0]:
        raise ValueError(
            f"grad must be (P,) with P = {fisher.shape[0]}, got {tuple(grad.shape)}"
        )
    if damping < 0.0:
        raise ValueError(f"damping must be >= 0, got {damping}")
    # Delegates to the shared eps -> 0 collapse kernel; bit-identical at this damping.
    delta: Array = regularized_solve(fisher, grad, eps=damping)
    return delta

Related APIsomnibias.curvature.regularized_solve↗

Other public import paths

• omnibias.curvature.damped_solve
• omnibias.curvature.natural_gradient.damped_solve

FUNCTION

glm_loss_gradient

omnibias.curvature.glm_loss_gradientSource ↗

glm_loss_gradient(X: Array, Y: Array, W: Array, beta: Array, c: Array, b: Array, *, activation: str='tanh', family: str='bernoulli') -> Array

Flat (P,) gradient of the GLM negative log-likelihood.

For the one-layer natural parameter eta_n = b + sigma(W x_n + beta) . c the exponential-family NLL gradient is (1/B) sum_n (A'(eta_n) - y_n) g_n with A' the GLM mean (link inverse) and g_n the closed-form per-sample parameter gradient. The residual A'(eta_n) - y_n is the canonical GLM working residual; pair this with omnibias.curvature.glm_fisher.glm_fisher for a Fisher-scoring step.

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.
activationkeyword-only	strDefault: 'tanh'	See the operation contract and implementation below.
familykeyword-only	strDefault: 'bernoulli'	See the operation contract and implementation below.

ReturnsArray

Explicit validation & failure cases

• When family not in _FAMILY_LOG_PARTITION, raises ValueError(f'unknown GLM family {family!r}; choose from {sorted(_FAMILY_LOG_PARTITION)}').

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

Executed regression · glm loss gradient rejects unknown family

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_glm_loss_gradient_rejects_unknown_family() -> None:
    X, W, beta, c, b, _ = _problem()
    Y = jnp.zeros(X.shape[0])
    with pytest.raises(ValueError, match="unknown GLM family"):
        glm_loss_gradient(X, Y, W, beta, c, b, family="nope")

python -m pytest "packages/omnibias-curvature/tests/test_natural_gradient.py::test_glm_loss_gradient_rejects_unknown_family" -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 glm_loss_gradient(
    X: Array,
    Y: Array,
    W: Array,
    beta: Array,
    c: Array,
    b: Array,
    *,
    activation: str = "tanh",
    family: str = "bernoulli",
) -> Array:
    r"""Flat ``(P,)`` gradient of the GLM negative log-likelihood.

    For the one-layer natural parameter ``eta_n = b + sigma(W x_n + beta) . c`` the
    exponential-family NLL gradient is ``(1/B) sum_n (A'(eta_n) - y_n) g_n`` with
    ``A'`` the GLM mean (link inverse) and ``g_n`` the closed-form per-sample
    parameter gradient. The residual ``A'(eta_n) - y_n`` is the canonical GLM
    working residual; pair this with :func:`omnibias.curvature.glm_fisher.glm_fisher`
    for a Fisher-scoring step.
    """
    if family not in _FAMILY_LOG_PARTITION:
        raise ValueError(
            f"unknown GLM family {family!r}; "
            f"choose from {sorted(_FAMILY_LOG_PARTITION)}"
        )
    spec = get_activation(activation)
    base = _FAMILY_LOG_PARTITION[family]
    n_batch = X.shape[0]

    def per_sample(x_n: Array) -> tuple[Array, Array]:
        eta = b + spec.forward(W @ x_n + beta) @ c
        g = one_layer_param_grad(x_n, W, beta, c, b, activation)
        return eta, g

    etas, gs = jax.vmap(per_sample)(X)  # (B,), (B, P)
    mu = etas if base is None else _log_partition_derivative(etas, base=base, order=1)
    g_loss: Array = (gs.T @ (mu - Y)) / n_batch
    return g_loss

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

Other public import paths

• omnibias.curvature.glm_loss_gradient
• omnibias.curvature.natural_gradient.glm_loss_gradient

FUNCTION

glm_natural_gradient_step

omnibias.curvature.glm_natural_gradient_stepSource ↗

glm_natural_gradient_step(X: Array, Y: Array, W: Array, beta: Array, c: Array, b: Array, *, activation: str='tanh', family: str='bernoulli', learning_rate: float=1.0, damping: float=0.001) -> tuple[Array, Array, Array, Array]

One natural-gradient (Fisher-scoring / IRLS) step on a one-layer GLM field.

Updates (b, c, beta, W) via theta <- theta - lr (F + damping I)^{-1} g with the closed-form GLM Fisher F (omnibias.curvature.glm_fisher.glm_fisher) and the GLM NLL gradient g (glm_loss_gradient). For family="gaussian" this is Gauss-Newton and, at damping = 0, equals omnibias.curvature.one_layer.mse_newton_step.

Returns (b_new, c_new, beta_new, W_new) with the input shapes.

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.
activationkeyword-only	strDefault: 'tanh'	See the operation contract and implementation below.
familykeyword-only	strDefault: 'bernoulli'	See the operation contract and implementation below.
learning_ratekeyword-only	floatDefault: 1.0	See the operation contract and implementation below.
dampingkeyword-only	floatDefault: 0.001	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 · gaussian realizable quadratic convergence

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_realizable_quadratic_convergence() -> None:
    X, Wt, bt, ct, b0t, rng = _problem(D=2, H=3, B=80, seed=9)
    eta = _field_eta(X, Wt, bt, ct, b0t)  # teacher; zero-residual target
    W = Wt + 0.1 * jnp.asarray(rng.normal(size=Wt.shape))
    beta = bt + 0.1 * jnp.asarray(rng.normal(size=bt.shape))
    c = ct + 0.1 * jnp.asarray(rng.normal(size=ct.shape))
    b = b0t + 0.1
    g0 = float(jnp.linalg.norm(glm_loss_gradient(X, eta, W, beta, c, b, family="gaussian")))
    for _ in range(15):
        b, c, beta, W = glm_natural_gradient_step(
            X, eta, W, beta, c, b, family="gaussian", learning_rate=1.0, damping=1e-10
        )
    gN = float(jnp.linalg.norm(glm_loss_gradient(X, eta, W, beta, c, b, family="gaussian")))
    assert gN < g0 * 1e-4

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

Read the implementation

def glm_natural_gradient_step(
    X: Array,
    Y: Array,
    W: Array,
    beta: Array,
    c: Array,
    b: Array,
    *,
    activation: str = "tanh",
    family: str = "bernoulli",
    learning_rate: float = 1.0,
    damping: float = 1e-3,
) -> tuple[Array, Array, Array, Array]:
    r"""One natural-gradient (Fisher-scoring / IRLS) step on a one-layer GLM field.

    Updates ``(b, c, beta, W)`` via
    ``theta <- theta - lr (F + damping I)^{-1} g`` with the closed-form GLM Fisher
    ``F`` (:func:`omnibias.curvature.glm_fisher.glm_fisher`) and the GLM NLL
    gradient ``g`` (:func:`glm_loss_gradient`). For ``family="gaussian"`` this is
    Gauss-Newton and, at ``damping = 0``, equals
    :func:`omnibias.curvature.one_layer.mse_newton_step`.

    Returns ``(b_new, c_new, beta_new, W_new)`` with the input shapes.
    """
    g_loss = glm_loss_gradient(
        X, Y, W, beta, c, b, activation=activation, family=family
    )
    fisher = glm_fisher(X, W, beta, c, b, activation=activation, family=family)
    theta = pack_params(b, c, beta, W)
    theta_new = natural_gradient_step(
        theta, g_loss, fisher, learning_rate=learning_rate, damping=damping
    )
    return unpack_params(theta_new, H=W.shape[0], D=W.shape[1])

Related APIsomnibias.curvature.glm_fisher.glm_fisher↗omnibias.curvature.glm_loss_gradient↗omnibias.curvature.natural_gradient_step↗omnibias.curvature.one_layer.pack_params↗omnibias.curvature.one_layer.unpack_params↗

Other public import paths

• omnibias.curvature.glm_natural_gradient_step
• omnibias.curvature.natural_gradient.glm_natural_gradient_step

FUNCTION

natural_gradient_step

omnibias.curvature.natural_gradient_stepSource ↗

natural_gradient_step(theta: Array, grad: Array, fisher: Array, *, learning_rate: float=1.0, damping: float=0.001) -> Array

Preconditioned parameter update theta - lr (F + damping I)^{-1} grad.

The metric-aware (natural-gradient) counterpart of vanilla gradient descent: descending along F^{-1} grad instead of grad makes the step invariant to smooth reparametrisations. theta and grad are flat (P,) vectors and fisher the (P, P) Fisher information.

Parameters

Name	Type & default	Contract
thetapositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
gradpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
fisherpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
learning_ratekeyword-only	floatDefault: 1.0	See the operation contract and implementation below.
dampingkeyword-only	floatDefault: 0.001	See the operation contract and implementation below.

ReturnsArray

Executed regression · natural gradient step formula

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_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::test_natural_gradient_step_formula" -q

Read the implementation

def natural_gradient_step(
    theta: Array,
    grad: Array,
    fisher: Array,
    *,
    learning_rate: float = 1.0,
    damping: float = 1e-3,
) -> Array:
    r"""Preconditioned parameter update ``theta - lr (F + damping I)^{-1} grad``.

    The metric-aware (natural-gradient) counterpart of vanilla gradient descent:
    descending along ``F^{-1} grad`` instead of ``grad`` makes the step invariant
    to smooth reparametrisations. ``theta`` and ``grad`` are flat ``(P,)`` vectors
    and ``fisher`` the ``(P, P)`` Fisher information.
    """
    delta = damped_solve(fisher, grad, damping=damping)
    out: Array = theta - learning_rate * delta
    return out

Related APIsomnibias.curvature.damped_solve↗

Other public import paths

• omnibias.curvature.natural_gradient_step
• omnibias.curvature.natural_gradient.natural_gradient_step

← Package overviewRunnable learning examples →