Convex Potential Flows: Universal Probability Distributions with Optimal Transport and Convex Optimization
Authors: Chin-Wei Huang
Ricky T. Q. Chen
Christos Tsirigotis
Aaron Courville
Venue: ICLR 2021 Poster (2021)
Keywords: Normalizing flows
generative models
variational inference
invertible neural networks
universal approximation
optimal transport
convex optimization

Abstract: Flow-based models are powerful tools for designing probabilistic models with tractable density. This paper introduces Convex Potential Flows (CP-Flow), a natural and efficient parameterization of invertible models inspired by the optimal transport (OT) theory. CP-Flows are the gradient map of a strongly convex neural potential function. The convexity implies invertibility and allows us to resort to convex optimization to solve the convex conjugate for efficient inversion. To enable maximum likelihood training, we derive a new gradient estimator of the log-determinant of the Jacobian, which involves solving an inverse-Hessian vector product using the conjugate gradient method. The gradient estimator has constant-memory cost, and can be made effectively unbiased by reducing the error tolerance level of the convex optimization routine. Theoretically, we prove that CP-Flows are universal density approximators and are optimal in the OT sense. Our empirical results show that CP-Flow performs competitively on standard benchmarks of density estimation and variational inference.

Published as a conference paper at ICLR 2021
CONVEX POTENTIAL FLOWS:
UNIVERSAL PROBABILITY DISTRIBUTIONS WITH OPTIMAL TRANSPORT AND CONVEX OPTIMIZATION
Chin-Wei Huang
Ricky T. Q. Chen
University of Montreal & Mila chin-wei. @ukang@umonteede
University of Toronto & Vector Institute xtqichen@cs. toronto. edu
Christos Tsirigotis
Aaron Courville
University of Montreal & Mila
University of Montreal, Mila CIFAR Fellow aaron. Boros.courviiieeeontttee.
socistos.tsirigon rea. ca
ABSTRACT
Flow-based models are powerful tools for designing probabilistic models with tractable density. This paper introduces Convex Potential Flows (CP-Flow), natural and efficient parameterization of invertible models inspired by the opti- mal transport (OT) theory. CP-Flows are the gradient map of strongly convex neural potential function. The convexity implies invertibility and allows us to re- sort convex optimization solve the convex conjugate for efficient inversion. To enable maximum likelihood training, we derive a new gradient estimator of the log- determminnnn of the Jacobian, which involves solving an inverse- Hessian vector product using the conjugate gradient method. The gradient estimator has constant- memory cost, and can made effectively unbiased by reducing the error tolerance level of the convex optimization routine. Theoretically, we prove that CP-Flows are universal density approximators and are optimal in the OT sense. Our empiri- cal results show that CP-Flow performs competitively standard benchmarks of density estimation and variational inference.
INTRODUCTION
Normalizing flows (Dinh et al., 2014; Rezende & Mohamed, 2015) have recently gathered much interest within the machine learning community, ever since its recent breakthrough in modelling high dimensional image data (Dinh et al., 2017; Kingma Dhariwal, 2018). They are characterized by invertible mapping that reshape the distribution of its input data into simpler more complex one. To enable efficient training, numerous tricks have been proposed to impose structural constraints its parameterization, such that the density of the model can be tractably computed.
We ask the following question: "what is the natural way to parameterize a normalizing flow?" To gain a bit more intuition, we start from the one-dimension case. If function f R R is continuous, it is invertible (injective onto its image) if and only if it is strictly monotonic. This means that if are only allowed move the probability mass continuously without flipping the order of the particles, then we can only rearrange them by changing the distance in between.
In this work, we seek to generalize the above intuition of monotone rearrangement in ID. We do so by motivating the parameterization of normalizing flows from an optimal transport perspective, which allows us to define some notion of rearrangement cost (Villani, 2008). It turns out, if we want the output of flow to follow some desired distribution, under mild regularity conditions, we can characterize the unique optimal mapping by a convex potential (Brenier, 1991). In light of this, propose to parameterize normalizing flows by the gradient map of a (strongly) convex potential. Owing to this theoretical insight, the proposed method is provably universal and optimal; this means the proposed flow family can approximate arbitrary distributions and requires the least amount of transport cost. Furthermore, the parameterization with convex potentials allows us to formulate model inversion and gradient estimation convex optimization problems. As such,

---

