bayesmithDocumentation0.10.0 source

bayesmith.optimize.certify

Source-derived signatures and contracts. See the reference index for coverage and conventions.

Module contract
How far an iterative optimiser is from the minimum, and when to ask.

Everything here takes a real-valued objective ``f`` over a pytree of arrays
and nothing else: no model, no data, no parameter names. It answers two
questions a block-coordinate or alternating optimiser cannot answer from its
own iterates.

**Where am I?** :func:`decrement` computes the Newton decrement
``lambda = sqrt(g^T H^-1 g)`` at a point, for ``g`` and ``H`` the gradient and
Hessian of ``f`` there. Near a minimum ``f`` is quadratic and
``f - f* = lambda**2 / 2``, so ``lambda`` is the distance to the minimum in
units of the curvature's own standard deviation -- which does not grow with
the number of terms in ``f`` and does not care which coordinates are slow.
A tolerance on the CHANGE of ``f`` between iterations answers neither: it is
relative to ``|f|``, and the contraction it can be extrapolated at is the
fastest mode still moving, so a slow mode hidden under a fast one passes.

**Can I trust the number?** The Hessian is never formed from the model --
every product is ``jax.jvp`` of ``jax.grad`` -- and the linear solve behind
``H^-1 g`` is inexact. :class:`Decrement` therefore carries an upper bound,
not only an estimate. With ``r = g - H x`` the true residual of the computed
``x`` and ``mu`` a LOWER bound on the smallest eigenvalue of ``H``,

    lambda  <=  (R + sqrt(R**2 + 4 g^T x)) / 2,      R = |r| / sqrt(mu)

by Cauchy-Schwarz in the ``H^-1`` inner product (``lambda**2 - g^T x = g^T
H^-1 r <= lambda R``), inflated by ``1 / sqrt(1 - eps kappa)`` for the
rounding of ``g`` and ``H`` themselves. :meth:`Decrement.certifies` reads
that bound, so an inexact solve can only make a verdict refuse.

**Everything turns on that lower bound, and only a proof will do.** A
:class:`Decrement` says where its floor came from, and
:meth:`Decrement.certifies` answers ``True`` for two of the three:

* ``"dense"`` -- the formed Hessian's own smallest eigenvalue, computed (up
  to :data:`DENSE_MAX` parameters). It also decides positive definiteness.
* ``"supplied"`` -- the caller's ``floor=``, which is a CLAIM the caller must
  be able to prove. A model with proper Gaussian priors can: ``H = J^T N^-1 J
  + P`` is bounded below by the prior precision ``P`` whenever the map is
  linear in every parameter jointly and the noise does not depend on it.
* ``"probe"`` -- a :data:`PROBE_STEPS`-step Lanczos probe of ``H``, whose
  Ritz interval is an ESTIMATE and not a bound. ``theta[0] - beta |s[0]|``
  says how far ``theta[0]`` is from SOME eigenvalue, not from the smallest,
  so where the Krylov space never reaches the bottom of the spectrum the
  probe's floor can sit above it: measured over 480 random spectra, 5 in 240
  did in float32, the worst by a factor 7.4. A probe therefore never
  certifies here. It still refuses -- non-positive curvature, or a bound
  already outside the caller's limit -- which is what it is kept for.

``"none"`` is the fourth value: no usable floor at all -- the dense path's
smallest eigenvalue was not positive, or the probe's interval did not clear
zero -- which is a refusal by itself.

**Getting there.** :func:`polish` takes Newton steps -- a Steihaug-truncated
conjugate-gradient direction, an Armijo backtracking line search, and each
step kept only if it lowers ``f`` -- which is what removes a first-order
optimiser's step-size floor before the decrement is asked.

**When to ask.** A decrement costs a gradient and a solve, so
:func:`settled`, :func:`gap_step` and :func:`change_program` are the cheap
tests that pick the iterations worth spending one on. They schedule; none of
them certifies.

**Where this came from, and what was checked before it moved.** Lifted from
``rheplicant.inference.certify`` at ``e1acb6f`` (T-004 half B), byte for byte:
the module took plain callables and pytrees and imported nothing from that
package, so there was nothing to adapt. Its 50 unit tests came with it, as
``tests/test_certify.py``. The acceptance test that drives a whole estimate to
its MAP stays upstream, because it needs that package's models.

It arrived having been through five independent reviews there. What they
established, recorded here because a reader of THIS repository cannot see
them: the bound's derivation holds; both proven floors were attacked
row by row (over four model families at 1200 latents, every claimed lower
bound sat below the true smallest eigenvalue); a zero gradient and a singular
Hessian refuse on both the dense and the conjugate-gradient path; and 23
mutants were each killed by a named assertion. The probe's unsoundness as a
floor is the measurement in the list above and is why :meth:`Decrement.
certifies` reads :data:`PROVEN_FLOORS` rather than the floor's value.

