pyzag.stochastic

Tools for converting deterministc models implemented in pytorch to stochastic models

Besides the model-conversion machinery (MapNormal, HierarchicalStatisticalModel), this module hosts the Pyro side of Gauss-Newton preconditioning: PyroGaussNewtonOptim and PreconditionedSVI let pyzag.preconditioning drive an SVI fit. They live here rather than next to the curvature engine because they must subclass Pyro types at class-definition time; pyzag.preconditioning itself stays importable without Pyro.

class pyzag.stochastic.MapNormal(cov: float, loc_suffix: str = '_loc', scale_suffix: str = '_scale')

Bases: object

A map between a deterministic torch parameter and a two-scale normal distribution

Parameters:

cov – coefficient of variation used to define the scale priors

Keyword Arguments:

  • sep (str) – seperator character in names

  • loc_suffix – suffix to add to parameter name to give the upper-level distribution for the scale

  • scale_suffix – suffix to add to the parameter name to give the lower-level distribution for the scale

class pyzag.stochastic.HierarchicalStatisticalModel(base: Module, parameter_mapper: MapNormal, noise_prior: Tensor, update_mask: bool = False)

Bases: PyroModule

Converts a torch model over to being a Pyro-based hierarchical statistical model

Parameters:

  • base (torch.nn.Module) – base torch module

  • parameter_mapper (MapParameter) – mapper class describing how to convert from Parameter to Distribution

  • noise_prior (float) – scale prior for white noise

Keyword Arguments:

update_mask (bool) – if True, update the mask to remove samples that are not valid

forward(*args: Tensor, results: Tensor | None = None, weights: Tensor | None = None, **kwargs) → Tensor

Call the base forward with the appropriate args.

Parameters:

*args – arguments forwarded to the underlying model. At least one must be a tensor so the batch shape can be inferred.

Keyword Arguments:

  • results (torch.tensor or None) – results to condition on.

  • weights (torch.tensor or None) – weights on the results; defaults to ones.

training: bool
class pyzag.stochastic.GaussianResidual(flat, blocks, plates, obs=None)

Bases: object

The MAP residual of a Gaussian-family Pyro model, with its site structure.

For a Delta guide the negative log-joint is 0.5 * ||r||^2 + sum(log scale), so Gauss-Newton applies exactly to the 0.5 * ||r||^2 part. This object carries that r plus enough structure for a curvature estimator to exploit the model’s plates.

flat

the flat residual, differentiable w.r.t. the guide’s parameters.

Type:

Tensor

blocks

site name -> slice into flat.

Type:

dict

plates

site name -> tuple of enclosing plate names.

Type:

dict

obs_name

name of the observed site, if any.

Type:

str or None

obs_shape

unflattened shape of the observed block.

Type:

torch.Size

obs_scale

the observed site’s scale (the noise level).

Type:

Tensor or None

obs_plate_axis(plate_name)

Positive axis of plate_name within the observed batch shape.

property prior_names

Names of the latent (unobserved) blocks, in trace order.

prior_flat()

The latent blocks concatenated, differentiably.

pyzag.stochastic.gaussian_map_residual(model, guide, *args, **kwargs)

Assemble the MAP residual of a Gaussian-family Pyro model from a trace.

Runs the guide, replays the model under it, and turns every Gaussian-family sample site into whitened residual rows – (value - loc) / scale for Normal, value / scale for HalfNormal – including the observed site. Any poutine scale / mask on a site is folded into its rows, so the resulting least-squares objective matches the weighting the ELBO uses.

Gauss-Newton deliberately ignores the sum(log scale) log-normalizer terms of the log-joint: they are not of least-squares form, and dropping them is the standard Gauss-Newton approximation, not an oversight. The residual must include the prior rows, though – the likelihood alone has exactly zero gradient w.r.t. hierarchical hyper-parameters, so a likelihood-only curvature would silently leave them unpreconditioned.

Parameters:

  • model – the Pyro model (e.g. a HierarchicalStatisticalModel).

  • guide – the guide, typically AutoDelta – a point-mass guide is what makes the objective a deterministic least-squares problem.

  • *args – forwarded to both model and guide.

  • **kwargs – forwarded to both model and guide.

Returns:

GaussianResidual

pyzag.stochastic.hierarchical_gn_diagonal(curvature, gres, param_names, *, plate_name='samples', nsub=8, generator=None, validate=True)

Structure-aware diag(J^T J) for a hierarchical MAP problem.

A hierarchical model makes the generic row-subsampling estimator useless: with one latent block per plate member, drawing nsub random rows out of ntime * nmember touches at most nsub members and leaves every other member with zero curvature, while the N / nsub rescaling badly distorts the ones it does hit. This estimator instead splits the residual and uses the method each part deserves:

  • prior rows carry no forward model, so their whole Jacobian comes back from a single batched reverse pass – exact, and effectively free;

  • observed rows, plate-local parameters use one cotangent per slice across the plate, which is exact per member at nmember times fewer sweeps (see _obs_group_cotangents());

  • observed rows, the shared noise scale are analytic: the residual is (y - f) / eps, so d r_k / d eps = -r_k / eps and the whole column norm follows from r alone, with no sweep at all.