Published as a conference paper at ICLR 2021
make use of existing tools from the convex optimization literature to cheaply and efficiently estimate all quantities of interest.
In terms of the benefits of parameterizing flow gradient field, the convex potential is an R + R function, which is different from most existing discrete-time flows which are I + R. This makes CP-Flow relatively compact. It is also arguably easier design convex architecture, as we do not need to satisfy constraints such as orthogonality or Lipschitzness; the latter two usually require a direct or an iterative reparameterization of the parameters. Finally, it is possible incor- porate additional structure such as equivariance (Cohen & Welling, 2016; Zaheer et al., 2017) into the flow's parameterization, making CP-Flow more flexible general purpose density model.
2 BACKGROUND: NORMALIZING FLOWS AND OPTIMAL TRANSPORT
Normalizing flows are characterized by a differentiable, invertible neural network such that the probability density of the network's output can be computed conveniently using the change-of- variable formula
0f(x) dx
05 (v)
PY (f(x)) = px
(y Px (x))
(1)
By
where the Jacobian determinant term captures the local expansion or contraction of the density near (resp. y) induced by the mapping f (resp. f) and px is the density of a random variable X. The invertibility requirement has led to the design of many special neural network parameterizations such triangular maps, ordinary differential equations, orthogonality or Lipschitz constraints.
Universal Flows For general learning framework be meaningful, model needs to be flexible enough to capture variations in the data distribution. In the context of density modeling, this cor- responds to the model's capability represent arbitrary probability distributions of interest. Even though there exists a long history of literature on universal approximation capability of deep neural networks (Cybenko, 1989: Lu al., 2017; Lin & Jegelka, 2018), invertible neural networks gener- ally have limited expressivity and cannot approximate arbitrary functions. However, for the purpose of approximating a probability distribution, it suffices to show that the distribution induced by normalizing flow is universal.
Among many ways to establish distributional universality of flow based methods (e.g. Huang et al. 2018; 2020b; Teshima et al. 2020; Kong & Chaudhuri 2020), one particular approach is to approx- imate a deterministic coupling between probability measures. Given a pair of probability densities px and py. deterministic coupling is mapping y such that g(X) py if X ~ px. We seek to find a coupling that is invertible, or at least can be approximated by invertible mappings.
Optimal Transport Let c(z, y) be a cost function. The Monge problem (Villani, 2008) pertains finding the optimal transport map g that realizes the minimal expected cost
inf Ex-px [c(X,ū(X))] Jo(xx, "=((X)>
(2)
When the second moments of X and Y are both finite, and X is regular enough (e.g. having density), then the special case of c(x, y) = || - w||² has an interesting solution, a celebrated theorem due Brenier (1987; 1991):
Theorem (Brenier's Theorem. Theorem 1.22 of Santambrogio (2015)). Let p., v be probability measures with finite second moment, and assume μ has a Lebesgue density px- Then there exists a convex potential G such that the gradient map g == VG (defined up to null set) uniquely solves the Monge problem in eq. (2) with the quadratic cost function c(z, y) || 2
Some recent works are also inspired by Brenier's theorem and utilize a convex potential to param- eterize a critic model, starting from Taghvaei & Jalali (2019), and further built upon by Makkuva et al. (2019) who parameterize generator with convex potential and concurrently by Korotin et al. (2019). Our work sets itself apart from these prior works that it is entirely Ikelihood-based minimizing the (empirical) KI. divergence as opposed to approximate optimal transport cost.
2

---

Published as a conference paper at ICLR 2021
(a)
(b)
(c)
(d)
Figure 1: Illustration of Convex Potential Flow. (a) Data drawn from mixture of Gaussians. (b) Learned
convex potential (c) Mesh grid distorted by the gradient map the convex potential f VF. (d) Encoding
of the data via the gradient map = fx)) Notably, the encoding is the value of the gradient of the convex potential. When the curvature of the potential function is locally flat, gradient values are small and this results contraction towards the origin.
CONVEX POTENTIAL FLOWS
Given strictly convex potential F, we can define an injective map (invertible from its image) via its gradient f = F since the Jacobian of f is the Hessian matrix of F, and is thus positive definite. In this section, we discuss the parameterization of the convex potential F (3.1), and then address gradient estimation for CP-Flows (3.2). We examine the connection to other parameterization of normalizing flows (3.3), and finally rigorously prove universality in the next section.
3.1 MODELING
Input Convex Neural Networks We use L(x) to denote linear layer, and L (x) to denote linear layer with positive weights. We use the (fully) input-convex neural network (ICNN, Amos al. (2017)) parameterize the convex potential, which has the following form
F(z) L (s(=k)) + Lx-1(x)
24 _ (s(=+11)) + Lx(x)
21 = L1
where s is non-decra,,,ng, convex activation function. In this work, we use softplus-type activation functions, which rich family of activation functions that can be shown uniformly approximate the ReLU activation. See Appendix B for details.
Invertibility and Inversion Procedure If the ac-
Algorithm 1 Inverting CP-Flow.
tivation is twice differentiable, then the Hessian HF is positive semi-definite. We can make strongly convex by adding quadratic term F, (x) g |1 + F(x), such that HF. - al > O. This means the gradient map fa V.. is injective onto its image. Furthermore. it is surjective since for any y E Rd, the potential æ + (x) y has a unique minimizer' satisfying the first order condi- tion (x) = y. due to the strong convexity and differentiability. We refer to this invertible mapping fa as the convex potential flow, the CP-Flow. The above discussion also implies we can plug in a black-box convex solver to invert the gradient map fa. which summarize in Algorithm 1. Inverting batch of independent inputs is simple as summing the convex potential over all inputs: since all of the entries of the scalar in the mini- batch are independent of each other, computing the gradient all I's wrt all x's amounts computing the gradient of the summation of I's wrt all r's. Due to the convex nature of the prohlem, wide selection of algorithms can be used with convergence guarantees (Nesterov, 1998). In practice, use the L-BFGS algorithm (Byrd et al., 1995) as our CvxSolver.
1: procedure INVERT(F, CvxSolver)
Initialize x < y
3:
def closure(:
4:
Compute loss: 1 t F(x) yT return
6: x <- ccxxolver(ccosuree.
return
'The minimizer xt corresponds to the gradient map of the conver conjugate of the potential. See Appendix for formal discussion.
3

---

Published as a conference paper at ICLR 2021
Estimating Log Probability Following equation (1), computing the log density for CP-Flows requires taking the log determinant of a symmetric positive definite Jacobian matrix (as it is the Hessian of the potential). There exists numerous works estimating spectral densities (e.g. Tal- Ezer & Kosloff, 1984; Silver & Röder, 1994; Han et al., 2018; Adams et al., 2018), of which this quantity is special case. See Lin et al. (2016) for overview of methods that only require access to Hessian-vecee vector products. Hessian- vector products (hvp) are cheap to compute with reverse-mode automatic differentiation (Baydin al., 2017), which does not require constructing the full Hessian matrix and has the same asymptotic cost as evaluating Fa.
In particular, the log determinant can be rewritten the form of a generalized trace tr log H. Chen et al. (2019a) limit the spectral norm (i.e. eigenvalues) of H and directly use the Taylor expansion of the matrix logarithm. Since our has unbounded eigenvalues, we use more complex algorithm designed for symmetric matrices, the stochastic Lanczos quadrature (SLQ; Ubaru et al., 2017). At the core of SLQ is the Lanezos method, which computes m eigenvalues of H by first constructing symmetric tridiagonal matrix T E _= and computing the eigenvalues of T. The Lanczos procedure only requires Hessian-vector products, and it can be combined with a stochastic trace estimator to provide stochastic estimate of our log probability. We chose SLQ because it has shown theoretically and empirically to have low variance (Ubaru et al., 2017).
3.2 (1)-MORYYY UNBIASED V Vlogder H ESTIMATOR
We would also like to have an estimator for the gradi- ent of the log determinant to enable variants of stochastic gradient descent for optimization. Unfortunately, directly backpropagating through the log determinant estimator is not ideal. Two major drawbacks of directly differenti- ating through SLQ are that it requires (i) diffcrentiating through an eigendecomposition routine and (ii) storing all Hessian-vecte products in memory (see fig. 2). Problem (i) is more specific to SLQ. because the gradient of an eigendecompoosttin is not defined when the eigenvalues are not unique (Seeger et al., 2017). Consequently, we have empirically observed that differentiating through SLQ can be unstable, frequently resulting NaNs due to the eigendecomposition. Problem (ii) will hold true for other algorithms that also estimate log det H with Hessian deestan-vvector products, and generally the only difference is that different numerical routine would need to be differenti- ated through. Due to these problems, we do not differentiate through SLQ, but we still use it as an efficient method for monitoring training progress.
Memory Usage (GB)
81.6
80
SLQ (m=50)
SLQ (m=20)
SLQ (m=10)
CG Backprop
32.2
20
11.3
3.2
Figure 2: Memory for training CIFAR- I0.
Instead, it is possible to construct an alternative formulation of the gradient as the solution of a con- vex optimization problem, foregoing the necessity of differentiating through an estimation routine of the log determinant. We adapt the gradient formula from Chen et al. (2019a, Appendix C) to the context of convex potentials. Using Jacobi's s formula" and the adjugate representation of the matrix inversel, for any invertible matrix H with parameter f, we have the following identity:
8 @ log det H det H au det H 20
det IT tr (adj(H)
0H i 1h
[
8H 80
E.
8H 80
Notably, in the last equality, used the Hutchinson trace estimator (Hutchinson, 1989) with Rademacher random vector v, leading to a (0(1)-memory unbiased Monte Carlo gradient estimator.
Computing the quantity u - in eq. (3) by constructing and inverting the full Hessian requires d calls to an automatic differentiation routine and is too costly for our purposes. However, we can recast this quantity as the solution of quadratic optimization problem
arg min
- 2
H - v 2
(4)
which has the unique minimizer 2" - since H is symmetric positive definite.
4

---

Published as a conference paper at ICLR 2021
We use the conjugate gradient (CG) method, which is specifically designed for solving the unconstrained optimization problems in eq. (4) with symmetric positive definite H. It uses only Hessian-vector products and is straightforward to parallelize. Conjugate gradient is guaranteed to return the exact solution 2* within d itera- tions, and the error of the approximation is known to converge exponentially fast ||~TH 2=I1æ I 2|= = II where z" is the estimate after m iterations. The rate of convergence 7 < relates to the condition number of H. For more details, see Nocedal & Wright (2006, Ch. 5). In practice, we terminate CG when || - vlloc < is satisfied for some user-controlled tolerance. Empirically, we find that stringent tolerance values are unnecessary for stochastic optimization (see appendix F).
Algorithm 2 Surrogate training objective.
1: procedure SURROGATEOOICCF CG)
2: Obtain the gradient f(x) å Fa)())
3: Sample Rademacher random vector
4: def hvp(e):
5:
return v" & f(x)
6:
2 2 stop gradien int (CG(hvp,
7: return hvp(z)
Estimating the full quantity in eq. (3) is then simply a matter of computing and differentiating a scalar quantity (a surrogate objective) involving another Hessian-vector product: d 49 ((=*)) Ho), where only H is differentiated through (since 2Th only used to approximate v H as a modifier of the gradient). We summarize this procedure in Algorithm 2. Similar to inversion, the hvp can also be computed in batch by summing over the data index, since all entries are independent.
3.3 CONNECTION TO OTHER NORMALIZING FLOWS
Residual Flow For a = 1, the gradient map fi resembles the residual flow (Behrmann et al., 2019; Chen et al., 2019a). They require the residual block equivalent to our gradient map f f-to be contractive (with Lipschitz constant strictly smaller than 1) as sufficient condition for invertibility. In contrast, we enforce invertibility by using strongly convex potentials, which guarantees that the inverse of our flow is globally unique. With this, we do not pay the extra compute cost for having to satisfy Lipschitz constraints using methods such as spectral normalization (Miyato et al., 2018). Our gradient estimator is also derived similarly to that of Chen et al. (2019a), though we have the benefit of using well-studied convex optimization algorithms for computing the gradients.
Sylvester Flow By restricting the architecture of our ICNN to one hidden layer, we can also re- cover a form similar to Sylvester Flows. For a 1-hidden layer ICNN (K = 1) and a = 1, we have Fi 1 | + L (s(L)z) 21)) + L2(x). Setting the weights of Le to zero, we have
fz Fi = x W diag(w3 )s (Wix+++ b2).
(5)
We notice the above form bears a close resemblance to the Sylvester normalizing flow (Van Den Berg et al., 2018) (with Q R and R from Van Den Berg et al. (2018) 8) being equal to WTT. diag(w;) ) and I, respectively). For the Sylvester flow be invertible, they require that R and R be triangular and Q be orthogonal, which is a computationally costly procedure. This orthogonality constraint also implies that the number of hidden units cannot exceed d. This restriction orthogonal matrices and one hidden layer are for applying Sylvester's 's determinant identity. In contrast, we do not require our weight matrices to be orthogonal, and we can use any hidden width and depth for the ICNN.
Sigmoidal Flow
Let s be the softplus activation function and = s'. Then for the 1-dimensional case (d = 1) and a = 0 (without the residual connection). we have
Ix Po(x) = E wi j=1
= Σ |wi, w (Iwi, sign ( const. (6)
which is equivalent to the sigmoidal flow of Huang et al. (2018) up to rescaling (since the weighted sum is no longer a convex sum) and a constant shift, and is monotone due to the positive weights. This correspondence is not surprising since a differentiable function is convex if and only if its derivative is monotonically non-decreasing. It also means we can parameterize an increasing func- tion as the derivative of convex function, which opens up autoregressive normalizing flows (Kingma et al., 2016; Huang et al., 2018; Müller et al., 2019; Jaini et al., 2019; Durkan et al., 2019; Wehenkel & Louppe, 2019).
direction for parameterizing
5

---

Published as a conference paper at ICLR 2021
Flows with Potential Parameterization Inspired by connections between optimal transport and continuous normalizing flows, some works (Zhang et al., 2018; Finlay et al., 2020a; Onken et al., 2020) have proposed to parameterize continuous-time transformations by taking the gradient of a scalar potential. They do not strictly require the potential to be convex since it is guaranteed to be invertible in the infinitesimal setting of continuous normalizing flows (Chen al., 2018). There exist works (Yang & Karniadakis, 2019; Finlay et al., 2020b; Onken et al., 2020) that have applied the theory of optimal transport to regularize continuous-time flows have low transport cost. In contrast, we connect optimal transport with discrete-timg time normalizing flows, and CP-Flow guaranteed by construction to converge pointwise the optimal mapping between distributions without explicit regularization (see Section 4).
4 THEORETICAL ANALYSES
As explained in Section 2, the parameterization of CP-Flow is inspired by the Brenier potential. So naturally we would hope show that (1) CP-Flows are distributionally universal, and that (2) the learned invertible map is optimal the sense of the average squared distance the input travels E|||=- - [=)]] Proofs of statements made in this section can be found Appendices and D.
To show (1), our first step is to show that ICNNs can approximate arbitrary convex functions. However, convergence of potential functions does not generally imply convergence of the gradi- ent fields. A classic example is the sequence F, sin(nx)/ / and the corresponding derivatives fn = cos(nx) F, - as n + 3o but fn does not. Fortunately, convexity allows us to control the variation of the gradient map (since the derivative of a convex function is monotone), our second step of approximation holds.
R be a proper convex convex function. Assume Fn + G. Then for almost every x E TR'", G is differentiable and fn(x) VF, - g(r).
Theorem 2. Let Fn Rd + R be differentiable convex functions and G : -
Combining these two steps and Brenier's theorem, we show that CP-Flow with softplus-type type activa- tion function is distributionally universal.
Theorem (Universality). Given random variables X μ and Y V, with μ being absolutely continuous w.r.. the Lebesgue measure, there exists a sequence of ICNN Fn with a softplus-tye type activation, such that V F2 Y distribution.
N.B. In the theorem we do not require the second moment to be finite, as for arbitrary random variables we can apply the standard truncation technique and redistribute the probability mass so that the new random variables are almost surely bounded. For probability measures with finite second moments, we indeed use the gradient map of ICNN to approximate the optimal transport map corresponding to the Brenier potential. In the following theorem, we show that the optimal transport map is the only such mapping that we can approximate we match the distributions.
Theorem 4 (Optimality). Let G be the Brenier potential of X P and Y ~ and let Fa be convergent sequence of differentiable, convex potentials, such that F, o X - Y in distribution. Then F,, converges almost surely to VG.
The theorem states that in practice, even if we optimize according to some loss that traces the convergence in distribution, our model is still able to recover the optimal transport map, as if were optimizing according to the transport cost. This allows us to estimate optimal transport maps without solving the constrained optimization in (2). See Seguy al. (2018) for some potential applications of the optimal transport map, such as domain adaptation or domain translation.
5 EXPERIMENT
We use CP-Flow to perform density estimation (RHS (1)) and variational inference (LHS of (1)) to assess its approximation capability, and the effectiveness of the proposed gradient estimator. All the details of experiments can be found in Appendix E. Code is available at https///ittub.com/CCV Huang/CP-Flow.

