API

MGVI

An implementation of the Metric Gaussian Variational Inference algorithm.

struct MGVI.BacktrackingLineSearch

Simple Armijo backtracking line search.

Unlike the line searches from LineSearches it is free of scalar control flow, so a NewtonCG using it (with suitable target functions) is suitable for program tracing and compilation, e.g. via Reactant.

Constructors:

• '''BacktrackingLineSearch(; fields...)'''

• c::Float64: sufficient-decrease constant Default: 0.0001

• ρ::Float64: step reduction factor Default: 0.5

• maxsteps::Int64: maximum number of step reductions Default: 20

struct GeoVIConfig

geoVI algorithm configuration.

Fields:

• linsolver: Linear solver to use, must be suitable for positive-definite operators
• linsolver_opts: Linear solver options
• optimizer: Optimization solver for the KL minimization
• optimizer_opts: Optimization solver options
• sampling_optimizer: MGVI.NewtonCG used for the nonlinear per-sample solves

linsolver must be a solver supported by LinearSolve, linsolver_opts is passed to LinearSolve.solve as keyword arguments. optimizer and optimizer_opts behave as in MGVIConfig.

struct GeoVIPreparedStep

A geoVI step prepared via geovi_prepare for a fixed model, data, number of residual samples and parameter dimensionality.

Run steps with geovi_step(prepared::GeoVIPreparedStep, center).

struct GeoVISampler

Generates zero-centered samples from the geoVI posterior approximation around a given expansion point.

geoVI locally isometrizes the Fisher metric M(ξ) = J'ℐJ + I via the embedding emb(ξ) = (euclidean_coords(model(ξ)), ξ): with the embedding Jacobian C = ∂emb₁/∂ξ frozen at the expansion point ξ̄, each sample solves the nonlinear least-squares problem

min_ξ ½ ‖ g̃(ξ) - t ‖²,   g̃(ξ) = (ξ - ξ̄) + C'(emb₁(ξ) - emb₁(ξ̄))

with target t ~ N(0, M(ξ̄)), starting from the linear (MGVI) residual M(ξ̄)⁻¹ t. For forward models that are linear in ξ this reduces exactly to MGVI residual sampling.

Constructor:

GeoVISampler(
    f_model::Function, center_point::Vector{<:Real}, linear_solver,
    sampling_optimizer::MGVI.NewtonCG, context::MGVIContext;
    linear_solver_opts::NamedTuple = (;)
)

linear_solver must be a solver supported by LinearSolve.

Call MGVI.sample_geovi_residuals(s::GeoVISampler, n::Integer) to generate 2n samples (each target t is solved for both signs ±t).

struct MGVIConfig

MGVI algorithm configuration.

Fields:

• linsolver: Linear solver to use, must be suitable for positive-definite operators
• linsolver_opts: Linear solver options
• optimizer: Optimization solver to use
• optimizer_opts: Optimization solver options

linsolver must be a solver supported by LinearSolve or MGVI.MatrixInversion. Use MatrixInversion only for low-dimensional problems. linsolver_opts is passed to LinearSolve.solve as keyword arguments.

optimizer may be MGVI.NewtonCG() or an optimization algorithm supported by OptimizationBase or Optim. optimizer_opts is algorithm-specific.

MGVIContext(rng::AbstractRNG, ad::AutoDiffOperators.ADSelector)

Specifies the linear operator type, RNG and automatic differentiation backend to be used by MGVI operations.

Note: When using a Zygote-based ADSelector, also load DistributionsAD to make forward models that construct Distributions objects Zygote-differentiable.

struct MGVI.MGVIKLTarget <: Function

The optimization target of mgvi_step: the sampled KL estimator (mean negative non-normalized log-posterior over the samples) as a function of the variational mean.

Construct with mgvi_kl_target.

struct MGVIPreparedStep

An MGVI step prepared via mgvi_prepare for a fixed model, data, number of residual samples and parameter dimensionality.

