pyzag.preconditioning

Two ways to use the Gauss-Newton curvature of a calibration residual, sharing one estimator. Pick by whether the curvature drifts over the fit:

lever	what it does	optimizers	refreshes?
GaussNewtonPreconditioner	reshapes the gradient each step	SGD-family only	yes
gauss_newton_rescalers()	changes coordinates once	any, incl. Adam	no

Adam and its relatives divide each coordinate by its own running gradient RMS, so they are exactly invariant to gradient preconditioning – the preconditioner rejects them rather than silently doing nothing. Use the reparametrization lever with those.

Gauss-Newton curvature for least-squares model calibration: two levers.

When calibrating a model against data by minimizing a least-squares loss L(theta) = 1/2 r(theta)^T W r(theta) (r the residual, W a diagonal weight), the raw gradient is badly scaled whenever the parameters have heterogeneous magnitudes or sensitivities – which slows first-order optimizers. The usual fix, pyzag.reparametrization, rescales each parameter by a hand-picked range (ub - lb); that works but demands prior knowledge of every parameter’s range.

This module offers data-driven alternatives that need no ranges, both built on the Gauss-Newton curvature H = J^T W J (J = dr/dtheta) and both sharing one estimator, CurvatureEstimator. It never forms the dense Jacobian: H comes from a small subsample of residual rows, one reverse-mode sweep each.

Lever 1 – preconditioning (GaussNewtonPreconditioner), a wrapper around a torch optimizer that reshapes the gradient every step:

theta <- optimizer_update( (H + lam * diag(H))^{-1} @ grad )

It can refresh H as the fit moves, and at lr=1 with SGD the update is exactly the damped Gauss-Newton step, so there is nothing to tune. It requires an optimizer whose step is proportional to its gradient – see the note below.

Lever 2 – static reparametrization (gauss_newton_rescalers() + pyzag.reparametrization.CurvatureRescale), which estimates H once and changes coordinates by 1 / sqrt(diag H):

scalers = gauss_newton_rescalers(model.named_parameters(), residual_closure)
Reparameterizer(scalers)(model)
opt = torch.optim.Adam(model.parameters(), lr=...)   # any optimizer

Because it is a reparametrization the optimizer’s own state moves into the scaled coordinates with the metric, so it works with any optimizer, Adam included. It cannot refresh, though: re-scaling mid-run would invalidate a stateful optimizer’s moments.

Which to reach for:

curvature over the fit	lever
stable	reparametrization – one estimate, any optimizer
drifts	preconditioning – refreshes, SGD-family only

Both are alternatives to a hand-picked RangeRescale, and lever 2 additionally separates the two jobs a range width does today: it takes the step scale from the data, leaving lb / ub free to be honest bounds rather than a compromise between bounding and conditioning.

Damping in lever 1 is Marquardt (lam * diag(H)), invariant to parameter scales.

The gain ratio – observed loss reduction over the reduction the quadratic model predicted – drives two independent mechanisms, and it is worth keeping them apart:

knob	question	action	cost
rho	is the cached H stale?	recompute H	nsub sweeps
lam_adapt	is the step too long?	adapt lam, reject bad step	free

rho=None turns off refreshing (compute H once and reuse it) but leaves damping active – which matters more in that configuration, since refreshing is no longer available as a corrective. Damping costs nothing extra: it reuses loss values the training loop already has.

Important

The base optimizer’s step must be proportional to the gradient it is given – SGD, with or without momentum. That proportionality is what makes the update the damped Gauss-Newton step.

Adam and its relatives do not qualify. They divide each coordinate by its own running gradient RMS, so multiplying the gradient by a diagonal leaves their step exactly unchanged; wrapping one is a no-op, not a weaker effect. The constructor rejects them for that reason. To condition a problem you intend to optimize with Adam, change coordinates instead – pyzag.reparametrization composes with Adam precisely because a reparametrization also moves the optimizer’s own state into the new space, which a preconditioner applied to the gradient cannot do.

Typical use (no parameter ranges required):

