PaperlayerView exact revision on arXiv

Source-linked AI summary

Decision-Making with Auto-Encoding Variational Bayes

Romain Lopez, Pierre Boyeau, Nir Yosef, Michael I. Jordan, Jeffrey Regier

arXiv:2002.07217v3stat.MLcs.AIcs.LG

TL;DR

Using the variational distribution as a posterior surrogate can bias decision-relevant estimates. This paper develops a three-step procedure that fits models and multiple approximate posteriors, then combines them with multiple importance sampling; experiments show it outperforms several alternatives.

  • Problem

    Using VAE variational distributions for Bayesian decisions can produce biased or high-variance posterior-expectation estimates, making them unreliable proposal distributions.

  • Method

    The paper fits a generative model, learns several approximate posteriors with alternative objectives, and combines them as proposals through multiple importance sampling.

  • Results

    The proposed three-step procedure outperforms VAE, IWAE, WW, χ-VAE, and IWAE combined with annealed importance sampling in real-world decision-making examples.

  • Takeaways & Limitations

    The best generative model does not necessarily yield the best decisions, whereas combining distinct approximate posteriors can improve decision-making performance.

  • Takeaways & Limitations

    The theoretical analysis focuses on probabilistic principal component analysis, and its error bounds are too loose for some worst-case comparisons between objectives.

Abstract

from arXiv · show

To make decisions based on a model fit with auto-encoding variational Bayes (AEVB), practitioners often let the variational distribution serve as a surrogate for the posterior distribution. This approach yields biased estimates of the expected risk, and therefore leads to poor decisions for two reasons. First, the model fit with AEVB may not equal the underlying data distribution. Second, the variational distribution may not equal the posterior distribution under the fitted model. We explore how fitting the variational distribution based on several objective functions other than the ELBO, while continuing to fit the generative model based on the ELBO, affects the quality of downstream decisions. For the probabilistic principal component analysis model, we investigate how importance sampling error, as well as the bias of the model parameter estimates, varies across several approximate posteriors when used as proposal distributions. Our theoretical results suggest that a posterior approximation distinct from the variational distribution should be used for making decisions. Motivated by these theoretical results, we propose learning several approximate proposals for the best model and combining them using multiple importance sampling for decision-making. In addition to toy examples, we present a full-fledged case study of single-cell RNA sequencing. In this challenging instance of multiple hypothesis testing, our proposed approach surpasses the current state of the art.

1 Introduction

The introduction argues that using the variational distribution for VAE-based decisions can fail because posterior expectations may be biased or high-variance and the learned model may not match the data-generating process. It proposes fitting and selecting a model, then learning separate proposal distributions for decision-making, with theoretical and empirical analysis in probabilistic PCA.

  • Motivation: VAE-based applications increasingly require Bayesian decisions, but exact posterior functionals are unavailable after fitting with the ELBO.Existing approaches substitute the variational distribution, use it for importance sampling, or discard it for alternative decision procedures.
  • Problem: Using the variational distribution can produce biased or high-variance posterior-expectation estimates, while the learned VAE model may still differ from the real data-generating process.Substitution typically causes bias; importance sampling commonly suffers high variance, and model misspecification creates a separate decision-making problem.
  • Proposed framework: The proposed three-step procedure fits models with objectives such as VAE, IWAE, WW, or χ-VAE, selects the best model using a metric, and then learns decision-specific proposals with the model fixed.The χ-VAE variant minimizes the χ2 divergence for fixed pθ.
  • Probabilistic PCA analysis: In probabilistic PCA, the analysis recovers that underdispersed proposals cause severe importance-sampling error and shows that overdispersion can worsen existing variational-Bayes biases during model learning.These theoretical findings are also confirmed empirically.

2 Background