Limits. The decrement is a statement about the quadratic model of ``f`` at
the point, so where the curvature changes over one unit of distance it is
local. Above :data:`DENSE_MAX` parameters with no ``floor``, nothing here
certifies: the probe is a heuristic, and a direction of negative curvature
that neither it nor the Krylov space meets is not seen. In a precision
whose epsilon times ``kappa`` approaches one, the gradient and the Hessian
are rounding and nothing here recovers them; the status says so instead.
RTOL MAXITER DENSE_MAX BATCH PROBE_STEPS PROVEN_FLOORS STATUS_SAID Decrement real_size decrement_program decrement POLISH_ITERATIONS POLISH_CG_MAXITER POLISH_CG_TOL ARMIJO ARMIJO_HALVINGS polish SETTLED_CHANGES OBJECTIVE_FLOOR_EPS RESOLUTION_EPS SOLVE_TOL_STEP SOLVE_TOL_FLOOR settled effective_tol solve_tol_floor change_program GapState gap_step

RTOL

Defined in bayesmith.optimize.certify, line 107.

RTOL: dict[int, float] = {4: 0.0001, 8: 1e-08}

MAXITER

Defined in bayesmith.optimize.certify, line 113.

MAXITER: int = 1000

DENSE_MAX

Defined in bayesmith.optimize.certify, line 129.

DENSE_MAX: int = 1024

BATCH

Defined in bayesmith.optimize.certify, line 130.

BATCH: int = 16

PROBE_STEPS

Defined in bayesmith.optimize.certify, line 136.

PROBE_STEPS: int = 32

PROVEN_FLOORS

Defined in bayesmith.optimize.certify, line 150.

PROVEN_FLOORS: frozenset[str] = frozenset({FLOOR_DENSE, FLOOR_SUPPLIED})

STATUS_SAID

Defined in bayesmith.optimize.certify, line 153.

STATUS_SAID: dict[int, str] = {CONVERGED: 'reached its residual', UNREACHED: 'did not reach its residual within its iteration cap', NONCONVEX: 'met a direction of non-positive curvature, so the point is not near a minimum', NONFINITE: 'was not finite'}

Decrement

Defined in bayesmith.optimize.certify, line 163.

class Decrement
estimate: float
distance: float
lambda2: float
residual: float
reach: float
kappa: float
products: int
dense: bool
floor_source: str
status: int
One Newton decrement: what it measured, and how well.

Attributes:
    estimate: ``sqrt(g^T x)`` for the computed ``x``, in units of the
        curvature's standard deviation -- the distance to the minimum of
        the objective's quadratic model, as the solve left it, and a
        lower bound on it.
    distance: the upper bound on that distance, ``inf`` when the residual
        or the curvature cannot bound it. This is what a verdict reads.
    lambda2, residual, reach, kappa: ``g^T x``; the true residual's size
        relative to ``g``; that residual divided by the square root of
        the curvature floor, which is the ``R`` of the bound; and the
        condition number the rounding allowance was computed from.
    products: Hessian-vector products spent, the residual's own and any
        spectral probe's included.
    dense: whether the Hessian was formed and solved directly.
    floor_source: ``"dense"``, ``"supplied"``, ``"probe"`` or ``"none"``
        -- where the curvature the bound divides by came from, and
        whether it is a proof (see the module docstring).
    status: one of :data:`CONVERGED`, :data:`UNREACHED`,
        :data:`NONCONVEX`, :data:`NONFINITE`.

Decrement.proven

proven(self) -> bool
Whether the floor the bound rests on is a proof and not an estimate.

Decrement.certifies

certifies(self, limit: float) -> bool
Whether the point is PROVEN within ``limit`` of the minimum.

Three things, and a probed floor fails the first however good the
other two look: the floor is a proof, the solve converged, and the
upper bound it allows is inside ``limit``.

real_size

Defined in bayesmith.optimize.certify, line 213.

real_size(template: Any) -> int
Real degrees of freedom in a pytree: a complex leaf counts twice.

What :data:`DENSE_MAX` is compared against, so a caller can ask the same
question before paying for anything.

decrement_program

Defined in bayesmith.optimize.certify, line 223.

decrement_program(objective: Callable[[Any], jax.Array], template: Any, *, floor: float | None=None, limit: float | None=None) -> Callable[[Any], tuple[jax.Array, ...]]
``values -> (lambda2, rho, reach, products, status, kappa)``, jitted once.

