FlowVAT: Normalizing Flow Variational Inference with Affine-Invariant Tempering

TL;DR

FlowVAT introduces affine‑invariant tempering for normalizing‑flow variational inference, tempering both the base distribution and the posterior while conditioning the flow on temperature $T$ to yield a single, temperature‑conditional model that generalizes across $T$. This approach preserves modes found at higher temperatures when sampling at $T=1$ and reduces mode‑seeking behavior without requiring annealing schedules, enabling more reliable inference for multimodal, high‑dimensional posteriors. Across 2D, 10D, and 20D synthetic multimodal targets and the eight‑schools benchmark, FlowVAT achieves richer mode recovery and competitive ELBOs, with tempered evidence estimates stabilizing for moderate temperatures. The method offers a promising step toward fully automatic black‑box variational inference for complex posteriors, supported by an inductive bias from affine‑invariant tempering and the generalization capabilities of overparameterized networks.

Abstract

Multi-modal and high-dimensional posteriors present significant challenges for variational inference, causing mode-seeking behavior and collapse despite the theoretical expressiveness of normalizing flows. Traditional annealing methods require temperature schedules and hyperparameter tuning, falling short of the goal of truly black-box variational inference. We introduce FlowVAT, a conditional tempering approach for normalizing flow variational inference that addresses these limitations. Our method tempers both the base and target distributions simultaneously, maintaining affine-invariance under tempering. By conditioning the normalizing flow on temperature, we leverage overparameterized neural networks' generalization capabilities to train a single flow representing the posterior across a range of temperatures. This preserves modes identified at higher temperatures when sampling from the variational posterior at $T = 1$, mitigating standard variational methods' mode-seeking behavior. In experiments with 2, 10, and 20 dimensional multi-modal distributions, FlowVAT outperforms traditional and adaptive annealing methods, finding more modes and achieving better ELBO values, particularly in higher dimensions where existing approaches fail. Our method requires minimal hyperparameter tuning and does not require an annealing schedule, advancing toward fully-automatic black-box variational inference for complicated posteriors.

Figures (11)

• Figure 1: Comparison of different variational inference methods on 2D multimodal targets. The top row shows the target distribution and results from FlowVAT (this work), target-only tempering, traditional annealing, and AdaAnn. The bottom row includes normalizing flow VI and fine-tuned versions of each method, where additional training epochs are run at either $T=1$ for the annealed methods or $T\in[0.95, 1.5]$ for the conditional tempering methods including FlowVAT. ELBO values with standard deviations are shown for each approach. We can see that our approach, FlowVAT, performs better than all other methods.
• Figure 2: Estimated evidence via importance sampling (see \ref{['eq:IS_evidence']}) for the 2D multimodal target posterior along with $1\sigma$ errorbars. As the true posterior is well approximated by the variational posterior, we do not expect a strong temperature dependence.
• Figure 3: The transformation represented by the normalizing flow at various temperatures for the FlowVAT model (top), or target-only tempering (bottom). These plots are produced by transforming a regular grid using the normalizing flow models; as such, an identity transformation would be represented by a regular grid with a grid-spacing of 2. Models without fine-tuning are used to produce this diagram as we are including temperatures outside of the fine-tuning temperature range.
• Figure 4: Estimated evidence via importance sampling (see \ref{['eq:IS_evidence']}) for instances of the 10D (left) and 20D (right) multimodal target posteriors along with $1\sigma$ errorbars. As the true posterior is well approximated by the variational posterior in the 10D case, there is no temperature dependence. In the more challenging 20D case, we can see that the estimated evidence converges for $T\geq1.3$.
• Figure 5: Comparison of variational inference methods on the eight schools model. Top row: Marginal histograms of $\theta_1$. Bottom row: Contour plots of the joint distribution between $\theta_1$ and $\tau$; additional plots are shown in \ref{['sec:eight_schools_plots']}. Contours represent the 68%, 95%, and 99% credible regions, with thick gray lines for the reference posterior and dashed red lines for each variational approximation.