API REFERENCE/omnibias-curvature

omnibias.curvature.torch.sharpness

0.1.0a2AlphaPython >=3.10AGPL / commercial

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

import omnibias.curvature.torch.sharpness

Module contract & mathematical conventions

Exact-curvature sharpness for **multi-layer torch** networks.

This is the torch, arbitrary-depth counterpart of omnibias.curvature.sharpness (which materialises the exact one-layer parameter Hessian in JAX). For a deep network the full parameter Hessian is too large to form, so the curvature functionals are computed **matrix-free** from exact **Hessian-vector products** (HVPs):

• hvp -- one exact HVP Hv via reverse-over-reverse autograd (torch.autograd.grad with create_graph=True twice). Because the omnibias torch activations are plain differentiable ops (no custom autograd.Function with a hand-written backward), this double-backward is the *exact* Hessian action -- and where the net uses the Riccati activations, the \sigma', \sigma'', \sigma''' that enter it are the bit-stable closed-form tower.
• top_eigenvalue -- \lambda_{\max}(H) by (two-phase) power iteration on the HVP; converges to the algebraically largest eigenvalue even when H is indefinite.
• hutchinson_trace / hutchinson_frobenius_sq -- unbiased Rademacher estimators of \operatorname{tr}(H)=\sum_i\lambda_i and \lVert H\rVert_F^2=\sum_i\lambda_i^2.
• sharpness_aware_loss / sam_objective -- differentiable training objectives; the penalty gradient rides on \sigma''' in closed form (reverse-over-reverse-over-reverse), so it is exact -- no finite-difference sharpness, no ascent-step approximation.

The API takes a scalar loss (with an autograd graph) and the params it depends on (e.g. list(model.parameters())), so it drops into any training loop -- an MLP, a PINN (JetMLP), or the omnibias.score.flow CNF velocity field.

Honesty

• The HVP and dense_hessian are **exact** (autograd), matching torch.autograd.functional.hessian to floating point.
• top_eigenvalue is exact *in the power-iteration limit* (a few dozen iterations suffice for a well-separated top eigenvalue).
• hutchinson_* are **unbiased stochastic** estimators -- exact in expectation, with variance ~ 1/n_samples. For a *reported* number on a small net, use dense_hessian + the matrix helpers instead.
• This measures / regularises curvature; it does **not** assert that flat minima always generalise (that is problem-dependent).

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 →

CLASS

HessianOperator

omnibias.curvature.torch.HessianOperatorSource ↗

class HessianOperator

Matrix-free exact Hessian of loss w.r.t. params.

Computes the first-order gradient once (with create_graph=True) and exposes repeated HVPs v \mapsto Hv against the retained graph.

Parameters

Name	Type & default	Contract
losspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
paramspositional-or-keyword	ParamsRequired	See the operation contract and implementation below.

Executed regression · exact sam fixed lam reports lam eff

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_exact_sam_fixed_lam_reports_lam_eff() -> None:
    """With ``lam_auto=False`` the telemetry reports the constant ``lam`` (backward-compat)."""
    net, x, y, params = _problem(seed=0)
    opt = ExactSAM(params, lr=1e-2, lam=1e-3)
    opt.step(lambda: _mse(net, x, y))
    assert opt._lam_eff == pytest.approx(1e-3)

python -m pytest "packages/omnibias-curvature/tests/test_torch_exact_sam.py::test_exact_sam_fixed_lam_reports_lam_eff" -q

Fields & methods

__init__method

__init__(self, loss: Tensor, params: Params) -> None

See the class contract above and the exact method implementation below.

Parameters

Name	Type & default	Contract
losspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
paramspositional-or-keyword	ParamsRequired	See the operation contract and implementation below.

Returns: None

Validation & failure cases

• When not self.params, raises ValueError('params is empty').

Executed regression · exact sam fixed lam reports lam eff

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_exact_sam_fixed_lam_reports_lam_eff() -> None:
    """With ``lam_auto=False`` the telemetry reports the constant ``lam`` (backward-compat)."""
    net, x, y, params = _problem(seed=0)
    opt = ExactSAM(params, lr=1e-2, lam=1e-3)
    opt.step(lambda: _mse(net, x, y))
    assert opt._lam_eff == pytest.approx(1e-3)

python -m pytest "packages/omnibias-curvature/tests/test_torch_exact_sam.py::test_exact_sam_fixed_lam_reports_lam_eff" -q

Method implementation

    def __init__(self, loss: Tensor, params: Params) -> None:
        self.params: list[Tensor] = [p for p in params]
        if not self.params:
            raise ValueError("params is empty")
        self.grads: tuple[Tensor, ...] = torch.autograd.grad(
            loss, self.params, create_graph=True,
        )

View source ↗

hvpmethod

hvp(self, vec: Sequence[Tensor], *, create_graph: bool=False) -> Vector

Exact H @ vec as a list matching params shapes.

Parameters

Name	Type & default	Contract
vecpositional-or-keyword	Sequence[Tensor]Required	See the operation contract and implementation below.
create_graphkeyword-only	boolDefault: False	See the operation contract and implementation below.

Returns: Vector

Executed regression · exact sam fixed lam reports lam eff

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_exact_sam_fixed_lam_reports_lam_eff() -> None:
    """With ``lam_auto=False`` the telemetry reports the constant ``lam`` (backward-compat)."""
    net, x, y, params = _problem(seed=0)
    opt = ExactSAM(params, lr=1e-2, lam=1e-3)
    opt.step(lambda: _mse(net, x, y))
    assert opt._lam_eff == pytest.approx(1e-3)

python -m pytest "packages/omnibias-curvature/tests/test_torch_exact_sam.py::test_exact_sam_fixed_lam_reports_lam_eff" -q

Method implementation

    def hvp(self, vec: Sequence[Tensor], *, create_graph: bool = False) -> Vector:
        """Exact ``H @ vec`` as a list matching ``params`` shapes."""
        dot = _flat_dot(self.grads, vec)
        hv = torch.autograd.grad(
            dot, self.params, retain_graph=True, create_graph=create_graph,
        )
        return list(hv)

View source ↗

grad_normproperty

grad_norm(self) -> Tensor

See the class contract above and the exact method implementation below.

Returns: Tensor

Executed regression · exact sam ascent mode reduces loss

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_exact_sam_ascent_mode_reduces_loss() -> None:
    net, x, y, params = _problem(seed=3)

    def closure() -> torch.Tensor:
        return _mse(net, x, y)

    f0 = float(closure().detach())
    opt = ExactSAM(params, lr=1e-2, mode="ascent", rho=0.05, iters=15, seed=0)
    for _ in range(60):
        opt.step(closure)
    assert float(closure().detach()) < f0

python -m pytest "packages/omnibias-curvature/tests/test_torch_exact_sam.py::test_exact_sam_ascent_mode_reduces_loss" -q

Method implementation

    def grad_norm(self) -> Tensor:
        return _flat_norm(self.grads)

View source ↗

paramsattribute

params: list[Tensor] = [p for p in params]

Public instance attribute assigned by __init__. The source assignment is params = [p for p in params]; it describes instance state, not a literal constructor default.

View source ↗

gradsattribute

grads: tuple[Tensor, ...] = torch.autograd.grad(loss, self.params, create_graph=True)

Public instance attribute assigned by __init__. The source assignment is grads = torch.autograd.grad(loss, self.params, create_graph=True); it describes instance state, not a literal constructor default.

View source ↗

Read the implementation

class HessianOperator:
    r"""Matrix-free exact Hessian of ``loss`` w.r.t. ``params``.

    Computes the first-order gradient once (with ``create_graph=True``) and
    exposes repeated HVPs :math:`v \mapsto Hv` against the retained graph.
    """

    def __init__(self, loss: Tensor, params: Params) -> None:
        self.params: list[Tensor] = [p for p in params]
        if not self.params:
            raise ValueError("params is empty")
        self.grads: tuple[Tensor, ...] = torch.autograd.grad(
            loss, self.params, create_graph=True,
        )

    def hvp(self, vec: Sequence[Tensor], *, create_graph: bool = False) -> Vector:
        """Exact ``H @ vec`` as a list matching ``params`` shapes."""
        dot = _flat_dot(self.grads, vec)
        hv = torch.autograd.grad(
            dot, self.params, retain_graph=True, create_graph=create_graph,
        )
        return list(hv)

    @property
    def grad_norm(self) -> Tensor:
        return _flat_norm(self.grads)

Other public import paths

• omnibias.curvature.torch.HessianOperator
• omnibias.curvature.torch.sharpness.HessianOperator

FUNCTION

curvature_sharpness

omnibias.curvature.torch.curvature_sharpnessSource ↗

curvature_sharpness(loss: Tensor, params: Params, *, measure: str='frobenius', n_samples: int=8, iters: int=30, generator: torch.Generator | None=None, differentiable: bool=False) -> Tensor

Matrix-free curvature sharpness of loss at the current params.

measure is "trace" / "frobenius" (Hutchinson) or "top_eig" (power iteration).

Parameters

Name	Type & default	Contract
losspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
paramspositional-or-keyword	ParamsRequired	See the operation contract and implementation below.
measurekeyword-only	strDefault: 'frobenius'	See the operation contract and implementation below.
n_sampleskeyword-only	intDefault: 8	See the operation contract and implementation below.
iterskeyword-only	intDefault: 30	See the operation contract and implementation below.
generatorkeyword-only	torch.Generator | NoneDefault: None	See the operation contract and implementation below.
differentiablekeyword-only	boolDefault: False	See the operation contract and implementation below.

ReturnsTensor

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():
    net, X, Y, params = _small_problem()
    with pytest.raises(ValueError, match="unknown sharpness measure"):
        S.curvature_sharpness(_mse(net, X, Y), params, measure="nope")

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

Regression example · all measures finite on deep net

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_all_measures_finite_on_deep_net():
    net, X, Y, params = _small_problem(in_dim=3, hidden=8, depth=3, n=24, seed=11)
    P = sum(p.numel() for p in params)
    assert P > 100  # genuinely multi-layer / larger than one-layer
    for measure in ("trace", "frobenius", "top_eig"):
        val = S.curvature_sharpness(_mse(net, X, Y), params, measure=measure,
                                    n_samples=8, iters=30,
                                    generator=torch.Generator().manual_seed(0))
        assert torch.isfinite(val)

python -m pytest packages/omnibias-curvature/tests/test_torch_sharpness.py -k test_all_measures_finite_on_deep_net -q

Read the implementation

def curvature_sharpness(
    loss: Tensor, params: Params, *,
    measure: str = "frobenius",
    n_samples: int = 8, iters: int = 30,
    generator: torch.Generator | None = None,
    differentiable: bool = False,
) -> Tensor:
    """Matrix-free curvature sharpness of ``loss`` at the current ``params``.

    ``measure`` is ``"trace"`` / ``"frobenius"`` (Hutchinson) or ``"top_eig"``
    (power iteration).
    """
    if measure == "trace":
        return hutchinson_trace(loss, params, n_samples=n_samples,
                                generator=generator, differentiable=differentiable)
    if measure == "frobenius":
        return hutchinson_frobenius_sq(loss, params, n_samples=n_samples,
                                       generator=generator, differentiable=differentiable)
    if measure == "top_eig":
        return top_eigenvalue(loss, params, iters=iters,
                              generator=generator, differentiable=differentiable)
    raise ValueError(
        f"unknown sharpness measure {measure!r}; choose from {list(_MEASURES)}"
    )

Related APIsomnibias.curvature.torch.hutchinson_frobenius_sq↗omnibias.curvature.torch.hutchinson_trace↗omnibias.curvature.torch.top_eigenvalue↗

Other public import paths

• omnibias.curvature.torch.curvature_sharpness
• omnibias.curvature.torch.sharpness.curvature_sharpness

FUNCTION

dense_hessian

omnibias.curvature.torch.dense_hessianSource ↗

dense_hessian(loss: Tensor, params: Params) -> Tensor

Materialise the exact (P, P) parameter Hessian by HVPs on the basis.

Only for **small** models (P up to a few hundred): it runs P backward passes. Handy as an oracle for the matrix-free estimators and for exact sharpness on tiny nets.

Parameters

Name	Type & default	Contract
losspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
paramspositional-or-keyword	ParamsRequired	See the operation contract and implementation below.

ReturnsTensor

Executed regression · hvp matches dense action

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_hvp_matches_dense_action():
    net, X, Y, params = _small_problem(seed=1)
    Hd = S.dense_hessian(_mse(net, X, Y), params)
    v = S._rand_like(params, generator=torch.Generator().manual_seed(3))
    Hv = S.hvp(_mse(net, X, Y), params, v)
    Hv_flat = torch.cat([h.reshape(-1) for h in Hv]).detach()
    v_flat = torch.cat([vi.reshape(-1) for vi in v])
    assert float((Hv_flat - Hd @ v_flat).abs().max()) < 1e-10

python -m pytest "packages/omnibias-curvature/tests/test_torch_sharpness.py::test_hvp_matches_dense_action" -q

Regression example · dense hessian matches torch func 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_dense_hessian_matches_torch_func_hessian():
    net, X, Y, params = _small_problem()
    Hd = S.dense_hessian(_mse(net, X, Y), params)
    flat_loss, theta0 = _flat_loss_fn(net, X, Y)
    Ho = torch.func.hessian(flat_loss)(theta0)
    assert Hd.shape == Ho.shape == (theta0.numel(), theta0.numel())
    assert torch.allclose(Hd, Hd.T, atol=1e-9), "dense Hessian not symmetric"
    assert float((Hd - Ho).abs().max()) < 1e-7

python -m pytest packages/omnibias-curvature/tests/test_torch_sharpness.py -k test_dense_hessian_matches_torch_func_hessian -q

Read the implementation

def dense_hessian(loss: Tensor, params: Params) -> Tensor:
    r"""Materialise the exact ``(P, P)`` parameter Hessian by HVPs on the basis.

    Only for **small** models (``P`` up to a few hundred): it runs ``P`` backward
    passes. Handy as an oracle for the matrix-free estimators and for exact
    sharpness on tiny nets.
    """
    op = HessianOperator(loss, params)
    sizes = [p.numel() for p in op.params]
    total = int(sum(sizes))
    rows = []
    for i in range(total):
        # basis vector e_i as a list-of-tensors
        e: Vector = []
        j = i
        for p, n in zip(op.params, sizes, strict=True):
            block = torch.zeros(n, dtype=p.dtype, device=p.device)
            if 0 <= j < n:
                block[j] = 1.0
            j -= n
            e.append(block.reshape(p.shape))
        hv = op.hvp(e, create_graph=False)
        rows.append(torch.cat([h.reshape(-1) for h in hv]).detach())
    return torch.stack(rows)

Related APIsomnibias.curvature.torch.HessianOperator↗

Other public import paths

• omnibias.curvature.torch.dense_hessian
• omnibias.curvature.torch.sharpness.dense_hessian

FUNCTION

hessian_eigenvalue_extremes

omnibias.curvature.torch.hessian_eigenvalue_extremesSource ↗

hessian_eigenvalue_extremes(loss: Tensor, params: Params, *, iters: int=30, generator: torch.Generator | None=None) -> tuple[float, float]

(lambda_min, lambda_max) of the parameter Hessian by power iteration.

Parameters

Name	Type & default	Contract
losspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
paramspositional-or-keyword	ParamsRequired	See the operation contract and implementation below.
iterskeyword-only	intDefault: 30	See the operation contract and implementation below.
generatorkeyword-only	torch.Generator | NoneDefault: None	See the operation contract and implementation below.

Returnstuple[float, float]

Executed regression · power iteration matches dense top eigenvalue

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_power_iteration_matches_dense_top_eigenvalue():
    net, X, Y, params = _small_problem(seed=2)
    Hd = S.dense_hessian(_mse(net, X, Y), params)
    ev = torch.linalg.eigvalsh(0.5 * (Hd + Hd.T))
    lam = float(S.top_eigenvalue(_mse(net, X, Y), params, iters=80,
                                 generator=torch.Generator().manual_seed(0)))
    assert abs(lam - float(ev[-1])) < 1e-5
    lo, hi = S.hessian_eigenvalue_extremes(_mse(net, X, Y), params, iters=120,
                                           generator=torch.Generator().manual_seed(0))
    assert abs(hi - float(ev[-1])) < 1e-5
    assert abs(lo - float(ev[0])) < 1e-3  # min converges slower; looser

python -m pytest "packages/omnibias-curvature/tests/test_torch_sharpness.py::test_power_iteration_matches_dense_top_eigenvalue" -q

Read the implementation

def hessian_eigenvalue_extremes(
    loss: Tensor, params: Params, *,
    iters: int = 30, generator: torch.Generator | None = None,
) -> tuple[float, float]:
    """``(lambda_min, lambda_max)`` of the parameter Hessian by power iteration."""
    op = HessianOperator(loss, params)
    v1 = _power_eigvec(op, iters=iters, generator=generator, shift=0.0, x0=None)
    mu = float(_rayleigh(op, v1))
    v2 = _power_eigvec(op, iters=iters, generator=generator, shift=mu, x0=None)
    eta = float(_rayleigh(op, v2))
    return (min(mu, eta), max(mu, eta))

Related APIsomnibias.curvature.torch.HessianOperator↗

Other public import paths

• omnibias.curvature.torch.hessian_eigenvalue_extremes
• omnibias.curvature.torch.sharpness.hessian_eigenvalue_extremes

FUNCTION

hessian_frobenius_sq

omnibias.curvature.torch.hessian_frobenius_sqSource ↗

hessian_frobenius_sq(H: Tensor) -> Tensor

\lVert H\rVert_F^2=\sum_i\lambda_i^2.

Parameters

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

ReturnsTensor

Executed regression · exact sam penalty lowers sharpness vs ablation

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_exact_sam_penalty_lowers_sharpness_vs_ablation() -> None:
    """Same init + data + base step: ``lam>0`` ends at a strictly lower exact ``||H||_F^2``
    than the ``lam=0`` ablation, while still reducing the data loss."""
    net_reg, x, y, p_reg = _problem(seed=1)
    net_base, _, _, p_base = _problem(seed=1)  # identical init and data (same seed)

    def train(net: JetMLP, params: list[torch.Tensor], lam: float) -> None:
        opt = ExactSAM(params, lr=5e-3, lam=lam, measure="frobenius", n_samples=6, seed=0)
        for _ in range(150):
            opt.step(lambda: _mse(net, x, y))

    f0 = _frobenius_sq(net_base, x, y, p_base)
    l0 = float(_mse(net_base, x, y).detach())
    train(net_base, p_base, lam=0.0)
    train(net_reg, p_reg, lam=3e-2)

    frob_base = _frobenius_sq(net_base, x, y, p_base)
    frob_reg = _frobenius_sq(net_reg, x, y, p_reg)
    loss_reg = float(_mse(net_reg, x, y).detach())

    assert loss_reg < l0  # the penalised run still trained
    assert frob_reg < frob_base  # ... to a strictly flatter minimum
    assert frob_base <= f0 + 1.0  # sanity: the ablation did not blow curvature up

python -m pytest "packages/omnibias-curvature/tests/test_torch_exact_sam.py::test_exact_sam_penalty_lowers_sharpness_vs_ablation" -q

Read the implementation

def hessian_frobenius_sq(H: Tensor) -> Tensor:
    r""":math:`\lVert H\rVert_F^2=\sum_i\lambda_i^2`."""
    return (H * H).sum()

Other public import paths

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

FUNCTION

hessian_top_eigenvalue

omnibias.curvature.torch.hessian_top_eigenvalueSource ↗

hessian_top_eigenvalue(H: Tensor) -> Tensor

\lambda_{\max}(H) via a symmetric eigendecomposition.

Parameters

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

ReturnsTensor

Executed regression · sharpness matches dense eigensolve and grows with beta

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_matches_dense_eigensolve_and_grows_with_beta() -> None:
    torch = pytest.importorskip("torch")
    from omnibias.struct.torch import chain_hessian, chain_sharpness

    trellis = random_chain(3)
    e = torch.tensor(trellis.emissions)
    trn = torch.tensor(trellis.transitions)
    st = torch.tensor(trellis.start)
    top = float(chain_sharpness(e, trn, 4.0, start=st))
    ref = float(np.linalg.eigvalsh(chain_hessian(e, trn, 4.0, start=st).numpy())[-1])
    assert abs(top - ref) < 1e-9
    # sharper relaxation at higher beta
    assert float(chain_sharpness(e, trn, 8.0, start=st)) > float(chain_sharpness(e, trn, 1.0, start=st))

python -m pytest "packages/omnibias-struct/tests/test_curvature.py::test_sharpness_matches_dense_eigensolve_and_grows_with_beta" -q

Read the implementation

def hessian_top_eigenvalue(H: Tensor) -> Tensor:
    r""":math:`\lambda_{\max}(H)` via a symmetric eigendecomposition."""
    Hs = 0.5 * (H + H.transpose(-1, -2))
    return torch.linalg.eigvalsh(Hs)[-1]

Other public import paths

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

FUNCTION

hessian_trace

omnibias.curvature.torch.hessian_traceSource ↗

hessian_trace(H: Tensor) -> Tensor

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

Parameters

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

ReturnsTensor

Read the implementation

def hessian_trace(H: Tensor) -> Tensor:
    r""":math:`\operatorname{tr}(H)=\sum_i\lambda_i`."""
    return torch.trace(H)

Other public import paths

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

FUNCTION

hutchinson_frobenius_sq

omnibias.curvature.torch.hutchinson_frobenius_sqSource ↗

hutchinson_frobenius_sq(loss: Tensor, params: Params, *, n_samples: int=8, generator: torch.Generator | None=None, differentiable: bool=False) -> Tensor

Unbiased Rademacher estimator of \lVert H\rVert_F^2.

Uses \mathbb E\lVert Hv\rVert^2 = \operatorname{tr}(H^\top H) for a Rademacher probe v.

Parameters

Name	Type & default	Contract
losspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
paramspositional-or-keyword	ParamsRequired	See the operation contract and implementation below.
n_sampleskeyword-only	intDefault: 8	See the operation contract and implementation below.
generatorkeyword-only	torch.Generator | NoneDefault: None	See the operation contract and implementation below.
differentiablekeyword-only	boolDefault: False	See the operation contract and implementation below.

ReturnsTensor

Executed regression · exact sam fixed lam reports lam eff

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_exact_sam_fixed_lam_reports_lam_eff() -> None:
    """With ``lam_auto=False`` the telemetry reports the constant ``lam`` (backward-compat)."""
    net, x, y, params = _problem(seed=0)
    opt = ExactSAM(params, lr=1e-2, lam=1e-3)
    opt.step(lambda: _mse(net, x, y))
    assert opt._lam_eff == pytest.approx(1e-3)

python -m pytest "packages/omnibias-curvature/tests/test_torch_exact_sam.py::test_exact_sam_fixed_lam_reports_lam_eff" -q

Regression example · hutchinson trace and frobenius unbiased

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_hutchinson_trace_and_frobenius_unbiased():
    net, X, Y, params = _small_problem(seed=4)
    Hd = S.dense_hessian(_mse(net, X, Y), params)
    tr_exact = float(torch.trace(Hd))
    fr_exact = float((Hd * Hd).sum())
    gen = torch.Generator().manual_seed(0)
    tr_est = float(S.hutchinson_trace(_mse(net, X, Y), params, n_samples=2000, generator=gen))
    fr_est = float(S.hutchinson_frobenius_sq(_mse(net, X, Y), params, n_samples=2000, generator=gen))
    # Unbiased estimators; at 2000 probes (seed-fixed => deterministic) the
    # relative error is well under 5%.
    assert abs(tr_est - tr_exact) < 0.05 * (1 + abs(tr_exact))
    assert abs(fr_est - fr_exact) < 0.05 * (1 + abs(fr_exact))

python -m pytest packages/omnibias-curvature/tests/test_torch_sharpness.py -k test_hutchinson_trace_and_frobenius_unbiased -q

Read the implementation

def hutchinson_frobenius_sq(
    loss: Tensor, params: Params, *,
    n_samples: int = 8, generator: torch.Generator | None = None,
    differentiable: bool = False,
) -> Tensor:
    r"""Unbiased Rademacher estimator of :math:`\lVert H\rVert_F^2`.

    Uses :math:`\mathbb E\lVert Hv\rVert^2 = \operatorname{tr}(H^\top H)` for a
    Rademacher probe ``v``.
    """
    op = HessianOperator(loss, params)
    acc = torch.zeros((), dtype=op.params[0].dtype, device=op.params[0].device)
    for _ in range(n_samples):
        v = _rand_like(op.params, generator=generator, rademacher=True)
        hv = op.hvp(v, create_graph=differentiable)
        acc = acc + _flat_dot(hv, hv)
    return acc / n_samples

Related APIsomnibias.curvature.torch.HessianOperator↗

Other public import paths

• omnibias.curvature.torch.hutchinson_frobenius_sq
• omnibias.curvature.torch.sharpness.hutchinson_frobenius_sq

FUNCTION

hutchinson_trace

omnibias.curvature.torch.hutchinson_traceSource ↗

hutchinson_trace(loss: Tensor, params: Params, *, n_samples: int=8, generator: torch.Generator | None=None, differentiable: bool=False) -> Tensor

Unbiased Rademacher estimator of \operatorname{tr}(H).

Parameters

Name	Type & default	Contract
losspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
paramspositional-or-keyword	ParamsRequired	See the operation contract and implementation below.
n_sampleskeyword-only	intDefault: 8	See the operation contract and implementation below.
generatorkeyword-only	torch.Generator | NoneDefault: None	See the operation contract and implementation below.
differentiablekeyword-only	boolDefault: False	See the operation contract and implementation below.

ReturnsTensor

Executed regression · exact sam reduces loss for every measure

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_exact_sam_reduces_loss_for_every_measure(measure: str) -> None:
    net, x, y, params = _problem(seed=0)

    def closure() -> torch.Tensor:
        return _mse(net, x, y)

    f0 = float(closure().detach())
    opt = ExactSAM(params, lr=1e-2, lam=1e-3, measure=measure, n_samples=4, iters=15, seed=0)
    last = None
    for _ in range(60):
        last = opt.step(closure)
    assert last is not None and torch.isfinite(last)
    assert float(closure().detach()) < f0

python -m pytest "packages/omnibias-curvature/tests/test_torch_exact_sam.py::test_exact_sam_reduces_loss_for_every_measure[trace]" -q

Regression example · hutchinson trace and frobenius unbiased

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_hutchinson_trace_and_frobenius_unbiased():
    net, X, Y, params = _small_problem(seed=4)
    Hd = S.dense_hessian(_mse(net, X, Y), params)
    tr_exact = float(torch.trace(Hd))
    fr_exact = float((Hd * Hd).sum())
    gen = torch.Generator().manual_seed(0)
    tr_est = float(S.hutchinson_trace(_mse(net, X, Y), params, n_samples=2000, generator=gen))
    fr_est = float(S.hutchinson_frobenius_sq(_mse(net, X, Y), params, n_samples=2000, generator=gen))
    # Unbiased estimators; at 2000 probes (seed-fixed => deterministic) the
    # relative error is well under 5%.
    assert abs(tr_est - tr_exact) < 0.05 * (1 + abs(tr_exact))
    assert abs(fr_est - fr_exact) < 0.05 * (1 + abs(fr_exact))

python -m pytest packages/omnibias-curvature/tests/test_torch_sharpness.py -k test_hutchinson_trace_and_frobenius_unbiased -q

Read the implementation

def hutchinson_trace(
    loss: Tensor, params: Params, *,
    n_samples: int = 8, generator: torch.Generator | None = None,
    differentiable: bool = False,
) -> Tensor:
    r"""Unbiased Rademacher estimator of :math:`\operatorname{tr}(H)`."""
    op = HessianOperator(loss, params)
    acc = torch.zeros((), dtype=op.params[0].dtype, device=op.params[0].device)
    for _ in range(n_samples):
        v = _rand_like(op.params, generator=generator, rademacher=True)
        hv = op.hvp(v, create_graph=differentiable)
        acc = acc + _flat_dot(v, hv)
    return acc / n_samples

Related APIsomnibias.curvature.torch.HessianOperator↗

Other public import paths

• omnibias.curvature.torch.hutchinson_trace
• omnibias.curvature.torch.sharpness.hutchinson_trace

FUNCTION

hvp

omnibias.curvature.torch.hvpSource ↗

hvp(loss: Tensor, params: Params, vec: Sequence[Tensor]) -> Vector

One exact Hessian-vector product H @ vec (see HessianOperator).

Parameters

Name	Type & default	Contract
losspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
paramspositional-or-keyword	ParamsRequired	See the operation contract and implementation below.
vecpositional-or-keyword	Sequence[Tensor]Required	See the operation contract and implementation below.

ReturnsVector

Executed regression · hvp matches dense action

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_hvp_matches_dense_action():
    net, X, Y, params = _small_problem(seed=1)
    Hd = S.dense_hessian(_mse(net, X, Y), params)
    v = S._rand_like(params, generator=torch.Generator().manual_seed(3))
    Hv = S.hvp(_mse(net, X, Y), params, v)
    Hv_flat = torch.cat([h.reshape(-1) for h in Hv]).detach()
    v_flat = torch.cat([vi.reshape(-1) for vi in v])
    assert float((Hv_flat - Hd @ v_flat).abs().max()) < 1e-10

python -m pytest "packages/omnibias-curvature/tests/test_torch_sharpness.py::test_hvp_matches_dense_action" -q

Regression example · hvp exact through inner derivative loss

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_hvp_exact_through_inner_derivative_loss():
    torch.manual_seed(0)
    net = JetMLP(1, 4, 1, depth=2, base="tanh").double()
    x = torch.linspace(-1.0, 1.0, 8, dtype=torch.float64).reshape(-1, 1).requires_grad_(True)
    params = [p for p in net.parameters() if p.requires_grad]

    def residual_loss():
        u = net(x)
        (du,) = torch.autograd.grad(u.sum(), x, create_graph=True)  # du/dx, keep graph
        r = du - u  # toy first-order PDE residual du/dx = u
        return (r ** 2).mean()

    # Matrix-free HVP == dense Hessian action, even though the loss contains an
    # inner spatial derivative (so this is a third-order autograd overall).
    Hd = S.dense_hessian(residual_loss(), params)
    v = S._rand_like(params, generator=torch.Generator().manual_seed(2))
    Hv = S.hvp(residual_loss(), params, v)
    Hv_flat = torch.cat([h.reshape(-1) for h in Hv]).detach()
    v_flat = torch.cat([vi.reshape(-1) for vi in v])
    assert float((Hv_flat - Hd @ v_flat).abs().max()) < 1e-9

    # The differentiable penalty backprops to finite, parameter-shaped grads.
    obj = S.sharpness_aware_loss(residual_loss(), params, lam=1e-3, measure="frobenius",
                                 n_samples=2, generator=torch.Generator().manual_seed(3))
    net.zero_grad()
    obj.backward()
    for p in params:
        assert p.grad is not None and p.grad.shape == p.shape
        assert torch.isfinite(p.grad).all()

python -m pytest packages/omnibias-curvature/tests/test_torch_sharpness_integration.py -k test_hvp_exact_through_inner_derivative_loss -q

Read the implementation

def hvp(loss: Tensor, params: Params, vec: Sequence[Tensor]) -> Vector:
    """One exact Hessian-vector product ``H @ vec`` (see :class:`HessianOperator`)."""
    return HessianOperator(loss, params).hvp(vec, create_graph=False)

Related APIsomnibias.curvature.torch.HessianOperator↗

Other public import paths

• omnibias.curvature.torch.hvp
• omnibias.curvature.torch.sharpness.hvp

FUNCTION

sam_objective

omnibias.curvature.torch.sam_objectiveSource ↗

sam_objective(loss: Tensor, params: Params, *, rho: float=0.05, iters: int=30, generator: torch.Generator | None=None) -> Tensor

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

Parameters

Name	Type & default	Contract
losspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
paramspositional-or-keyword	ParamsRequired	See the operation contract and implementation below.
rhokeyword-only	floatDefault: 0.05	See the operation contract and implementation below.
iterskeyword-only	intDefault: 30	See the operation contract and implementation below.
generatorkeyword-only	torch.Generator | NoneDefault: None	See the operation contract and implementation below.

ReturnsTensor

Read the implementation

def sam_objective(
    loss: Tensor, params: Params, *,
    rho: float = 0.05, iters: int = 30,
    generator: torch.Generator | None = None,
) -> Tensor:
    r"""Ascent-free "SAM done right": :math:`L + \text{gap}_\rho`."""
    return loss + sam_sharpness_gap(
        loss, params, rho=rho, iters=iters, generator=generator, differentiable=True,
    )

Related APIsomnibias.curvature.torch.sam_sharpness_gap↗

Other public import paths

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

FUNCTION

sam_sharpness_gap

omnibias.curvature.torch.sam_sharpness_gapSource ↗

sam_sharpness_gap(loss: Tensor, params: Params, *, rho: float=0.05, iters: int=30, generator: torch.Generator | None=None, differentiable: bool=False) -> Tensor

Exact second-order SAM inner-max gap \rho\lVert\nabla L\rVert + \tfrac12\rho^2\max(\lambda_{\max},0).

Upper-bounds \max_{\lVert\varepsilon\rVert\le\rho}L(\theta+\varepsilon)-L(\theta) to O(\rho^3); the linear term is what classic SAM estimates with one ascent step, the curvature term is the exact-Hessian correction.

Parameters

Name	Type & default	Contract
losspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
paramspositional-or-keyword	ParamsRequired	See the operation contract and implementation below.
rhokeyword-only	floatDefault: 0.05	See the operation contract and implementation below.
iterskeyword-only	intDefault: 30	See the operation contract and implementation below.
generatorkeyword-only	torch.Generator | NoneDefault: None	See the operation contract and implementation below.
differentiablekeyword-only	boolDefault: False	See the operation contract and implementation below.

ReturnsTensor

Executed regression · exact sam ascent mode reduces loss

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_exact_sam_ascent_mode_reduces_loss() -> None:
    net, x, y, params = _problem(seed=3)

    def closure() -> torch.Tensor:
        return _mse(net, x, y)

    f0 = float(closure().detach())
    opt = ExactSAM(params, lr=1e-2, mode="ascent", rho=0.05, iters=15, seed=0)
    for _ in range(60):
        opt.step(closure)
    assert float(closure().detach()) < f0

python -m pytest "packages/omnibias-curvature/tests/test_torch_exact_sam.py::test_exact_sam_ascent_mode_reduces_loss" -q

Regression example · sam gap upper bounds sampled worst case

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():
    net, X, Y, params = _small_problem(in_dim=1, hidden=2, depth=1, n=8, seed=9)
    rho = 0.08
    gap2 = float(S.sam_sharpness_gap(_mse(net, X, Y), params, rho=rho, iters=100,
                                     generator=torch.Generator().manual_seed(0)))
    assert gap2 > 0.0

    rng = np.random.default_rng(0)
    flats = [p.reshape(-1) for p in params]
    total = sum(f.numel() for f in flats)
    worst = 0.0
    with torch.no_grad():
        L0 = float(_mse(net, X, Y))
        for _ in range(1500):
            d = rng.normal(size=total)
            d = rho * d / np.linalg.norm(d)
            d_t = torch.tensor(d, dtype=torch.float64)
            idx = 0
            for f in flats:
                k = f.numel()
                f += d_t[idx:idx + k]
                idx += k
            worst = max(worst, float(_mse(net, X, Y)) - L0)
            idx = 0
            for f in flats:
                k = f.numel()
                f -= d_t[idx:idx + k]
                idx += k
    assert gap2 >= worst - 1e-3, f"gap2={gap2} under-shoots sampled worst={worst}"
    assert gap2 <= 2.0 * worst + 1e-6

python -m pytest packages/omnibias-curvature/tests/test_torch_sharpness.py -k test_sam_gap_upper_bounds_sampled_worst_case -q

Read the implementation

def sam_sharpness_gap(
    loss: Tensor, params: Params, *,
    rho: float = 0.05, iters: int = 30,
    generator: torch.Generator | None = None,
    differentiable: bool = False,
) -> Tensor:
    r"""Exact second-order SAM inner-max gap
    :math:`\rho\lVert\nabla L\rVert + \tfrac12\rho^2\max(\lambda_{\max},0)`.

    Upper-bounds :math:`\max_{\lVert\varepsilon\rVert\le\rho}L(\theta+\varepsilon)-L(\theta)`
    to :math:`O(\rho^3)`; the linear term is what classic SAM estimates with one
    ascent step, the curvature term is the exact-Hessian correction.
    """
    op = HessianOperator(loss, params)
    gnorm = op.grad_norm if differentiable else op.grad_norm.detach()
    v_top, lam = _top_eigvec_and_value(op, iters=iters, generator=generator)
    if differentiable:
        hv = op.hvp(v_top, create_graph=True)
        lam = _flat_dot(v_top, hv)
    lam_pos = torch.clamp(lam, min=0.0)
    return rho * gnorm + 0.5 * rho * rho * lam_pos

Related APIsomnibias.curvature.torch.HessianOperator↗

Other public import paths

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

FUNCTION

sharpness_aware_loss

omnibias.curvature.torch.sharpness_aware_lossSource ↗

sharpness_aware_loss(loss: Tensor, params: Params, *, lam: float=0.01, measure: str='frobenius', n_samples: int=8, iters: int=30, generator: torch.Generator | None=None) -> Tensor

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

Differentiable end-to-end; the penalty gradient rides on the closed-form \sigma''' via reverse-over-reverse-over-reverse autograd.

