Comparison of Affine and Rational Quadratic Spline Coupling and Autoregressive Flows through Robust Statistical Tests

Abstract

Normalizing flows have emerged as a powerful brand of generative models, as they not only allow for efficient sampling of complicated target distributions but also deliver density estimation by construction. We propose here an in-depth comparison of coupling and autoregressive flows, both based on symmetric (affine) and non-symmetric (rational quadratic spline) bijectors, considering four different architectures: real-valued non-Volume preserving (RealNVP), masked autoregressive flow (MAF), coupling rational quadratic spline (C-RQS), and autoregressive rational quadratic spline (A-RQS). We focus on a set of multimodal target distributions of increasing dimensionality ranging from 4 to 400. The performances were compared by means of different test statistics for two-sample tests, built from known distance measures: the sliced Wasserstein distance, the dimension-averaged one-dimensional Kolmogorov–Smirnov test, and the Frobenius norm of the difference between correlation matrices. Furthermore, we included estimations of the variance of both the metrics and the trained models. Our results indicate that the A-RQS algorithm stands out both in terms of accuracy and training speed. Nonetheless, all the algorithms are generally able, without too much fine-tuning, to learn complicated distributions with limited training data and in a reasonable time of the order of hours on a Tesla A40 GPU. The only exception is the C-RQS, which takes significantly longer to train, does not always provide good accuracy, and becomes unstable for large dimensionalities. All algorithms were implemented using TensorFlow2 and TensorFlow Probability and have been made available on GitHub.

1. Introduction

The modern data science revolution has opened a great window of opportunities for scientific and societal advancement. In particular, machine learning (ML) technologies are being applied in a wide variety of fields from finance to astrophysics. It is thus crucial to carefully study the capabilities and limitations of ML methods in order to ensure their systematic usage. This is particularly pressing when applying ML to scientific research, for instance in a field such as high-energy physics (HEP), where one often deals with complicated high-dimensional data and high levels of precision are needed.

In this paper, we focus on normalizing flows (NFs) [1,2,3,4], a class of neural density estimators that for one, offers a competitive approach to generative models, such as generative adversarial networks (GANs) [5] and variational autoencoders (VAEs) [6,7], for the generation of synthetic data and, for another, opens up a wide range of applications due to its ability to directly perform density estimation. Even though we have in mind applications of NFs to HEP, in this paper, we remain agnostic with respect to the applications and only performed a general comparative study of the performances of coupling and autoregressive NFs when used to learn high-dimensional multi-modal target distributions. Nevertheless, it is worth mentioning some of the applications of NFs to HEP can also be extended to several other fields of scientific research.

While applications of the generative direction of NFs is rather obvious in a field such as HEP, which bases its foundations on Monte Carlo simulations, it is interesting to mention some of the possible density estimation applications. The ability to directly learn the likelihood, or the posterior in a Bayesian framework, has applications ranging from analysis, inference, reinterpretation, and preservation to simulation-based likelihood-free inference [8,9,10,11,12,13], unfolding of HEP analyses [14], generation of effective priors for Bayesian inference [15,16,17,18,19,20], systematic uncertainty estimation and parametrization, generation of effective proposals for sequential Monte Carlo [21,22,23,24,25,26,27,28], numerical integration based on importance sampling algorithms [29,30,31,32], and probabilistic programming applied to fast inference [33,34].

The basic principle behind NFs is to perform a series of invertible bijective transformations on a simple base probability density function (PDF) to approximate a complicated PDF of interest. The optimal parameters of the transformations, often called “bijectors”, are derived from training neural networks (NNs) that directly take the negative log-likelihood of the true data computed with the NF distribution as the loss function. As it turns out, PDFs are everywhere in HEP: from the likelihood function of an experimental or a phenomenological result to the distribution that describes a particle-collision process. Thus, NFs have found numerous applications in HEP: they have been used for numerical integration and event generation [35,36,37,38,39,40,41,42], anomaly detection [43,44,45], detector unfolding [46,47], etc.

The growing interest in NFs implies the urgency of testing state-of-the-art architectures against complex data to ensure their systematic usability and to assess their expected performances. The purpose of this work was then to evaluate the performance of NFs against generic complicated distributions of increasing dimensionality. By performing this study, we aimed to make a concrete step forward in the general understanding of the realistic performances and properties of NFs, especially in high-precision scenarios. This work comprises a substantial upgrade with respect to our early study of Ref. [48], as we now have included more NF architectures, extended the dimensionality of the distributions, and significantly improved the testing strategy.

Our strategy was the following. We implemented in Python, using TensorFlow2 with TensorFlow Probability, four of the most commonly used NF architectures of the coupling and autoregressive type: real-valued non-volume preserving (RealNVP) [49], masked autoregressive flow (MAF) [50], coupling rational quadratic spline (C-RQS) [51], and autoregressive rational quadratic spline (A-RQS) [51].

We tested these NF architectures considering correlated mixture of Gaussian (CMoG) multi-modal distributions with dimensionalities ranging from 4 to 400. We also performed a small-scale hyperparameter scan, explicitly avoiding the fine-tuning of the models and generating the best result for each NF architecture and target distribution.

The performances were measured by means of different test statistics for two-sample testing built from known distance measures: the sliced Wasserstein distance, the dimension-averaged one-dimensional Kolmogorov–Smirnov statistic, and the Frobenius norm of the difference between correlation matrices. The analyses were performed by comparing a test sample, drawn from the original distribution, with an NF-generated one. Moreover, all test-statistics calculations were cross-validated, and an error was assigned both to the evaluation procedure, with repeated calculations of the metrics on different instances of the test and NF-generated samples, and to the training procedure, with repeated calculations on models trained with different instances of the training sample.

This paper is organized as follows. In Section 2, we describe the concept of NFs in more detail, focusing on the coupling and autoregressive types. In Section 3, we introduce the specific NF architectures under investigation. In Section 4, we present the metrics used in our analysis, and in Section 5, we discuss our results. Finally, we provide our concluding remarks in Section 6, with emphasis on the several prospective research avenues that we plan to follow.

2. Normalizing Flows

Normalizing Flows are made of series of bijective, continuous, and invertible transformations that map a simple base PDF to a more complicated target PDF. The purpose of NFs is to estimate the unknown underlying distribution of some data of interest and to allow for the generation of samples approximately following the same distribution. Since the parameters of both the base distribution and the transformations are known, one can $generate$ samples from the target distribution by drawing samples from the base distribution and then applying the proper transformation. This is known as the generative direction of the flow. Furthermore, since the NF transformations are invertible, one can also obtain the probability density of the true samples via inverse transformations from the target to the base PDF. This is known as the normalizing direction of the flow. It is called “normalizing” because the base distribution is often Gaussian even though this is not a requirement, and this is also the origin of the name “normalizing flows”.

The basic idea behind NFs is the change of variable formula for a PDF. Let $X,Y\in {\mathbb{R}}^D$ be random variables with PDFs $p_X,p_Y:{\mathbb{R}}^D\to \mathbb{R}$. Let us define a bijective map $\mathbf{g}:X\to Y$, with inverse $\mathbf{f}={\mathbf{g}}^{-1}$. The two densities are then related by the well known formula

$$p_Y\left(y\right)=p_X\left({\mathbf{g}}^{-1}\left(y\right)\right)\mid det{\mathrm{J}}_g{\mid }^{-1}=p_X\left(\mathbf{f}\left(y\right)\right)\mid det{\mathrm{J}}_f\mid$$

(1)

where ${\mathrm{J}}_f={\displaystyle \frac{\partial \mathbf{f}}{\partial y}}$ is the Jacobian of $\mathbf{f}\left(y\right)$, and ${\mathrm{J}}_g={\displaystyle \frac{\partial \mathbf{g}}{\partial x}}$ is the Jacobian of $\mathbf{g}\left(x\right)$. (Throughout the paper we always interpret X as the base distribution and Y as the target distribution, i.e., the data. We also always model flows in the generative direction, from base to data).