Grouped cotangents would double-count a parameter shared across the plate, so shared parameters are excluded from the swept estimate and handled by the analytic term. Hierarchical hyper-parameters are shared but touch the likelihood only through the prior, so the prior block already covers them; validate checks that on the first call rather than assuming it.

Parameters:
Keyword Arguments:

  • plate_name (str) – the plate whose members own private latents.

  • nsub (int) – number of plate-slices to sample from the observed block.

  • generator (torch.Generator, optional) – RNG for that subsample.

  • validate (bool) – check the shared-parameter assumption (one extra sweep).

Returns:

(diag, nsweeps) – the length-p curvature diagonal and the number of reverse-mode sweeps it cost.

Return type:

tuple

class pyzag.stochastic.PyroGaussNewtonOptim(optim_constructor, optim_args, residual_closure, *, plate_name='samples', nsub=8, rho=0.25, lam=0.01, lam_adapt=True, min_refresh_interval=1, generator=None, on_refresh=None, curvature_fn=None, check_optimizer=True)

Bases: PyroOptim

Gauss-Newton preconditioning for Pyro SVI.

A stock pyro.optim.PyroOptim builds one torch optimizer per parameter, so it can never see the cross-parameter structure a preconditioner is made of. This subclass instead builds a single optimizer over every parameter at once, and rescales the gradient SVI just computed by the Gauss-Newton curvature before handing it over. SVI accepts it because it only checks isinstance(optim, PyroOptim).

Two details of SVI shape the design. Parameters are created lazily, on the first step, so the optimizer and the curvature engine are built on first call rather than in __init__. And SVI.step has already run backward and dropped the graph by the time the optimizer is invoked, so a refresh cannot reuse it – residual_closure re-evaluates the model to get a differentiable residual. That cost is paid only on refresh steps.

Because pyro hands over an unordered set of parameters, they are sorted by param-store name to give the packed vector a stable, reproducible layout.

Parameters:

  • optim_constructor – a torch optimizer class or factory, as for pyro.optim.PyroOptim.

  • optim_args (dict) – its keyword arguments. Unlike the base class this must be a plain dict – a per-parameter callable cannot describe a single optimizer covering all parameters.

  • residual_closure (callable) – returns the residual at the current parameters, differentiable w.r.t. them. Returning a GaussianResidual (from gaussian_map_residual()) selects the structure-aware hierarchical estimator; a plain tensor falls back to generic row subsampling.

Keyword Arguments:

  • plate_name (str) – plate whose members own private latents.

  • nsub (int) – plate-slices (or rows, generically) sampled per refresh.

  • rho (float or None) – gain-ratio refresh threshold; None computes the curvature once and reuses it.

  • lam (float) – initial Marquardt damping factor.

  • lam_adapt (bool) – Levenberg-Marquardt damping adaptation; see pyzag.preconditioning.GaussNewtonCurvature.adapt_damping().

  • min_refresh_interval (int) – minimum steps between refreshes.

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

  • on_refresh (callable, optional) – refresh callback, see pyzag.preconditioning.GaussNewtonCurvature.

  • curvature_fn (callable, optional) – full override, called as curvature_fn(curvature, residual, param_names) and returning (H, nsweeps).

  • check_optimizer (bool) – reject a base optimizer that is invariant to gradient preconditioning. SVI’s usual choice, ClippedAdam, is one of those – see pyzag.preconditioning.check_optimizer_respects_gradient_scale().

record_loss(loss)

Tell the optimizer the objective value at the current parameters.

The gain-ratio staleness test needs the loss, and SVI.step does not pass it to the optimizer. PreconditionedSVI calls this automatically; with a plain SVI call it yourself before each svi.step, or the curvature is computed once and never refreshed.

recover_from_failed_evaluation()

Back out the last step after the model failed to evaluate.

Called by PreconditionedSVI when the ELBO itself raises: with SVI the failure happens inside loss_and_grads, before the optimizer is ever invoked, so recovery has to be driven from the training loop.

property recorded_loss

The most recent loss passed to record_loss(), or None.

get_state()

Serializable state of the single inner optimizer.

set_state(state_dict)

Stage state to be loaded when the inner optimizer is built.

class pyzag.stochastic.PreconditionedSVI(model, guide, optim, loss, loss_and_grads=None, num_samples=0, num_steps=0, **kwargs)

Bases: SVI

SVI that reports each step’s loss to its optimizer.

Identical to the base class except that the loss is handed to the optimizer before the update. PyroGaussNewtonOptim needs the objective value at the current parameters to run its gain-ratio staleness test, and the stock SVI.step computes exactly that but returns it to the caller instead of passing it down. Optimizers without a record_loss method are unaffected.

step(*args, **kwargs)

Take one SVI step, recording the pre-update loss on the optimizer.

A model driven outside its valid region raises here, inside loss_and_grads – before the optimizer runs, so the optimizer cannot react on its own. If it can back the last step out and damp harder, do that and report the step as nan rather than letting the whole fit die on one bad point.