---

achieve high Published as a conference paper at ICLR 2021
data: x
feaf(2)
fcp(1)
16
Figure 4: Approximating optimal transport map via maximum likelihood (minimizing KL divergenee). In the first figure the left show the data in dimensions. The datapoints are colored according to their horizontal values (#1). The flows Sia and Sap are trained to transform the data into a standard Gaussian prior. figures on the right, we plot the expected quadratic transportation cost versus the KL divergence for different numbers of dimensionality. During training the KL is minimized, so the curves read from the right to the left.
Despite the universal property, having a poor parameterization can lead to difficulties in optimization and limit the effective expressivity of the model. We propose archi- tectural enhancement of ICNN, defined as follows (note the change notation: instead of writing the pre-activations 2, we use h to denote the activated units):
ICNN Architecture
Foug (x) L (kx) + (x+1
h4 == concat (lhg, 1)
h4 := s(L (hx-1) + Lx(r)) 1 1 s(Lx aug (I))
(7)
where half of the hidden units are directly connected to the input, so the gradient would have some form of skip connection. We call this the input-augmented ICNN. Unless otherwise stated, the input-augmented ICNN as the default architecture.
5.1 TOY EXAMPLES
Having distributional universality for single flow layer means that expressiveness without composing too many flows. We demonstrate this by fitting the den- sity on some toy examples taken from Papa- makarios et al. (2017) and Behrmann al. (2019). We compare with the masked au- toregressive flow (MAF, Papamakarios (2017)) and the neural autoregressive flow (NAF, (Huang al., 2018)). Results are pre- sented in fig. 3. We try to match the network size for each data. All models fit the first data well. As affine couplings cannot split probabil- ity mass, MAF fails to fit to the second and third d Although the last dataset is intrinsi- cally harder to fit (as NAF, another universal density model, also fails to fit it well), the proposed method still manages to learn the correct density with high fidelity.
Data
MAF
NAF
CP-Flow
Figure 3: Learning toy densities.
5.2 APPROXIMATING OPTIMAL COUPLING
As predicted by Theorem 4, CP-Flow is guaranteed converge to the optimal coupling minimizing the expected quadratic cost. We empirically verify it by learning the Gaussian density and com- paring the expected quadratic distance between the input and output of the flow against J between the Gaussian data and the standard Gaussian prior (as there is closed-form expression). In fig. 4, we see that the transport cost gets closer to the optimal value when the learned density
Behhraan et al. (2019) demonstrates many flow layers. But here we restrict the number of flow layers to he 3 5.
potentially improve the affine coupling models by composing

---

Published as a conference paper at ICLR 2021
Model
POWER
GAS
HEPMASS
MINIBOONE
BSDS300
Real NVP (Dinh et al., 2017)
-0.17
-8.33
18.71
13.55
-153.28
FFJORD (Grathwohl al., 2018)
-0.46
-8.59
14.92
10.43
-15..40
MADE
(Germain et al., 2015)
3.08
-3.56
20.98
15.59
-148.85
MAF (Papamakarios et al., 2017)
-0.24
-10.08
17.70
11.75
-155.69
TAN (Oliva et al., 2018)
-0.48
-11.19
15.12
11.01
-157.03
NAF (Huang et al., 2018)
-0.62
-11.96
15.09
8.86
-157.73
CP-Flow (Ours)
-0.52
-10.36
16.93
10.58
-154.99
Table I:
Average test negative log-likelihood (in nats) of tabular datasets in Papamakarios et al. (2017) for density estimation models (lower is better). Standard deviation is presented in the appendix E.4.
MNIST
CIFAR-10
Model
Bits/dim
N. params
Bits/dim
N. params
Real NVP (Dinh et al., 2017)
1.05
N/A
3,49
N/A
Glow (Kingma & Dhariwal, 2018)
N/A
3.35
44.0M
RQ-NSF (Durkan et al., 2019)
3.38
11.8
Residual Flow (Chen et al., 2019a)
0.97
16.6M
3.28
25.2M
Coupling Block Ablation
1.02
3.1M
3,58
2.9M
Residual Block Ablation
2.9M
3,46
3.1M
CP-Flow (Ours)
1.02
2.9M
3,40
1.9M
Table 2: Negative oogllikelhhod (in bits) on held-out test data (lower is better). (2019). Obtained from official open source code.
"Taken from Durkan al.
approaches the data distribution (measured by the KL divergence). We compare against the linear inverse autoregressive flow (Kingma al., 2016), which has the capacity represent the multivari- ate Gaussian density, yet it does not learn the optimal coupling.
5.3 DENSITY ESTIMATION
We demonstrate the efficacy of our model and the proposed gradient estimator by performing density estimation on the standard benchmarks.
Tabular Data We the datasets preprocessed by Papamakarios et al. (2017). In table 1, report average negative log-likelihood estimates evaluated on held-out test sets, for the best hyper- parameters found via grid search. The search was focused on the number of flow blocks, the width and depth of the ICNN potentials. See appendix E.4 for details. Our models perform competi- tively against alternative approaches the literature. We also perform an ablation on the CG error tolerance and ICNN architectures in appendix F.
Image Data Next, we apply CP-Flow to model the density of standard image datasets, MNIST and CIFAR-I0. For this, use convolutional layers in place of fully connected layers. Prior works have had to use large architectures, with many flow blocks composed together, resulting in large number of parameters to optimize. While we also compose multiple blocks of CP-Flows, we find that CP-Flow can perform relatively well with fewer number of parameters (table 2). Notably, we achieve comparable bits per dimension to Neural Spline Flows (Durkan et al., 2019) work promoting fewer parameters-whin having using around 16% number of parameters.
As prior works use different architectures with widely varying hyperparameters, we perform a more careful ablation study using coupling (Dinh al., 2014; 2017) and invertible residual blocks (Chen et al.. 2019a). We replace each of our flow blocks with the corresponding baseline. We find that on CIFAR-10, 10, the baseline flow models do not perform nearly well as CP-Flow. We believe this may be because CP-Flows are universal with just one flow block, whereas coupling and invertible residual blocks limited in expressivity Lipschitz-onnsraiine..
8

---

Published as a conference paper at ICLR 2021
5.4 AMORTIZING ICNN FOR VARIATIONAL INFERENCE
Normalizing flows also allow us to employ a larger,
more flexible family of distributions for variational inference (Rezende & MMhaame, 2015). We repli- cate the experiment conducted in Van Den Berg et al. (2018) to enhance the variational autoencoder (Kingma Welling, 2013). For inference amor- tization, we use the partially input convex neural network from Amos et al. (2017), and use the out- of the encoder as the additional input for condi- tioning. As table 3 shows, the performance of CP- Flow is close to the best reported in Van Den Berg et al. (2018) without changing the experiment setup. This shows that the convex potential parameteriza- tion along with the proposed gradient estimator can learn to perform accurate amortized inference.
Gaussian
4.53
104.28
110.80
Planar
4.40
102.65
109.66
IAF Sylvester
4.47
102.41
111.58
4.45
99.00
104.62
CP-Flow (vanilla)
4.47
102.06
106.53
CP-Flow (aug)
4.45
100.82
105.17
Table 3: Negative ELBO of VAE (lower is better). Standard deviation reported in appendix E.6.
Also, we show that replacing the vanilla ICNN with the input-augmentedd ICNN leads improve- ment of the likelihood estimates.
6 CONCLUSION
We propose a new parameterization of normalizing flows using the gradient map of a convex po- tential. We make connections to the optimal transport theory to show that the proposed flow is universal density model, and leverage tools from convex optmiization to enable efficient training and model inversion. Experimentally, we show that the proposed method works reasonably well when evaluated on standard benchmarks.
Furthermore, we demonstrate that the performance can improved by designing better ICNN ar- chitectures. We leave the exploration for better ICNN and convolutional ICNN architecture to improve density estimation and generative modeling for future research.
ACKNOWLEDCEMENTS
We would like to acknowledge the Python community (Van Rossum Drake Jr, 1995; Oliphant, 2007) for developing the tools that enabled this work, including numpy (Oliphant, 2006; Van Der Walt et al., 2011; Walt et al.. 2011; Harris et al., 2020), PyTorch (Paszke et al., 2019), Mat- plotlib (Hunter, 2007), seaborn (Waskom al., 2018), pandas (McKinney, 2012), and SciPy (Jones et al., 2014).
REFERENCES
Ryan P Adams, Jeffrey Pennington, Matthew Johnson, Jamie Smith, Yaniv Ovadia, Brian Patton, and James Saunderson. Estimating the spectral density of large implicit matrices. arXiv preprint arXiv:1802.0345)) 2018.
Brandon Amos, Lei Xu, and Zico Kolter. Input convex neural networks. In International Confer- ence on Machine Learning, pp. 146-155, 2017.
Atilim Günes Baydin, Barak Pearlmutter, Alexey Andreyevich Radul, and Jeffrey Mark Siskind. Automatic differentiation in machine learning: survey. The Journal of Machine Learning Re- search, (811)5555--6337,2011
Jens Behrmann, Will Grathwohl, Ricky T. Q. Chen, David Duvenaud, and Hörn-ke Jacobsen. Invertible residual networks. International Conference on Machine Learning, pp. 573-582, 2019.
Jens Behrmann, Paul Vicol, Kuan-Chieh Wang, Roger Grosse, and Jöm-Henrik Jacobsen. Un- derstanding and mitigating exploding inverses in invertible neural networks. arXiv preprint a7X1v:2006.003377. 2020.

