Estimating Complex Densities using Two-Stage Normalizing Flows

TL;DR

This paper introduces a Two-Stage Normalizing Flows framework that effectively approximates and samples from complex, intractable distributions by leveraging partial information and heterogeneous data sources, enhancing inference in scientific applications.

Contribution

The paper presents a novel two-stage normalizing flows method that combines component densities learned from samples with analytical terms to model complex distributions.

Findings

Accurately recovers nonlinear target structures

Provides stable and flexible density approximations

Effective in real-world astronomy application

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.