romtools.workflows.inverse.vi_drivers

Variational inference drivers for Gaussian posterior approximations.

This module provides derivative-free variational inference (VI) routines for black-box forward models. An approximate posterior is represented by a Gaussian variational family, and the optimizer updates the variational parameters to maximize the evidence lower bound (ELBO).

Theory

run_vi solves a Gaussian variational inference problem by maximizing

\[\max_{\zeta}\; \mathbb{E}_{\theta \sim q(\theta;\zeta)} \left[ \log p(y,\theta) - \log q(\theta;\zeta) \right],\]

where \(q(\theta;\zeta)\) is the variational search distribution and \(\zeta\) denotes its parameters. In this implementation, the variational family is Gaussian in either diagonal form or fixed-correlation multivariate form, with the optimizer state stored as the variational mean and log standard deviation.

For black-box forward models, the code uses the score-function (REINFORCE/log-likelihood-trick) estimator instead of reparameterization gradients through the PDE model:

\[\nabla_{\zeta} \mathbb{E}_{\theta \sim q(\theta;\zeta)} \left[\mathcal{L}(\theta;\zeta)\right] = \mathbb{E}_{\theta \sim q(\theta;\zeta)} \left[ \mathcal{L}(\theta;\zeta)\, \nabla_{\zeta}\log q(\theta;\zeta) \right],\]

with

\[\mathcal{L}(\theta;\zeta) = \log p(y,\theta) - \log q(\theta;\zeta).\]

The expectation is approximated with Monte Carlo or randomized quasi-Monte Carlo samples from the current variational distribution. Because only \(\nabla_{\zeta}\log q\) is needed, the forward model itself is treated as derivative-free.

When optimizer_method="adam", the selected score gradient is passed to a stateful Adam ascent update. By default, the raw score gradient is preconditioned with the damped mean-field Gaussian Fisher matrix before Adam, matching the ABRIS setup; the default initial learning rate is \(0.1/d\), where \(d\) is the parameter dimension.

When optimizer_method="newton", the routine also forms a second-order score-function estimator for curvature:

\[\nabla_{\zeta}^2 \mathbb{E}_{\theta \sim q(\theta;\zeta)} \left[\mathcal{L}(\theta;\zeta)\right] = \mathbb{E}_{\theta \sim q(\theta;\zeta)} \left[ \mathcal{L}(\theta;\zeta) \left( \nabla_{\zeta}\log q\,\nabla_{\zeta}\log q^{\top} + \nabla_{\zeta}^2 \log q \right) \right].\]

This is the curvature model used by the Newton update path. The implementation supports a metric rescaling via newton_metric and either diagonal or full projected Hessian solves via newton_hessian_type. The newton_regularization_strategy option selects a fixed eigenvalue tolerance, a tolerance proportional to the Hessian spectral norm, or "per_parameter" direction-wise damping. All strategies form curvature from \(-H\) and replace eigenvalues strictly below the tolerance with the reciprocal of newton_fallback_learning_rate (default 0.01). Eigenvalues at or above the tolerance are retained. newton_regularization controls that tolerance. Before damping and line-search scaling, replaced directions take gradient steps with this learning rate. The separate newton_additive_regularization adds \(\lambda I\) for the "absolute" and "hessian_norm" strategies, and adds \(\lambda\,\mathrm{diag}(|H_{ii}|+\epsilon)\) for "per_parameter", where \(\epsilon\) is newton_regularization_epsilon. For the natural metric, regularization is evaluated after metric rescaling. The newton_curvature_strategy option controls stochastic coupling: the default "same_sample" reuses the gradient samples, "independent" evaluates a separate batch (whose size is set by newton_hessian_num_samples), and "lagged" preconditions with the previous exponentially averaged Hessian. The lagged averaging weight is set by newton_hessian_averaging_factor. When newton_adaptive_regularization is enabled, a rejected step decreases the next step size and increases additive regularization by newton_regularization_increase_factor. An accepted step with lower ELBO instead increases the next step size and divides additive regularization by newton_regularization_decrease_factor. Both controls stay unchanged after an accepted step whose ELBO does not drop. Additive regularization is bounded by its configured initial value and newton_regularization_max_multiplier times that value.

Variance Reduction

score_function_entropy_strategy="analytic" (the default) estimates the log-joint score contribution and adds the Gaussian entropy gradient exactly. The "joint" strategy instead keeps \(-\log q\) inside the sampled ELBO integrand. The analytic strategy can have lower variance far from the optimum; with a suitable baseline, the joint strategy has zero per-sample contribution when the posterior belongs to the variational family. Both choices are score-function estimators, not pathwise estimators.

The gradient estimator optionally supports a baseline through baseline_method. In particular, the leave-one-out option subtracts a sample mean baseline that is independent of the current score term, preserving unbiasedness while reducing Monte Carlo variance:

\[\widehat g_{\mathrm{LOO}} = \frac{1}{N}\sum_{i=1}^{N} \left( \mathcal{L}(\theta^{(i)};\zeta) - b^{(-i)} \right) \nabla_{\zeta}\log q(\theta^{(i)};\zeta),\]

where \(b^{(-i)}\) is the mean ELBO over all samples except \(\theta^{(i)}\).

Gaussian Structure

For Gaussian priors and Gaussian variational families, the ELBO integrand decomposes into

\[\log p(y,\theta) = \log p(y\mid\theta) + \log p_0(\theta),\]

so the forward-model dependence enters through the log-likelihood term, while the prior and variational-density terms are handled analytically. This is why the routine can remain derivative-free with respect to the forward model while still using gradient- and Hessian-based updates in the variational parameters.