Parameters

Name	Type & default	Contract
losspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
paramspositional-or-keyword	ParamsRequired	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.
n_sampleskeyword-only	intDefault: 8	See the operation contract and implementation below.
iterskeyword-only	intDefault: 30	See the operation contract and implementation below.
generatorkeyword-only	torch.Generator | NoneDefault: None	See the operation contract and implementation below.

ReturnsTensor

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():
    D, H_teacher, H_student, n_train = 3, 3, 6, 30
    Xtr, Ytr, Xte, Yte = _teacher_student(17, D, H_teacher, n_train, noise=0.1)

    torch.manual_seed(100)
    net_plain = JetMLP(D, H_student, 1, depth=2, base="tanh").double()
    net_flat = JetMLP(D, H_student, 1, depth=2, base="tanh").double()
    net_flat.load_state_dict(net_plain.state_dict())  # same init

    p_plain = _train(net_plain, Xtr, Ytr, lam=0.0, steps=200, lr=8e-3, seed=1)
    p_flat = _train(net_flat, Xtr, Ytr, lam=3e-3, steps=200, lr=8e-3, seed=1)

    def report(net, params):
        tr = float(_mse(net, Xtr, Ytr).detach())
        te = float(_mse(net, Xte, Yte).detach())
        curv = float(S.top_eigenvalue(_mse(net, Xtr, Ytr), params, iters=60,
                                      generator=torch.Generator().manual_seed(0)))
        return tr, te, curv

    tr_p, te_p, curv_p = report(net_plain, p_plain)
    tr_f, te_f, curv_f = report(net_flat, p_flat)

    assert tr_p < 0.05 and tr_f < 0.08, (tr_p, tr_f)
    assert curv_f < 0.75 * curv_p, f"curv flat={curv_f:.3f} vs plain={curv_p:.3f}"
    # exact-curvature regularisation also improved generalisation here
    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_torch_sharpness.py::test_sharpness_aware_training_finds_flatter_minimum" -q