---

Published as a conference paper at ICLR 2021
Yann Brenier. Décomposition polaire et réarrangement monotone des champs de vecteurs. CR Acad. Sci. Paris Sér. I Math., 305:805- 005:805-808, 1987.
Yann Brenier. Polar factorization and monotone rearrangement of vector-valued functions. Commu- nications on pure and applied mathematics,, 44(4):375 1991.
Richard H Byrd, Peihuang Lu, Jorge Nocedal, and Ciyou Zhu. A limited memory algorithm for bound constrained optimization. SIAM Journal on scientific computing, 16(5):1190-1208, 1995.
Ricky Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. In Advances in neural information processing systems, pp. 6571-6583, 2018.
Ricky T. Q. Chen, Jens Behrmann, David Duvenaud, and Jörn -orn-Henrik Jacobsen. Residual flows for invertible generative modeling. In Advances in Neural Information Processing Systems, pp- 1916,9926, 2019a.
Yize Chen, Yuanyuan Shi, and Baosen Zhang. Optimal control via neural networks: A convex approach. In International Conference on Learning Representations, 2019b.
Mircea Cimpoi, Subhransu Maji, lasonas Kokkinos, Sammy Mohamed, and Andrea Vedaldi. De- scribing textures in the wild. In Conference on Computer Vision and Pattern Recognition, 2014.
Taco Cohen and Max Welling. Group equivariant convolutional networks. In International confer-
ence on machine learming. pp. 2990-2990 2016.
George Cybenko. Approximation by superpositions of sigmoidal function. Mathematics of control, signals and systems, 2(4):303,314, 1989.
Laurent Dinh, David Krueger, and Yoshua Bengio. Nice: Non-linear independent components esti- mation. arXiv preprint arXiv:1410.8516, 2014.
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real nvp. In International Conference Learning Representations, 2017.
Vincent Dumoulin, Ethan Perez, Nathan Schucher, Florian Strub, Harm de Vries, Aaron Courville, and Yoshua Bengio. Feature-wise transformations. Distill, 3(7)el11, 2018.
Conor Durkan, Artur Bekasov, lain Murray, and George Papamakarios. Neural spline flows. In Advances in Neural Information Processing Systems, pp. 7511-7522, 2019.
Chris Finlay, Augusto Gerolin, Adam Oberman, and Aram-Alexandre Pooladian. Learning nor- malizing flows from entropy-kantorovice potentials. arXiv preprint arXiv:2006.0660 2020a.
Chris Finlay, Hennik Jacobsen, Levon Nurbekyan, and Adam Oberman. How train your neural ode: the world of jacobian and kinetic regularization. In International Conference on Machine Learning, 2020b.
Mathieu Germain, Karol Gregor, Iain Murray, and Hugo Larochelle. Made: Masked autoencoder for distribution estimation. In Proceedings of the 32nd International Conference Machine Learning, volume 37 of Proceedings of Machine Learning Research, pp. 881-889, PMLR, 07-09 Jul 2015.
Ian Goodfellow, David Warde-Farley, Mehdi Mirza, Aaron Courville, and Yoshua Bengio. Maxout networks. In International conference on machine learning, pp. 1319-1327. PMLR, 2013.
Will Grathwohl, Ricky T. Q. Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud. Ffjord: Free-form continuous dynamics for scalable reversible generative models. In International Con- ference on Learning Representations, 2018.
David Ha, Andrew Dai, and Quoe V Le. Hypernetworks. arXiv preprint arXiv: 0609.090900......... 2016.
Insu Han, Haim Avron, and Jinwoo Shin. Stochastic chebyshev gradient descent for spectral opti- mization. In Advances in Neural Information Processing Systems, pp. 7386-7396, 2018.
10

