pyzag.curvature

Gauss-Newton curvature extraction, shared by everything that needs H.

H = J^T W J for a residual r(theta), formed without ever materializing the dense Jacobian: each sampled row costs one reverse-mode sweep. Two consumers build on this, and they want quite different things from it:

Keeping the extraction here means one implementation of the row sampling, the cotangent grouping and the packing, rather than one per consumer.

class pyzag.curvature.CurvatureEstimator(parameters, *, mode='diag', nsub=None, sample_indices=None, cotangents=None, weights=None, generator=None, sampling='random')

Form the Gauss-Newton curvature H = J^T W J of a residual.

The estimation half of Gauss-Newton, with no notion of an optimizer, a training loop, or a cached value that might go stale. It packs a parameter list into a flat vector, takes reverse-mode products against a residual, and turns a chosen set of rows into H. Two rather different consumers need exactly that and nothing more:

  • GaussNewtonCurvature adds caching, staleness and damping, to precondition a training loop step by step;

  • gauss_newton_rescalers() calls estimate() exactly once, to build a static reparametrization.

Parameters:

parameters (iterable of Tensor) – leaf tensors with requires_grad=True. Their shapes and order are frozen here.

Keyword Arguments:

  • mode (str) – "diag" (default) or "full".

  • nsub (int or None) – residual rows to subsample, one reverse sweep each. None (default) uses every row, which is exact. Subsampling is an opt-in trade: see estimate() for what it costs you. Ignored if sample_indices or cotangents is given.

  • sample_indices (array-like of int, optional) – explicit rows to sample.

  • cotangents (callable or sequence, optional) – custom cotangents; see estimate() for the precondition they carry.

  • weights (Tensor, optional) – diagonal of W, broadcastable to the residual shape. None means W = I.

  • generator (torch.Generator, optional) – RNG for reproducible subsampling.

  • sampling (str) –

    "random" (default) draws uniformly without replacement; "stratified" draws one row from each of nsub equal blocks of the flattened residual.

    Measured on the NEML2 calibration the two are equivalent: both give a usable scale above nsub ~ 32 and both fail below it, because a parameter whose sensitivity is concentrated in a small part of the residual can be missed either way. Stratified is offered because it cannot alias against a (ntime, nbatch) layout the way an explicit stride can, but it is not a substitute for a large enough nsub.

    Watch the scale, not H. 1 / sqrt(H) amplifies a near-zero entry without bound, so ||H - H_exact|| badly understates the damage: at nsub=8 that norm is 17% off while the scale it implies is wrong by twelve orders of magnitude.

estimate(residual)

Return (H, nsweeps), sampling rows as configured.

With a custom cotangents sequence each cotangent costs one sweep, but the vector it returns is the sum of the rows it selects. That sum is the correct contribution only when those rows have disjoint parameter support – otherwise the estimate picks up cross terms. Valid for block-diagonal Jacobians (one cotangent per time step across a pyro.plate of independent specimens); invalid for parameters shared across the grouped rows, which must be handled separately.

Custom cotangents also carry their own weighting: put sqrt(w_k) in entry k rather than 1, since (sum_k sqrt(w_k) J_k)^2 equals sum_k w_k J_k^2 when the supports are disjoint. They are used as given, with no n / nsampled extrapolation – a custom grouping describes its own coverage.

property nparam

Total number of scalar parameters p.

pack(grads)

Flatten a per-parameter gradient tuple to a length-p vector, substituting zeros for None (the allow_unused case).

pack_grads()

Flatten the parameters’ current .grad to a length-p vector.

theta()

Flat, detached copy of the current parameter values.

vjp(residual, cotangent)

One reverse-mode sweep: J^T cotangent, packed to a length-p vector.

weights_flat(residual)

The diagonal of W, flattened to match the flattened residual.

write_grads(vec)

Overwrite every parameter’s .grad with the matching slice of vec. Any previously accumulated gradient is discarded.