Bayesian decision-making uses a model and posterior distribution to choose optimal actions, while variational methods approximate posterior expectations when direct access is unavailable. Alternative divergence-based objectives motivate inference procedures, but existing error bounds are too loose for direct worst-case comparisons.

  • Bayesian decision-making: Bayesian decision-making combines a model and posterior distribution to make optimal decisions.For observation x, the optimal action is defined through a posterior expectation of a loss-dependent function.
  • Variational autoencoders: VAEs learn a generative model by using qφ(z ∣x) to approximate the posterior pθ(z ∣x) while maximizing the ELBO.The ELBO and reverse KL variational gap decompose the evidence, and jointly maximizing the ELBO yields a reverse-KL-minimizing variational distribution.
  • Approximate inference: When the posterior is inaccessible, practitioners can estimate posterior expectations using plugin estimators or self-normalized importance sampling from qφ(z ∣x).Importance sampling uses weights w(x,z) := pθ(x,z)/qφ(z∣x).
  • Approximate inference: Error bounds motivate alternative inference procedures based on forward KL or χ2 divergence, including WW and χ-VAE.Forward-KL bounds motivate WW, while χ2-divergence bounds motivate investigation of χ-VAE.
  • Limitations: Existing upper bounds are too loose to directly compare the worst-case performance of χ-VAE and WW over a function class.The comparison limitation is discussed in further detail in Appendix C.

3 Theoretical analysis for pPCA

The section uses tractable pPCA to analyze how variational objectives affect importance-sampling efficiency, posterior approximation, and model selection. Its results favor posterior proposals with larger variance and motivate combining proposals for downstream estimation.

  • Analysis setup: Theoretical analysis specializes to pPCA because posterior inference is tractable, while Gaussian posterior behavior offers intuition for broader practical models.The authors use pPCA as a precise special case and connect its relevance to the success of Gaussian approximations such as Laplace’s method.
  • Importance sampling: Theorem 1 bounds importance-sampling error when particle count scales as n = β exp{∆KL(pθ ∥qφ)}, linking sample efficiency to posterior–proposal divergence.The bound assumes the posterior second moment of f(z) is bounded by κ and requires β > log t∗(x).
  • Importance sampling: The spectral norm ∥A(x)∥2 controls log-weight concentration, and its asymmetric minimum in one dimension favors overestimating rather than underestimating posterior variance.The analysis indicates that larger values of D(x) can be more suitable for importance sampling than smaller values.
  • Model selection: With few particles, variational Bayes can bias model parameters toward regions where the variational bound is tighter, making proposal choice relevant for model selection.For one particle, IWELBO tightness equals ELBO tightness and is measured by the reverse KL divergence; higher D(x) increases parameter-space sensitivity.
  • Empirical analysis: In pPCA experiments, IWAE, WW, and χ-VAE outperform VAE on held-out exact likelihood, while MIS proposals achieve the best posterior-expectation MAE.The three-step IWAE-MIS procedure significantly outperforms all single-proposal methods, whereas WW proposals perform worse than expected.

4 Classification-based decision theory

This section evaluates rejection-based classification on semi-supervised MNIST using the M1+M2 generative model. Although plugin classification achieves 95%–97% accuracy, SNIS performs around 60%, consistent with diagnostics indicating large proposal-estimation errors.

  • Experimental setup: The MNIST task predicts whether an image’s label belongs to a chosen digit subset while permitting rejection of ambiguous cases.The experiment uses labels 0–8 and treats label 9 as unlabelled.
  • Contribution: This is, to the authors’ knowledge, the first evaluation of semi-supervised generative models in a rejection-based decision-making scenario.The evaluation uses an SNIS estimator in this setting.
  • Model: The semi-supervised M1+M2 model represents the class with discrete latent variable c and additional variation with latent variables u and z.The model is used to evaluate rejection-based decisions in a generative classification setting.
  • Variational fitting: Reverse-KL variational fitting can assign zero probability to many classes because it covers only one distributional mode.The text notes that alternative divergences, including WW and χ-VAE, can remedy overconfidence in the VAE-trained M1+M2 model.
  • Results: 95%–97% plugin accuracy contrasts with around 60% SNIS accuracy across methods for labels 1 through 8.The plugin estimator performs well, whereas SNIS performs poorly in the reported classification comparison.
  • Results: PSIS estimates greater than one for all algorithms indicate that the variational distribution may cause large estimation error when used as a proposal.The results include AUPRC for classifying nines and goodness-of-fit metrics, and IWAE and WW learn the best generative model in terms of IWELBO.

5 Multiple testing and differential gene expression