Regression example · hvp exact through inner derivative loss

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_hvp_exact_through_inner_derivative_loss():
    torch.manual_seed(0)
    net = JetMLP(1, 4, 1, depth=2, base="tanh").double()
    x = torch.linspace(-1.0, 1.0, 8, dtype=torch.float64).reshape(-1, 1).requires_grad_(True)
    params = [p for p in net.parameters() if p.requires_grad]

    def residual_loss():
        u = net(x)
        (du,) = torch.autograd.grad(u.sum(), x, create_graph=True)  # du/dx, keep graph
        r = du - u  # toy first-order PDE residual du/dx = u
        return (r ** 2).mean()

    # Matrix-free HVP == dense Hessian action, even though the loss contains an
    # inner spatial derivative (so this is a third-order autograd overall).
    Hd = S.dense_hessian(residual_loss(), params)
    v = S._rand_like(params, generator=torch.Generator().manual_seed(2))
    Hv = S.hvp(residual_loss(), params, v)
    Hv_flat = torch.cat([h.reshape(-1) for h in Hv]).detach()
    v_flat = torch.cat([vi.reshape(-1) for vi in v])
    assert float((Hv_flat - Hd @ v_flat).abs().max()) < 1e-9

    # The differentiable penalty backprops to finite, parameter-shaped grads.
    obj = S.sharpness_aware_loss(residual_loss(), params, lam=1e-3, measure="frobenius",
                                 n_samples=2, generator=torch.Generator().manual_seed(3))
    net.zero_grad()
    obj.backward()
    for p in params:
        assert p.grad is not None and p.grad.shape == p.shape
        assert torch.isfinite(p.grad).all()

