Bernstein Flows for Flexible Posteriors in Variational Bayes

TL;DR

Bernstein Flow Variational Inference (BF-VI) introduces a flexible, black-box variational framework that uses Bernstein-polynomial transformation models to approximate complex, multi-parameter posteriors. By combining transformation-models with triangular/masked autoregressive maps, BF-VI yields accurate posterior approximations in low dimensions and often outperforms NF-based VI in higher dimensions. The approach is demonstrated across single- and multi-parameter Bayesian models and is further applied to a semi-structured melanoma problem that jointly models image data and interpretable tabular components, enabling uncertainty quantification for interpretable parameters. While tail underestimation can occur in high dimensions due to optimization and KL asymmetry, BF-VI proves to be a practical, scalable tool for uncertainty quantification in complex Bayesian settings and semi-structured modeling.

Abstract

Variational inference (VI) is a technique to approximate difficult to compute posteriors by optimization. In contrast to MCMC, VI scales to many observations. In the case of complex posteriors, however, state-of-the-art VI approaches often yield unsatisfactory posterior approximations. This paper presents Bernstein flow variational inference (BF-VI), a robust and easy-to-use method, flexible enough to approximate complex multivariate posteriors. BF-VI combines ideas from normalizing flows and Bernstein polynomial-based transformation models. In benchmark experiments, we compare BF-VI solutions with exact posteriors, MCMC solutions, and state-of-the-art VI methods including normalizing flow based VI. We show for low-dimensional models that BF-VI accurately approximates the true posterior; in higher-dimensional models, BF-VI outperforms other VI methods. Further, we develop with BF-VI a Bayesian model for the semi-structured Melanoma challenge data, combining a CNN model part for image data with an interpretable model part for tabular data, and demonstrate for the first time how the use of VI in semi-structured models.

Figures (14)

• Figure 1: Overview of the transformation model. A: shows the bijective transformation function $g: Z \rightarrow \theta$ (or its approximation $f_\text{BP}$) mapping between B: a predefined latent density $f_Z$ and C: a potentially complex posterior (or its variational distribution).
• Figure 2: Bernoulli experiment. Left panel: Comparison of the analytical posterior for the parameter $\pi$ in the Bernoulli model ${Y \sim \mathop{\mathrm{Ber}}\nolimits(\pi)}$ with variational distributions achieved via Gaussian-VI and BF-VI with BP order $M=1,10,50$. Right panel: The dependence of the divergence $\mathop{\mathrm{KL}}\nolimits\infdivx{q_\lambda(w)}{p(w \mid D)}\xspace$ on $M$ for 20 runs.
• Figure 3: Cauchy experiment: comparison of MCMC posterior distribution of the parameter $\xi$ in the Cauchy model $Y\sim \text{Cauchy}(\xi, \gamma$) and the variational distributions estimated via Gaussian-VI or BF-VI ($M=2,6,10,30,50$). For the BF-VI method, the curves are overlays of 10 independent runs.
• Figure 4: Toy linear regression example visualization of the posterior. The model has four parameters: two regression coefficients $\beta_1$ and $\beta_2$, the intercept $\mu_0$, and the standard derivation $\sigma$. Samples from the true posterior resulting from MCMC (red) are overlayed with samples from the BF-VI approximation (blue).
• Figure 5: Posterior predictive distribution of the non-linear regression model $(Y \mid\text{\boldmath$x$}) \sim N(\mu(\text{\boldmath$x$}),\sigma=0.2)$ where the conditional mean is modeled by a BNN using MCMC, MF-Gaussian-VI, or BF-FI.