Skip to main content
API REFERENCE/omnibias-curvature

omnibias.curvature.glm_fisher

0.1.0a2AlphaPython >=3.10AGPL / commercial

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

import omnibias.curvature.glm_fisher
Module contract & mathematical conventions

Multi-parameter Fisher information for one-layer GLM fields (JAX).

Two complementary objects lift the *scalar* Fisher information A''(theta) (the exact one-dimensional Fisher information) to many parameters:

  • fisher_information_metric -- the exponential-family **Fisher-Rao metric in natural coordinates**: the (d, d) diagonal matrix diag(A''(eta_1), ..., A''(eta_d)) for d coordinates sharing the scalar log-partition base. This is the direct multivariate generalisation of the 1-D Fisher information; hand it to omnibias-geometry (see omnibias.geometry.jax.ops.exponential_family_fisher_metric) to obtain its Christoffel symbols and curvature.
  • glm_fisher -- the (P, P) Fisher information of the one-layer field **parameters** theta = (b, c, beta, W) under a GLM observation model whose natural parameter is the field output eta_n = f_theta(x_n). It reuses the closed-form per-sample gradient omnibias.curvature.one_layer.one_layer_param_grad and weights each outer product by the GLM variance A''(eta_n):

.. math::

F(\theta) = \frac{1}{B} \sum_n A''(\eta_n)\, g_n g_n^\top, \qquad g_n = \nabla_\theta \eta_n .

This is the Fisher-scoring / IRLS information matrix and generalises omnibias.curvature.one_layer.mse_gauss_newton_fisher (the Gaussian family, A'' = 1).

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

fisher_information_metric​

omnibias.curvature.fisher_information_metricSource ↗
fisher_information_metric(eta: Array | float, *, base: str='softplus') -> Array

Exponential-family Fisher-Rao metric diag(A''(eta_k)) in natural coords.

eta carries the natural parameters on its last axis (shape (..., d)); the returned metric has shape (..., d, d) and is diagonal with the GLM variance A''(eta_k) of the scalar log-partition base on the diagonal. For d == 1 this is the 1-D Fisher information embedded as a (1, 1) matrix.

Parameters

NameType & defaultContract
etapositional-or-keywordArray | floatRequiredSee the operation contract and implementation below.
basekeyword-onlystrDefault: 'softplus'See the operation contract and implementation below.
ReturnsArray
Read the implementation
def fisher_information_metric(
    eta: Array | float, *, base: str = "softplus"
) -> Array:
    r"""Exponential-family Fisher-Rao metric ``diag(A''(eta_k))`` in natural coords.

    ``eta`` carries the natural parameters on its last axis (shape ``(..., d)``);
    the returned metric has shape ``(..., d, d)`` and is diagonal with the GLM
    variance ``A''(eta_k)`` of the scalar log-partition ``base`` on the diagonal.
    For ``d == 1`` this is the 1-D Fisher information embedded as a ``(1, 1)``
    matrix.
    """
    v = _log_partition_derivative(eta, base=base, order=2)
    d = v.shape[-1]
    eye = jnp.eye(d, dtype=v.dtype)
    out: Array = v[..., :, None] * eye
    return out
Other public import paths
  • omnibias.curvature.fisher_information_metric
  • omnibias.curvature.glm_fisher.fisher_information_metric
FUNCTION

glm_fisher​

omnibias.curvature.glm_fisher.glm_fisherSource ↗
glm_fisher(X: Array, W: Array, beta: Array, c: Array, b: Array, *, activation: str='tanh', family: str='bernoulli') -> Array

(P, P) GLM Fisher information of the one-layer field parameters.

The field output eta_n = b + sigma(W x_n + beta) . c is the natural parameter of the chosen exponential family ("bernoulli" / "poisson" / "gaussian"). Returns F = (1/B) sum_n A''(eta_n) g_n g_n^T with g_n the closed-form per-sample parameter gradient -- a positive-semidefinite Fisher matrix.

Parameters

NameType & defaultContract
Xpositional-or-keywordArrayRequiredSee the operation contract and implementation below.
Wpositional-or-keywordArrayRequiredSee the operation contract and implementation below.
betapositional-or-keywordArrayRequiredSee the operation contract and implementation below.
cpositional-or-keywordArrayRequiredSee the operation contract and implementation below.
bpositional-or-keywordArrayRequiredSee the operation contract and implementation below.
activationkeyword-onlystrDefault: 'tanh'See the operation contract and implementation below.
familykeyword-onlystrDefault: 'bernoulli'See the operation contract and implementation below.
ReturnsArray
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
Read the implementation
def glm_fisher(
    X: Array,
    W: Array,
    beta: Array,
    c: Array,
    b: Array,
    *,
    activation: str = "tanh",
    family: str = "bernoulli",
) -> Array:
    r"""``(P, P)`` GLM Fisher information of the one-layer field parameters.

    The field output ``eta_n = b + sigma(W x_n + beta) . c`` is the natural
    parameter of the chosen exponential ``family`` (``"bernoulli"`` /
    ``"poisson"`` / ``"gaussian"``). Returns
    ``F = (1/B) sum_n A''(eta_n) g_n g_n^T`` with ``g_n`` the closed-form
    per-sample parameter gradient -- a positive-semidefinite Fisher matrix.
    """
    spec = get_activation(activation)
    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)
    w = _glm_weights(etas, family)  # (B,)
    fisher: Array = (gs.T * w) @ gs / n_batch
    return fisher
Related APIsomnibias.curvature.one_layer_param_grad↗omnibias.jax.get_activation↗
← Package overviewRunnable learning examples →