---

Published as a conference paper at ICLR 2021
Charles R Harris, K Jarrod Millman, Stéfan J van der Walt, Ralf Gommers, Pauli Virtanen, David Cournapeau, Eric Wieser, Julian Taylor. Sebastian Berg, Nathaniel J Smith, et al. Array program- ming with numpy. Nature, 585(7825):357-332, 2020.
Chin-War Chin-Wei Huang, David Krueger, Alexandre Lacoste, and Aaron Courville. Neural autoregressive flows. In International Conference on Machine Leurning, pp. 2078-2087, 2018.
Chin-Wat Wei Huang, Laurent Dinh, and Aaron Courville. Augmented normalizing flows: Bridging the gap between generative flows and latent variable models. arXiv preprint arXiv:2002.07001, 2020a.
Chin-Wei Huang, Laurent Dinh, and Aaron Courville. Solving ode with universal flows: Approx- imation theory for flow-based models. In ICLR 2020 Workshop on Integration of Deep Neural Models and Differential Equations, 2020b.
Gao Huang. Zhuang Liu, Laurens Van Der Maaten, and Kilian Weinberger. Densely connected convolutional networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 4700-4708, 2017.
John Hunter. Matplotlib: 2d graphics environment. Computing in science & engineering, 9(3): 90, 2007.
Michael F Hutchinson. stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Communications in Statistix-Simullation and Computation, 18(3):1059-100 1989.
Priyank Jaini, Kira A Selby, and Yaoliang Yu. Sum-of polynomial flow. In International Conference on Machine Learning, pp. 3009-3018, 2019.
Eric Jones, Travis Oliphant, and Pearu Peterson. {SciPy]: Open source scientific tools for [Python]. 2014.
Diederik Kingma and Max Welling. 1rXiv:1312.6114, 2013.
Auto-encoding variational bayes.
arXiv preprint
Durk P Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible Ixl convolutions. In Advances in neural information processing systems, pp. 10215-10224, 2018.
Durk P Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling. Im- proved variational inference with inverse autoregressive flow. In Advances in neural information processing systems, pp. 4743-4751, 2016.
Zhifeng Kong and Kamalika Chaudhuri. The expressive power of a class of normalizing flow mod- els. arXiv preprint arXiv::006..0 00000000003022110000 2020.
Alexander Korotin, Vage Egiazarian, Arip Asadulaev, Alexander Safin, and Evgeny Burnaev. Wasserstein-2 generative networks. arXiv preprint arXiv: 2019.
Hongzhou Lin and Stefanie Jegelka. Resnet with one-neuron hidden layers is a universal approxi- mator. In Advances in neural information processing systems, pp. 6169-6178. 2018
Lin Lin, Yousef Saad, and Chao Yang. Approximating spectral densities of large matrices. SIAM review, 58(1):34-6) 65, 2016.
Zhou Lu, Hongming Pu, Feicheng Wang. Zhiqiang Hu, and Liwei Wang. The expressive power of neural networks: A view from the width. In Advances in neural information processing systems, pp. 6231-62333 2017.
Ashok Vardhan Makkuva, Amirhossein Taghvaci, Sewoong Oh, and Jason D Lee. Optimal transport mapping via input convex neural networks. arXiv preprint arXiv: 10962, 2019.
Wes McKinney. Python for data analysis: Data wrangling with Pandas, NumPy, and IPython." O' Reilly Media, Inc.", 2012.
11