python -m pytest packages/omnibias-curvature/tests/test_torch_sharpness_integration.py -k test_hvp_exact_through_inner_derivative_loss -q

Read the implementation

def sharpness_aware_loss(
    loss: Tensor, params: Params, *,
    lam: float = 1e-2, measure: str = "frobenius",
    n_samples: int = 8, iters: int = 30,
    generator: torch.Generator | None = None,
) -> Tensor:
    r"""Curvature-regularised loss :math:`L + \lambda\,\mathcal S(\nabla^2 L)`.

    Differentiable end-to-end; the penalty gradient rides on the closed-form
    :math:`\sigma'''` via reverse-over-reverse-over-reverse autograd.
    """
    penalty = curvature_sharpness(
        loss, params, measure=measure, n_samples=n_samples, iters=iters,
        generator=generator, differentiable=True,
    )
    return loss + lam * penalty

Related APIsomnibias.curvature.torch.curvature_sharpness↗

Other public import paths

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

FUNCTION

top_eigenvalue

omnibias.curvature.torch.top_eigenvalueSource ↗

top_eigenvalue(loss: Tensor, params: Params, *, iters: int=30, generator: torch.Generator | None=None, differentiable: bool=False) -> Tensor

\lambda_{\max}(H) (algebraically largest) by power iteration.

