An Optimal Transport-Based Generative Model for Bayesian Posterior Sampling

1. Problem Statement

The paper addresses the fundamental challenge in Bayesian inference: sampling from posterior distributions $\pi_n(\theta)$ when the normalizing constant (marginal likelihood) is intractable.

• The Bottleneck: Standard methods like Markov Chain Monte Carlo (MCMC) suffer from slow mixing and high correlation in high-dimensional or complex models. Variational Inference (VI) scales well but often provides biased estimates and lacks exact sampling guarantees.
• The Goal: Develop a method to learn a deterministic map $T^*$ that pushes forward a simple reference distribution $\mu$ (e.g., a standard Gaussian) to the target posterior $\pi_n$. Once learned, this map allows for the generation of an arbitrary number of independent posterior samples efficiently, bypassing the need for Markov chain simulation.

2. Methodology

The core contribution is a Generative Model based on Optimal Transport (OT) that learns a specific class of transport maps constrained by OT theory.

A. Theoretical Foundation

Instead of learning an arbitrary transport map (which leads to non-unique solutions and instability), the authors constrain the map $T^*$ to be the Optimal Transport Map in the Monge formulation:
$$T^* = \arg \min_{T \in \mathcal{T}} \mathbb{E}_{X \sim \mu} [\|T(X) - X\|_2^2]$$
subject to $T_\# \mu = \pi_n$.

• Uniqueness & Parsimony: By Brenier's theorem, if $\mu$ is absolutely continuous, the optimal map is unique and takes the form $T^* = \nabla u^*$, where $u^*$ is a convex function. This ensures the map is the "most parsimonious" transformation, preserving geometric structure and enabling unique multivariate ranks.
• Optimization Objective: Directly enforcing the pushforward constraint is computationally intractable. Instead, the authors minimize the Kullback-Leibler (KL) divergence between the transported distribution and the target:
  $$\min_{T \in \mathcal{T}} d_{KL}(T_\# \mu \parallel \pi_n)$$
  Using Lemma 1, this is reformulated into a tractable objective that avoids computing the intractable normalizing constant of $\pi_n$:
  $$\max_{T \in \mathcal{T}} \mathbb{E}_{\mu} [\log(\pi_n \circ T) + \log |\det(J_T)|]$$

B. Architecture for Different Parameter Spaces

The paper proposes specific structural constraints for the class of maps $\mathcal{T}$ based on the nature of the parameters $\theta$:

1. Continuous Parameters:

   • The map is defined as the gradient of a convex function $u$.
   • Approximation: The convex potential $u$ is approximated by a maximum of local convex potentials: $u(x) = \max_{k \in [L]} u_k(x)$.
   • Each local potential $u_k$ is a sum of "convex units" (integrals of increasing bounded functions like Tanh or Softsign).
   • This Max-of-Convex structure allows the model to naturally handle multimodal distributions by assigning different local potentials to different modes.

2. Mixed Discrete-Continuous Parameters:

   • For models with latent discrete variables (e.g., mixture models), the map is non-invertible in the traditional sense.
   • Theorem 1: The authors derive that the optimal potential takes the form $u^*(x) = \kappa \max_{k} (\langle x^{(1)}, b_k \rangle + \phi_k(x^{(2)}))$, where $x^{(1)}$ is the discrete embedding and $x^{(2)}$ is continuous.
   • The discrete component is selected via an argmax (similar to a multi-class Probit model), and the continuous component is mapped via the gradient of the selected local potential.

3. Specialized Adaptations:

   • Nearly Gaussian Posteriors: For regular parametric models where the posterior is asymptotically normal (Bernstein-von Mises theorem), the authors restrict $\mathcal{T}$ to affine maps, drastically reducing computation while maintaining statistical accuracy.
   • Latent Variable Models: To handle the exponential growth of latent states (e.g., in Gaussian Mixture Models), they introduce a Mean-Field Transport class. This decomposes the global potential into a sum of local potentials per observation, reducing complexity from exponential to linear in the sample size.

3. Key Contributions

1. Structural Uniqueness: Unlike previous transport map approaches (e.g., triangular maps or general normalizing flows) that may suffer from non-uniqueness or coordinate-ordering dependencies, this method enforces the OT constraint, guaranteeing a unique, interpretable solution.
2. Multimodal Handling: The "Max-of-Convex" architecture explicitly addresses the difficulty of learning multimodal distributions, a known weakness of standard Normalizing Flows (NFs) and Triangular maps.
3. Mixed-Variable Framework: The paper extends OT theory to mixed discrete-continuous spaces, providing a rigorous formulation for Bayesian latent variable models where standard invertible flows fail.
4. New Inferential Tools: Because the map is an OT map (gradient of a convex function), it preserves the ordering of center-outward quantiles. This enables:
   • Construction of multivariate credible regions.
   • Definition of Bayesian p-values based on geometric extremity.
   • Visualization of joint uncertainty via sign curves and rank contours.

4. Results

The authors evaluate their method against Transport Monte Carlo (TMC) and various Normalizing Flows (Planar, MAF, NSF) across several benchmarks:

• Multimodal Mixtures (Simulation): In high-dimensional Gaussian mixture models ($d=20, K=10$), the proposed method significantly outperformed competitors in terms of 2-Wasserstein distance ($W_2$). TMC and NFs failed to capture all modes or lost correlation structure, while the OT map accurately preserved the geometry of disconnected clusters.
• Bayesian Logistic Regression: The method achieved credible intervals and predictive accuracy comparable to MCMC (Gibbs sampling) but with independent samples. It successfully identified significant covariates and provided accurate uncertainty quantification.
• Gaussian Mixture Models (Latent Variables): In a hierarchical GMM setting, the "Mean-Field" OT approach successfully estimated both the cluster means and the discrete latent assignments, a task where standard NFs are inapplicable.
• Real-World Data (Yeast Dataset): Applied to gene functional classification, the method matched MCMC in variable selection and predictive accuracy (76.8% test accuracy) while being computationally efficient.
• Visualization: The method successfully generated center-outward quantile contours for non-convex (banana-shaped) and multimodal distributions, preserving directional information that other methods distorted.

5. Significance

This paper bridges the gap between Optimal Transport theory and Bayesian computation.

• Efficiency: It offers a "one-shot" sampling mechanism. Once the map is trained, generating millions of independent samples is computationally trivial, unlike MCMC which requires sequential sampling.
• Interpretability: The structural constraints (convexity, monotonicity) provide a geometric interpretation of the posterior, allowing for novel exploratory analysis tools (ranks, quantile contours) that are not available with black-box generative models.
• Robustness: By explicitly modeling the geometry of the target (via the max-of-potentials structure), it overcomes the "mode collapse" and "disconnected support" issues that plague many modern generative models in statistical settings.

In summary, the authors propose a theoretically grounded, structurally constrained generative model that serves as a powerful, efficient, and interpretable alternative to MCMC and Variational Inference for complex Bayesian inference problems.