opt = torch.optim.SGD(model.parameters(), lr=1.0)
pre = GaussNewtonPreconditioner(opt, model.parameters(), rho=0.25, nsub=8)
for _ in range(niter):
    loss = pre.step(lambda: model(time, temperature, loading) - data)

The closure returns the residual vector r(theta) (differentiable w.r.t. the parameters); the preconditioner computes the gradient and curvature from it. At lr=1 with SGD the update is exactly the damped Gauss-Newton step.

Forming H lives in pyzag.curvature (CurvatureEstimator), shared by both levers. GaussNewtonCurvature extends it with the caching, staleness and damping a training loop needs. GaussNewtonPreconditioner drives a plain torch loop; pyzag.stochastic.PyroGaussNewtonOptim drives Pyro SVI, where the training loop belongs to somebody else.

Scope / notes:

• Physical bounds are intentionally not handled here – a preconditioner only rescales the step direction. Enforce bounds separately (e.g. a projection after step or a bounds-only pyzag.reparametrization).

• Non-finite residuals or preconditioned gradients (which stiff models can produce) cause the update to be skipped with a warning, never silently.

• For mode="full" use nsub (or len(sample_indices)) at least the number of parameters, otherwise the sampled H is rank deficient.

class pyzag.preconditioning.GaussNewtonCurvature(parameters, *, rho=0.25, lam=0.01, lam_adapt=True, min_refresh_interval=1, on_refresh=None, nsub=8, **estimator_kwargs)

A CurvatureEstimator plus the state a training loop needs.

Caches H, judges when it has gone stale (should_refresh()), and runs the Levenberg-Marquardt damping schedule (adapt_damping()). A driver calls begin_step(), gain_ratio(), adapt_damping(), should_refresh(), refresh(), precondition() and note_step() in that order.

Keyword Arguments:

• rho (float or None) – gain-ratio threshold; None never refreshes after the first step (a fixed preconditioner).

• lam (float) – initial Marquardt damping factor.

• lam_adapt (bool) – adapt lam from the gain ratio – grow it when a step underdelivers, shrink it when a step delivers, reject a step that increased the loss. See adapt_damping().

• min_refresh_interval (int) – steps that must reuse the cached curvature before it may be refreshed again. The default 1 allows a refresh at most every other step; pass 0 to allow one on every step.

• on_refresh (callable, optional) – called on every refresh as on_refresh(step=..., gain_ratio=..., n_refresh=...).

Every other keyword is forwarded to CurvatureEstimator.

property H

a length-p vector (diag) or p x p matrix (full), or None before the first refresh.

Type:

The cached curvature

adapt_damping(loss_v, gain_ratio)

Levenberg-Marquardt damping update, driven by the gain ratio.

Gauss-Newton drops the sum(r_i * grad^2 r_i) term of the true Hessian, and that term is not always small. For a Gaussian scale parameter in unconstrained coordinates (sigma = exp(u)) the dropped term is r^2 – exactly as large as the Gauss-Newton term r^2 itself – so the undamped step is a systematic factor-of-two overshoot in every scale coordinate. Damping is what absorbs that: on a step whose observed reduction falls short of the quadratic model’s prediction, lam grows and the step shrinks towards gradient descent; on a step that delivers, lam decays back towards the full Newton step.

Rejection keys off the observed loss increase, not off gain_ratio. The two are not interchangeable: gain_ratio() returns a -1.0 sentinel when the quadratic model predicted an increase, which says the curvature is stale, not that the step was bad – the actual loss may well have fallen. Reading that sentinel as a regression would back out perfectly good steps.

Parameters:

• loss_v (float) – objective at the current parameters.

• gain_ratio (float or None) – from gain_ratio().

Returns:

True if the previous step should be rejected – the loss went up, so the caller should call undo_last_step().

Return type:

bool

begin_step()

Advance the step counter. Returns the new 0-based step index.

gain_ratio(loss_v)

Gain ratio of the PREVIOUS step: observed vs. predicted loss reduction.

