Source-linked AI summary
Yes, but Did It Work?: Evaluating Variational Inference
Yuling Yao, Aki Vehtari, Daniel Simpson, Andrew Gelman
TL;DR
VI approximations can be difficult to evaluate because they may fail for structural, optimization, or tail-related reasons, while ELBO values do not provide a reliable trust threshold. The paper proposes PSIS for assessing and improving full-posterior estimates and VSBC for evaluating average point-estimate calibration. PSIS uses ˆk to diagnose posterior discrepancy, while VSBC uses calibration symmetry to reveal bias, with conclusions bounded by the data-generating model and average-case setting.
Problem
Variational approximations can be difficult to assess because approximation and optimization failures are possible, while ELBO values are not a reliable measure of fit.
Method
The paper proposes PSIS to assess and adjust full-posterior approximations and VSBC to assess average calibration of VI point estimates.
Results
PSIS can expose severe posterior mismatch, including ˆk = 9.8 for ADVI in a cancer-classification example, while the Gaussian approximation could not recover some bimodal coefficients.
Takeaways & Limitations
PSIS and VSBC provide complementary diagnostics for deciding whether VI’s full posterior or point estimates are reliable for the model and data setting being assessed.
Takeaways & Limitations
VSBC evaluates average performance on model-generated data, so passing does not guarantee good performance for a particular real-data realization or when the model is misspecified.
Abstract
from arXiv · showhide
While it's always possible to compute a variational approximation to a posterior distribution, it can be difficult to discover problems with this approximation. We propose two diagnostic algorithms to alleviate this problem. The Pareto-smoothed importance sampling (PSIS) diagnostic gives a goodness of fit measurement for joint distributions, while simultaneously improving the error in the estimate. The variational simulation-based calibration (VSBC) assesses the average performance of point estimates.
1. Introduction
Variational inference is widely used for scalable Bayesian inference, but its approximate posterior can be difficult to assess because theoretical guarantees are limited and common diagnostics are inadequate. The paper proposes PSIS and a point-estimate calibration diagnostic to evaluate different aspects of VI quality.
- VI formulation: VI fixes a family of approximate densities and selects the member minimizing KL divergence to the true posterior, equivalently maximizing the ELBO.This formulation avoids direct posterior expectations or sampling.
- Why diagnostics are needed: VI approximations can fail because of optimization issues, restrictive approximation families, posterior asymmetry, or KL divergence’s under-penalization of light tails.These failure modes motivate an explicit diagnostic algorithm.
- Why diagnostics are needed: ELBO values are difficult to interpret and compare across approximations because of an unknown reparameterization-dependent constant and their uninterpretable scale.Consequently, ELBO alone does not indicate when a variational posterior is trustworthy.
- Proposed diagnostics: The paper proposes diagnostics for the quality of the entire variational posterior and for the average bias of a point estimate under correct model specification.The two diagnostics target complementary aspects of approximation quality.
- Proposed diagnostics: PSIS uses a generalized Pareto diagnostic to estimate discrepancy between the approximate and true distributions while improving importance-sampling estimates.Its continuous estimated shape parameter, ˆk, provides the diagnostic measurement.
- Proposed diagnostics: The second diagnostic evaluates the median or modal point estimate’s average behavior and can reveal systemic bias or calibration problems in univariate functionals.It can monitor the magnitude of bias while the diagnostic is computed.
2. Is the Joint Distribution Good Enough?
The joint-distribution diagnostic uses importance sampling with Pareto smoothing to trade bias against variance and uses ˆk to assess how closely VI approximates the full posterior. Its interpretation supports retaining, tuning, or replacing VI, while marginal diagnostics and approximation-family changes require caution.
- Importance sampling trade-offs: Plain VI has low variance but unknown bias, whereas importance sampling is consistent with small self-normalization bias but may have large or infinite variance.This motivates seeking an intermediate bias-variance trade-off.
- PSIS procedure: PSIS fits a generalized Pareto distribution to the largest importance ratios, replaces them with fitted expected values, and truncates weights at the raw maximum.The remaining weights are unchanged, producing smoothed weights for lower-variance estimation.
- PSIS procedure: PSIS has smaller mean square error than plain importance sampling and truncated importance sampling.The method is presented as Bayesian importance sampling with a prior on the largest ratios.
- Interpreting ˆk: A large ˆk indicates finite-sample PSIS failure and therefore a large estimation error for the VI approximation.The diagnostic remains useful even when VI is not being used as a PSIS proposal.
- Interpreting ˆk: If ˆk < 0.5, PSIS has a fast convergence rate and the variational approximation is considered close enough to the true density.Further PSIS use is recommended to improve estimation.
- Interpreting ˆk: If 0.5 < ˆk < 0.7, the approximation is imperfect but useful, with practically acceptable finite-sample convergence and Monte Carlo error.PSIS is again recommended to shrink errors.
- Interpreting ˆk: If ˆk > 0.7, PSIS convergence becomes impractically slow, so researchers should tune VI or turn to exact MCMC.Finite-sample ˆk may exceed 1 even though the theoretical k is smaller than 1.
- Reparameterization: The ˆk diagnostic is invariant to smooth reparameterization because the target-to-proposal density-ratio distribution is unchanged.However, changing the approximation family changes the computation and can change PSIS diagnostics.
3. Assessing the Average Performance of the Point Estimate
VSBC evaluates the average calibration of VI point estimates across datasets simulated from the model, complementing PSIS’s fixed-dataset full-posterior assessment. It uses symmetry of marginal calibration probabilities to identify bias, but its conclusions are limited to model-generated data and average behavior.
- Purpose: VSBC assesses the calibration of a VI posterior’s center or point estimate, even when the full variational posterior is a poor approximation.This complements PSIS, which evaluates the full posterior distribution.
- Procedure: The procedure simulates multiple datasets from the model, fits VI to each dataset, and computes marginal calibration probabilities.These probabilities are then examined across replications and parameter dimensions.
- Procedure: VSBC tests whether marginal calibration probabilities are symmetric around 0.5 using histogram inspection or a Kolmogorov-Smirnov test.Asymmetry indicates bias in the corresponding variational marginal.
- Relationship to PSIS: Unlike PSIS’s fixed-data assessment, VSBC measures average point-estimate calibration over datasets that could be generated by the model.The paper recommends using the two diagnostics as appropriate because they operate under different paradigms.
- Limitations: VSBC can target one-dimensional marginals or functionals, but passing on average does not guarantee good performance for a particular dataset or misspecified real data.The authors recommend conservative interpretation: failure predicts poor performance on the model, whereas passing is not a guarantee for real data.
- Interpretation: Under symmetry assumptions, an unbiased VI estimate yields symmetric VSBC p-values, while positive or negative bias produces right- or left-skewed p-values.The symmetry of the true posterior is stronger than necessary in practice for this result.
4. Applications
Across linear, logistic, hierarchical, and horseshoe-regression examples, PSIS and VSBC expose approximation errors that convergence measures or predictive density can miss. PSIS also improves estimates when the approximation is usable, while reparameterization can reduce errors in hierarchical models.
- PSIS and VSBC diagnostics apply to mean-field Gaussian ADVI implementations across the examples.
- Linear Regression: With M = 1000 replications, VSBC found β1 and β2 marginally unbiased but σ overestimated in Bayesian linear regression.The σ p-value distribution was right-skewed, while the coefficient distributions were not rejected by the two-sided KS-test.
- Linear Regression: Changing the relative ELBO stopping threshold from 10^-5 to 10^-2 increased PSIS ˆk from 0.61 to 4.4, revealing sensitivity that default convergence could miss.When ˆk < 0.7, PSIS adjustment eventually made estimation errors negligible in the example.
- Logistic Regression: As logistic-regression design correlation increased, PSIS ˆk rose and lpd became misleading because prediction improved while the mean-field approximation worsened.The lpd discrepancy showed a sharp increase around ˆk = 0.7.
- Logistic Regression: Across first- and second-moment RMSEs, PSIS adjustment reduced errors, whereas no estimation was considered reliable when ˆk > 0.7.Plain importance sampling had larger RMSE because of higher variance.
- Hierarchical and Horseshoe Models: In hierarchical and horseshoe models, VI errors depended on parameterization and posterior shape: reparameterization reduced eight-school biases, while ADVI missed bimodality and produced biased shrinkage estimates.The centered eight-school parameterization had joint ˆk = 1.00, whereas the non-centered version had ˆk = 0.64; VSBC identified contrasting τ biases across parameterizations.
5. Discussion
The discussion emphasizes that PSIS is local and can miss unseen posterior modes, while VSBC can diagnose potential missing modes through prior-based simulation. The proposed diagnostics also support reparameterization choices, but extensions to modified variational objectives remain future work.
- PSIS can fail when variational samples do not cover all posterior modes, making its estimated k insensitive to unseen modes.The authors characterize PSIS as a local diagnostic and illustrate failure for a two-mode posterior.
- Multiple over-dispersed initializations can detect bimodality, while KL(p, q) offers another divergence-based diagnostic computable through PSIS.
- VSBC samples true parameters from the prior and, unless the prior is too restrictive, its p-value can diagnose a potential missing mode.
- PSIS diagnostics can guide reparameterization by comparing k values, although finding an optimal parameterization may be as difficult as finding the true posterior.
- The authors identify applying diagnostics to modified approximations and jointly optimizing divergence order and α as future research directions.
A.1. Proof to Proposition 1: Marginal ˆk in PSIS diagnostic
The proposition establishes that a finite α-moment condition for the joint density ratio implies the corresponding condition for each marginal ratio. The proof uses conditional distributions under the true posterior and variational approximation.
- If the joint ratio p(θ)/q(θ) has a finite α-moment for some α > 1, then each marginal ratio p(θ_i)/q(θ_i) also has a finite α-moment.
- The proof assumes, without loss of generality, that the parameter space is R^K, using a smooth transformation otherwise.
- For each margin θ_i, the proof defines conditional distributions of the remaining parameters under p and q given θ_i.
A.2. Proof to Proposition 2: Symmetry in VSBC-Test
Under symmetric variational and true posteriors, unbiased VI produces symmetric VSBC p-values, while positive or negative bias produces right or left skewness. The proof connects this result to prior-predictive simulation and the unbiasedness relation.
- With symmetric q and p, an unbiased VI estimate makes the VSBC p-value distribution symmetric.
- Positive or negative VI bias produces respectively right- or left-skewed VSBC p-value distributions.
- The proposition writes q(θ|y) to emphasize that the variational approximation depends on observed data.
- Sampling θ^(0) from the prior and y from the likelihood makes θ^(0) a draw from the true posterior conditional on observed y.
- The proof uses symmetry of q, the unbiasedness condition, and symmetry of p in its successive equalities.
- A positive-bias condition changes the relevant equality into a less-than relation, yielding the corresponding skewness direction.
B.1. Linear and Logistic Regressions
The regression experiments evaluate VSBC and PSIS diagnostics for ADVI, including sensitivity to stopping time and predictor correlation. They use simulated linear and logistic regression settings with repeated replications and posterior-sample-based diagnostics.
- Linear regressions: The linear-regression simulations use n = 10000 observations and d = 100 regressors with β_i ∼ N(0, 1) and σ ∼ gamma(0.5, 0.5).The corresponding Stan implementation is shown in Figure IX.
- Convergence and diagnostics: ADVI can be sensitive to the stopping time because Monte Carlo evaluation of the objective function produces substantial uncertainty.Stan treats ADVI as converged when the running average or median relative ELBO norm change falls below the default tolerance of 0.01.
- VSBC evaluation: The VSBC experiment evaluates ADVI with 1000 replications using a conservative tolerance of 10^-4 and learning rate η = 0.05.The predictor matrix is fixed across replications, and β1, β2, and log σ are preregistered before testing.
- Logistic regressions: The logistic-regression simulations are presented as simulation results in Figures 3 and 4 of the main paper, with the Stan implementation shown in Figure X.
- PSIS evaluation: The PSIS experiment varies predictor correlation in a two-regressor design and compares k-hat with log predictive density from 100 independent test data.The reference posterior is obtained with NUTS in Stan using four chains and 3000 iterations per chain.
B.2. Eight-School Model
The eight-school example models school-specific treatment effects with a Bayesian hierarchical normal model and compares centered and non-centered parameterizations. Its VSBC analysis finds average unbiasedness for θ1 but bias for τ under the stated prior.
- Model setup: The eight-school study analyzes treatment effects from eight schools using a Bayesian hierarchical normal model with shared hyperparameters μ and τ.Each school separately reported its treatment-effect mean and standard deviation, while θ_j denotes the underlying treatment effect for school j.
- Parameterizations: The centered and non-centered parameterizations provide alternative Stan implementations of the eight-school model.The centered form produces strong dependency between τ and θ, whereas the non-centered form extracts that dependency.
- PSIS evaluation: The PSIS evaluation uses 10^5 posterior samples, with the marginal k-hat computed from the NUTS density.That density is typically unavailable for more complicated problems in practice.
- VSBC evaluation: The VSBC test uses 1000 replications and preregisters θ1 and log τ before evaluation.
- VSBC result: Under a prior that is essentially flat over the interesting likelihood region, VSBC indicates that VI estimates θ1 are on average unbiased while τ is biased in both parameterizations.The result depends on the choice of prior, and the corresponding p-value skewness is shown in Figure XIII.
B.3. Cancer Classification Using Horseshoe Priors
The cancer-classification experiment applies a regularized horseshoe prior to a high-dimensional leukemia dataset with far more predictors than observations. Its very large k-hat indicates that VI is far from the true posterior, so PSIS adjustment is not performed.
- Data and motivation: The leukemia dataset contains 72 observations and 7129 standardized features, with binary outcomes modeled using logistic regression.The dataset is high-dimensional, with far more predictors than observations.
- Prior specification: The analysis uses a regularized horseshoe prior because only a few predictors are expected to have distinguishable coefficients and many predictors are correlated.The regularized horseshoe is described as a generalization of the horseshoe prior.
- Prior specification: The intercept has a N(0, 10) prior and is not shrunk, while the global shrinkage scale is set to τ0 = 2.The Stan implementation is summarized in Figure XIV.
- Inference setup: The reference NUTS fit uses four chains with 3000 iterations each, and β1834 is selected because it has the largest posterior mean and a bimodal posterior with a spike at zero.ADVI uses the same parameterization with learning rate η = 0.1 and tolerance tol rel obj = 0.001.
- Diagnostic result: An extremely large k-hat indicates that VI is far from the true posterior, so no PSIS adjustment is conducted.The k-hat estimate uses 10^4 posterior samples.