Variational Inference

Sampling-based methods (MCMC, HMC, SMC) approximate the posterior with a set of draws. Variational Inference (VI) instead picks a tractable family of distributions $q_{\varphi }$ and optimizes $\varphi$ so that $q_{\varphi }$ is as close as possible to the true posterior, by maximizing the Evidence Lower BOund (ELBO):

$$\text{ELBO}(\varphi )={\mathbb{E}}_{z\sim q_{\varphi }}\left[\log p(x,z)-\log q_{\varphi }(z)\right]\le \log p(x)$$

This trades exactness for speed: a converged VI fit is one optimization run, not thousands of MCMC iterations — at the cost of being only as good as the chosen family $q_{\varphi }$ lets it be.

The model

/// `mu ~ Normal(0, 1)`; observe `y ~ Normal(mu, 0.5)` at the fixed value 1.5.
fn model() -> Model<f64> {
    sample(addr!("mu"), Normal::new(0.0, 1.0).unwrap())
        .bind(|mu| observe(addr!("y"), Normal::new(mu, 0.5).unwrap(), 1.5).map(move |_| mu))
}

The same conjugate Normal-Normal model as the SMC and ABC tutorials, with observe used normally this time — VI, unlike ABC, does need to score the model's log-density.

Building a guide and running the optimizer

    let mut rng = StdRng::seed_from_u64(11);

    // A real-supported latent gets a Normal factor -- the family that can
    // represent this model's posterior exactly (finding FG-17: guide families
    // are chosen to match the latent's support instead of a one-size-fits-all
    // Normal that would mismatch bounded latents).
    let mut guide = MeanFieldGuide::new();
    guide.add_latent(addr!("mu"), Support::Real, 0.0);

    let config = VIConfig {
        n_iterations: 800,
        n_samples_per_iter: 32,
        base_learning_rate: 0.3,
        ..VIConfig::default()
    };
    let result = optimize_meanfield_vi_with_config(&mut rng, model, guide, &config);

A few things worth calling out:

• guide.add_latent(addr, Support::Real, init) picks the variational family from the latent's declared support instead of defaulting every latent to an unconstrained Normal — a Normal guide on a strictly-positive or unit-interval latent would propose out-of-support values whose model log-density is $-\infty$, collapsing the ELBO.
• Both location and scale are optimized. Each variational factor has two free parameters (e.g. mu and log_sigma for a Normal factor, the latter unconstrained via a log transform for positivity); optimize_meanfield_vi_with_config updates both by gradient ascent, not just the location.
• Gradients are common-random-numbers finite differences. Fugue models are plain Rust closures with no autodiff, so the ELBO gradient is estimated by central finite differences with the +ε/-ε evaluations sharing an RNG seed — this cancels Monte Carlo noise in the difference, leaving only the (much smaller) finite-difference bias.
• A Robbins-Monro step schedule and an ELBO-plateau convergence test mean the optimizer can stop before n_iterations — check result.converged and result.iterations.

Reading the result

    let VariationalParam::Normal { mu, log_sigma } = result
        .guide
        .params
        .get(&addr!("mu"))
        .expect("guide has a factor for mu")
    else {
        panic!("Support::Real latent must produce a Normal factor");
    };
    let fitted_sigma = log_sigma.exp();
    println!("Fitted q(mu) = Normal({mu:.4}, {fitted_sigma:.4})");
    println!("Exact posterior = Normal({POSTERIOR_MEAN}, {POSTERIOR_SD:.4})");

Checking against ground truth

    // Mean-field VI is exact for this model family (see module docs), so a
    // healthy optimizer run should land close to the true posterior. The
    // tolerances are set from repeated runs at this seed/config: comfortably
    // above run-to-run stochastic-gradient noise, but tight enough that
    // regressing either the location or (log-space) scale update of finding
    // FG-04 -- which previously left the scale un-optimized entirely -- would
    // fail this assertion (an un-optimized scale stays at its ~1.0 init,
    // several sigma away from the target 0.4472).
    assert!(
        (*mu - POSTERIOR_MEAN).abs() < 0.15,
        "VI fitted mean {mu} too far from exact posterior mean {POSTERIOR_MEAN}"
    );
    assert!(
        (fitted_sigma - POSTERIOR_SD).abs() < 0.15,
        "VI fitted sd {fitted_sigma} too far from exact posterior sd {POSTERIOR_SD}"
    );

Run it yourself:

cargo run --example vi_inference

Next

• Back to Sequential Monte Carlo or Approximate Bayesian Computation to compare methods on the same target.
• Advanced Inference overview for a summary table of when to reach for each method.