``objective`` is a real function of a pytree shaped like ``template``; a
complex leaf enters as its real and imaginary parts, so ``n`` counts real
degrees of freedom. Build this once per shape and call it at each point:
it is the expensive half, and :func:`decrement` is the reading of its
output.

Up to :data:`DENSE_MAX` parameters the Hessian is formed from ``n``
products and solved by an eigendecomposition after scaling by its
diagonal, and the curvature floor the bound needs is that scaled matrix's
own smallest eigenvalue. Above it, ``H x = g`` is conjugate gradients
from zero, unpreconditioned, and the floor is either ``floor`` -- a
verified lower bound on the smallest eigenvalue of ``H``, which is what
makes a certificate there sound -- or, failing that, a
:data:`PROBE_STEPS`-step Lanczos probe of ``H``.

``limit`` is the distance the caller will certify. Given one, the
iteration stops as soon as the verdict cannot change: when the bound is
inside ``limit``, or when ``g^T x`` alone (which only grows) is already
outside it. Without one it stops at a relative residual of
:data:`RTOL`. Either way the residual returned is ``|g - H x|``
recomputed with one more product, because a recursive residual drifts
from the true one in finite precision, and it is the true one the bound
needs.

decrement

Defined in bayesmith.optimize.certify, line 315.

decrement(objective: Callable[[Any], jax.Array], at: Any, *, program: Callable[[Any], tuple[jax.Array, ...]] | None=None, floor: float | None=None, limit: float | None=None) -> Decrement
The Newton decrement of ``objective`` at ``at``, with its error bound.

``program`` is a :func:`decrement_program` built earlier for this shape;
without one this builds (and compiles) a program for this call alone, and
``floor`` and ``limit`` are passed to it.

The bound is ``(R + sqrt(R**2 + 4 lambda2)) / 2`` for ``R`` the residual
measured against the curvature floor, inflated by ``1 / sqrt(1 - eps
kappa)``: a relative perturbation ``eps`` of ``H`` moves ``g^T H^-1 g`` by
up to ``eps kappa`` of itself, and ``g`` and every Hessian product are
rounded. That second term is negligible in an 8-byte float and refuses
any Hessian with ``kappa`` near ``1 / eps`` in a 4-byte one, where the
decrement's digits are rounding.

POLISH_ITERATIONS

Defined in bayesmith.optimize.certify, line 515.

POLISH_ITERATIONS: int = 3

POLISH_CG_MAXITER

Defined in bayesmith.optimize.certify, line 523.

POLISH_CG_MAXITER: int = 50

POLISH_CG_TOL

Defined in bayesmith.optimize.certify, line 528.

POLISH_CG_TOL: float = 1e-05

ARMIJO

Defined in bayesmith.optimize.certify, line 534.

ARMIJO: float = 0.0001

ARMIJO_HALVINGS

Defined in bayesmith.optimize.certify, line 538.

ARMIJO_HALVINGS: int = 30

polish

Defined in bayesmith.optimize.certify, line 586.

polish(objective: Callable[[Any], jax.Array], at: Any, *, iterations: int=POLISH_ITERATIONS) -> tuple[Any, jax.Array]
Newton steps on ``objective``, each kept only if it helps.

Returns ``(point, kept)``, ``kept`` counting the steps that were taken --
zero means the point came back unchanged.

A first-order optimiser restarted at each outer iteration has a precision
floor set by its step size rather than by the problem: Adam with zeroed
moments takes a first bias-corrected step of ``step_size *
sign(gradient)`` however small the gradient is, and the resulting map has
a fixed point a fraction of a step size from the optimum -- measured at
2881 curvature standard deviations on a two-parameter power law fitted
over 4096 channels, at every inner step count, shrinking with the step
size and not with the iterations.

A Newton step uses the curvature, so its length goes to zero with the
gradient and the iterate lands on the optimum wherever the first-order
steps left it. The direction is :func:`_steihaug`'s: conjugate gradients
on Hessian-vector products, no Hessian formed, stopped at negative
curvature. The step along it is backtracked (halved up to
:data:`ARMIJO_HALVINGS` times) until Armijo's condition holds,
``f(x + t p) - f(x) <= ARMIJO * t * g.p`` with ``g.p < 0`` and the result
finite; otherwise the iteration keeps its point. The decrease is compared
as a DIFFERENCE, because ``f(x) + small`` rounds to ``f(x)`` when ``|f|``
is large: measured, ``after <= before`` accepted a 4-byte-float jump from
``y = 3`` to ``y = -97.9`` on a potential offset by 1e7, where no step
could be resolved at all.