---

Published as a conference paper at ICLR 2021
Takeru Miyato, Toshiki Kataoka, Masanori Koyama, and Yuichi Yoshida. Spectral normalization for generative adversarial networks. arXiv preprint arXiv:1802.05957 2018.
Thomas Müller, Brian McWilliams, Fabrice Rousselle, Markus Gross, and Jan Novák. Neural im- portance sampling. ACM Transactions on Graphics (TOG), 38(5):1-19, 2019.
Yurii Nesterov. Introductory lectures convex programming volume i: Basic course. Lecture notes, 3(4):5, 1998.
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Ng. Reading digits in natural images with unsupervised feature learning. 2011.
Jorge Nocedal and Stephen Wright. Numerical optimization. Springer Science & Business Media, 2006.
Travis E Oliphant. A guide to NumPy, volume 1. Trelgol Publishing USA, 2006.
Travis E Oliphant. Python for scientific computing. Computing in Science & Engineering, 9(3): 10-20, 2007.
Junier Oliva, Avinava Dubey, Manzil Zaheer, Barnabas Poczos, Ruslan Salakhutdinov, Eric Xing. and Jeff Schneider. Transformation autoregressive networks. In Proceedings of the 35th In- ternational Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 3898-3907. PMLR, 10-15 Jul 2018.
Derek Onken, Samy Wu Fung. Xingjian Li, and Lars Ruthotto. Ot-flow: Fast and accurate continu- ous normalizing flows via optimal transport. arXiv preprint arXiv:2006.00/ 000044 2020.
George Papamakarios, Theo Pavlakou, and lain Murray. Masked autoregressive flow for density estimation. In Advances in Neural Information Processing Systems, pp. 2338-2347, 2017.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high- performance deep learning library. In Advances in neural information processing systems, pp. 8026-8037, 2019.
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In Interna- tional Conference on Machine Learning, pp. 1530-1538, 2015.
R Tyrrell Rockafellar. Convex analysis. Number 28. Princeton university press, 1970.
Ludger Rüschendorf and Svetlozar Rachev. characterization of random variables with minimum 12-distance. Journal of multivariate analysis, 32(1):48-54 1990.
Filippo Santambrogio. Optimal transport for applied mathematicians. Birkäuser, NY, 55(58-63):94, 2015.
Matthias Seeger, Asmus Hetzel, Zhenwen Dai, Eric Meissner, and Neil D Lawrence, differentiating linear algebra. arXiv preprini 1rXiv:1710.08717.1 2017
Auto-
Vivien Seguy, Bharath Bhushan Damodaran, Remi Flamary, Nicolas Courty, Antoine Rolet, and Mathieu Blondel. Large scale optimal transport and mapping estimation. In International Con- ference on Learning Representations, 2018.
RN Silver and Röder. Densities states mega-dimensional hamiltonian matrices. International Journal of Modern Physics C, 5(04):735-753, 1994.
Amirhossein Taghvaci and Amin Jalali. 2- 2-wasserstein approximation restrricd convex potentials with application improved training for gans. arXiv preprint arXiv: 2019.
Hillel Tal-Ezer and Ronnie Kosloff. An accurate and efficient scheme for propagating the time dependent schrödinger equation. The Journal chemical physics, 8199)3397733771 1984.
12

---