Returns None when there is no anchored quadratic model, or when the prediction is too small to divide by. Only meaningful when the previous step’s quadratic model predicted a non-negligible change; at/near convergence pred_red -> 0 and the ratio is 0/0 noise, which must not be mistaken for staleness.

note_refresh(gain_ratio, refreshed)

Record this step’s refresh decision.

Takes the decision rather than being one of a pair of methods, so a caller cannot handle one branch and forget the other.

note_step(g, dtheta, loss_v, theta=None, opt_state=None)

Anchor the quadratic model at the step that just completed.

dtheta must be the actual parameter displacement the optimizer produced, not the proposed direction, so the gain ratio accounts for the base optimizer’s learning rate and momentum.

precondition(g)

Return (H + lam * diag(H))^{-1} g, Marquardt-damped and scale-invariant.

refresh(residual, w_flat=None)

Re-estimate H from residual and cache it.

w_flat is accepted for backwards compatibility and ignored – the weights come from weights.

reject_failed_evaluation(optimizer)

Back out the last step after the model could not be evaluated at all.

adapt_damping() compares two loss values, so it cannot react to a point where the objective does not exist – a stiff forward model driven outside its valid region raises instead of returning a number, and the exception escapes before any comparison is possible. That is simply the limiting case of a step that was too long, so it is handled the same way, with a harder damping bump because the evidence is stronger.

Returns:

True if a good point was available to fall back to. When False there is nothing to undo – the very first evaluation failed – and the caller should let the error propagate rather than pretend to recover.

Return type:

bool

request_refresh()

Force a refresh on the next step.

reset_anchor()

Drop the cached quadratic model (its anchor point is no longer valid).

set_H(H, sweeps=0)

Install a curvature computed elsewhere – e.g. a structure-aware estimator that assembles H from several blocks, each estimated by the method that suits it – counting it as a refresh.

Parameters:

• H (Tensor) – length-p vector in "diag" mode, p x p matrix in "full" mode.

• sweeps (int) – number of reverse-mode sweeps it cost, for accounting.

should_refresh(gain_ratio)

Whether the curvature must be (re)computed on this step. With rho=None it is computed once and never refreshed.

undo_last_step(optimizer)

Restore the parameters and the optimizer state to before the last step.

Restoring the parameters alone is not enough: a stateful optimizer’s momentum buffer still holds the rejected direction and would immediately re-apply it, so the same step is proposed, rejected, and re-proposed forever while lam ratchets to its ceiling. The optimizer state has to rewind with the parameters.

class pyzag.preconditioning.GaussNewtonPreconditioner(optimizer, parameters, *, mode='diag', nsub=8, sample_indices=None, cotangents=None, rho=0.25, lam=0.01, lam_adapt=True, weights=None, min_refresh_interval=1, generator=None, on_refresh=None, check_optimizer=True)

Wrap a torch optimizer and precondition its gradient with the Gauss-Newton curvature, refreshed on a gain-ratio trigger.

Parameters:

• optimizer (torch.optim.Optimizer) – the base optimizer to wrap. Its param_groups must cover exactly parameters.

• parameters (iterable of Parameter) – the parameters being calibrated.

Keyword Arguments:

• mode (str) – "diag" (default) preconditions with the diagonal of H; "full" uses the full p x p matrix.

• nsub (int) – number of residual rows to subsample when estimating H (one reverse-mode sweep each). Ignored if sample_indices is given.

• sample_indices (array-like of int, optional) – explicit rows of the flattened residual to sample (overrides nsub). Use this to sample representatively across conditions; a poorly chosen stride can alias with batch structure and starve some parameters, so when in doubt leave it None and let the estimator draw nsub rows at random.

• cotangents (callable or sequence, optional) – custom cotangent vectors, overriding nsub and sample_indices. Each costs one sweep and may select a group of rows – valid only when those rows have disjoint parameter support. See GaussNewtonCurvature.refresh().