With differentiable=True returns the Rayleigh quotient at the (detached) top eigenvector, whose gradient is the exact Hellmann-Feynman derivative v^\top (\partial_\theta H)\, v.

Parameters

Name	Type & default	Contract
losspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
paramspositional-or-keyword	ParamsRequired	See the operation contract and implementation below.
iterskeyword-only	intDefault: 30	See the operation contract and implementation below.
generatorkeyword-only	torch.Generator | NoneDefault: None	See the operation contract and implementation below.
differentiablekeyword-only	boolDefault: False	See the operation contract and implementation below.

ReturnsTensor

Executed regression · exact sam reduces loss for every measure

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_exact_sam_reduces_loss_for_every_measure(measure: str) -> None:
    net, x, y, params = _problem(seed=0)

    def closure() -> torch.Tensor:
        return _mse(net, x, y)

    f0 = float(closure().detach())
    opt = ExactSAM(params, lr=1e-2, lam=1e-3, measure=measure, n_samples=4, iters=15, seed=0)
    last = None
    for _ in range(60):
        last = opt.step(closure)
    assert last is not None and torch.isfinite(last)
    assert float(closure().detach()) < f0

python -m pytest "packages/omnibias-curvature/tests/test_torch_exact_sam.py::test_exact_sam_reduces_loss_for_every_measure[top_eig]" -q