The scVI case study evaluates Bayesian differential-expression decisions by ranking genes while controlling posterior expected FDR. χ-VAE-MIS provides the best FDR estimates, outperforming single-proposal alternatives and χ-VAE with AIS.

  • FDR control: The procedure selects the largest gene set whose posterior expected FDR remains below a practitioner-specified target.FDR-estimation quality is measured by mean absolute error between ground-truth and posterior expected FDR.
  • Results: 0.85 AUPRC is achieved by the VAE objective, whereas all other objective combinations exceed 0.95 for differential-expression gene ranking.VAE performs worst on held-out log-likelihood, while χ-VAE performs best.
  • Results: VAE-based proposals have large FDR-estimation error, which becomes worse with the plugin estimator; other models yield significantly lower error.The comparison covers proposals based on the fitted scVI models.
  • Results: χ-VAE-MIS achieves the best FDR estimates, improving over every single-proposal alternative and over χ-VAE combined with AIS.The study uses AIS with 500 steps and 100 prior samples to approximate the fitted model’s posterior.

6 Discussion

The paper proposes a theoretically motivated three-step procedure for VAE-based decision-making that outperforms several alternatives in real-world experiments. Its main cost is roughly threefold training overhead, while remaining robust to posterior-collapse mitigations and distinct from related loss-adaptation approaches.

  • Contributions: The three-step procedure uses derived self-normalized importance-sampling error and variational-Bayes bias analyses, and outperforms VAE, IWAE, WW, χ-VAE, and IWAE with AIS.The procedure is theoretically motivated by analyses on the pPCA model and evaluated in important real-world examples.
  • Computational cost: Training three VAEs with different loss functions adds roughly a constant factor of three in computational overhead versus a standard VAE.In pPCA, one VAE takes 12 seconds, while steps one and two together take 53 seconds on a single-NVIDIA-GeForce machine.
  • Robustness: Cyclical annealing and lagging inference networks improve VAE performance against posterior collapse, but the proposed three-step procedure still outperforms them.Posterior collapse occurs when the variational distribution equals the prior and the data likelihood is independent of latent variables.
  • Related approaches: Losscalibrated inference and amortized Monte Carlo integration are not directly applicable because decision-loss-specific ELBO adaptation biases learning pθ.Both approaches are presented as complementary frameworks for fitting distributions tailored to losses or particular posterior expectations.

Broader impact

The proposed method may improve data-based decisions across biological, supply-chain, and personalized-medicine applications. Its neural-network foundation can fit non-simple parametric data accurately but introduces both benefits and risks through its black-box nature.

  • The method enables practitioners to make better decisions based on data.
  • Potential applications include differential expression, supply-chain inventory and pricing, and personalized medicine treatments.
  • Neural networks can fit data accurately when the distribution is not simple or parametric, but their black-box nature brings benefits and risks.

Appendices … B Chi-VAEs

The appendices provide theoretical proofs, concentration results, experimental details, and a χ-VAE variant motivated by decision-making. The χ-VAE alternates IWELBO model updates with χ2-divergence-based inference updates and may use Student’s t proposals to improve posterior coverage.

  • Appendices: The appendices cover proofs, posterior-expectation bounds, bivariate Gaussian derivations, pPCA and MNIST experiments, and further model-specific discussion.Appendix B presents χ-VAE details, while Appendices A and C–H provide the listed theoretical, analytical, experimental, and methodological materials.
  • A.1 Proof of Lemma 1: The pPCA posterior is expressed in closed form using Gaussian conditioning, with M = W^TW + σ2I.The proof introduces M and derives the tractable posterior before analyzing the log-density ratio.
  • A.1 Proof of Lemma 1: Under the posterior, the log importance-weight ratio is a quadratic function of an isotropic Gaussian and is therefore shown to be sub-exponential.The Gaussian reparameterization enables concentration analysis of the log-ratio.
  • A.1 Proof of Lemma 1: The log-ratio concentration is centered around the forward KL divergence ΔKL(pθ(z ∣x) ∥ qφ(z ∣x)).The proof applies the quadratic-Gaussian concentration lemma to obtain bounds around this mean.
  • A.2 Proof of Lemma 2: The quadratic-Gaussian concentration proof diagonalizes A by an orthogonal eigendecomposition and establishes sub-exponential parameters for the centered quadratic form.The moment-generating function is finite for λ < 1/(4||A||2), yielding parameters involving ||A||2.
  • A.3 Proof of Theorem 1: Theorem 1 requires n = β exp{ΔKL(pθ ∥ qφ)} importance-sampling particles, for β > log t*(x), when the posterior second moment of f(z) is bounded by κ.This gives a sufficient sample-size condition for the importance-sampling estimator.
  • B Chi-VAEs: The proposed χ-VAE is an incremental WW variant that combines CHIVI, WW, and a reparameterized Student’s t variational posterior.The authors describe this combination as a novel variant not previously encountered in the literature.
  • B Chi-VAEs: χ-VAE alternates IWELBO model updates with CUBO-based inference updates targeting χ2 divergence, while Student’s t posteriors are expected to improve posterior coverage when Gaussian proposals fail.For fixed qφ, model parameters follow IWELBO gradients; for fixed pθ, inference parameters minimize an equivalent CUBO objective.