• rho (float or None) – gain-ratio threshold. H is refreshed when the observed loss reduction of the previous step is below rho times the value the cached curvature predicted – i.e. when the curvature model has gone stale. Larger rho refreshes more eagerly. Set to None to never refresh after the first step: a fixed preconditioner, the cheapest option and often sufficient when the curvature profile is stable (e.g. a diagonal), where it acts as a data-driven analogue of a static per-parameter rescaling.

• lam (float) – initial Marquardt damping factor; the solve uses H + lam * diag(H).

• lam_adapt (bool) – Levenberg-Marquardt damping adaptation – see GaussNewtonCurvature.adapt_damping(). On by default because the second-order term Gauss-Newton drops is not always small.

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

• min_refresh_interval (int) – number of steps that must reuse the cached curvature before it may be refreshed again – a hard floor on refresh cost. The default 1 therefore allows a refresh at most every other step; pass 0 to allow one on every step.

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

• on_refresh (callable, optional) – called whenever the curvature is refreshed, as on_refresh(step=<0-based step index>, gain_ratio=<float or None>, n_refresh=<count>). The refreshed step indices are also recorded on self.refresh_steps (useful for annotating plots with the recompute iterations).

• check_optimizer (bool) – reject a base optimizer that is invariant to gradient preconditioning (Adam and relatives), which would make the whole wrapper a silent no-op. See check_optimizer_respects_gradient_scale().

property lam

Marquardt damping factor.

property mode

"diag" or "full".

property n_refresh_sweeps

Total reverse-mode sweeps spent refreshing the curvature.

property n_refreshes

Number of curvature refreshes.

property n_steps

Number of completed updates.

property parameters

The parameters being calibrated.

property refresh_steps

Step indices at which the curvature was refreshed.

property rho

Gain-ratio refresh threshold (None = fixed preconditioner).

step(residual_closure)

Evaluate the residual, precondition the gradient, and step the wrapped optimizer. Returns the scalar loss (float) at the current parameters, or nan if the residual was non-finite and the update was skipped.

Parameters:

residual_closure (callable) – returns the residual r(theta) as a tensor that is differentiable w.r.t. parameters.

zero_grad(*args, **kwargs)

Delegate to the wrapped optimizer (provided for drop-in familiarity).

pyzag.preconditioning.apply_preconditioned_update(curvature, g, optimizer, loss_v=None, theta_before=None)

Precondition g, install it as .grad, and let optimizer step.

The tail every driver shares, whoever owns the training loop. The optimizer is handed (H + lam)^-1 g – the damped Gauss-Newton direction – and its own rule (learning rate, momentum) scales it. The gradient is overwritten, not accumulated.

This requires an optimizer whose step is proportional to the gradient it is given; see check_optimizer_respects_gradient_scale().

Parameters:

• curvature (GaussNewtonCurvature) – the engine holding the parameters.

• g (Tensor) – the raw packed gradient.

• optimizer (torch.optim.Optimizer) – applies the update.

• loss_v (float, optional) – objective at the current parameters. When given, the quadratic model is anchored here for the next step’s gain-ratio test; when None the trigger stays inert.

• theta_before (Tensor, optional) – parameter vector before this step, recorded so a step that turns out to increase the loss can be backed out. Read from the parameters when omitted.

Returns:

whether the update was applied. False means the curvature produced a non-finite direction, so the parameters were left untouched and a refresh is forced for the next step.

Return type:

bool

pyzag.preconditioning.check_optimizer_is_memoryless(optimizer)

Warn if optimizer carries state that damping cannot reach.

Levenberg-Marquardt damping assumes the step is a function of the current damped gradient: raise lam, get a shorter step. A momentum buffer breaks that – it contributes its own accumulated velocity, so a step that overshoots keeps overshooting no matter how hard lam is raised. In practice the run enters a reject/restore cycle and lam ratchets to its ceiling while the loss stalls. Restoring the optimizer state alongside the parameters does not help, because the restored state is what holds the offending velocity.

pyzag.preconditioning.check_optimizer_respects_gradient_scale(optimizer)

Raise if optimizer would annihilate the preconditioner.