Let us now consider a set of parameters $\left\{\varphi \right\}$ characterizing the chosen base density $p_X$ (typically the mean vector and covariance matrix of a multivariate Gaussian) and parametrize the map $\mathbf{g}$ by another set of parameters $\left\{\theta \right\}$. One can then perform a maximum likelihood estimation of the parameters $\Phi =\{\varphi ,\theta \}$ given some measured data $\mathcal{D}={\left\{y^I\right\}}_{I=1}^N$ distributed according to the unknown PDF $p_y$. The log-likelihood of the data is given by the following expression:

$$\begin{array}{c}logp(\mathcal{D}\mid \Phi )={\displaystyle \sum _{I=1}^{N}log{p}_Y(y^I\mid \Phi )}\hfill \\ ={\displaystyle \sum _{I=1}^{N}log{p}_X({\mathbf{f}}_{\theta }\left(y^I\right)\mid \theta ,\varphi )+log\mid det{\mathrm{J}}_f\mid }\hfill \end{array}$$

(2)

where we make the dependence of $\mathbf{f}$ on $\theta$ explicit through the notation ${\mathbf{f}}_{\theta }$. Then, the best estimate of the parameters $\Phi$ is given by

$$\widehat{\Phi }=\underset{\Phi }{arg\; max}logp(\mathcal{D}\mid \Phi )$$

(3)

Once the parameters $\widehat{\Phi }$ have been estimated from the data, the approximated target distribution can be sampled by applying the generative map $\mathbf{g}$ to samples obtained from the base PDF. The normalizing direction $\mathbf{f}$ can instead be used to perform density evaluation by transforming the new data of interest into sample generated by the base PDF, which is easier to evaluate.

Beside being invertible, the map $\mathbf{g}$ should satisfy the following properties:

• It should be sufficiently expressive to appropriately model the target distribution;
• It should be computationally efficient, meaning that both $\mathbf{f}$ (for training, this means computing the likelihood) and $\mathbf{g}$ (for generating samples), as well as their Jacobian determinants, must be easily calculable.

The composition of invertible bijective functions is also an invertible bijective function. Thus, $\mathbf{g}$ can be generalized to a set of $N_t$ transformations as $\mathbf{g}=g_{N_t}\circ g_{N_t-1}\circ \dots g_1$ with inverse $\mathbf{f}=f_1\circ \dots f_{N_t-1}\circ f_{N_t}$ and $det{\mathrm{J}}_f={\prod }_{n=1}^{N_t}det{\mathrm{J}}_{f_n}$, where each $f_n=g_n^{-1}$ depends on a $y_n$ intermediate random variable. This is a standard strategy to increase the flexibility of the overall transformation.

Typically, but not mandatorily, NF models are implemented using NNs to determine the parameters of the bijectors. The optimal values are obtained by minimizing a loss function corresponding to minus the log-likelihood defined as in Equation (2). This is a natural approach when the samples from the target density are available but the density itself cannot be evaluated, a common occurrence in fields, such as HEP, that heavily rely on Monte Carlo simulations. (Approaches beyond maximum likelihood, which use different loss functions, have also been considered in the literature, such as in Refs [52,53,54,55,56,57]. In this paper we always use the maximum likelihood approach and minus the log-likelihood as loss function.) This makes the models extremely flexible, with a usually stable training, at the cost of a potentially large number of parameters. Nonetheless, the flow transformation must be carefully designed; for instance, even if a given map and its inverse, with their respective Jacobians, are computable, one direction might be more efficient than the other, leading to models that favor sampling over evaluation (and training) or vice versa. Among the wide and growing variety of NF architectures available (see Ref. [58] for an overview), we focus in this work on coupling [4] and autoregressive flows [59], arguably the most widely used implementations of NFs, particularly in HEP.

2.1. Coupling Flows

Coupling flows, originally introduced in Ref. [4], are made of stacks of so-called coupling layers, in which each sample with dimension D is partitioned into two samples A and B with dimensions d and $D-d$, respectively. The parameters of the bjiector transforming the sample A are modeled by a NN that uses B as input, effectively constructing the $p\left(y_d\right|x_{d-D})$ conditional probability distributions. At each coupling layer in the stack, different partitionings are applied, usually by shuffling the dimensions before partitioning, so that all dimensions are properly transformed.

In other words, starting from a disjoint partition of a random variable $Y\in {\mathbb{R}}^D$ such that $(y^A,y^B)\in {\mathbb{R}}^d\times {\mathbb{R}}^{D-d}$ and a bijector $\mathbf{h}(·;\theta ):{\mathbb{R}}^d\to {\mathbb{R}}^d$, a coupling layer maps $\mathbf{g}:X\to Y$ as follows:

$$\begin{array}{cccc}& y^A\hfill & =\hfill & \mathbf{h}(x^A;\Theta \left(x^B\right)),\hfill \\ & y^B\hfill & =\hfill & x^B\hfill \end{array}$$

(4)

where the parameters $\theta$ are defined by a generic function $\Theta \left(x^B\right)$ only defined on the ${\mathbb{R}}^{D-d}$ partition, generally modeled by an NN. The function $\Theta \left(x^B\right)$ is called a conditioner, while the bijectors $\mathbf{h}$ and $\mathbf{g}$ are called coupling function and coupling flow, respectively. The necessary and sufficient condition for the coupling flow $\mathbf{g}$ to be invertible is that the coupling function $\mathbf{h}$ is invertible. In this case, the inverse transformation is given by

$$\begin{array}{cc}& x^A={\mathbf{h}}^{-1}(y^A;\Theta \left(x^B\right)),\hfill \\ & x^B=y^B\hfill \end{array}$$

(5)

Notice that despite the presence of a NN, whose inverse is unknown, to parametrize the conditioner, the invertibility of $\mathbf{h}$ is guaranteed by the fact that such a conditioner is a function of the unchanged dimensions only. The Jacobian of $\mathbf{g}$ is then a two-block triangular matrix. The dimensions $\{1:d\}$ are given by the Jacobian of $\mathbf{h}$, and the dimensions $\{d:D\}$ is the identity matrix. Thus, the Jacobian determinant is simply the following:

$$det{\mathrm{J}}_g=\prod _{i=1}^d{\displaystyle \frac{\partial h_i}{\partial x_i^A}}.$$

(6)

Note that the choice of the partition is arbitrary. The most common choice is to split the dimensions in half, but other partitions are possible [58]. Obviously, even when dimensions are halved, the way in which they are halved is not unique and can be implemented in several different ways, for instance through random masks, or through random shuffling before dividing the first and last half.

2.2. Autoregressive Flows

Autoregressive flows, first introduced in Ref. [59], can be viewed as a generalization of coupling flows. Now, the transformations of each dimension i are modeled by an autoregressive DNN according to the previously transformed dimensions of the distribution, resulting in the $p\left(y_i\right|y_{1:i-1})$ conditional probability distributions, where $y_{1:i-1}$ is a shorthand notation to indicate the list of variables $y_1,\dots y_{i-1}$. After each autoregressive layer, the dimensions are permuted to ensure the expressivity of the bijections over the full dimensionality of the target distribution.

Let us consider a bijector $\mathbf{h}(·;\theta ):\mathbb{R}\to \mathbb{R}$, parametrized by $\theta$. We can define an autoregressive flow function $\mathbf{g}$ such that

$$\begin{array}{cccc}& y_1\hfill & =\hfill & x_1,\hfill \\ & y_i\hfill & =\hfill & \mathbf{h}(x_i;{\Theta }_i\left(y_{1:i-1}\right)),i=2,\dots ,D\hfill \end{array}$$

(7)

The resulting Jacobian of $\mathbf{g}$ is again a triangular matrix, whose determinant is easily computed as

