Skip to content

virgil.svi

variational fits a guide (an approximate posterior) to the posterior of numpyro_model with numpyro's SVI, taking the same arguments as fit and starting from a fit and its Gauss–Newton covariance. See Variational inference.

Variational inference with numpyro's SVI, started from a fit.

Provisional

May change in any 0.4.x release without a deprecation cycle.

variational takes the same arguments as fit and numpyro_model, and fits a guide (an approximate posterior) to the posterior by maximizing the evidence lower bound (ELBO) with numpyro's SVI. The guide works in the unconstrained coordinates NUTS samples, and starts at the Laplace approximation: its centre at a FitResult and its width from gauss_newton_mass. The default guide is a normalizing flow that bends that Gaussian to follow curved ridges and skewed scales; a Gaussian guide is close to the Laplace approximation, for which see laplace_cov. Every guide covers one mode.

See Variational inference for the choice of guide and how to check the result. References: ADVI (Kucukelbir et al. 2017, https://arxiv.org/abs/1603.00788), the block neural autoregressive flow (De Cao, Titov & Aziz 2019, https://arxiv.org/abs/1904.04676), the inverse autoregressive flow (Kingma et al. 2016, https://arxiv.org/abs/1606.04934) and Pareto-smoothed importance sampling (Vehtari et al. 2024, https://arxiv.org/abs/1507.02646).

VariationalResult dataclass

The result of variational.

Attributes:

Name Type Description
samples dict

Draws from the guide, num_samples per site, in the model's own parameters and under the site names NUTS uses (MCMC.get_samples()), deterministic sites included.
losses ndarray

The loss (the negative ELBO, a Monte Carlo estimate) at each step.
guide AutoGuide

The fitted guide; guide.sample_posterior(key, params, sample_shape=(n,)) draws more.
params dict

The guide's fitted parameters.
converged bool

Whether the loss stopped falling: the mean loss over the last window steps is not below that of the window steps before it by more than twice their combined standard error. Necessary, not sufficient.
info dict

guide (its name), steps, elbo (minus the mean loss over the last window; None for a point-mass guide, whose loss is given as log_joint), dense_start (whether the width started from the Gauss–Newton covariance), khat (the Pareto k̂ of the importance weights p/q of num_samples draws: below 0.7 the guide is good enough to reweight; None if not computed or for a point-mass guide), time (seconds, compilation included) and dtype.

samples instance-attribute

losses instance-attribute

guide instance-attribute

params instance-attribute

converged instance-attribute

info instance-attribute

__init__(samples, losses, guide, params, converged, info)

variational(model, priors, data, regularizers=(), *, noise=None, likelihoods=(), start=None, guide='bnaf', steps=3000, optimizer=None, learning_rate=0.003, num_particles=8, num_samples=2000, key=None, flat_coordinates=False, dense_start=True, init_scale=0.1, window=None, psis=True, dtype='float64', **options)

Approximate the posterior with a guide fitted by numpyro's SVI.

Maximizes the ELBO with Adam on the posterior of numpyro_model. The guide starts at start with the Gauss–Newton width, or width init_scale given regularizers, dense_start=False or start=None (then at the priors' medians). Guide choice and checks: Variational inference.

Parameters:

Name Type Description Default
model

As for numpyro_model; see Common arguments.

 required
priors

As for numpyro_model; see Common arguments.

 required
data

As for numpyro_model; see Common arguments.

 required
regularizers

As for numpyro_model; see Common arguments.

 required
noise

As for numpyro_model; see Common arguments.

 required
likelihoods

As for numpyro_model; see Common arguments.

 required
start FitResult or dict

Where the guide starts: a fit or its values (default: the priors' medians, width init_scale).

 None
guide str or callable

"bnaf" (default), "iaf", "mvn" or "laplace"; or a callable guide(model, init_loc_fn=...) returning a numpyro autoguide. The flows need two or more parameters; the dense guides raise ValueError above 4096 coordinates (use laplace_samples for images).

 'bnaf'
steps int

Steps (3000), draws per ELBO estimate (8), draws returned (2000).

 3000
num_particles int

Steps (3000), draws per ELBO estimate (8), draws returned (2000).

 3000
num_samples int

Steps (3000), draws per ELBO estimate (8), draws returned (2000).

 3000
optimizer optional

Optimizer (default optax.adam), its step size (3e-3), the PRNG key (default PRNGKey(0)), flat_coordinates as for numpyro_model, and the precision (default "float64", as for fit).

 None
learning_rate optional

Optimizer (default optax.adam), its step size (3e-3), the PRNG key (default PRNGKey(0)), flat_coordinates as for numpyro_model, and the precision (default "float64", as for fit).

 None
key optional

Optimizer (default optax.adam), its step size (3e-3), the PRNG key (default PRNGKey(0)), flat_coordinates as for numpyro_model, and the precision (default "float64", as for fit).

 None
flat_coordinates optional

Optimizer (default optax.adam), its step size (3e-3), the PRNG key (default PRNGKey(0)), flat_coordinates as for numpyro_model, and the precision (default "float64", as for fit).

 None
dtype optional

Optimizer (default optax.adam), its step size (3e-3), the PRNG key (default PRNGKey(0)), flat_coordinates as for numpyro_model, and the precision (default "float64", as for fit).

 None
dense_start optional

Start the width at the Gauss–Newton covariance (default True; skipped, with a warning, given regularizers). init_scale (default 0.1) is the width of the coordinates it does not cover, or of all of them without it.

 True
init_scale optional

Start the width at the Gauss–Newton covariance (default True; skipped, with a warning, given regularizers). init_scale (default 0.1) is the width of the coordinates it does not cover, or of all of them without it.

 True
window optional

Convergence-test window (default steps // 10); whether to compute the Pareto k̂.

 None
psis optional

Convergence-test window (default steps // 10); whether to compute the Pareto k̂.

 None
**options

Fixed error terms and reject_unphysical, as for numpyro_model.

 {}

Returns:

Type Description
VariationalResult

samples (dict of draws), losses, guide, params, converged and info; warns if the loss had not settled.

Examples:

>>> vi = variational(scene, priors, data, start=result)