Published as a conference paper at ICLR 2021
Takeshi Teshima, Isao Ishikawa, Koichi Tojo, Kenta Oono, Masahiro Ikeda, and Masashi Sugiyama. Coupling- based invertible neural networks are universal diffeomorphism approximators. arXiv preprint arXiv:2006.11 100.: 2020.
Shashanka Ubaru, Jie Chen, and Yousef Saad. Fast estimation of tr(f(A)) via Stochastic Lanczos Quadrature. SIAM Journal on Matrix Analysis and Applications, 38(4): 10755-99,,,, 2017.
Rianne Van Den Berg, Leonard Hasenclever, Jakub M Tomezak, and Max Welling. Sylvester nor- malizing flows for variational inference. In 34th Conference on Uncertainty Artificial In- telligence 2018, UAI 2018, pp. 393-402. Association For Uncertainty in Artificial Intelligence (AUAI), 2018.
Stefan Van Der Walt, S Chris Colbert, and Gael Varoquaux. The numpy array: a structure for efficient numerical computation. Computing in Science & Engineering, 13(2):22, 2011.
Guido Van Rossum and Fred L Drake Jr. Python reference manual. Centrum voor Wiskunde en Informatica Amsterdam, 1995.
Cédric Villani. Optimal transport: old and new, volume 338. Springer Science & Business Media, 2008.
Stéfan van der Walt, S Chris Colbert, and Gael Varoquaux. The numpy array: a structure for efficient numerical computation. Computing in science & engineering, 13(2):22-30, 2011.
Michael Waskom, Olga Botvinník, Drew O'Kane, Paul Hobson, Joel Ostblom, Saulius Lukauskas, David C Gemperline, Tom Augspurger, Yaroslav Halchenko, John B. Cole, Jordi Warmenhoven, Julian de Ruiter, Cameron Pye. Stephan Hoyer, Jake Vanderplas, Santi Villalba, Gero Kunter, Eric Quintero, Pete Bachant, Marcel Martin, Kyle Meyer, Alistair Miles, Yoav Ram, Thomas Brunner, Tal Yarkoni, Mike Lee Williams, Constantine Evans, Clark Fitzgerald, Brian, and Adel Qalieh. mwaskom/seaborn: v0.9.0 (july 2018), July 2018. URL https::/doo httpa/////ooo.///////////
5281/zenode 1313201.
Antoine Wehenkel and Gilles Louppe. Unconstrained monotonic neural networks. In Advances in Neural Information Processing Systems, pp. 1555-1555, 2019.
Liu Yang and George Em Karniadakis. Potential flow generator with l2 optimal transport regularity for generative models. arXiv preprint arXiv: 1462, 2019.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander Smola. Deep sets. Advances in neural information processing systems, pp. 3391- 1401,2017
Linfeng Zhang, Lei Wang, al. Monge-ampère flow for generative modeling. arXiv preprint 1809. 10188, 2018.
Bolei Zhou, Agata Lapedriza, Aditya Khosla, Aude Oliva, and Antonio Torralba. Places: A 10 mil- lion image database for scene recognition. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2017.
13


