[Paper Review] Piecewise Normalizing Flows

H. T. J. Bevins, Will Handley+1 more|arXiv (Cornell University)|May 4, 2023

Machine Learning and Data Classification4 citations

TL;DR

This paper proposes piecewise normalizing flows (PNF), which improve multi-modal density estimation by clustering target data into modes before training separate normalizing flows per cluster. By aligning each cluster's topology with the standard normal base distribution, PNF achieves higher emulation accuracy than prior methods like learned rejection sampling (Stimper et al., 2022), especially under early stopping, while enabling parallel training and better expressivity for complex distributions.

ABSTRACT

Normalizing flows are an established approach for modelling complex probability densities through invertible transformations from a base distribution. However, the accuracy with which the target distribution can be captured by the normalizing flow is strongly influenced by the topology of the base distribution. A mismatch between the topology of the target and the base can result in a poor performance, as is typically the case for multi-modal problems. A number of different works have attempted to modify the topology of the base distribution to better match the target, either through the use of Gaussian Mixture Models (Izmailov et al., 2020; Ardizzone et al., 2020; Hagemann & Neumayer, 2021) or learned accept/reject sampling (Stimper et al., 2022). We introduce piecewise normalizing flows which divide the target distribution into clusters, with topologies that better match the standard normal base distribution, and train a series of flows to model complex multi-modal targets. We demonstrate the performance of the piecewise flows using some standard benchmarks and compare the accuracy of the flows to the approach taken in Stimper et al. (2022) for modelling multi-modal distributions. We find that our approach consistently outperforms the approach in Stimper et al. (2022) with a higher emulation accuracy on the standard benchmarks.

Motivation & Objective

• To address the poor performance of standard normalizing flows on multi-modal distributions due to topological mismatch between base and target distributions.
• To improve flow expressivity and training efficiency by dividing the target distribution into topologically simpler clusters.
• To develop a differentiable, bijective framework that maintains invertibility while enabling parallel training across clusters.
• To demonstrate superior performance compared to state-of-the-art methods such as learned rejection sampling (Stimper et al., 2022) on benchmark multi-modal distributions.

Proposed method

• The method begins by clustering target samples using a differentiable clustering algorithm, assigning each sample to one of k clusters.
• A separate normalizing flow (e.g., MAF) is trained for each cluster, using a standard normal base distribution per cluster.
• Each flow learns the invertible transformation from the standard normal distribution to the conditional distribution of the target within its assigned cluster.
• The overall model is differentiable and bijective when the cluster assignment is known, enabling sampling and density evaluation across the full target distribution.
• The architecture allows parallel training of individual flows, reducing overall training time and improving scalability.
• The approach is evaluated using standard benchmarks and physical data sets, with performance compared to single-flow and learned-base methods.

Experimental results

Research questions

• RQ1 Can dividing a multi-modal target distribution into clusters improve the accuracy of normalizing flow modeling compared to standard flows with a single base distribution?
• RQ2 How does the performance of piecewise normalizing flows compare to learned rejection sampling (Stimper et al., 2022) in terms of density estimation accuracy and training efficiency?
• RQ3 To what extent does the piecewise approach reduce the formation of spurious 'bridges' between modes that plague standard normalizing flows?
• RQ4 Does the use of clustering improve training convergence and reduce computational cost, especially under early stopping?
• RQ5 Can the method maintain differentiability and bijectivity while enabling parallel training of multiple flows?

Key findings

• The piecewise normalizing flow approach consistently outperforms the method of Stimper et al. (2022) in emulation accuracy on standard multi-modal benchmarks, particularly under early stopping.
• The PNF method reduces the formation of spurious bridges between modes, as visually confirmed in Figure 1, by aligning each cluster’s topology with the Gaussian base distribution.
• The method enables parallel training of individual flows per cluster, significantly improving training efficiency and scalability.
• On physical data sets (HEPMASS, MINIBOONE, POWER, GAS), the PNF performs competitively with standard MAFs, though the silhouette score favors the minimum tested number of clusters (k=2), suggesting limited multi-modality in these data sets.
• The approach maintains differentiability and bijectivity when cluster assignment is fixed and the clustering process is differentiable, preserving the core advantages of normalizing flows.
• The method shows promise for mitigating data-driven selection bias in machine learning by enabling user-defined cluster weighting for data generation.