Regression example · power iteration matches dense top eigenvalue

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_power_iteration_matches_dense_top_eigenvalue():
    net, X, Y, params = _small_problem(seed=2)
    Hd = S.dense_hessian(_mse(net, X, Y), params)
    ev = torch.linalg.eigvalsh(0.5 * (Hd + Hd.T))
    lam = float(S.top_eigenvalue(_mse(net, X, Y), params, iters=80,
                                 generator=torch.Generator().manual_seed(0)))
    assert abs(lam - float(ev[-1])) < 1e-5
    lo, hi = S.hessian_eigenvalue_extremes(_mse(net, X, Y), params, iters=120,
                                           generator=torch.Generator().manual_seed(0))
    assert abs(hi - float(ev[-1])) < 1e-5
    assert abs(lo - float(ev[0])) < 1e-3  # min converges slower; looser

python -m pytest packages/omnibias-curvature/tests/test_torch_sharpness.py -k test_power_iteration_matches_dense_top_eigenvalue -q

Read the implementation

def top_eigenvalue(
    loss: Tensor, params: Params, *,
    iters: int = 30, generator: torch.Generator | None = None,
    differentiable: bool = False,
) -> Tensor:
    r""":math:`\lambda_{\max}(H)` (algebraically largest) by power iteration.

    With ``differentiable=True`` returns the Rayleigh quotient at the (detached)
    top eigenvector, whose gradient is the exact Hellmann-Feynman derivative
    :math:`v^\top (\partial_\theta H)\, v`.
    """
    op = HessianOperator(loss, params)
    v_top, lam = _top_eigvec_and_value(op, iters=iters, generator=generator)
    if not differentiable:
        return lam
    hv = op.hvp(v_top, create_graph=True)
    return _flat_dot(v_top, hv)

Related APIsomnibias.curvature.torch.HessianOperator↗

Other public import paths

• omnibias.curvature.torch.top_eigenvalue
• omnibias.curvature.torch.sharpness.top_eigenvalue

← Package overviewRunnable learning examples →