$$det{\mathrm{J}}_g=\prod _{i=1}^D{\displaystyle \frac{\partial h_i}{\partial x_i}}.$$

(8)

where $\partial h_i/\partial x_i$ are the diagonal terms of the Jacobian.

Given that the structure of the bijector is similar to that of the coupling flow, also, in this case, the bijector is referred to as a coupling function. Note that ${\Theta }_j$ can also be alternatively determined with the precedent untransformed dimensions of X [59] such that

$$\begin{array}{cccc}& y_1\hfill & =\hfill & x_1,\hfill \\ & y_i\hfill & =\hfill & \mathbf{h}(x_i;{\Theta }_i\left(x_{1:i-1}\right)),i=2,\dots ,D\hfill \end{array}$$

(9)

The choice of variables used to model the conditioner may depend on whether the NF is intended for sampling or density estimation. In the former case, $\Theta$ is usually chosen to be modeled from the base variable X so that the transformations in the generative direction would only require one forward pass through the flow. The transformations in the normalizing direction would instead require D iterations trough the autoregressive architecture. This case is referred to as inverse autoregressive flow [59] (Notice that in Ref. [58], parametrizing the flow in the normalizing direction (the opposite of our choice), apparently uses the inverse of our formulas for direct and inverse flows. Our notation (and nomenclature) is consistent with that of Ref. [50]) and corresponds to the transformations in Equation (9). Conversely, in the case of density estimation, it is convenient to parametrize the conditioner using the target variable Y since transformations would be primarily in the normalizing direction. This case is referred to as direct autoregressive flow and corresponds to the transformations in Equation (7). In any case, when training the NFs, one always needs to perform the normalizing transformations to estimate the log-likelihood of the data, as in Equation (2). In our study, we only consider the direct autoregressive flow described by Equation (7).

3. Architectures

In the previous section, we describe NFs, focusing on the two most common choices for parametrizing the bijector $\mathbf{g}$ in terms of the coupling function $\mathbf{h}$. The only missing ingredient to make NFs concrete, remains the explicit choice of $\mathbf{h}$. For this study, we have chosen four of the most popular implementations of coupling and autoregressive flows: the real-valued non-volume preserving (RealNVP) [49], the masked autoregressive flow (MAF) [50], and the coupling and autoregressive rational-quadratic neural spline flows (C-RQS and A-RQS) [51]. (Reference [51] refers to coupling and autoregressive RQS flows as RQ-NSF (C) and RQ-NSF (AR), where RQ-NSF stands for rational-quadratic neural spline flow, and A and C for autoregressive and coupling, respectively). We discuss them in turn in the following subsections and give additional details about our specific implementation in Appendix A.1, Appendix A.2, Appendix A.3 and Appendix A.4.

3.1. The RealNVP

The RealNVP [49] is a type of coupling flow whose coupling functions $\mathbf{h}$ are affine functions with the following form:

$$\begin{array}{ccccc}& y_i\hfill & =\hfill & x_i,\hfill & i=1,\dots ,d,\hfill \\ & y_i\hfill & =\hfill & x_i{e}^{s_{i-d}\left(x_{1:d}\right)}+t_{i-d}\left(x_{1:d}\right),\hfill & i=d+1,\dots ,D\hfill \end{array}$$

(10)

where the s and t functions, defined on ${\mathbb{R}}^d\to {\mathbb{R}}^{D-d}$, respectively correspond to the scale and translation transformations modeled by NNs. The product in Equations (10) is intended elementwise for each i so that, $x_{d+1}$ is multiplied by $s_1$, $x_{d+2}$ by $s_2$ and so on up to $x_D$, which is multiplied by $s_{D-d}$. The Jacobian of this transformation is a triangular matrix with diagonal $\mathrm{diag}({\mathbb{I}}_d,\mathrm{diag}(exp\left(s_{i-d}\left(x_{1:d}\right)\right)))$ with $i=d+1,\dots ,D$, so that its determinant is independent of t and is simply given by

$$det\mathrm{J}=\prod _{i=1}^{D-d}{e}^{s_i\left(x_{1:d}\right)}=exp\left({\displaystyle \sum _{i=1}^{D-d}{s}_i\left(x_{1:d}\right)}\right)$$

(11)

The inverse of Equation (10) is given by

$$\begin{array}{ccccc}& x_i\hfill & =\hfill & y_i,\hfill & i=1,\dots ,d,\hfill \\ & x_i\hfill & =\hfill & {\displaystyle \left(y_i-t_{i-d}\left(y_{1:d}\right)\right)e^{-s_{i-d}\left(y_{1:d}\right)},}\hfill & i=d+1,\dots ,D\hfill \end{array}$$

(12)

A crucial property of the affine transformation (10) is that its inverse (12) is again an affine transformation depending only on s and t, and not on their inverse. This implies that the s and t functions can be arbitrarily complicated (indeed they are parametrized by a DNN), still leaving the RealNVP equally efficient in the forward (generative) and backward (normalizing) directions.

3.2. The MAF

The MAF algorithm was developed starting from the masked autoencoder for distribution estimation (MADE) [60] approach for implementing an autoregressive neural network through layers masking (see Appendix A.2).

In the original MAF implementation [50], the bijectors are again affine functions described as

$$\begin{array}{cccc}& y_1\hfill & =\hfill & x_1,\hfill \\ & y_i\hfill & =\hfill & x_i{e}^{s_{i-1}\left(y_{1:i-1}\right)}+t_{i-1}\left(y_{1:i-1}\right),i=2,\dots ,D\hfill \end{array}$$

(13)

The functions s and t are now defined on ${\mathbb{R}}^{D-1}\to {\mathbb{R}}^{D-1}$. The determinant of the Jacobian is simply

$$det\mathrm{J}=\prod _{i=1}^{D-1}{e}^{s_i\left(y_{1:i}\right)}=exp\left({\displaystyle \sum _{i=1}^{D-1}{s}_i\left(y_{1:i}\right)}\right)$$

(14)

and the inverse transformation is

$$\begin{array}{cccc}& x_1\hfill & =\hfill & y_1,\hfill \\ & x_i\hfill & =\hfill & \left(y_i-t_{i-1}\left(y_{1:i-1}\right)\right)e^{-s_{i-1}\left(y_{1:i-1}\right)},i=2,\dots ,D.\hfill \end{array}$$

(15)

As in the case of the RealNVP, the affine transformation guarantees that the inverse transformation only depends on s and t and not on their inverse, allowing for the choice of arbitrarily complicated functions without affecting computational efficiency.

3.3. The RQS Bijector

The bijectors in a coupling or masked autoregressive flow are not restricted to affine functions. It is possible to implement more expressive transformations as long as they remain invertible and computationally efficient. This is the case of the so-called rational-quadratic neural spline flows [51].