Run steps with mgvi_step(prepared::MGVIPreparedStep, center).

struct MGVIResult

State resulting from mgvi_step.

Fields:

• samples: The samples drawn by MGVI
• mnlp: The mean of the negative non-normalized log-posterior over the samples
• info: Additional information given by the linear solver and optimization algorithm.

struct MatrixInversion

Solve linear systems by direct matrix inversion.

Note: Will instantiate implicit matrices/operators in memory explicitly.

struct NewtonCG

Inexact-Newton optimizer with a matrix-free conjugate-gradient inner solver.

Each step solves the Newton system approximately, with the cg residual target set by the Eisenstat–Walker forcing sequence ‖r‖ <= min(1/2, √‖∇f‖) * ‖∇f‖, and finishes with a line search along the resulting direction.

Constructors:

• '''NewtonCG(; fields...)'''

• steps::Int64: number of NewtonCG steps Default: 4

• absdelta::Float64: stop early when a NewtonCG step decreases the target by no more than this (in uncompiled operation only); if positive, cg iterations also stop once their quadratic-energy decrease becomes negligible compared to the previous NewtonCG step's improvement Default: 0.0

• cg_maxiter::Int64: maximum number of cg iterations per NewtonCG step, 0 chooses automatically based on the problem dimensionality Default: 0

• cg_steps_traced::Int64: number of cg iterations per NewtonCG step when compiled via program tracing, which requires a fixed iteration count for now Default: 5

• linesearcher::Any: LineSearcher that will be used after cg iterations are finished Default: StrongWolfe{Float64}()

struct MGVI.PDLinMapWithChol{T} <: LinearMaps.LinearMap{T}

A LinearMap that stores both a map and the lower-triangular map of its Cholesky decomposition.

struct ResidualSampler

Generates zero-centered samples from the posterior's covariance approximated by the Fisher information.

This sampler uses Conjugate Gradients to iteratively invert the Fisher information, never instantiating the covariance in memory explicitly.

The Fisher information in canonical coordinates and Jacobian of the coordinate transformation are provided as arguments.

Constructor:

ResidualSampler(
    f_model::Function, center_point::Vector{<:Real}, linear_solver, context::MGVIContext;
    linear_solver_opts::NamedTuple = (;)
)

linear_solver must be a solver supported by LinearSolve or MGVI.MatrixInversion. Use MatrixInversion only for low-dimensional problems. linear_solver_opts is passed to LinearSolve.solve as keyword arguments.

Call MGVI.sample_residuals(s::ResidualSampler[, n::Integer]) to generate a single or n samples.

MGVI._batched_cg(
    apply_A, B::AbstractMatrix{<:Real}, max_iter::Integer, rtol::Real
)

Solve A * X == B column-wise via conjugate-gradient iterations, with apply_A(P) applying the positive-definite A to the columns of P.

Iterates until the residual norms of all columns have been reduced by the factor rtol, but at most max_iter times. Free of array mutation and, when compiled via program tracing, of scalar control flow (max_iter unrolled cg iterations then, masked to no-ops for converged columns), so it is suitable for compilation e.g. via Reactant, unlike dynamically terminated solvers.

MGVI.euclidean_coords(d::Distributions.Distribution)

Compute coordinates of the parameters of d in which the Fisher information metric of d's distribution family is the identity.

The Jacobian C of euclidean_coords with respect to the flat parameters of d (see MGVI.flat_params) satisfies C'C == fisher_information(d) - exactly for distribution families with flat Fisher geometry, up to curvature terms otherwise (e.g. for Normal, where the cross terms between mean and scale vanish only at μ == 0).

Used by geoVI (see geovi_step) to construct local isometries.

MGVI.fisher_information(distribution::Distributions.Distribution)

Get the fisher information matrix/operator (as a LinearMap) for the given distribution.