C Limitations of standard results for posterior statistics estimators … E.3 Model details and neural networks architecture

The appendix explains why standard importance-sampling error bounds are inadequate for algorithm selection, then provides analytical derivations and implementation details for the Gaussian and pPCA experiments. It specifies the estimators, divergence-based approximations, simulation setup, dataset split, model parameterization, and optimization procedure.

  • C Limitations of standard results for posterior statistics estimators: Standard upper bounds on importance-sampling error are not useful for comparing algorithms because they can be uninformative or loose.The appendix identifies two reasons: infinite χ2 divergence can invalidate an upper bound, and loose bounds can leave the preferred algorithm unclear.
  • C Limitations of standard results for posterior statistics estimators; C.1 Proof of Proposition 1: The plugin estimator appears suboptimal because its bias does not vanish with infinite samples, whereas the SNIS bound may be uninformative when χ2 divergence is infinite.This can occur for a VAE, so the bounds do not immediately determine which estimator performs better.
  • C Limitations of standard results for posterior statistics estimators: Pinsker-based bounds can make both the VAE and WW algorithm minimize an upper bound on plugin mean-square error, leaving the preferred method unclear.The bias can be further bounded using the square root of either ∆KL(pθ ∥qφ) or ∆KL(qφ ∥pθ).
  • D Analytical derivations in the bivariate Gaussian setting: The bivariate Gaussian analysis characterizes approximations by optimizing reverse KL, forward KL, or χ2 divergence while assuming the variational mean is correct.The derivation uses Gaussian conditioning, a closed-form matrix inverse, and the multivariate Gaussian KL expression.
  • D.1 Reverse KL; D.2 Forward KL: The reverse-KL and forward-KL sections derive the solutions to their respective Gaussian optimization problems.Both sections state that the solution follows from similar calculations using the Gaussian KL and matrix-inverse expressions.
  • D.3 Chi-square divergence: The χ2-divergence result is reported from prior derivations for exponential-family and multivariate Gaussian distributions, specifically using the desired quantity for α = −1.The appendix relies on existing results rather than re-deriving the expression.
  • D.4 Importance-weighted variational inference: For IWVI, the low-dimensional IWELBO is estimated with naive Monte Carlo using 10,000 samples, and λ1 and λ2 are optimized numerically by Nelder–Mead.Most quantities are not available in closed form.

E.4 Additional results … F.4 Estimation of posterior expectations for the M1+M2 model

Additional experiments show that proposal choice, particle count, and posterior-collapse mitigation affect estimation and decision quality, while the MNIST study reports similar accuracies and derives posterior estimators for M1+M2. The proposed three-step procedure performs best on MNIST, whereas the SNIS estimator may underperform because it omits the classifier.

  • E.4 Additional results: For pPCA, VAE variational proposals yield poor importance-weighted estimates for most models, while χ and MIS proposals achieve the best PSIS levels.PSIS exceeds 0.7 for IWAE, WW, and χ model combinations using the VAE variational distribution.
  • E.5 Results with an increased number of particles: With more particles, IWAE’s held-out likelihood worsens, whereas WW benefits more, learns the best generative model, and achieves lower MAE than IWAE.The comparison is reported on pPCA simulated data.
  • E.6 Benchmarking for posterior collapse methods: Posterior-collapse alleviation improves the VAE baseline’s held-out log-likelihood and MAE, but lagging inference networks improve held-out likelihood by only 2% and remain below other baselines.Extended studies of posterior-collapse effects are left for future work.
  • F.1 Dataset: The MNIST experiment uses a 50% training and 50% test split.The dataset is MNIST.
  • F.2 Model details and neural networks architecture: The models use convolutional and fully connected neural networks with SELU nonlinearities, dropout rate 0.1, 100 Adam epochs, learning rate 0.001, and batch size 512.The variational encoder has three convolutional layers and two fully connected layers; other distributions use one hidden layer of size 128.
  • F.3 Additional results: MNIST models have relatively similar accuracy, but the three-step procedure applied to IWAE provides the best accuracy levels.IWAE is identified as the best generative model in this comparison.
  • F.4 Estimation of posterior expectations for the M1+M2 model: For M1+M2, posterior expectations can be estimated by naive Monte Carlo or by SNIS using qφ(z,u ∣x,c) as an importance-sampling proposal.SNIS estimates pθ(c,x), normalizes across labels to recover pθ(c ∣x), and marginalizes latent variables z,u.
  • F.4 Estimation of posterior expectations for the M1+M2 model: The SNIS estimator omits qφ(c ∣z), so it may underperform the plugin estimator because that classifier is trained with a classification loss on labeled data.This limitation motivates expecting lower performance from SNIS than from the plugin estimator.

