Estimating Complex Densities using Two-Stage Normalizing Flows

TL;DR

A Two-Stage Normalizing Flows framework for approximating and sampling from intractable target distributions, able to accurately recover complex, highly nonlinear target structures using only partial information about the target density, is proposed.

Abstract

In many scientific applications, the target probability distribution cannot be evaluated in closed form or sampled from directly. Instead, it can often be decomposed into multiple components, some of which are accessible only through samples generated by simulators or external datasets, while others admit tractable mathematical expressions or are specified through statistical assumptions about variable relationships. Developing inference methods that coherently integrate these heterogeneous sources of information remains an open challenge. In this paper, we propose a Two-Stage Normalizing Flows framework for approximating and sampling from such distributions. The method first learns the densities of components for which only samples are available, and then combines the outputs with the analytically specified terms to reconstruct the full target distribution in a second stage. The resulting model enables both point-wise density evaluation and efficient generation of representative samples, without requiring direct access to the full target density or joint samples from the complete model. We assess the proposed approach through simulation studies in joint density inference and Bayesian hierarchical models with inaccessible likelihoods. The proposed framework is able to accurately recover complex, highly nonlinear target structures using only partial information about the target density, providing stable and flexible approximations in settings where standard modeling assumptions do not hold (or when complete access to the target distribution is not available). Analysis of a large scale astronomy application highlights interesting differences between our method and existing approaches. Our normalizing flows procedure offers a robust and flexible approach to inference for intractable target distributions across both simulated and real-world applications.

Figures (8)

• Figure 1: Illustration of the training process of a Normalizing Flows model. Starting from a simple base distribution (typically a standard normal), a sequence of invertible transformations gradually morphs the density into the complex shape of the target distribution.
• Figure 2: Comparison of the Stage 2 Normalizing Flow joint model with the ground-truth distribution. The top panel compares marginal fits, and the bottom panel compares bivariate fits.
• Figure 3: Marginal posterior estimates for each parameter under the proposed Normalizing Flows method and the Two-Stage Fully Bayesian procedure, compared with the true analytic marginals.
• Figure 4: Illustration of how the sampling prior displaces the software-generated sampler (a) and how the Two-Stage Fully Bayesian method fails to recover the posterior two-dimensional dependence structure and covariance (b,c). The estimates obtained by the Two-Stage Fully Bayesian method show strong dependence on the software-specific prior of the sampler, making the procedure sensitive to prior misspecification and less reliable when the prior shifts the generated draws away from the high-density regions of the target distribution.
• Figure 5: Selected posterior marginal profile comparison when the sampling prior is uniform in $\mathbb{R}^5$.