geovi_prepare(
    forward_model, data, n_residuals::Integer,
    center_proto::AbstractVector{<:Real},
    config::GeoVIConfig, context::MGVIContext;
    device = nothing
)

Prepare repeated geoVI steps for the given model and data and return a GeoVIPreparedStep, like mgvi_prepare does for MGVI steps.

If device (an MLDataDevices.AbstractDevice) is given, the step computation is set up on the device via HeterogeneousComputing.on_device. For a Reactant device it is compiled once and then reused for every subsequent step, which requires config.optimizer and config.sampling_optimizer to be NewtonCGs with MGVI.BacktrackingLineSearch linesearchers and an Enzyme-based context.ad.

The prepared step uses MGVI._batched_cg for the linear part of the residual sampling instead of config.linsolver; maxiters and reltol in config.linsolver_opts are honored.

geovi_sample(
    forward_model, data, n_residuals::Integer, center::AbstractVector{<:Real},
    config::GeoVIConfig, context::MGVIContext
)

Draws samples from the geoVI posterior approximation centered at center without performing an optimization step.

Returns a matrix whose 2 * n_residuals columns are the samples.

geovi_step(
    forward_model, data, n_residuals::Integer, center_init::AbstractVector{<:Real},
    config::GeoVIConfig, context::MGVIContext
)

Performs one geoVI step and returns a tuple (result::MGVIResult, updated_center::AbstractVector{<:Real}).

Like mgvi_step, but the samples are drawn from the geoVI approximation: the posterior is approximated by pulling a standard normal distribution back through a local isometry of the Fisher metric (see GeoVISampler), instead of by a multivariate normal distribution. This improves the approximation for forward models with significant nonlinearity at the posterior's scale.

Requires MGVI.euclidean_coords to be defined for the distribution type returned by forward_model.

Note: The prior is implicit, it is a standard (uncorrelated) multivariate normal distribution of the same dimensionality as center_init.

geovi_step(prepared::GeoVIPreparedStep, center::AbstractVector{<:Real})

Performs one geoVI step like geovi_step(forward_model, data, n_residuals, center, config, context), using a step prepared with geovi_prepare.

Returns a tuple (result::MGVIResult, updated_center).

mgvi_kl_target(forward_model, data, residual_samples::AbstractMatrix{<:Real})

Return the function of the variational mean that mgvi_step minimizes, given zero-centered residual samples (as matrix columns).

The result is a pure function of its argument and suitable for program tracing and compilation (e.g. via Reactant, with an Enzyme-based gradient), provided forward_model is.

mgvi_mvnormal_pushfwd_function(
    forward_model, data, config::MGVIConfig,
    center_point::AbstractVector{<:Real}, context::MGVIContext
)

Returns a function that pushes a multivariate normal distribution forward to the MGVI posterior approximation.

This currently instantiates the full Jacobian of the forward model as a matrix in memory, and so should not be used for very high-dimensional problems.

mgvi_prepare(
    forward_model, data, n_residuals::Integer,
    center_proto::AbstractVector{<:Real},
    config::MGVIConfig, context::MGVIContext;
    device = nothing
)

Prepare repeated MGVI steps for the given model and data and return a MGVIPreparedStep.

center_proto provides the shape and type of the (later) center points, its values are not used. The number of residual samples and the parameter dimensionality are fixed at preparation time.

If device (an MLDataDevices.AbstractDevice) is given, the step computation is set up on the device via HeterogeneousComputing.on_device. For a Reactant device it is compiled once and then reused for every subsequent step, which requires config.optimizer to be a NewtonCG with a MGVI.BacktrackingLineSearch linesearcher and an Enzyme-based context.ad.

The prepared step uses MGVI._batched_cg for residual sampling instead of config.linsolver; maxiters and reltol in config.linsolver_opts are honored.

mgvi_sample(
    forward_model, data, n_residuals::Integer, center::AbstractVector{<:Real},
    config::MGVIConfig, context::MGVIContext
)

Draws samples from the MGVI posterior approximation centered at center without performing an optimization step.