Gauss-Newton preconditioning replaces the gradient with (H + lam)^-1 g and relies on the optimizer’s step being proportional to what it is handed – that is what makes the update the damped Gauss-Newton step. Adam and its relatives normalize each coordinate by its own running gradient RMS, so multiplying the gradient by any fixed diagonal leaves their step exactly unchanged: wrapping one is a no-op, not a weaker effect.

That failure is silent and easy to mistake for a working configuration, so it is rejected outright rather than warned about. Use SGD (optionally with momentum), or reach for pyzag.reparametrization, which conditions the problem by changing coordinates and therefore composes with Adam.

pyzag.preconditioning.gauss_newton_rescalers(named_parameters, residual_closure, *, normalize=True, cond_max=1000000000000.0, bounds=None, offsets=None, **estimator_kwargs)

Build a static, curvature-derived reparametrization – the second lever.

Estimates H once at the current parameter values and returns one CurvatureRescale per parameter, scaled by 1 / sqrt(diag H). Feed the result straight to Reparameterizer:

scalers = gauss_newton_rescalers(model.named_parameters(), residual_closure)
Reparameterizer(scalers)(model)
opt = torch.optim.Adam(model.parameters(), lr=...)   # any optimizer

Use this rather than GaussNewtonPreconditioner when you want to keep an Adam-family optimizer (which is invariant to gradient preconditioning), or when the curvature is stable enough that one estimate serves the whole fit. Use the preconditioner instead when the curvature drifts, since only it can refresh. Cost here is a single estimate – nsub reverse sweeps – and nothing per step.

Parameters:

• named_parameters – (name, parameter) pairs, or a dict. Names must be the dotted paths Reparameterizer matches on.

• residual_closure (callable) – returns the residual at the current parameters, differentiable w.r.t. them.

Keyword Arguments:

• cond_max (float) – reject the estimate if max(H) / min(H) exceeds this. An under-sampled row set can leave a parameter’s curvature many orders of magnitude below the rest, and 1 / sqrt(H) then hands that parameter an essentially infinite step. The damage is invisible in H itself – on this calibration nsub=8 gives an H 17% off whose implied scale is wrong by twelve orders of magnitude – so the check has to be on the conditioning, not on the norm. Raise it only if your parameters genuinely differ that much in sensitivity.

• normalize (bool) – divide the scales by their geometric mean, so the learning rate sets the average physical step and the curvature sets only how that step is distributed across parameters. On by default, and you almost always want it: 1 / sqrt(H) carries units of theta / residual, so its absolute level is a property of the residual’s units rather than of the parameters. Measured on the NEML2 calibration it sits ~920x below the hand-picked range widths, which would freeze an optimizer tuned for those. Normalizing makes the learning rate mean the same thing it does for any other coordinate choice; the curvature still supplies the whole relative profile.

• bounds (dict, optional) – name -> (lb, ub) in natural units. Bounds are only bounds here – they do not influence the step scale, so they can be generous. Omitting them means no clamping at all, which is a real behavioural change if you are replacing a RangeRescale.

• offsets (dict, optional) – name -> offset, the natural value at scaled zero. Defaults to 0.

• **estimator_kwargs – forwarded to CurvatureEstimator (nsub, sample_indices, cotangents, weights, generator).

Returns:

name -> CurvatureRescale.

Return type:

dict

Raises:

ValueError – if any parameter has non-positive curvature. That means the residual does not depend on it (or not detectably), so there is no scale to derive and 1 / sqrt(H) is not defined – silently flooring it would hand back an arbitrary scale dressed up as a measurement.

The scaler

gauss_newton_rescalers() returns pyzag.reparametrization.CurvatureRescale objects, installed with the existing pyzag.reparametrization.Reparameterizer. See pyzag.reparametrization.

Preconditioning an SVI fit

The Pyro-facing half lives in pyzag.stochastic, because those classes subclass Pyro types and this module is deliberately importable without Pyro. Use pyzag.stochastic.PyroGaussNewtonOptim with pyzag.stochastic.PreconditionedSVI, and build the residual with pyzag.stochastic.gaussian_map_residual().