API

MGVI

An implementation of the Metric Gaussian Variational Inference algorithm.

struct MVGIConfig

MGVI clgorithm configuration.

Fields:

• linar_solver: Linear solver to use, must be suitable for positive-definite operators
• 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.

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

MGVIContext{OP=LinearMap}(rng::AbstractRNG, ad::AutoDiffOperators.ADSelector)

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

struct MGVIResult

State resulting from mgvi_step.

Fields:

• smples: The samples drawn by MVGI
• 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

Constructors:

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

• α::Float64: amount of previous NewtonCG improvement guarding the lower bound to the improvement between consecutive cg iterations from the second NewtonCG step on Default: 0.1

• steps::Int64: Number of total NewtonCG steps Default: 4

• i₀::Int64: maximum number of cg iterations in the first NewtonCG step Default: 5

• i₁::Int64: maximum number of cg iterations from the second NewtonCG step on Default: 50

• 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-tringangular 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 must be a solver supported by LinearSolve or MGVI.MatrixInversion. Use MatrixInversion only for low-dimensional problems.

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

MGVI.fisher_information(distribution::Distributions.Distribution)

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

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 Jabocian of the forward model as a matrix in memory, and so should not be used for very high-dimensional problems.

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

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}).

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 valuated 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