The spline bijectors are made of K bins, where in each bin one defines a monotonically increasing rational-quadratic function. The binning is defined on an interval $\mathbb{B}=[-B,B]$, outside of which the function is set to the identity transformation. The bins are defined by a set of $K+1$ coordinates ${\{(x_i^{\left(k\right)},y_i^{\left(k\right)})\}}_{k=0}^K$, called knots, strictly monotonically increasing between $\{(x_i^{\left(0\right)},y_i^{\left(0\right)})=(-B,-B)$ and $\{(x_i^{\left(K\right)},y_i^{\left(K\right)})=(B,B)$. We use the bracket index notation to denote knots’ coordinates, which are defined for each dimension of the vectors $x_i$ and $y_i$. It is possible to construct a rational-quadratic spline bijector with the desired properties with the following procedure [61].

Let us define the quantities

$$\begin{array}{c}{h}_i^{\left(k\right)}=x_i^{(k+1)}-x_i^{\left(k\right)},\hfill \\ {\Delta }_i^{\left(k\right)}=(y_i^{(k+1)}-y_i^{\left(k\right)})/h_i^{\left(k\right)}.\hfill \end{array}$$

(16)

Obviously, ${\Delta }_i^{\left(k\right)}$ represents the variation of $y_i$ with respect to the variation of $x_i$ within the k-th bin. Moreover, since we assumed strictly monotonically increasing coordinates, ${\Delta }_i^{\left(k\right)}$ is always positive or zero. We are interested in defining a bijector $\mathbf{g}\left(x_i\right)$ and mapping the $\mathbb{B}$ interval to itself, such that $\mathbf{g}(x_i^{\left(k\right)})=y_i^{\left(k\right)}$, and with derivatives $d_i^{\left(k\right)}=d{y}_i^{\left(k\right)}/d{x}_i^{\left(k\right)}$ satisfying the following conditions:

$$\begin{array}{cc}{d}_i^{\left(k\right)}=d_i^{(k+1)}=0\hfill & \mathrm{for}{\Delta }_i^{\left(k\right)}=0,\hfill \\ d_i^{\left(k\right)},d_i^{(k+1)}>0\hfill & \mathrm{for}{\Delta }_i^{\left(k\right)}>0\hfill \end{array}$$

(17)

Such condition is necessary and also sufficient, in the case of a rational quadratic function, to ensure monotonicity [61]. Moreover, for the boundary knots, we set $d_i^{\left(0\right)}=d_i^{\left(K\right)}=1$ to match the linear behavior outside the $\mathbb{B}$ interval.

For $x_i\in [x_i^{\left(k\right)},x_i^{(k+1)}]$, we define

$${\theta }_i=(x_i-x_i^{\left(k\right)})/h_i^{\left(k\right)}$$

(18)

such that ${\theta }_i\in [0,1]$. Then, for $x_i$ in each of the intervals $[x_i^{\left(k\right)},x_i^{(k+1)}]$ with $k=0,\dots ,K-1$, we define

$$y_i=P_i^{\left(k\right)}\left({\theta }_i\right)/Q_i^{\left(k\right)}\left({\theta }_i\right)$$

(19)

with the functions P and Q defined by

$$\begin{array}{ccc}{P}_i^{\left(k\right)}\left({\theta }_i\right)\hfill & =\hfill & {\Delta }_i^{\left(k\right)}{y}_i^{(k+1)}{\theta }_i^2+{\Delta }_i^{\left(k\right)}{y}_i^{\left(k\right)}{(1-{\theta }_i)}^2\hfill \\ & & +(y_i^{\left(k\right)}{d}_i^{(k+1)}+y_i^{(k+1)}{d}_i^{\left(k\right)}){\theta }_i(1-{\theta }_i),\hfill \\ Q_i^{\left(k\right)}\left({\theta }_i\right)\hfill & =\hfill & {\Delta }_i^{\left(k\right)}+(d_i^{(k+1)}+d_i^{\left(k\right)}-2{\Delta }_i^{\left(k\right)}){\theta }_i(1-{\theta }_i)\hfill \end{array}$$

(20)

The ratio in Equation (19) can then be written in the simplified form

$$y_i=y_i^{\left(k\right)}+{\displaystyle \frac{(y_i^{(k+1)}-y_i^{\left(k\right)})({\Delta }_i^{\left(k\right)}{\theta }_i^2+d_i^{\left(k\right)}{\theta }_i(1-{\theta }_i))}{{\Delta }_i^{\left(k\right)}+(d_i^{(k+1)}+d_i^{\left(k\right)}-2{\Delta }_i^{\left(k\right)}){\theta }_i(1-{\theta }_i)}}$$

(21)

The Jacobian ${\mathrm{J}}_{\mathbf{g}}=\partial y_i/\partial x_j$ is then diagonal, with entries given by

$${\displaystyle \frac{{({\Delta }_i^{\left(k\right)})}^2(d_i^{(k+1)}{\theta }_i^2+2{\Delta }_i^{\left(k\right)}{\theta }_i(1-{\theta }_i)+d_i^{\left(k\right)}{(1-{\theta }_i)}^2)}{{({\Delta }_i^{\left(k\right)}+(d_i^{(k+1)}+d_i^{\left(k\right)}-2{\Delta }_i^{\left(k\right)}){\theta }_i(1-{\theta }_i))}^2}}$$

(22)

for $i=1,\dots ,D$. The inverse of the transformation (19) can also be easily computed by solving the quadratic Equations (19) with respect to $x_i$.

In practice, B and K are hyperparameters, while ${\{(x_i^{\left(k\right)},y_i^{\left(k\right)})\}}_{k=0}^K$ and ${\{d_i^{\left(k\right)}\}}_{k=1}^{K-1}$ are $2(K+1)$ plus $K-1$ parameters, modeled by an NN, which determine the shape of the spline function. The different implementations of the RQS bijector, in the context of coupling and autoregressive flows, are determined by the way in which such parameters are computed. We briefly describe them in turn in the following two subsections.

3.4. The C-RQS

In the coupling flow case (C-RQS), one performs the usual partitioning of the D dimensions in the two sets composed of the first d and last $D-d$ dimensions. The first d dimensions are then kept unchanged $y_i=x_i$ for $i=1,\dots ,d$, while the parameters describing the RQS transformations of the other $D-d$ dimensions are determined from the inputs of the first d dimensions, denoted by $x_{1:d}$. Schematically, we could write

$$\begin{array}{ccc}& x_i^{\left(k\right)}\hfill & =x_i^{\left(k\right)}\left(x_{1:d}\right),\hfill \\ & y_i^{\left(k\right)}\hfill & =y_i^{\left(k\right)}\left(x_{1:d}\right),\hfill \\ & d_i^{\left(k\right)}\hfill & =d_i^{\left(k\right)}\left(x_{1:d}\right)\hfill \end{array}$$

(23)

for $i=d+1,\dots ,D$.

A schematic description of our implementation of the C-RQS is given in Appendix A.3.

3.5. The A-RQS

The RQS version of the MAF, which we call A-RQS, is instead obtained by leaving unchanged the first dimension $y_1=x_1$ and determining the parameters of the transformation of the i-th dimension from the output of all preceding dimensions, denoted by $y_{1:i-1}$. Schematically, this is given by

$$\begin{array}{ccc}& x_i^{\left(k\right)}\hfill & =x_i^{\left(k\right)}\left(y_{1:i-1}\right),\hfill \\ & y_i^{\left(k\right)}\hfill & =y_i^{\left(k\right)}\left(y_{1:i-1}\right),\hfill \\ & d_i^{\left(k\right)}\hfill & =d_i^{\left(k\right)}\left(y_{1:i-1}\right)\hfill \end{array}$$

(24)

for $i=2,\dots ,D$.

4. Non-Parametric Quality Metrics

We assessed the performance of our trained models using three distinct metrics: the dimension-averaged 1D Kolmogorov–Smirnov (KS) two-sample test statistic ${\overline{D}}_{y,z}$, the sliced Wasserstein distance (SWD) ${\overline{W}}_{y,z}$, and the Frobenius norm (FN) of the difference between the correlation matrices of two samples ${\parallel C\parallel }_F$. With a slight abuse of nomenclature, we refer to these three different distance measures with vanishing optimal value simply as KS, SWD, and FN, respectively. Each of these metrics served as a separate test statistic in a two-sample test, where the null hypothesis assumed that both samples originate from the same target distribution. For each metric, we established its distribution under the null hypothesis by drawing both samples from the target distribution. We then compared this distribution with the test-statistic calculated from a two-sample test between samples drawn from the target and NF distributions to assign each model a p-value for rejecting the null hypothesis.

To quantify the uncertainty on the test-statistics computed for the test vs NF-generated samples, we performed the tests 10 times using differently seeded target- and NF-generated samples. We calculated p-values based on the mean test statistic and its $\pm 1$ standard deviation.

For model comparison and to assess the uncertainty on the training procedure, we trained 10 instances of each model configuration, defined by a set of hyperparameter values. Rather than selecting the single best-performing instance to represent the best architecture, we averaged the performances across these instances and identified the architecture with the best average performance. After selecting this top-performing model, which we called the “best” model, we reported both its average and peak performances.

When computing the test-statistics distributions under the null hypothesis and evaluating each model’s p-values, we found that the discriminative power of the KS metric was larger than that of the SWD and FN ones. For this reason, we used the result of the KS-statistic to determine the best model and then showed results also for the other two statistics. Nevertheless, even though the best model could vary depending on the metric used, no qualitative difference in the conclusions would arise from choosing FN or SWD as the ranking metric; i.e., results were not identical but were consistent among the three metrics.

It is important to stress that despite our insistence on using non-parametric quality metrics, we actually know the target density, and we used this information for bootstrapping uncertainties and computing p-values. In real-world examples, the target density is generally not known, and depending on the number of available samples, our procedure for evaluation needs to be adapted or may end up being unusable. Nevertheless, this well-defined statistical approach is crucial for us since we aim to draw rather general conclusions, which strongly depend on the ability to estimate the uncertainties and should rely on robust statistical inference.

In the following, we briefly introduce the three aforementioned metrics. To do so, we employ the following notation: we indicate with N the number of D-dimensional points in each sample and use capital indices $I,J$ to run over N and lowercase indices $i,j$ to run over D (We warn the reader not to confuse the dimensionality D with the KS test-statistic $D_{x,y}$.). We also use Greek letters indices $\alpha ,\beta$ to run over slices (random directions).

• Kolmogorov–Smirnov test
  The KS test is a statistical test used to determine whether or not two 1D samples are drawn from the same unknown PDF. The null hypothesis is that both samples come from the same PDF. The KS test statistic is given by

  $$D_{y,z}={\mathbf{sup}}_x\mid F_y\left(x\right)-F_z\left(x\right)\mid ,$$

  (25)

  where $F_{y,z}\left(x\right)$ are the empirical distributions of each of the samples $\left\{y_I\right\}$ and $\left\{z_I\right\}$, and $\mathbf{sup}$ is the supremum function. For characterizing the performances of our results, we computed the KS test-statistic for each of the 1D marginal distributions along the D dimensions and took the average as follows:

  $${\overline{D}}_{y,z}={\displaystyle \frac{1}{D}}\sum _{i=1}^D{D}_{y,z}^i.$$

  (26)
  The actual test statistic that we consider in this paper is the scaled version of ${\overline{D}}_{y,z}$, given by

  $$t_{\mathrm{KS}}=\sqrt{{\displaystyle \frac{N}{2}}}{\overline{D}}_{y,z}$$

  (27)
  The $\sqrt{N/2}$ factor comes from the known $\sqrt{m·n/(m+n)}$ factor in the scaled KS statistic with different-sized samples of sizes m and n, respectively.
  Notice that even though the test statistic

  $$\sqrt{{\displaystyle \frac{m·n}{m+n}}}{D}_{y,z}$$

  (28)

  is asymptotically distributed according to the Kolmogorov distribution [62,63,64,65,66], the same is not true for our $t_{\mathrm{KS}}$ statistic due to correlations among dimensions. Nevertheless, our results seem to suggest that the asymptotic distribution of $t_{\mathrm{KS}}$ for large D (that means when the average is taken over many dimensions) has a reasonably universal behavior, translating into almost constant rejection lines (solid gray lines with different thicknesses) in the upper panels of Figure 1.
• Sliced Wasserstein Distance
  The SWD [67,68] is a distance measure for comparing two multidimensional distributions based on the 1D Wasserstein distance [69,70]. The latter distance between two univariate distributions is given as a function of their respective empirical distributions as follows

  $$W_{y,z}={\int }_{\mathbb{R}}dx\mid F_y\left(x\right)-F_z\left(x\right)\mid$$

  (29)
  Intuitively, the difference between the WD and the KS test statistic is that the latter considers the maximum distance, while the former is based on the integrated distance.
  Our implementation of the SWD is defined as follows. For each model and each dimensionality D, drew $2D$ random directions ${\widehat{v}}_{\alpha }^i$, with $i=1,\dots D$ and $\alpha =1,\dots ,2D$, uniformly distributed on the surface of the unit N-sphere (This can be done by normalizing the D-dimensional vector obtained by sampling each components from an independent standard normal distribution [71]). Given two D-dimensional samples $\left\{{\mathbf{y}}_I\right\}=\left\{y_I^i\right\}$ and $\left\{{\mathbf{z}}_I\right\}=\left\{z_I^i\right\}$, we considered the $2D$ projections

  $$\left\{y_I^{\alpha }\right\}=\left\{\sum _{i=1}^D{y}_I^i{\widehat{v}}_i^{\alpha }\right\},\left\{z_I^{\alpha }\right\}=\left\{\sum _{i=1}^D{z}_I^i{\widehat{v}}_i^{\alpha }\right\}$$

  (30)

  and computed the corresponding $2D$ Wasserstein distances as follows:

  $$W_{y,z}^{\alpha }={\int }_{\mathbb{R}}dx\mid F_{y^{\alpha }}\left(x\right)-F_{z^{\alpha }}\left(x\right)\mid$$

  (31)

  was

  $${\overline{W}}_{y,z}={\displaystyle \frac{1}{2D}}\sum _{\alpha =1}^{2D}{W}_{y,z}^{\alpha }.$$

  (32)
  In analogy with the scaled KS test statistic, we defined the scaled SWD test statistic as follows:

  $$t_{\mathrm{SWD}}=\sqrt{{\displaystyle \frac{N}{2}}}{\overline{W}}_{y,z}$$

  (33)
• Frobenius norm
  The FN of a matrix M is given by

  $${\parallel M\parallel }_F=\sqrt{{\displaystyle \sum _{i,j}}\mid m_{ij}{\mid }^2},$$

  (34)

  where $m_{ij}$ are the elements of M. By defining $C=C_y-C_z$, where $C_y,C_z$ are the two $N\times N$ correlation matrices of the samples $\left\{{\mathbf{y}}_I\right\}=\left\{y_I^i\right\}$ and $\left\{{\mathbf{z}}_I\right\}=\left\{z_I^i\right\}$, its FN, given by

  $${\parallel C\parallel }_F=\sqrt{{\displaystyle \sum _{i,j}}\mid c_{y,ij}-c_{z,ij}{\mid }^2}$$

  (35)

  is a distance measure between the two correlation matrices. In analogy with the previously defined test statistics, we defined the scaled FN test statistic as follows:

  $$t_{\mathrm{FN}}=\sqrt{{\displaystyle \frac{N}{2}}}{\displaystyle \frac{{\parallel C\parallel }_F}{D}},$$

  (36)

  where we also divided by the number of dimensions D to remove the approximately linear dependence of the FN distance on the dimensionality of the samples.

5. Testing the Normalizing Flows

We tested the four architectures discussed above on CMoG distributions defined as a mixture of $n=3$ components and $D=4,8,16,32,64,100,200,400$ dimensional multivariate Gaussian distributions with diagonal covariance matrices, parametrized by means randomly generated in the $[0,10]$ interval and standard deviations randomly generated in the $[0,1]$ interval (The values for the means and standard deviations were chosen so that the different components could generally be resolved). The components were mixed according to an n dimensional categorical distribution (with random probabilities). This meant that a different probability was assigned to each component, while different dimensions of the same component multivariate Gaussian were assigned the same probability. The resulting multivariate distributions had random order-one off-diagonal elements in the covariance matrix and multi-modal 1D marginal distributions (see, for illustration, Figure A1 and Figure A2).

In our analysis, we considered a training set of ${10}^5$ points, a validation set of $3\times {10}^4$ points, and a test set equal in size to the training set, with ${10}^5$ points. It is important to note that the chosen size of the test set corresponds to the most stringent condition for evaluating the NF models. This is because the NF cannot be expected to approximate the real target distribution more accurately than the uncertainty determined by the size of the training sample.

In practical terms, the most effective NF would be indistinguishable from the target distribution when tested on a sample size equivalent to the training set. In our analysis, we found that models tested with ${10}^5$ samples often led to rejection at the $2\sigma /3\sigma$ level, at least with the most powerful KS test. However, this should not be viewed as a poor outcome. Rather, it suggests that one needs to utilize a test set as large as the training set to efficiently discern the NF from the true model, while smaller samples are effectively indistinguishable from those generated with the target distribution.

An alternative approach, which we did not adopt due to computational constraints, involves calculating the sample size required to reject the null hypothesis at a given confidence level. This approach offers a different but equally valid perspective, potentially useful for various applications. Nevertheless, our approach was efficient for demonstrating that NFs can perform exceptionally well on high-dimensional datasets and for comparing, among each other, the performances of different NF architectures.

For each of the four different algorithms described above and for each value of N, we performed a small scan over some of the free hyperparameters. Details on the choice of the hyperparameters are reported in Appendix B. All models were trained on Nvidia A40 GPUs.

The performances of the best NF architectures are reported in Figure 1 and

Table A2. Values of the most relevant hyperparameters and metrics for the average best models obtained for the CMoG distributions. The number of training epochs and the training and prediction times are averages. The columns KS, Sliced WD, and Frobenius Norm contain the values of the corresponding test statistics shown in the left panels of Figure 1. The row corresponding to the best model for each dimension, that is, the one with the minimum KS test statistic, is shown in bold.

Results for Average Best Models
Hidden	# of	Algorithm	Spline	Trainable	KS	Sliced	Frobenius	# of	Training	Prediction
Layers	Bijec.		Knots	Parameters		WD	Norm	Epochs	Time (s)	Time (s)
4D
$3\times 128$	10	MAF	–	346,960	$1.5\pm 0.4$	$3.1\pm 0.7$	$0.9\pm 0.3$	442	5301	17
$3\times 256$	5	RealNVP	–	666,900	$1.6\pm 0.5$	$3.7\pm 0.9$	$1.0\pm 0.4$	616	8777	8
$\mathbf{3}\times \mathbf{128}$	2	A-RQS	8	91,064	$\mathbf{1.2}\pm \mathbf{0.1}$	$\mathbf{2.6}\pm \mathbf{0.4}$	$\mathbf{0.7}\pm \mathbf{0.2}$	$\mathbf{670}$	$\mathbf{7606}$	$\mathbf{54}$
$3\times 128$	5	C-RQS	8	196,710	$1.4\pm 0.2$	$3.3\pm 0.3$	$1.0\pm 0.5$	483	12,346	26
8D
$\mathbf{3}\times \mathbf{128}$	5	MAF	–	181,200	$\mathbf{1.3}\pm \mathbf{0.2}$	$\mathbf{3.3}\pm \mathbf{0.7}$	$\mathbf{0.8}\pm \mathbf{0.3}$	$\mathbf{713}$	$\mathbf{5757}$	$\mathbf{11}$
$3\times 128$	10	RealNVP	–	346,960	$2.2\pm 0.5$	$5.2\pm 1.1$	$1.4\pm 0.4$	340	8242	12
$3\times 128$	2	A-RQS	12	140,592	$1.6\pm 0.5$	$4.7\pm 2.1$	$1.0\pm 0.4$	477	5516	315
$3\times 128$	5	C-RQS	8	227,660	$2.4\pm 0.6$	$6.1\pm 1.9$	$1.6\pm 0.5$	294	$\mathrm{11,500}$	27
16D
$\mathbf{3}\times \mathbf{128}$	5	MAF	–	196,640	$\mathbf{1.6}\pm \mathbf{0.2}$	$\mathbf{3.7}\pm \mathbf{0.6}$	$\mathbf{7.1}\pm \mathbf{1.9}$	$\mathbf{479}$	$\mathbf{3410}$	$\mathbf{16}$
$3\times 128$	5	RealNVP	–	181,200	$2.6\pm 0.3$	$6.4\pm 1.3$	$12\pm 3$	366	4263	8
$3\times 128$	2	A-RQS	12	214,880	$1.8\pm 0.3$	$3.1\pm 0.5$	$2.2\pm 1.0$	327	3705	65
$3\times 128$	10	C-RQS	8	579,120	$2.0\pm 0.6$	$3.9\pm 1.3$	$4.6\pm 2.0$	558	64,056	53
32D
$\mathbf{3}\times \mathbf{128}$	5	MAF	–	227,520	$\mathbf{1.7}\pm \mathbf{0.3}$	$\mathbf{4.5}\pm \mathbf{0.9}$	$\mathbf{2.1}\pm \mathbf{0.6}$	$\mathbf{595}$	$\mathbf{4193}$	$\mathbf{24}$
$3\times 128$	10	RealNVP	–	393,280	$2.6\pm 0.5$	$6.8\pm 1.6$	$3.0\pm 0.8$	676	16,946	12
$3\times 128$	2	A-RQS	8	264,384	$1.9\pm 0.3$	$4.4\pm 1.0$	$1.7\pm 0.4$	375	4705	97
$3\times 256$	10	C-RQS	12	2,798,560	$2.0\pm 0.5$	$5.2\pm 1.7$	$2.3\pm 0.8$	750	75,606	83
64D
$3\times 128$	10	MAF	–	578,560	$2.1\pm 0.6$	$7\pm 3$	$1.3\pm 0.5$	537	5289	89
$3\times 256$	10	RealNVP	–	1,564,800	$2.3\pm 0.4$	$6.5\pm 1.7$	$1.5\pm 0.3$	711	13,871	16
$\mathbf{3}\times \mathbf{128}$	2	A-RQS	8	462,464	$\mathbf{1.5}\pm \mathbf{0.2}$	$\mathbf{4.0}\pm \mathbf{0.7}$	$\mathbf{0.9}\pm \mathbf{0.1}$	$\mathbf{523}$	$\mathbf{6197}$	$\mathbf{332}$
$3\times 256$	5	C-RQS	12	2,139,360	$2.4\pm 0.9$	$6.4\pm 2.7$	$2.6\pm 0.8$	813	36,608	53
100D
$3\times 128$	10	MAF	–	717,520	$2.2\pm 0.8$	$7\pm 4$	$2.1\pm 1.2$	778	7824	144
$3\times 256$	10	RealNVP	–	1,703,400	$3.2\pm 1.3$	$10\pm 6$	$2.9\pm 1.9$	991	23,500	19
$\mathbf{3}\times \mathbf{128}$	2	A-RQS	12	994,904	$\mathbf{1.4}\pm \mathbf{0.2}$	$\mathbf{4.0}\pm \mathbf{0.9}$	$\mathbf{1.0}\pm \mathbf{0.2}$	$\mathbf{588}$	$\mathbf{6219}$	$\mathbf{1027}$
200D
$3\times 128$	10	MAF	–	1,103,520	$3.7\pm 1.4$	$12\pm 6.0$	$3.4\pm 2.4$	612	6149	393
$3\times 256$	10	RealNVP	–	2,088,400	$6.3\pm 2.3$	$19\pm 11$	$5.4\pm 2.2$	1000	22,704	25
$\mathbf{3}\times \mathbf{128}$	2	A-RQS	12	1,923,504	$\mathbf{1.3}\pm \mathbf{0.2}$	$\mathbf{3.3}\pm \mathbf{1.1}$	$\mathbf{1.0}\pm \mathbf{0.5}$	$\mathbf{703}$	$\mathbf{9943}$	$\mathbf{4900}$
400D
$3\times 128$	10	MAF	–	1,875,520	$4.3\pm 1.8$	$14\pm 8$	$8\pm 6$	600	4612	1242
$3\times 256$	0	RealNVP	–	2,858,400	$8.4\pm 2.4$	$21\pm 11$	$24\pm 10$	824	23,705	38
$\mathbf{3}\times \mathbf{128}$	2	A-RQS	8	2,542,304	$\mathbf{2.2}\pm \mathbf{0.9}$	$\mathbf{6.9}\pm \mathbf{1.8}$	$\mathbf{11}\pm \mathbf{4}$	$\mathbf{796}$	$\mathbf{9970}$	$\mathbf{9738}$

and

Table A3. Values of the most relevant hyperparameters and metrics for the absolute best models obtained for the CMoG distributions. The number of training epochs and the training and prediction times are averages. The columns KS, Sliced WD, and Frobenius Norm contain the values of the corresponding-statistics shown in the right panels of Figure 1. The row corresponding to the best model for each dimension, that is, the one with the minimum KS test statistic, is shown in bold.

Results for Absolute Best Models
Hidden	# of	Algorithm	Spline	Trainable	KS	Sliced	Frobenius	# of	Training	Prediction
Layers	Bijec.		Knots	Parameters		WD	Norm	Epochs	Time (s)	Time (s)
4D
$\mathbf{3}\times \mathbf{128}$	10	MAF	–	346,960	$\mathbf{1.1}\pm \mathbf{0.1}$	$\mathbf{2.0}\pm \mathbf{0.5}$	$\mathbf{0.5}\pm \mathbf{0.1}$	$\mathbf{442}$	$\mathbf{5301}$	$\mathbf{17}$
$3\times 256$	5	RealNVP	–	666,900	$1.1\pm 0.2$	$2.3\pm 0.6$	$0.7\pm 0.3$	616	8777	8
$3\times 128$	2	A-RQS	8	91,064	$1.1\pm 0.2$	$2.2\pm 0.6$	$0.5\pm 0.2$	670	7606	54
$3\times 128$	5	C-RQS	8	196,710	$1.2\pm 0.2$	$2.8\pm 0.7$	$0.6\pm 0.2$	483	12,346	26
8D
$3\times 128$	5	MAF	–	181,200	$1.1\pm 0.2$	$2.7\pm 0.4$	$0.6\pm 0.1$	713	5757	11
$3\times 128$	10	RealNVP	–	346,960	$1.3\pm 0.2$	$3.2\pm 1.0$	$0.6\pm 0.1$	340	8242	12
$\mathbf{3}\times \mathbf{128}$	2	A-RQS	12	140,592	$\mathbf{1.1}\pm \mathbf{0.1}$	$\mathbf{2.6}\pm \mathbf{0.8}$	$\mathbf{0.6}\pm \mathbf{0.2}$	$\mathbf{477}$	$\mathbf{5516}$	$\mathbf{315}$
$3\times 128$	5	C-RQS	8	227,660	$1.4\pm 0.1$	$3.1\pm 0.8$	$1.0\pm 0.2$	294	11,500	27
16D
$3\times 128$	5	MAF	–	196,640	$1.3\pm 0.1$	$2.8\pm 0.9$	$2.4\pm 0.9$	479	3410	16
$3\times 128$	5	RealNVP	–	181,200	$2.2\pm 0.1$	$4.6\pm 0.6$	$6.3\pm 0.9$	366	4263	8
$3\times 128$	2	A-RQS	12	214,880	$1.4\pm 0.1$	$2.4\pm 0.5$	$1.1\pm 0.2$	327	3705	65
$\mathbf{3}\times \mathbf{128}$	10	C-RQS	8	579,120	$\mathbf{1.1}\pm \mathbf{0.1}$	$\mathbf{2.3}\pm \mathbf{0.3}$	$\mathbf{1.9}\pm \mathbf{0.3}$	$\mathbf{558}$	64,056	$\mathbf{53}$
32D
$\mathbf{3}\times \mathbf{128}$	5	MAF	–	227,520	$\mathbf{1.2}\pm \mathbf{0.1}$	$\mathbf{2.8}\pm \mathbf{0.7}$	$\mathbf{0.9}\pm \mathbf{0.2}$	$\mathbf{595}$	$\mathbf{4193}$	$\mathbf{24}$
$3\times 128$	10	RealNVP	–	393,280	$1.7\pm 0.2$	$4.0\pm 1.4$	$1.8\pm 0.5$	676	16,946	12
$3\times 128$	2	A-RQS	8	264,384	$1.3\pm 0.1$	$2.9\pm 0.4$	$1.2\pm 0.2$	375	4705	97
$3\times 256$	10	C-RQS	12	2,798,560	$1.4\pm 0.2$	$2.9\pm 0.5$	$1.2\pm 0.1$	750	75,606	83
64D
$3\times 128$	10	MAF	–	578,560	$1.5\pm 0.1$	$3.1\pm 0.7$	$0.8\pm 0.1$	537	5289	89
$3\times 256$	10	RealNVP	–	1,564,800	$1.8\pm 0.1$	$3.7\pm 0.8$	$1.1\pm 0.2$	711	13,871	16
$\mathbf{3}\times \mathbf{128}$	2	A-RQS	8	462,464	$\mathbf{1.1}\pm \mathbf{0.2}$	$\mathbf{3.1}\pm \mathbf{0.8}$	$\mathbf{0.6}\pm \mathbf{0.1}$	$\mathbf{523}$	$\mathbf{6197}$	$\mathbf{332}$
$3\times 256$	5	C-RQS	12	2,139,360	$1.5\pm 0.1$	$3.7\pm 0.5$	$1.8\pm 0.1$	813	36,608	53
100D
$3\times 128$	10	MAF	–	717,520	$1.3\pm 0.1$	$3.7\pm 0.7$	$0.8\pm 0.2$	778	7824	144
$3\times 256$	10	RealNVP	–	1,703,400	$1.9\pm 0.1$	$3.1\pm 0.5$	$1.4\pm 0.1$	991	23,500	19
$\mathbf{3}\times \mathbf{128}$	2	A-RQS	12	994,904	$\mathbf{1.1}\pm \mathbf{0.2}$	$\mathbf{2.8}\pm \mathbf{0.8}$	$\mathbf{0.8}\pm \mathbf{0.3}$	$\mathbf{588}$	$\mathbf{6219}$	$\mathbf{1027}$
200D
$3\times 128$	10	MAF	–	1,103,520	$2.2\pm 0.3$	$5.9\pm 1.3$	$1.0\pm 0.1$	612	6149	393
$3\times 256$	10	RealNVP	–	2,088,400	$4.0\pm 0.2$	$8.8\pm 1.2$	$2.9\pm 0.2$	1000	22,704	25
$\mathbf{3}\times \mathbf{128}$	2	A-RQS	12	1,923,504	$\mathbf{1.0}\pm \mathbf{0.1}$	$\mathbf{2.3}\pm \mathbf{0.7}$	$\mathbf{0.6}\pm \mathbf{0.1}$	$\mathbf{703}$	$\mathbf{9943}$	$\mathbf{4900}$
400D
$3\times 128$	10	MAF	–	1,875,520	$1.7\pm 0.2$	$13.2\pm 1.3$	$12.0\pm 0.7$	600	4612	1242
$3\times 256$	10	RealNVP	–	2,858,400	$5.3\pm 0.1$	$10.1\pm 0.2$	$8.3\pm 0.1$	824	23,705	38
$\mathbf{3}\times \mathbf{128}$	2	A-RQS	8	2,542,304	$\mathbf{1.2}\pm \mathbf{0.1}$	$\mathbf{3.3}\pm \mathbf{0.7}$	$\mathbf{2.7}\pm \mathbf{0.6}$	$\mathbf{796}$	$\mathbf{9970}$	$\mathbf{9738}$

.

Figure 1 shows the values of the three test statistics (vertical panels) for the average (left panels) and absolute (right panels) best models obtained with the four different architectures. The three gray lines with different thicknesses represent the values of the test statistics corresponding to $1\sigma$, $2\sigma$, and $3\sigma$ rejection (p-values of $0.68$, $0.95$, and $0.99$, respectively) of the null hypothesis that the two samples (test and NF-generated) are drawn from the same PDF. These rejection lines were obtained through ${10}^4$ pseudo-experiments. The curve for the best C-RQS models stops at 64D since the training becomes unstable and the model does not converge. The situation could likely be improved by adding regularization and by fine-tuning the hyperparameters. However, to allow for a fair comparison with the other architectures, where regularization and fine-tuning are not necessary for a reasonable convergence, we avoided pushing C-RQS beyond 64D. Also notice that the uncertainty shown in the point at 64D for the C-RQS is artificially very small since only a small fraction of the differently seeded runs converged. This uncertainty should therefore be considered unreliable.

All plots in Figure 1 include uncertainties. As already mentioned, the best model was chosen as the one with best architecture on average, and therefore, over 10 different trainings were performed with differently seeded training samples. For the selection of the best model, the left plots show the performances averaged over the 10 trainings, with error bands representing the corresponding standard deviations, while the right plots show the performances of the absolute best instance among the 10 trained replicas, with the error band representing the standard deviation over the 10 replicas generated for testing (test and NF-generated samples). In other words, we can say that the uncertainties shown in the left plots are the standard deviations due to repeated training, while the uncertainties shown in the right plots are the standard deviations due to repeated generation/evaluation (testing).

Figure 1 clearly highlights the distinct characteristics that establish the A-RQS as the top-performing algorithm:

• Its performances are almost independent of the data dimensionality;
• The average best model is generally not rejected at $3\sigma$ level when evaluated with a number of points equal to the number of training points;
• The absolute best model is generally not rejected at $2\sigma$ level when evaluated with a number of points equal to the number of training points;
• The uncertainties due to differently seeded training and testing are generally comparable, while for all other models, the uncertainty from training is generally much larger than the one from evaluation.

All values shown in Figure 1 are reported in Table A2 and Table A3 for the average and absolute best models, respectively. In the tables, we also show the total number of trainable parameters, the average number of epochs, training time, and prediction time. It is interesting to look at the training and prediction times. Indeed, while for the coupling flows, even though training time is much larger than prediction time, both times grow with a similar rate; for the autoregressive flows, the prediction time grows faster than does the training time, which is almost constant. This is because, as we have already mentioned, the MAF is a “direct flow”, very fast for density estimation and therefore for training (single pass through the flow) but slower for generation and therefore for testing (N passes through the flow, with N the dimensionality of the target distribution). Still, testing was reasonably fast, considering that each test actually consisted of 10 tests with three metrics and ${10}^5$ points per sample. All trainings/tests took less than a few hours (sometimes, especially in small dimensionality, a few minutes), which means that all models, expect the C-RQS in large dimensionalities, are fairly fast both in training and inference (Notice that even though the training/testing times do not go beyond a few hours, we trained and tested 10 replicas of four architectures in eight different dimensionalities (apart from C-RQS) and with a few different values of the hyperparameters for a total of about 1360 runs (see

Table A1. Hyperparameter values used in our analysis. The last row shows the total number of runs for each architecture, with the 10 replicas and the different dimensionalities being taken into account.

Hyperparameter Values
Hyperpar.	MAF	RealNVP	A-RQS	C-RQS
Number of	$5,10$	$5,10$	2	$5,10$
bijectors
Number of	$3\times 128$	$3\times 128$	$3\times 128$	$3\times 128$
hidden	$3\times 256$	$3\times 256$	$3\times 256$	$3\times 256$
layers
Number of	–	–	$8,12$	$8,12$
spline knots
Total number	320	320	320	400
of runs

). This took several months of GPU time, showing how resource demanding is to reliably estimate uncertainties of ML models, even in relatively simple cases.) Another interesting number in the tables is the total number of trainable parameters. Such number makes clear how the autoregressive architectures are more expressive than are the simple coupling ones, giving better results with relatively fewer parameters. It is also clear from the table that the improvement stepping from a simple linear affine bijector to a rational quadratic spline based on the same architecture is much larger for autoregressive architectures and less evident for the simplest coupling ones. The large number of parameters needed to obtain reasonable results from C-RQS may be the origin of its training instability at large dimensionality.

6. Conclusions and Outlook

Normalizing flows have shown many potential applications in a wide variety of research fields including HEP, both in their normalizing and generative directions. However, to ensure a standardized usage and to match the required precision, their application to high-dimensional datasets need to be properly evaluated. This paper makes a step forward in this direction by quantifying the ability of coupling and autoregressive flow architectures to model distributions of increasing complexity and dimensionality.

Our strategy consisted in performing statistically robust tests utilizing different figures of merits and including estimates of the variances induced both by the training and the evaluation procedures.

We focused on the most widely used NF architectures in HEP, the coupling (RealNVP and C-RQS) and the autoregressive flows (MAF and A-RQS), and compared them against generic multimodal distributions of increasing dimensionality.

As the main highlight, we found that the A-RQS is highly capable of precisely learning all the high-dimensional complicated distributions it was tested against, always within a few hours of training on a Tesla A40 GPU and with limited training data. Moreover, the A-RQS architecture, showed great generalization capabilities, obtaining almost constant results over a very wide range of dimensionalities, ranging from 4 to 400.

As for the other tested architectures, our results show that reasonably good results can be obtained with all of them but the C-RQS, which ended up being the least capable in generalizing to large dimensionality, with unstable and longer trainings, especially in high dimensionality.

Our analysis was performed implementing all architectures in TensorFlow2 with TensorFlow Probability using Python. The code is available in Ref. [72], while a general-purpose user-friendly framework for NFs in TensorFlow2 named NF4HEP is under development and can be found in Ref. [73]. Finally, a code for statistical inference and two-sample tests, implementing the metrics considered in this paper (and others) in TensorFlow2, is available in Ref. [74].

We stress that the intention of this study was to secure generic assessments of how NFs perform in high dimensions. For this reason, the target distributions were chosen independently of any particular experimentally driven physics dataset. An example of application to a physics dataset, in the direction of building an unsupervised DNNLikelihood [75], is presented in Ref. [76]. Nonetheless, these studies represent the firsts of a series to come. Let us briefly mention, in turn, the research directions we aim to follow starting from the present paper.

• Development of reliable multivariate quality metrics, including approaches based on machine learning [77,78]. We note the importance of performing statistically meaningful tests on generative models, ideally including uncertainty estimation. A thorough study of different quality metrics against high dimensional data is on its way. Moreover, new results [79,80,81] suggest that classifier-based two-sample tests have the potential to match the needs of the HEP community when paired with a careful statistical analysis. These tests can leverage different ML models to provide high flexibility and sensitivity together with short training times, especially when based on kernel methods [81]. On the other hand, further studies are needed to investigate their efficiency and scalability to high dimensions.
• Study of the dependence of the NF performances on the size of the training sample [82]. In the present paper, we always kept the number of training points to ${10}^5$. It is clear that such a number is fairly large in small dimensionality, such as $N=4$ dimensions, and undersized for large dimensionality, such as $N\ge 100$. It is important to study the performances of the considered NF architectures in the case of scarce or very abundant data and to assess the dependence of the final precision on the number of training samples. This can also be related to developing techniques to infer the uncertainty of the NF models.
• Studies on how NFs can be used for statistical augmentation. For instance, NFs can be used for normalizing direction to build effective priors and proposals to enhance (in terms of speed and convergence time) known sampling techniques, such as Markov chain Monte Carlo, whose statistical properties are well established.
• The ability to preserve and distribute pretrained NF-based models. A final issue that needs to be addressed to ensure a widespread use of NFs in HEP is the ability to preserve and distribute pretrained NF-based models. This is, for the time being, not an easy and standard task, and support from the relevant software developers in the community is crucial to achieving this goal.