Marginal Flow: a flexible and efficient framework for density estimation

TL;DR

Marginal Flow tackles the bottleneck in density estimation of balancing exact density evaluation with efficient sampling by marginalizing latent parameters. It defines $q(x|w)$ with latent $w$ and learns a distribution $q_theta(w)$ to form $q_theta(x)$ through Monte Carlo marginalization, implemented by a flexible neural mapping $f_theta$ from a base $p_base(z)$. The framework supports learning densities on lower-dimensional manifolds, multi-modal distributions, and conditioning, while enabling efficient training with forward KL, reverse KL, or other objectives and avoiding strict bijective constraints. Empirically, Marginal Flow is orders of magnitude faster than competing models for both sampling and exact density evaluation and demonstrates strong performance across synthetic densities, simulation-based inference benchmarks, Wishart PD-matrix mixtures, and manifold learning in image latent spaces. The work also shows how to extend the method to diverse mixture families (e.g., Dirichlet, Wishart) and provides resources for reproducibility, highlighting practical impact for probabilistic modeling, SBI, and geometric data analysis.

Abstract

Current density modeling approaches suffer from at least one of the following shortcomings: expensive training, slow inference, approximate likelihood, mode collapse or architectural constraints like bijective mappings. We propose a simple yet powerful framework that overcomes these limitations altogether. We define our model $q_θ(x)$ through a parametric distribution $q(x|w)$ with latent parameters $w$. Instead of directly optimizing the latent variables $w$, our idea is to marginalize them out by sampling $w$ from a learnable distribution $q_θ(w)$, hence the name Marginal Flow. In order to evaluate the learned density $q_θ(x)$ or to sample from it, we only need to draw samples from $q_θ(w)$, which makes both operations efficient. The proposed model allows for exact density evaluation and is orders of magnitude faster than competing models both at training and inference. Furthermore, Marginal Flow is a flexible framework: it does not impose any restrictions on the neural network architecture, it enables learning distributions on lower-dimensional manifolds (either known or to be learned), it can be trained efficiently with any objective (e.g. forward and reverse KL divergence), and it easily handles multi-modal targets. We evaluate Marginal Flow extensively on various tasks including synthetic datasets, simulation-based inference, distributions on positive definite matrices and manifold learning in latent spaces of images.

Figures (16)

• Figure 1: Motivation for marginalization: learned distribution and samples when optimizing directly the parameters ${\bm{w}}_i$ compared to resampling them from a learnable $q_\theta({\bm{w}})$, as in Marginal Flow.
• Figure 2: Marginal flow model diagram. Evaluating the modeled density $q_\theta({\bm{x}})$ (center) and sampling from $q_\theta({\bm{x}})$ (right) requires to first sample the parameters ${\bm{w}}_i$ (left).
• Figure 3: Runtime for sampling (left) and exact density evaluation (right) of 100 points.
• Figure 4: Toy example of density defined on (unknown) 1D manifold. (left) Training data consists of 1500 points. (center) Marginal Flow perfectly learns the density and discovers the correct manifold. Free-form Flow learns an incorrect manifold and is not able to embed the density in 2D space. (right) Flow Matching and Normalizing Flow learn the density but cannot account for a manifold.
• Figure 5: Toy example of multi-modal density learned by log-likelihood on 150 data points. For a fair comparison, all models use a uniform base distribution. Note that Marginal Flow is not a mixture model (for which this task would be trivial) since ${\bm{w}}_i$ are always resampled (see Figure \ref{['fig:motivation_marginalization']}).