A complex parameter set is returned unchanged: the Hessian-vector product
of a real function of complex arguments is not the operator this step
needs.

SETTLED_CHANGES

Defined in bayesmith.optimize.certify, line 665.

SETTLED_CHANGES: int = 2

OBJECTIVE_FLOOR_EPS

Defined in bayesmith.optimize.certify, line 678.

OBJECTIVE_FLOOR_EPS: int = 64

RESOLUTION_EPS

Defined in bayesmith.optimize.certify, line 688.

RESOLUTION_EPS: float = 4.0

SOLVE_TOL_STEP

Defined in bayesmith.optimize.certify, line 695.

SOLVE_TOL_STEP: float = 0.01

SOLVE_TOL_FLOOR

Defined in bayesmith.optimize.certify, line 696.

SOLVE_TOL_FLOOR: dict[int, float] = {4: 2.4e-07, 8: 1e-12}

settled

Defined in bayesmith.optimize.certify, line 699.

settled(trace: list[float], tol: float) -> bool
Whether the last :data:`SETTLED_CHANGES` changes of ``trace`` are within ``tol``.

``trace`` holds the objective after each iteration. Each change is
measured in both directions, ``|f[k] - f[k-1]|``, against ``tol *
max(|f[k]|, 1)``, so a rise larger than that is never settled; and two
changes rather than one, so an iteration that happens to land at the same
value from the other side of a minimum is not read as a fixed point.

effective_tol

Defined in bayesmith.optimize.certify, line 716.

effective_tol(tol: float, objective: jax.Array) -> float
``max(tol, OBJECTIVE_FLOOR_EPS * eps)`` for the objective's dtype.

solve_tol_floor

Defined in bayesmith.optimize.certify, line 722.

solve_tol_floor(objective: jax.Array) -> float
How far an inner solve's tolerance may be tightened, by precision.

:data:`SOLVE_TOL_FLOOR` for the precisions it names, and two machine
epsilons for any other -- a table lookup would raise on those, and a run
should not fail because its dtype is unusual.

change_program

Defined in bayesmith.optimize.certify, line 733.

change_program(resolution_eps: float=RESOLUTION_EPS) -> Callable
``(terms0, scales0, terms1, scales1) -> (decrease, resolution)``, jitted.

The change of an objective given as per-element TERMS, summed term by
term, and the rounding that change carries: ``resolution_eps * eps *
sqrt(sum (scale0 + scale1)**2)``, each term's two evaluations added as
independent errors.

Term by term because the difference of two totals is resolved only to
``eps |f|``, and ``|f|`` carries every constant the change does not:
near convergence the two evaluations of one term are within a factor of
two of each other, so their difference is exact in the working dtype
(Sterbenz), and the sum of small differences carries almost none of the
total's rounding.

GapState

Defined in bayesmith.optimize.certify, line 760.

class GapState
contraction: float | None = None
moved: bool = False
decrease: float | None = None
resolution: float | None = None
What :func:`gap_step` carries from one iteration to the next.

Attributes:
    contraction: the last contraction estimated from a decrease the
        arithmetic resolved, or ``None`` when there is none (none yet, or
        a rise since).
    moved: whether any decrease so far exceeded its resolution.
    decrease, resolution: the previous iteration's decrease and its
        resolution.

gap_step

Defined in bayesmith.optimize.certify, line 778.

gap_step(state: GapState, decrease: float, resolution: float, tol: float) -> tuple[GapState, bool, float | None, float | None]
One iteration of the gap PRE-SCREEN: ``(state, passed, gap, rho)``.

``decrease`` is ``D[k] = f[k-1] - f[k]`` and ``resolution`` bounds its
rounding. The contraction is ``max(D[k], 0) / D[k-1]``, clipped to 1,
re-estimated whenever ``D[k-1]`` exceeds its resolution (the last
estimate is kept otherwise, and a run that never moved beyond the
resolution uses 0); the gap is ``(max(D[k], 0) + resolution) / (1 - rho)``,
the decrease plus its geometric tail, and the iteration passes when that
is at most ``tol``. A rise the arithmetic resolves passes nothing and
forgets the contraction: an exact conditional minimisation cannot raise
the objective.

This is a screen and not a certificate. The ratio reads the FASTEST mode
still moving, so a slow mode with a gap left under a fast one's decreases
passes it -- measured at 0.5 to 1.0 curvature standard deviations on two
correlated pairs, and at 1.33 over random dense precisions. What
certifies is :func:`decrement`, which an outer loop computes on an
iteration that passes this screen, so that it pays one solve per
candidate stop and not one per iteration.

Type a concept or API name.