Reviews:
[{"review_id": "tcuPtxKO8N1", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\nThis paper proposes the flow based representation of a probability distribution so that the corresponding density remains tractable. In particular, the push-forward map that generates the desired distribution is characterized by the gradient of a strongly convex potential function. The invertability of the mapping as well as the Jacobian of the mapping is hence guaranteed by such a convexity property of the potential function. The proposed CP-flows are proved to be universal density approximators and are optimal in the OT (2-Wasserstein) sense.\n\nComment:\nIt seems to me that the proposed CP-flow model is quite similar to a set of recent works that use the ICNN to parameterized the convex Brenier potential of 2-Wasserstein. This limits the novelty of this paper. \nOn the other hand, I think one contribution of this paper is to study how the gradient of the log determinant of the Hessian of the potential function can be computed, which relies on a standard technique, the Hutchinson trace estimator.\nAnother contribution of this paper is to study universality of the CP-flow model as a density approximators. Since any convex function yields an affine-max representation which is a special case of the ICNN, the Brenier's theorem implies the universality.\n\nConcerns:\n1. Unlike the traditional flow models which consists of a sequence of simple transformations, it seems to me that CP-flow is not exactly a \"flow\" model as it only involves a single (but complex) transformation (which of course contains multiple layers). Correspondingly, the computation involved in the flow model (inversion, gradient evaluation) are substantially more expensive.\n2. The presentation of the paper focuses on the description of the model. However, I would suggest the authors to add a few paragraphs to briefly discuss some problems that takes the CP-flow model as a subroutine, e.g. variational inference or maximum likelihood loss. This would make the paper more self-contained and it will also justify the need to have a traceable density.  \n3. More on the writing of the paper. Some derivations to show that why (3) is important are appreciated, e.g. \n$det \\frac{\\partial f^{-1}(x)}{\\partial x} = \\frac{1}{det(\\nabla f(f^{-1}(x)))}$. Again, it would be good for a broader audience if the paper is more self-contained.\n4. How to efficiently implement the batch computation of the gradient of the model? What do you mean by \"Inverting a batch of independent inputs is as simple as summing the convex potential over all inputs\"?\n\n*********\nThe authors partially addressed my concerns. Therefore I raise my score to 6.", "extra_scores": {}, "extra_text": {}}, {"review_id": "kJm-bQTQYD", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "#### Summary and contributions\nThis paper presents CP-Flow, an alternative formulation of normalizing flows via the gradient of a strictly convex function. Concretely, an optimal transport map is constructed for the quadratic transport cost. This map can be constructed via the gradient of a strictly convex function owing to Brenier’s Theorem. The strict convexity of the potential function ensures that the gradient map is injective and surjective (i.e., invertible). This leads to relations with normalizing flows, invertible maps that can be trained via exact maximum likelihood. The authors borrow techniques from convex optimization to invert the gradient map and to compute the log-determinant of the Hessian (and its gradient), quantities that are required for density evaluation and gradient-based training. Using arguments from optimal transport literature, the authors prove the universality and optimality of the gradient maps when they are realized using input-convex neural networks (ICNNs). Experiments in various settings of density estimation and variational inference have been presented that demonstrate a competitive performance of CP-Flow against competing normalizing flow models.\n\n#### Strengths\nThis work makes a significant contribution to the body of literature on optimal transport-based generative models. By parameterizing the flow as the gradient of a convex potential function ($R^d \\rightarrow R$), this method improves the parameter efficiency of invertible models. Moreover, the network architectures, although restricted to ICNNs, can be much more flexible than in conventional normalizing flows where networks have to satisfy restrictive conditions to ensure invertibility. Theoretical results also demonstrate that ICNNs have the capacity to learn transport maps between arbitrary continuous probability measures. The theoretical results are supported by the empirical evaluation which, despite lacking scale, is sufficient to demonstrate the potential of the proposed method. \n\n#### Weaknesses\nThe primary weakness of this work is that some directly related works that utilize Brenier’s Theorem to learn transport maps via ICNNs have not been discussed in sufficient detail. Apart from flow-based setups, other works have also explored the construction of transport maps using ICNNs which have not been discussed (see for example [1, 2, 3]). Moreover, the experiments lack comparisons with these works or a discussion of why such comparisons are not possible.  Another recent related work is [4]. \nA potential concern is the efficiency of solving a quadratic optimization problem for computation of the gradient of the log-det of Hessian for each gradient step. Did the authors compare the runtimes of the different methods?\n\n[1] Korotin, Alexander, et al. \"Wasserstein-2 Generative Networks.\" arXiv preprint arXiv:1909.13082 (2019).    \n[2] Taghvaei, Amirhossein, and Amin Jalali. \"2-wasserstein approximation via restricted convex potentials with application to improved training for gans.\" arXiv preprint arXiv:1902.07197 (2019).    \n[3] Makkuva, Ashok Vardhan, et al. \"Optimal transport mapping via input convex neural networks.\" arXiv preprint arXiv:1908.10962 (2019).    \n[4] Finlay, Chris, et al. \"Learning normalizing flows from Entropy-Kantorovich potentials.\" arXiv preprint arXiv:2006.06033 (2020).\n\n\n#### Additional feedback\n\nSuggestions:\n- Include a proper discussion of related works that use ICNNs to parameterize the Brenier potential and compare with these methods experimentally, if possible.\n\nQuestions:\n- Is there a reason that the authors did not include a NSF ablation in Table 2?\n- Did the authors investigate the quality of generated samples for the image datasets (using FID, IS, etc.) as the poor quality of images is a known problem for flow-based models?\n-------\nThe authors have addressed most of my concerns. I believe the paper is a good contribution to the literature on normalizing flows; therefore, I firmly vote for acceptance. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "TBtSxId3Rn", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "#####################\n\nSummary and contributions:\n\nThe paper provides a novel approach to normalizing flows. It models a normalizing flow as the gradient map of a convex potential function.  The gradient estimation and model inversion, that are the computationally expensive part of arbitrary complex NFs, are formulated as cheap convex optimization problems. It proves that  the proposed CP-flow is a universal density approximator and also shows that it is an optimal transport (OT) map. \n\n#####################\n\nReasons for score:\n\nI vote for clear acceptance. The idea of convex potential flow sounds very interesting and the way the computational complexities are very useful for designing more advanced NFs. The paper is very well written. \n\nStrength:\n\n1. The paper is very well written and structured and is easy to read for a wide audience. It also provides a good review of main papers in the literature\n2. The motivation for using CP-flow is well explained. \n3. The claims are well supported by theoretical proofs and empirical studies.   \n4. The gradient estimation and model inversion, that are the computationally complex part of arbitrary complex NFs, are formulated as convex optimization problems that gain advantage of fast converging and cheap optimization algorithms and it also leverages the efficient Hessian-vector product computation.\n5. The experimental results show the proposed flow can perform competitively with much less number of parameters .\n\n#####################\n\nAdditional Feedback and Questions:\n\n1. After reading the paper it is not quite clear why we do need an optimality in Monge sense (Theorem 4) and what is the point of optimal transport in this work. It is worth expanding or adding more insights to the motivation given in the Introduction by the notion of rearrangement cost. \n2. Orthogonality or Lipschitzness constraints are mentioned without citation in introduction.\n3. What is so special about the soft-plus type nonlinearity in theorem 3, can we use other non-linearities such as the symmetric-log derived in [1] that are monotonic and differentiable by construction ?\n4. CP-flow looks more expressive than NAF in toy examples but why is it outperformed in density estimation? As a more insightful comparison, I suggest comparing the “transport cost” of the NAF with that of ICNN in section 5.2 (Figure 4) as NAF is outperforming ICNN in all the benchmarks and is universal.\n5. It is worth comparing the number of parameters of the optimal CP-flow in Table 1 with the available methods to have a better understanding of models’ flexibility. Maybe, CP-flow can achieve the SOTA if the number of parameters are normalized.\n6. I wonder if CP-flow can outperform the Residual flow which uses a similar gradient estimator routine if its number of parameters, in Table 2, is increased.\n7. Compared to RealNVP, CP-flow requires an optimization solution per each training update. Also, as noted in the paper, the computational cost of the CP-flow is less than residual flows as it saves the Spectral normalization process. So how the speed and convergence rate of CP-flow is compared against the benchmarks, assuming all are using the same hardware e.g GPUs (and cvx opt are implemented in GPU). It looks better to have a sense of it in one of the experiments. \n8. I am willing to see the randomly generated sample of the CP-flow especially to compare its local/global stability against the Residual flow due to its Lipschitz constraint as discussed in [2].\n9. A schematic architecture of the ICNN model helps better understand it.\n\nRef:\n[1] M. Karami, D. Schuurmans, Jascha Sohl-Dickstein, Daniel Duckworth, Laurent Dinh, “Invertible Convolutional Flow\", Advances in Neural Information Processing Systems (NeurIPS) 2019,\n\n[2] Behrmann, Jens, et al. \"Understanding and mitigating exploding inverses in invertible neural networks.\" arXiv preprint arXiv:2006.09347 (2020) ", "extra_scores": {}, "extra_text": {}}, {"review_id": "8DUVqvG-c17", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "### Summary\nThe authors introduce CP-Flows, a way to parameterize normalizing flows by constructing an input-convex neural net with softplus-type activation functions and considering its gradient as the flow. They add a quadratic term to ensure invertibility. Using convex optimization techniques their method only needs access to convex optimization solvers. They show that this architecture is universal (that is, starting from a measure $\\mu$, there is a sequence of CP-Flows converging weakly to a desired distribution $\\nu$). They also prove that the constructed flow converges pointwise to the optimal Brenier map for Euclidean cost. They perform a set of experiments on synthetic and real-world datasets, and show their method delivers its promises.\n\n### Strong/Weak Points\n+ The ideas described in the paper are simple and easy to understand.\n+ The paper is generally well-written (with exceptions detailed below).\n+ The ideas for computing the trace of log Hessian and its gradient are neat; however, I am not sure if they are novel and not present in the literature.\n+ All the results are asymptotic. The convergence results are weak: convergence to the optimal map (Theorem 4) is point-wise, and convergence of the distributions is weak (Theorem 3). \n+ There are lots of tricks used here and there (computing the trace of log Hessian and its gradient) to reduce computational complexity, but nothing explicit is computed theoretically nor presented experimentally. \n+ I am not sure if Theorem 2 and Theorem 4 are already known or not. I might not recall the reference, but I am pretty confident that they have been existing in the literature.\n+ In the experiments, looking at Table 1, NAF always outperfrom CP-Flow (and no explanation has given for why is it the case). The same is for Table 3, where Sylvester outperforms CP-Flow in every dataset.\n\nI have decided to give a 5 to this paper, as theoretically it does not add significantly to the current theory of OT and NFs. Experimentally, I feel that the performance is (probably) marginally better than the other methods if it is not worse.\n\n### Additional Feedback\nFixing the following errors improves readability of the paper:\n- Section 1, paragraph 3, line 4, \"the network\" has not been introduced and not clear where it is referring to.\n- same, line 7, universality and optimality is not clear (in what sense). It becomes clear later.\n- page 2, line 2, \"the network architecture\" still is dangling.\n- Section 2, two lines above \"Universal Flows\" does not make sense. \"ordinary differential equations, orthogonality or Lipschitz constants\" are not NN architectures.\n- same, two lines above \"Optimal Transport\" there is no reference why \"under very mild conditions a coupling exists\". \n- same, \"Optimal Transport\". The problem described is not Monge's problem. It's Monge's formulation for optimal transportation.\n- page 3, \"Invertibility and ...\" the activation changed from $s$ to $g$ suddenly.\n- page 4, lines 11, 12 are hard to understand. Why NaNs creep in?\n- page 4, (3), why the terms are in red? \n- page 5, \"Sigmoidal Flow\" activation is called $s$ while in the formula it is $\\sigma$.\n- same, \"Flows with Potential ...\" the last two lines, I do not see what the contrast is made against. \n- page 6, Section 5, I do not see why the notation changes from $z$ to $h$.\n- page 7, I do not know how to interpret the middle plots for $f_\\mathrm{iaf}(x)$ and $f_\\mathrm{cp}(x)$.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The main contribution of the paper is a novel parametrization of normalizing flows using ideas from optimal transport theory. This new parametrization allows viewing a normalizing flow as the gradient of a convex function, which leads to an efficient method for gradient estimation for likelihood estimation, using only access to convex optimization solvers. The paper is overall well-written and provides a clean analysis. Theoretical results from the paper are supported by experiments. The paper was overall viewed favorably by the reviewers. 
Decision: Accept (Poster)