G Analysis of alternate divergences for the M1+M2 model … H.2 Model details and neural networks architecture

The supplemental analysis characterizes alternate divergences for the M1+M2 model and documents the simulated single-cell dataset and scVI-based model used in experiments. It also identifies pathological uncertainty calibration and importance-sampling instability in these settings.

  • G Analysis of alternate divergences for the M1+M2 model: The M1+M2 model exhibits pathological behavior similar to the factor-analysis case, worsened by composite-loss fitting and linked to poorly calibrated neural-network uncertainties.
  • G Analysis of alternate divergences for the M1+M2 model: Alternate divergences decompose into similar closed-form expressions that depend on λ when inference for u and z is semi-exact.The derivations assume qφ(z ∣x)qφ(u ∣c,z) = pθ(z,u ∣c,x), with pθ(c ∣z) having complete support for fixed z.
  • G Analysis of alternate divergences for the M1+M2 model: Reverse-KL reduces to a constant plus the differential entropy of qφ(c ∣z), whose absolute value is bounded by log C.
  • G Analysis of alternate divergences for the M1+M2 model: Forward-KL and chi-square divergences become unbounded when qφ(c ∣z) is zero while pθ(c ∣z) is nonzero.Forward-KL contains pθ(c ∣z)log qφ(c ∣z), whereas chi-square contains pθ(c ∣z)^2/qφ(c ∣z).
  • H.1 Dataset: The simulated scRNA dataset contains N = 1000 cells and G = 100 genes, using two cell states with Poisson-lognormal counts and an 80% −20% train-test split.The two states share covariance, and each cell state is modeled categorically with parameter p.
  • H.1 Dataset: Ground-truth gene log fold changes were assigned among similar, up-regulated, and down-regulated statuses, with Gaussian centers 0, −1, and 1 and σ = 0.16.Gene-expression means for population a were sampled uniformly on (10,100).
  • H.2 Model details and neural networks architecture: The experiment uses a scVI variant in which latent z_n represents biological state, l_n represents library size, and counts follow a zero-inflated negative binomial model.The negative-binomial mean is l_n times the normalized neural-network output h_ng = f_w(z_n), which is deterministic conditional on z_n.
  • H.2 Model details and neural networks architecture: Importance-sampling weights for the original scVI model had high PSIS diagnostics, while high posterior variance for l_n caused numerical instabilities despite more efficient FDR control with alternate divergences.The authors describe the proposal distributions as poor on this dataset and removed priors to work around the problem.

H.3 Additional results

Additional results show that PSIS alone can mislead model and variational-distribution selection, whereas PRAUC and FDR calibration better distinguish reliable differential-expression approaches. IWAE and χ appear sturdy, while traditional VAE fails to estimate FDR reliably.

  • Differential expression: PSIS does not provide a complete picture for selecting a decent model/variational-distribution combination.On differential expression, appealing PSIS values can coexist with deceiving gene rankings.
  • Differential expression: VAE generative-model combinations offer appealing PSIS values but deceiving gene rankings, with PRAUC AUC = 0.94.This result indicates that strong PSIS values do not necessarily imply useful differential-expression rankings.
  • Differential expression: The match between FDR evolution and posterior expected FDR for IWAE and χ hints that they constitute sturdy approaches for differential-expression tasks.The comparison uses the FDR evolution and posterior expected FDR of gene rankings obtained by each model.
  • Differential expression: Traditional VAE fails to estimate FDR reliably.The assessment compares posterior expected FDR with ground-truth FDR for selecting genes with the highest DE probability.
Loading 2002.07217v3…