Returns a matrix whose 2 * n_residuals columns are the samples center ± residual.

mgvi_step(
    forward_model, data, n_residuals::Integer, center_init::AbstractVector{<:Real},
    config::MGVIConfig, context::MGVIContext
)

Performs one MGVI step and returns a tuple (result::MGVIResult, updated_center::AbstractVector{<:Real}).

The posterior distribution is approximated with a multivariate normal distribution. The covariance is approximated with the inverse Fisher information evaluated at center_init. Samples are drawn according to this covariance, which are then used to estimate and minimize the KL divergence between the true posterior and the approximation.

Note: The prior is implicit, it is a standard (uncorrelated) multivariate normal distribution of the same dimensionality as center_init.

Example

using Random, Distributions, MGVI
import LinearSolve, Zygote

context = MGVIContext(ADSelector(Zygote))

model(x::AbstractVector) = Normal(x[1], 0.2)
true_param = [2.0]
data = rand(model(true_param), 1)[1]
center = [1.3]

config = MGVIConfig(
    linsolver = LinearSolve.KrylovJL_CG(),
    optimizer = MGVI.NewtonCG()
)
n_residuals = 12
n_steps = 5

res, center = mgvi_step(model, data, n_residuals, center, config, context)
for i in 1:n_steps-1
    res, center = mgvi_step(model, data, n_residuals, center, config, context)
end

samples_from_est_covariance = res.samples

mgvi_step(prepared::MGVIPreparedStep, center::AbstractVector{<:Real})

Performs one MGVI step like mgvi_step(forward_model, data, n_residuals, center, config, context), using a step prepared with mgvi_prepare.

Returns a tuple (result::MGVIResult, updated_center).

pareto_diagnostic(log_ratios::AbstractVector{<:Real})

Run Pareto-smoothed importance sampling (PSIS) on unnormalized log importance ratios log p - log q.

Returns (; pareto_shape, weights, result) with the Pareto shape k̂, the normalized Pareto-smoothed importance weights and the full PSIS.PSISResult. k̂ diagnoses how well the variational distribution q covers the posterior p: k̂ < 0.5 is good, values up to 0.7 are usable, higher values indicate that q misses posterior mass.

pareto_diagnostic(
    forward_model, data, samples::AbstractMatrix{<:Real},
    center::AbstractVector{<:Real}, context::MGVIContext
)

PSIS diagnostic of the metric-Gaussian posterior approximation q = N(center, (J'ℐJ + I)⁻¹) for posterior samples (as matrix columns) drawn around center, as returned by mgvi_step.

The log-density of q enters only up to its sample-independent normalization, which cancels in the diagnostic and in the self-normalized weights. Exact for MGVI samples of a converged fit; for geoVI samples the per-sample nonlinear transport is ignored, making the diagnostic approximate.

MGVI.sample_geovi_residuals(s::GeoVISampler, n::Integer)

Generate 2n zero-centered geoVI residual samples, as columns of a matrix.

Each of the n targets is solved for both signs, sharing the underlying random draw, so columns 2i-1 and 2i form an antithetic pair (only approximately mirrored, due to the nonlinearity).

MGVI.sample_residuals(
    f_model, center::AbstractVector{<:Real},
    sample_n::AbstractMatrix{<:Real}, sample_η::AbstractMatrix{<:Real},
    ad::ADSelector;
    cg_max_iterations::Integer = 4 * length(center),
    cg_rtol::Real = sqrt(eps(Float64))
)

Generate zero-centered residual samples like sample_residuals(::ResidualSampler, n), but as a pure function of pre-drawn standard normal samples sample_n (data-parameter space, size m × k) and sample_η (parameter space, size n × k), solving all k linear systems together with MGVI._batched_cg.

Returns an n × k matrix of residual samples.

Free of RNG state and array mutation, so suitable for program tracing and compilation, e.g. via Reactant with an Enzyme-based ad.