[name=Theorem]thm [name=Proposition]prop

# Convex Potential Flows:
Universal Probability Distributions with
Optimal Transport and Convex Optimization

Chin-Wei Huang
Mila, University of Montreal
<EMAIL_ADDRESS>&Ricky T. Q. Chen
University of Toronto & Vector Institute
<EMAIL_ADDRESS>&Christos Tsirigotis
Mila, University of Montreal
<EMAIL_ADDRESS>&Aaron Courville
Mila, University of Montreal & CIFAR Fellow
<EMAIL_ADDRESS>
###### 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.

##  1 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 an invertible
mapping that can reshape the distribution of its input data into a simpler or
more complex one. To enable efficient training, numerous tricks have been
proposed to impose structural constraints on 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 a function $f:\mathbb{R}\rightarrow\mathbb{R}$ is
continuous, it is invertible (injective onto its image) if and only if it is
strictly monotonic. This means that if we are only allowed to 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 1D. 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 a 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, we 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 as convex optimization problems. As such, we
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 a flow as a gradient field, the
convex potential is an $\mathbb{R}^{d}\rightarrow\mathbb{R}$ function, which
is different from most existing discrete-time flows which are
$\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$. This makes CP-Flow relatively
compact. It is also arguably easier to design a 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 to incorporate additional structure such
as equivariance (Cohen & Welling, 2016; Zaheer et al., 2017) into the flow’s
parameterization, making CP-Flow a more flexible general purpose density
model.

##  2 Background: Normalizing Flows and Optimal Transport

Normalizing flows are characterized by a differentiable, invertible neural
network $f$ such that the probability density of the network’s output can be
computed conveniently using the change-of-variable formula

$\displaystyle p_{Y}(f(x))=p_{X}(x)\left|\frac{\partial f(x)}{\partial
x}\right|^{-1}\qquad\Longleftrightarrow\qquad
p_{Y}(y)=p_{X}(f^{-1}(y))\left|\frac{\partial f^{-1}(y)}{\partial y}\right|$
(1)

where the Jacobian determinant term captures the local expansion or
contraction of the density near $x$ (resp. $y$) induced by the mapping $f$
(resp. $f^{-1}$), and $p_{X}$ is the density of a random variable $X$. The
invertibility requirement has led to the design of many special neural network
parameterizations such as triangular maps, ordinary differential equations,
orthogonality or Lipschitz constraints.

#### Universal Flows

For a general learning framework to be meaningful, a model needs to be
flexible enough to capture variations in the data distribution. In the context
of density modeling, this corresponds to the model’s capability to 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 et al., 2017; Lin & Jegelka, 2018),
invertible neural networks generally 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
a 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 approximate a _deterministic coupling_ between
probability measures. Given a pair of probability densities $p_{X}$ and
$p_{Y}$, a deterministic coupling is a mapping $g$ such that $g(X)\sim p_{Y}$
if $X\sim p_{X}$. We seek to find a coupling that is invertible, or at least
can be approximated by invertible mappings.

#### Optimal Transport

Let $c(x,y)$ be a cost function. The _Monge problem_ (Villani, 2008) pertains
to finding the optimal transport map $g$ that realizes the minimal expected
cost

$\displaystyle J_{c}(p_{X},p_{Y})=\inf_{\widetilde{g}:\widetilde{g}(X)\sim
p_{Y}}\mathbb{E}_{X\sim p_{X}}[c(X,\widetilde{g}(X))]$ (2)

When the second moments of $X$ and $Y$ are both finite, and $X$ is regular
enough (e.g. having a density), then the special case of $c(x,y)=||x-y||^{2}$
has an interesting solution, a celebrated theorem due to Brenier (1987; 1991):
[Brenier’s Theorem, Theorem 1.22 of Santambrogio (2015)] Let $\mu,\nu$ be
probability measures with a finite second moment, and assume $\mu$ has a
Lebesgue density $p_{X}$. Then there exists a convex potential $G$ such that
the gradient map $g:=\nabla G$ (defined up to a null set) uniquely solves the
Monge problem in eq. 2 with the quadratic cost function $c(x,y)=||x-y||^{2}$.

Some recent works are also inspired by Brenier’s theorem and utilize a convex
potential to parameterize a critic model, starting from Taghvaei & Jalali
(2019), and further built upon by Makkuva et al. (2019) who parameterize a
generator with a convex potential and concurrently by Korotin et al. (2019).
Our work sets itself apart from these prior works in that it is entirely
likelihood-based, minimizing the (empirical) KL divergence as opposed to an
approximate optimal transport cost.

##  3 Convex Potential Flows

(a)

(b)

(c)

(d)

Figure 1: Illustration of Convex Potential Flow. (a) Data $x$ drawn from a
mixture of Gaussians. (b) Learned convex potential $F$. (c) Mesh grid
distorted by the gradient map of the convex potential $f=\nabla F$. (d)
Encoding of the data via the gradient map $z=f(x)$. 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
in a contraction towards the origin.

Given a strictly convex potential $F$, we can define an injective map
(invertible from its image) via its gradient $f=\nabla 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 a linear layer, and $L^{+}(x)$ to denote a linear
layer with positive weights. We use the (fully) input-convex neural network
(ICNN, Amos et al. (2017)) to parameterize the convex potential, which has the
following form

$F(x)=L_{K+1}^{+}(s(z_{K}))+L_{K+1}(x)\qquad
z_{k}:=L_{k}^{+}(s(z_{k-1}))+L_{k}(x)\qquad z_{1}:=L_{1}(x)$

where $s$ is a non-decreasing, convex activation function. In this work, we
use softplus-type activation functions, which is a rich family of activation
functions that can be shown to uniformly approximate the ReLU activation. See
Appendix A for details.

Algorithm 1 Inverting CP-Flow.

1:procedure Invert($F,y,\texttt{CvxSolver}$)

2: Initialize $x\leftarrow y$

3: def closure():

4: Compute loss: $l\leftarrow F(x)-y^{\top}x$

5: return $l$

6: $x\leftarrow\texttt{CvxSolver}(\texttt{closure},x)$

7: return $x$

#### Invertibility and Inversion Procedure

If the activation $s$ is twice differentiable, then the Hessian $H_{F}$ is
positive semi-definite. We can make it strongly convex by adding a quadratic
term $F_{\alpha}(x)=\frac{\alpha}{2}||x||^{2}_{2}+F(x)$, such that
$H_{F_{\alpha}}\succeq\alpha I\succ 0$. This means the gradient map
$f_{\alpha}=\nabla F_{\alpha}$ is injective onto its image. Furthermore, it is
surjective since for any $y\in\mathbb{R}^{d}$, the potential $x\mapsto
F_{\alpha}(x)-y^{\top}x$ has a unique minimizer111The minimizer $x^{*}$
corresponds to the gradient map of the _convex conjugate_ of the potential.
satisfying the first order condition $\nabla F_{\alpha}(x)=y$, due to the
strong convexity and differentiability. We refer to this invertible mapping
$f_{\alpha}$ as the _convex potential flow_ , or the CP-Flow. The above
discussion also implies we can plug in a black-box convex solver to invert the
gradient map $f_{\alpha}$, which we summarize in Algorithm 1. Inverting a
batch of independent inputs is as simple as summing the convex potential over
all inputs: since all of the entries of the scalar $l$ in the minibatch are
independent of each other, computing the gradient all $l$’s wrt all $x$’s
amounts to computing the gradient of the summation of $l$’s wrt all $x$’s. Due
to the convex nature of the problem, a wide selection of algorithms can be
used with convergence guarantees (Nesterov, 1998). In practice, we use the _L-
BFGS_ algorithm (Byrd et al., 1995) as our CvxSolver.

#### 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 on 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 a special case.
See Lin et al. (2016) for an overview of methods that only require access to
Hessian-vector products. Hessian-vector products (hvp) are cheap to compute
with reverse-mode automatic differentiation (Baydin et al., 2017), which does
not require constructing the full Hessian matrix and has the same asymptotic
cost as evaluating $F_{\alpha}$.

In particular, the log determinant can be rewritten in the form of a
generalized trace $\operatorname{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 $H$ has unbounded eigenvalues, we use a
more complex algorithm designed for symmetric matrices, the _stochastic
Lanczos quadrature_ (SLQ; Ubaru et al., 2017). At the core of SLQ is the
Lanczos method, which computes $m$ eigenvalues of $H$ by first constructing a
symmetric tridiagonal matrix $T\in\mathbb{R}^{m\times m}$ 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
a 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 ${\mathcal{O}}(1)$-Memory Unbiased $\nabla\log\det H$ estimator

Figure 2: Memory for training CIFAR-10.

We would also like to have an estimator for the _gradient_ 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
differentiating through SLQ are that it requires (i) differentiating through
an eigendecomposition routine and (ii) storing all Hessian-vector products in
memory (see fig. 2). Problem (i) is more specific to SLQ, because the gradient
of an eigendecomposition 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 in NaNs due
to the eigendecomposition. Problem (ii) will hold true for other algorithms
that also estimate $\log\det H$ with Hessian-vector products, and generally
the only difference is that a different numerical routine would need to be
differentiated through. Due to these problems, we do not differentiate through
SLQ, but we still use it as an efficient method for monitoring training
progress.

Instead, it is possible to construct an alternative formulation of the
gradient as the solution of a convex 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 formula∗ and the
adjugate representation of the matrix inverse†, for any invertible matrix $H$
with parameter $\theta$, we have the following identity:

(3)

Notably, in the last equality, we used the Hutchinson trace estimator
(Hutchinson, 1989) with a Rademacher random vector $v$, leading to a
${\mathcal{O}}(1)$-memory, unbiased Monte Carlo gradient estimator.

Computing the quantity $v^{\top}H^{-1}$ 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 a quadratic optimization problem

$\arg\min_{z}\left\\{\frac{1}{2}z^{\top}Hz-v^{\top}z\right\\}$ (4)

which has the unique minimizer $z^{*}=H^{-1}v$ since $H$ is symmetric positive
definite.

Algorithm 2 Surrogate training objective.

1:procedure SurrogateObj($F,x,\texttt{CG}$)

2: Obtain the gradient $\smash{f(x)\triangleq\nabla_{x}F(x)}$

3: Sample Rademacher random vector $r$

4: def hvp($v$):

5: return $v^{\top}\frac{\partial}{\partial x}f(x)$

6:
$z\leftarrow\texttt{stop\\_gradient}\left(\texttt{CG}(\texttt{hvp},r)\right)$

7: return $\texttt{hvp}(z)^{\top}r$

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 $z^{*}$ within $d$ iterations, and the error of the
approximation is known to converge exponentially fast $||z^{m}-z^{*}||_{H}\leq
2\gamma^{m}||z^{0}-z^{*}||_{H}$, where $z^{m}$ is the estimate after $m$
iterations. The rate of convergence $\gamma<1$ relates to the condition number
of $H$. For more details, see Nocedal & Wright (2006, Ch. 5). In practice, we
terminate CG when $||Hz^{m}-v||_{\infty}<\tau$ is satisfied for some user-
controlled tolerance. Empirically, we find that stringent tolerance values are
unnecessary for stochastic optimization (see appendix E).

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: $\frac{d}{d\theta}\left((z^{m})^{\top}Hv\right)$,
where only $H$ is differentiated through (since $z^{m}$ is only used to
approximate $v^{\top}H^{-1}$ 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 $\alpha=1$, the gradient map $f_{1}$ resembles the residual flow (Behrmann
et al., 2019; Chen et al., 2019a). They require the residual block—equivalent
to our gradient map $f$—to be _contractive_ (with Lipschitz constant strictly
smaller than 1) as a 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
recover a form similar to Sylvester Flows. For a 1-hidden layer ICNN ($K=1$)
and $\alpha=1$, we have
$F_{1}=\frac{1}{2}||x||^{2}_{2}+L_{2}^{+}(s(L_{1}x))+L_{2}(x)$. Setting the
weights of $L_{2}$ to zero, we have

$f_{1}(x)=\nabla_{x}F_{1}(x)=x+W_{1}^{\top}\text{diag}({w_{2}^{+}})s^{\prime}(W_{1}x+b_{1}).$
(5)

We notice the above form bears a close resemblance to the Sylvester
normalizing flow (Van Den Berg et al., 2018) (with ${\bm{Q}}$, ${\bm{R}}$ and
$\widetilde{{\bm{R}}}$ from Van Den Berg et al. (2018) being equal to
$W_{1}^{\top}$, $\text{diag}(w_{2}^{+})$ and $I$, respectively). For the
Sylvester flow to be invertible, they require that ${\bm{R}}$ and
$\widetilde{{\bm{R}}}$ be triangular and ${\bm{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 to
orthogonal matrices and one hidden layer are for applying Sylvester’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 $\sigma=s^{\prime}$. Then for
the 1-dimensional case ($d=1$) and $\alpha=0$ (without the residual
connection), we have

$\frac{\partial}{\partial
x}F_{0}(x)=\sum_{j=1}w_{1,j}w_{2,j}^{+}\sigma(w_{1,j}x+b_{1,j})=\sum_{j=1}|w_{1,j}|w_{2,j}^{+}\sigma(|w_{1,j}|x+\operatorname{sign}(w_{1,j})b_{1,j})+\text{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 function as the derivative of a convex function, which opens up
a new direction for parameterizing 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).

#### Flows with Potential Parameterization

Also inspired by Brenier’s theorem, Zhang et al. (2018) parameterize the
transform by taking the gradient of a potential, which can be thought of as
the linearized Monge-Ampère equation. They do not require the potential to be
convex since it is guaranteed to be invertible in the infinitesimal setting of
continuous normalizing flows (Chen et al., 2018). There exist other recent
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 to have low transport cost. In contrast, CP-Flow is guaranteed by
construction to converge pointwise to the optimal mapping between
distributions without explicit regularization (see Section 4). Another recent
work by Finlay et al. (2020a) is based off of entropy-regularized optimal
transport, which motivates the parameterization of a continuous normalizing
flow using scalar potentials.

##  4 Theoretical Analyses

As explained in Section 2, the parameterization of CP-Flow is inspired by the
Brenier potential. So naturally we would hope to show that (1) CP-Flows are
distributionally universal, and that (2) the learned invertible map is optimal
in the sense of the average squared distance the input travels
$\mathbb{E}[||x-f(x)||^{2}]$. Proofs of statements made in this section can be
found in Appendices B and C.

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 gradient fields. A classic example is the
sequence $F_{n}=\sin(nx)/\sqrt{n}$ and the corresponding derivatives
$f_{n}=\cos(nx)\sqrt{n}$: $F_{n}\rightarrow 0$ as $n\rightarrow\infty$ but
$f_{n}$ does not. Fortunately, convexity allows us to control the variation of
the gradient map (since the derivative of a convex function is monotone), so
our second step of approximation holds.

[] Let $F_{n}:\mathbb{R}^{d}\rightarrow\mathbb{R}$ be differentiable convex
functions and $G:\mathbb{R}^{d}\rightarrow\mathbb{R}$ be a proper convex
function. Assume $F_{n}\rightarrow G$. Then for almost every
$x\in\mathbb{R}^{d}$, $G$ is differentiable and $f_{n}(x):=\nabla
F_{n}(x)\rightarrow\nabla G(x)=:g(x)$.

Combining these two steps and Brenier’s theorem, we show that CP-Flow with
softplus-type activation function is distributionally universal.

[Universality] Given random variables $X\sim\mu$ and $Y\sim\nu$, with $\mu$
being absolutely continuous w.r.t. the Lebesgue measure, there exists a
sequence of ICNN $F_{n}$ with a softplus-type activation, such that $\nabla
F_{n}\circ X\rightarrow Y$ in 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 if
we match the distributions.

[Optimality] Let $G$ be the Brenier potential of $X\sim\mu$ and $Y\sim\nu$,
and let $F_{n}$ be a convergent sequence of differentiable, convex potentials,
such that $\nabla F_{n}\circ X\rightarrow Y$ in distribution. Then $\nabla
F_{n}$ converges almost surely to $\nabla G$.

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 we were optimizing according to the
transport cost. This allows us to estimate optimal transport maps without
solving the constrained optimization in (2). See Seguy et 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 of (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 D.

#### ICNN Architecture

Despite the universal property, having a poor parameterization can lead to
difficulties in optimization and limit the effective expressivity of the
model. We propose an architectural enhancement of ICNN, defined as follows
(note the change in notation: instead of writing the pre-activations $z$, we
use $h$ to denote the activated units):

$\displaystyle F^{aug}(x)$ $\displaystyle:=L_{K+1}^{+}(h_{K})+L_{K+1}(x)$
$\displaystyle
h_{k}:=\texttt{concat}([\widetilde{h}_{k},h_{k}^{aug}])\qquad\widetilde{h}_{k}$
$\displaystyle:=s(L_{k}^{+}(h_{k-1})+L_{k}(x))\qquad
h_{k}^{aug}=s(L_{k}^{aug}(x))$ (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, we use the input-augmented ICNN as
the default architecture.

###  5.1 Toy examples

DataMAFNAFCP-Flow

Figure 3: Learning toy densities.

Having distributional universality for a single flow layer means that we can
achieve high expressiveness without composing too many flows. We demonstrate
this by fitting the density on some toy examples taken from Papamakarios et
al. (2017) and Behrmann et al. (2019). We compare with the masked
autoregressive flow (MAF, Papamakarios et al. (2017)) and the neural
autoregressive flow (NAF, (Huang et al., 2018)). Results are presented 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 probability mass, MAF fails to fit
to the second and third datasets222Behrmann et al. (2019) demonstrates one can
potentially improve the affine coupling models by composing many flow layers.
But here we restrict the number of flow layers to be 3 or 5.. Although the
last dataset is intrinsically 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.

###  5.2 Approximating optimal coupling

data: $x$$z=f_{iaf}(x)$$z=f_{cp}(x)$

$d=8$$d=16$

Figure 4: Approximating optimal transport map via maximum likelihood
(minimizing KL divergence). In the first figure on the left we show the data
in 2 dimensions. The datapoints are colored according to their horizontal
values ($x_{1}$). The flows $f_{iaf}$ and $f_{cp}$ are trained to transform
the data into a standard Gaussian prior. In the 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.

As predicted by Theorem 4, CP-Flow is guaranteed to converge to the optimal
coupling minimizing the expected quadratic cost. We empirically verify it by
learning the Gaussian density and comparing the expected quadratic distance
between the input and output of the flow against $J_{{||x-y||}^{2}}$ between
the Gaussian data and the standard Gaussian prior (as there is a closed-form
expression). In fig. 4, we see that the transport cost gets closer to the
optimal value when the learned density approaches the data distribution
(measured by the KL divergence). We compare against the linear inverse
autoregressive flow (Kingma et al., 2016), which has the capacity to represent
the multivariate Gaussian density, yet it does not learn the optimal coupling.

###  5.3 Density estimation

Model | Power | Gas | Hepmass | Miniboone | BSDS300
---|---|---|---|---|---
Real NVP (Dinh et al., 2017) | -0.17 | -8.33 | 18.71 | 13.55 | -153.28
FFJORD (Grathwohl et al., 2018) | -0.46 | -8.59 | 14.92 | 10.43 | -157.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 1: 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 section D.4.

We demonstrate the efficacy of our model and the proposed gradient estimator
by performing density estimation on the standard benchmarks.

#### Tabular Data

We use the datasets preprocessed by Papamakarios et al. (2017). In table 1, we
report average negative log-likelihood estimates evaluated on held-out test
sets, for the best hyperparameters found via grid search. The search was
focused on the number of flow blocks, the width and depth of the ICNN
potentials. See section D.4 for details. Our models perform competitively
against alternative approaches in the literature. We also perform an ablation
on the CG error tolerance and ICNN architectures in appendix E.

#### Image Data

Next, we apply CP-Flow to model the density of standard image datasets, MNIST
and CIFAR-10. For this, we 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 a 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)—another work promoting fewer parameters—while 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
et 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, the baseline flow models do not perform nearly as 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 are limited in
expressivity or Lipschitz-constrained.

| 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) | 1.06 | N/A |  | 3.35 | 44.0M†
RQ-NSF (Durkan et al., 2019) | — | — |  | 3.38 | 11.8M†
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 | 1.04 | 2.9M |  | 3.46 | 3.1M
CP-Flow (Ours) | 1.02 | 2.9M |  | 3.40 | 1.9M
Table 2: Negative log-likelihood (in bits) on held-out test data (lower is
better). †Taken from Durkan et al. (2019). ‡Obtained from official open source
code.

###  5.4 Amortizing ICNN for Variational Inference

| Freyfaces | Omniglot | Caltech
---|---|---|---
Gaussian | 4.53 | 104.28 | 110.80
Planar | 4.40 | 102.65 | 109.66
IAF | 4.47 | 102.41 | 111.58
Sylvester | 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 section D.6.

Normalizing flows also allow us to employ a larger, more flexible family of
distributions for variational inference (Rezende & Mohamed, 2015). We
replicate the experiment conducted in Van Den Berg et al. (2018) to enhance
the variational autoencoder (Kingma & Welling, 2013). For inference
amortization, we use the partially input convex neural network from Amos et
al. (2017), and use the output of the encoder as the additional input for
conditioning. 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 parameterization along with the
proposed gradient estimator can learn to perform accurate amortized inference.
Also, we show that replacing the vanilla ICNN with the input-augmented ICNN
leads to improvement of the likelihood estimates.

##  6 Conclusion

We propose a new parameterization of normalizing flows using the gradient map
of a convex potential. We make connections to the optimal transport theory to
show that the proposed flow is a universal density model, and leverage tools
from convex optimization 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 be improved by designing
better ICNN architectures. We leave the exploration for a better ICNN and
convolutional ICNN architecture to improve density estimation and generative
modeling for future research.

## Acknowledgements

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), Matplotlib (Hunter,
2007), seaborn (Waskom et al., 2018), pandas (McKinney, 2012), and SciPy
(Jones et al., 2014).

## References

  * Adams et al. (2018) Ryan P Adams, Jeffrey Pennington, Matthew J Johnson, Jamie Smith, Yaniv Ovadia, Brian Patton, and James Saunderson.  Estimating the spectral density of large implicit matrices.  _arXiv preprint arXiv:1802.03451_ , 2018.
  * Amos et al. (2017) Brandon Amos, Lei Xu, and J Zico Kolter.  Input convex neural networks.  In _International Conference on Machine Learning_ , pp. 146–155, 2017.
  * Baydin et al. (2017) Atılım Günes Baydin, Barak A Pearlmutter, Alexey Andreyevich Radul, and Jeffrey Mark Siskind.  Automatic differentiation in machine learning: a survey.  _The Journal of Machine Learning Research_ , 18(1):5595–5637, 2017.
  * Behrmann et al. (2019) Jens Behrmann, Will Grathwohl, Ricky T. Q. Chen, David Duvenaud, and Jörn-Henrik Jacobsen.  Invertible residual networks.  In _International Conference on Machine Learning_ , pp. 573–582, 2019.
  * Behrmann et al. (2020) Jens Behrmann, Paul Vicol, Kuan-Chieh Wang, Roger Grosse, and Jörn-Henrik Jacobsen.  Understanding and mitigating exploding inverses in invertible neural networks.  _arXiv preprint arXiv:2006.09347_ , 2020.
  * Brenier (1987) Yann Brenier.  Décomposition polaire et réarrangement monotone des champs de vecteurs.  _CR Acad. Sci. Paris Sér. I Math._ , 305:805–808, 1987\.
  * Brenier (1991) Yann Brenier.  Polar factorization and monotone rearrangement of vector-valued functions.  _Communications on pure and applied mathematics_ , 44(4):375–417, 1991.
  * Byrd et al. (1995) 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.
  * Chen et al. (2018) Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud.  Neural ordinary differential equations.  In _Advances in neural information processing systems_ , pp. 6571–6583, 2018.
  * Chen et al. (2019a) Ricky T. Q. Chen, Jens Behrmann, David K Duvenaud, and Jörn-Henrik Jacobsen.  Residual flows for invertible generative modeling.  In _Advances in Neural Information Processing Systems_ , pp. 9916–9926, 2019a.
  * Chen et al. (2019b) Yize Chen, Yuanyuan Shi, and Baosen Zhang.  Optimal control via neural networks: A convex approach.  In _International Conference on Learning Representations_ , 2019b.
  * Cimpoi et al. (2014) Mircea Cimpoi, Subhransu Maji, Iasonas Kokkinos, Sammy Mohamed, and Andrea Vedaldi.  Describing textures in the wild.  In _Conference on Computer Vision and Pattern Recognition_ , 2014\.
  * Cohen & Welling (2016) Taco Cohen and Max Welling.  Group equivariant convolutional networks.  In _International conference on machine learning_ , pp. 2990–2999, 2016.
  * Cybenko (1989) George Cybenko.  Approximation by superpositions of a sigmoidal function.  _Mathematics of control, signals and systems_ , 2(4):303–314, 1989.
  * Dinh et al. (2014) Laurent Dinh, David Krueger, and Yoshua Bengio.  Nice: Non-linear independent components estimation.  _arXiv preprint arXiv:1410.8516_ , 2014.
  * Dinh et al. (2017) Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio.  Density estimation using real nvp.  In _International Conference on Learning Representations_ , 2017.
  * Dumoulin et al. (2018) Vincent Dumoulin, Ethan Perez, Nathan Schucher, Florian Strub, Harm de Vries, Aaron Courville, and Yoshua Bengio.  Feature-wise transformations.  _Distill_ , 3(7):e11, 2018.
  * Durkan et al. (2019) Conor Durkan, Artur Bekasov, Iain Murray, and George Papamakarios.  Neural spline flows.  In _Advances in Neural Information Processing Systems_ , pp. 7511–7522, 2019.
  * Finlay et al. (2020a) Chris Finlay, Augusto Gerolin, Adam M Oberman, and Aram-Alexandre Pooladian.  Learning normalizing flows from entropy-kantorovich potentials.  _arXiv preprint arXiv:2006.06033_ , 2020a.
  * Finlay et al. (2020b) Chris Finlay, Jörn-Henrik Jacobsen, Levon Nurbekyan, and Adam M Oberman.  How to train your neural ode: the world of jacobian and kinetic regularization.  In _International Conference on Machine Learning_ , 2020b.
  * Germain et al. (2015) Mathieu Germain, Karol Gregor, Iain Murray, and Hugo Larochelle.  Made: Masked autoencoder for distribution estimation.  In _Proceedings of the 32nd International Conference on Machine Learning_ , volume 37 of _Proceedings of Machine Learning Research_ , pp. 881–889. PMLR, 07–09 Jul 2015.
  * Goodfellow et al. (2013) 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.
  * Grathwohl et al. (2018) 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 Conference on Learning Representations_ , 2018.
  * Ha et al. (2016) David Ha, Andrew Dai, and Quoc V Le.  Hypernetworks.  _arXiv preprint arXiv:1609.09106_ , 2016.
  * Han et al. (2018) Insu Han, Haim Avron, and Jinwoo Shin.  Stochastic chebyshev gradient descent for spectral optimization.  In _Advances in Neural Information Processing Systems_ , pp. 7386–7396, 2018.
  * Harris et al. (2020) 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 programming with numpy.  _Nature_ , 585(7825):357–362, 2020.
  * Huang et al. (2018) Chin-Wei Huang, David Krueger, Alexandre Lacoste, and Aaron Courville.  Neural autoregressive flows.  In _International Conference on Machine Learning_ , pp. 2078–2087, 2018.
  * Huang et al. (2020a) Chin-Wei Huang, Laurent Dinh, and Aaron Courville.  Augmented normalizing flows: Bridging the gap between generative flows and latent variable models.  _arXiv preprint arXiv:2002.07101_ , 2020a.
  * Huang et al. (2020b) Chin-Wei Huang, Laurent Dinh, and Aaron Courville.  Solving ode with universal flows: Approximation theory for flow-based models.  In _ICLR 2020 Workshop on Integration of Deep Neural Models and Differential Equations_ , 2020b.
  * Huang et al. (2017) Gao Huang, Zhuang Liu, Laurens Van Der Maaten, and Kilian Q Weinberger.  Densely connected convolutional networks.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pp. 4700–4708, 2017.
  * Hunter (2007) John D Hunter.  Matplotlib: A 2d graphics environment.  _Computing in science & engineering_, 9(3):90, 2007.
  * Hutchinson (1989) Michael F Hutchinson.  A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines.  _Communications in Statistics-Simulation and Computation_ , 18(3):1059–1076, 1989.
  * Jaini et al. (2019) Priyank Jaini, Kira A Selby, and Yaoliang Yu.  Sum-of-squares polynomial flow.  In _International Conference on Machine Learning_ , pp. 3009–3018, 2019.
  * Jones et al. (2014) Eric Jones, Travis Oliphant, and Pearu Peterson.  $\\{$SciPy$\\}$: Open source scientific tools for $\\{$Python$\\}$.  2014\.
  * Kingma & Welling (2013) Diederik P Kingma and Max Welling.  Auto-encoding variational bayes.  _arXiv preprint arXiv:1312.6114_ , 2013.
  * Kingma & Dhariwal (2018) Durk P Kingma and Prafulla Dhariwal.  Glow: Generative flow with invertible 1x1 convolutions.  In _Advances in neural information processing systems_ , pp. 10215–10224, 2018.
  * Kingma et al. (2016) Durk P Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling.  Improved variational inference with inverse autoregressive flow.  In _Advances in neural information processing systems_ , pp. 4743–4751, 2016.
  * Kong & Chaudhuri (2020) Zhifeng Kong and Kamalika Chaudhuri.  The expressive power of a class of normalizing flow models.  _arXiv preprint arXiv:2006.00392_ , 2020.
  * Korotin et al. (2019) Alexander Korotin, Vage Egiazarian, Arip Asadulaev, Alexander Safin, and Evgeny Burnaev.  Wasserstein-2 generative networks.  _arXiv preprint arXiv:1909.13082_ , 2019.
  * Lin & Jegelka (2018) Hongzhou Lin and Stefanie Jegelka.  Resnet with one-neuron hidden layers is a universal approximator.  In _Advances in neural information processing systems_ , pp. 6169–6178, 2018.
  * Lin et al. (2016) Lin Lin, Yousef Saad, and Chao Yang.  Approximating spectral densities of large matrices.  _SIAM review_ , 58(1):34–65, 2016.
  * Lu et al. (2017) 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–6239, 2017.
  * Makkuva et al. (2019) Ashok Vardhan Makkuva, Amirhossein Taghvaei, Sewoong Oh, and Jason D Lee.  Optimal transport mapping via input convex neural networks.  _arXiv preprint arXiv:1908.10962_ , 2019.
  * McKinney (2012) Wes McKinney.  _Python for data analysis: Data wrangling with Pandas, NumPy, and IPython_.  ” O’Reilly Media, Inc.”, 2012.
  * Miyato et al. (2018) Takeru Miyato, Toshiki Kataoka, Masanori Koyama, and Yuichi Yoshida.  Spectral normalization for generative adversarial networks.  _arXiv preprint arXiv:1802.05957_ , 2018.
  * Müller et al. (2019) Thomas Müller, Brian McWilliams, Fabrice Rousselle, Markus Gross, and Jan Novák.  Neural importance sampling.  _ACM Transactions on Graphics (TOG)_ , 38(5):1–19, 2019.
  * Nesterov (1998) Yurii Nesterov.  Introductory lectures on convex programming volume i: Basic course.  _Lecture notes_ , 3(4):5, 1998.
  * Netzer et al. (2011) Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y Ng.  Reading digits in natural images with unsupervised feature learning.  2011\.
  * Nocedal & Wright (2006) Jorge Nocedal and Stephen Wright.  _Numerical optimization_.  Springer Science & Business Media, 2006.
  * Oliphant (2006) Travis E Oliphant.  _A guide to NumPy_ , volume 1.  Trelgol Publishing USA, 2006.
  * Oliphant (2007) Travis E Oliphant.  Python for scientific computing.  _Computing in Science & Engineering_, 9(3):10–20, 2007.
  * Oliva et al. (2018) Junier Oliva, Avinava Dubey, Manzil Zaheer, Barnabas Poczos, Ruslan Salakhutdinov, Eric Xing, and Jeff Schneider.  Transformation autoregressive networks.  In _Proceedings of the 35th International Conference on Machine Learning_ , volume 80 of _Proceedings of Machine Learning Research_ , pp. 3898–3907. PMLR, 10–15 Jul 2018.
  * Onken et al. (2020) Derek Onken, Samy Wu Fung, Xingjian Li, and Lars Ruthotto.  Ot-flow: Fast and accurate continuous normalizing flows via optimal transport.  _arXiv preprint arXiv:2006.00104_ , 2020.
  * Papamakarios et al. (2017) George Papamakarios, Theo Pavlakou, and Iain Murray.  Masked autoregressive flow for density estimation.  In _Advances in Neural Information Processing Systems_ , pp. 2338–2347, 2017.
  * Paszke et al. (2019) 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.
  * Rezende & Mohamed (2015) Danilo Rezende and Shakir Mohamed.  Variational inference with normalizing flows.  In _International Conference on Machine Learning_ , pp. 1530–1538, 2015.
  * Rockafellar (1970) R Tyrrell Rockafellar.  _Convex analysis_.  Number 28. Princeton university press, 1970.
  * Rüschendorf & Rachev (1990) Ludger Rüschendorf and Svetlozar T Rachev.  A characterization of random variables with minimum l2-distance.  _Journal of multivariate analysis_ , 32(1):48–54, 1990.
  * Santambrogio (2015) Filippo Santambrogio.  Optimal transport for applied mathematicians.  _Birkäuser, NY_ , 55(58-63):94, 2015.
  * Seeger et al. (2017) Matthias Seeger, Asmus Hetzel, Zhenwen Dai, Eric Meissner, and Neil D Lawrence.  Auto-differentiating linear algebra.  _arXiv preprint arXiv:1710.08717_ , 2017.
  * Seguy et al. (2018) Vivien Seguy, Bharath Bhushan Damodaran, Remi Flamary, Nicolas Courty, Antoine Rolet, and Mathieu Blondel.  Large scale optimal transport and mapping estimation.  In _International Conference on Learning Representations_ , 2018.
  * Silver & Röder (1994) RN Silver and H Röder.  Densities of states of mega-dimensional hamiltonian matrices.  _International Journal of Modern Physics C_ , 5(04):735–753, 1994.
  * Taghvaei & Jalali (2019) Amirhossein Taghvaei and Amin Jalali.  2-wasserstein approximation via restricted convex potentials with application to improved training for gans.  _arXiv preprint arXiv:1902.07197_ , 2019.
  * Tal-Ezer & Kosloff (1984) Hillel Tal-Ezer and Ronnie Kosloff.  An accurate and efficient scheme for propagating the time dependent schrödinger equation.  _The Journal of chemical physics_ , 81(9):3967–3971, 1984.
  * Teshima et al. (2020) 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.11469_ , 2020.
  * Ubaru et al. (2017) 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):1075–1099, 2017.
  * Van Den Berg et al. (2018) Rianne Van Den Berg, Leonard Hasenclever, Jakub M Tomczak, and Max Welling.  Sylvester normalizing flows for variational inference.  In _34th Conference on Uncertainty in Artificial Intelligence 2018, UAI 2018_ , pp. 393–402. Association For Uncertainty in Artificial Intelligence (AUAI), 2018.
  * Van Der Walt et al. (2011) 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.
  * Van Rossum & Drake Jr (1995) Guido Van Rossum and Fred L Drake Jr.  _Python reference manual_.  Centrum voor Wiskunde en Informatica Amsterdam, 1995.
  * Villani (2008) Cédric Villani.  _Optimal transport: old and new_ , volume 338.  Springer Science & Business Media, 2008.
  * Walt et al. (2011) 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.
  * Waskom et al. (2018) Michael Waskom, Olga Botvinnik, 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://doi.org/10.5281/zenodo.1313201.
  * Wehenkel & Louppe (2019) Antoine Wehenkel and Gilles Louppe.  Unconstrained monotonic neural networks.  In _Advances in Neural Information Processing Systems_ , pp. 1545–1555, 2019.
  * Yang & Karniadakis (2019) Liu Yang and George Em Karniadakis.  Potential flow generator with $l_{2}$ optimal transport regularity for generative models.  _arXiv preprint arXiv:1908.11462_ , 2019.
  * Zaheer et al. (2017) Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander J Smola.  Deep sets.  In _Advances in neural information processing systems_ , pp. 3391–3401, 2017.
  * Zhang et al. (2018) Linfeng Zhang, Lei Wang, et al.  Monge-ampère flow for generative modeling.  _arXiv preprint arXiv:1809.10188_ , 2018.
  * Zhou et al. (2017) Bolei Zhou, Agata Lapedriza, Aditya Khosla, Aude Oliva, and Antonio Torralba.  Places: A 10 million image database for scene recognition.  _IEEE Transactions on Pattern Analysis and Machine Intelligence_ , 2017.

##  Appendix A Softplus Type Activation

In this section, we let $r(x)=\max(0,x)$ be the ReLU activation function.

######  Definition 1.

We say a function $s$ is of the softplus type if the following holds

  1. (a)

$s\geq r$

  2. (b)

$s$ is convex

  3. (c)

$|s(x)-r(x)|\rightarrow 0$ as $|x|\rightarrow\infty$

Note that a softplus-type activation function is necessarily continuous, non-
decreasing, and uniformly approximating ReLU in the following sense:

$|s(xa)/a-r(x)|\rightarrow 0$

uniformly for all $x\in\mathbb{R}$ as $a\rightarrow\infty$.

The following proposition characterizes a big family of softplus-type
functions, and establishes a close connection between softplus type functions
and probability distribution functions.

######  Proposition 1.

Let $p$ be a probability density function of a random variable with mean zero.
Then the convolution $s:=p*r$ is a softplus-type function. Moreover,
$s(x)=\int_{-\infty}^{x}F_{p}(y)dy$, where $F_{p}$ is the distribution
function of $p$, and $s$ is at least twice differentiable.

###### Proof.

We first prove a claim (i): $xF_{p}(x)\rightarrow 0$ as $x\rightarrow-\infty$.
First, for $x\leq 0$,

$0\geq xF_{p}(x)=\int_{-\infty}^{x}xp(y)dy\geq\int_{-\infty}^{x}yp(y)dy$

Since $1_{y\leq x}yp(y)\rightarrow 0$ as $x\rightarrow-\infty$ and $|1_{y\leq
x}yp(y)|\leq|yp(y)|$, which is integrable by assumption, the integral on the
RHS of the above goes to $0$ by the dominated convergence theorem.

We now show the identity. By definition, since $\int yp(y)dy=0$

$\displaystyle s(x)$ $\displaystyle=\int\max(x-y,0)p(y)dy=\int\max(x,y)p(y)dy$
$\displaystyle=\int_{-\infty}^{x}xp(y)dy+\int_{x}^{\infty}yp(y)dy=xF_{p}(x)+\int_{x}^{\infty}yp(y)dy$
(8)

where we’ve used claim (i) to evaluate $xF_{p}(x)$ as $x\rightarrow-\infty$.
On the other hand, integration by part implies

$\displaystyle\int_{-\infty}^{x}F_{p}(y)dy=yF_{p}(y)\big{|}_{-\infty}^{x}-\int_{-\infty}^{x}yp(y)dy=xF_{p}(x)+\int_{x}^{\infty}yp(y)dy$
(9)

Twice differentiability follows from the differentiability of $F_{p}$.

(a) Now since $r(x)$ is convex, Jensen’s inequality gives

$s(x)=\int r(y)p(x-y)dy=\mathbb{E}[r(y)]\geq r(\mathbb{E}[y])=r(x)$

(b) $s$ is convex because $s^{\prime}=F_{p}$ is non-decreasing.

(c) To show that $s$ ane $r$ are asymptotically the same, we notice the
integral on the RHS of (8) goes to $0$ as $|x|\rightarrow\infty$ (by the
dominated or monotone convergence theorem). It suffices to show $xF_{p}(x)$
goes to $0$ as $x\rightarrow-\infty$, which is just claim (i), and
$x-xF_{p}(x)>0$ goes to $0$ as $x\rightarrow\infty$. To show the latter, we
can rewrite $x-xF_{p}(x)=x(1-F_{p}(x))=\int_{x}^{\infty}xp(y)dy$, and the same
argument as the claim holds with a vanishing upper bound.

∎

When $p$ is taken to be the standard logistic density, the corresponding
$s=p*r$ is simply the regular softplus activation function. We list a few
other softplus-type functions in table 4 and visualize them in fig. 5. We
experimented with the Gaussian-softplus and the logistic-softplus.

| $p(y)$ | $s:=p*r$
---|---|---
Logistic | $\frac{\exp(-x)}{(1+\exp(-x))^{2}}$ | $\log\left(1+\exp(x)\right)$
Laplace | $\frac{e^{-|x|}}{2}$ | $r(x)+\frac{e^{-|x|}}{2}$
Gaussian | $\frac{e^{-\frac{x^{2}}{2}}}{\sqrt{2\pi}}$ | $\frac{{\sqrt{\frac{\pi}{2}}x\operatorname{erf}\left(\frac{x}{\sqrt{2}}\right)+e^{-\frac{x^{2}}{2}}+\sqrt{\frac{\pi}{2}}x}}{\sqrt{2\pi}}$
Table 4: Formula of some softplus-type functions. Figure 5: Softplus-type
functions.

##  Appendix B Universality Proof

#### Notation:

Given a convex set $\Omega\subseteq\mathbb{R}^{d}$, we let
${\mathcal{C}}(\Omega)$ denote the set of continuous functions on $\Omega$,
and ${\mathcal{C}}_{\times}(\Omega):=\\{f\in{\mathcal{C}}(\Omega):f\text{ is
convex}\\}$ denote the set of convex, continuous functions.

We first show that ICNNs with a suitable activation function are dense in
$C_{\times}$. A similar result can be found in Chen et al. (2019b), where they
use a different constructive proof: first show that piecewise maximum of
affine functions, i.e. the maxout unit (Goodfellow et al., 2013), can
approximate any convex function, and then represent maxout using ICNN. We
emphasize our construction is simpler (see proof of Proposition 3).

The following proposition proves that functions that are pointwise maximum of
affine functions, are a dense subset of $C_{\times}$.

######  Proposition 2.

Pointwise maximum of affine functions is dense in
${\mathcal{C}}_{\times}([0,1]^{d})$.

###### Proof.

Fix some $\epsilon>0$. Since $f\in{\mathcal{C}}_{\times}([0,1]^{d})$ is
uniformly continuous on $[0,1]^{d}$, there exists some $\delta>0$ such that
$|f(x)-f(y)|<\epsilon$ provided that $||x-y||<\delta$. Let $n$ be big enough
such that $2^{-n}<\delta$, and let ${\mathcal{X}}$ be the set of points whose
coordinates sit on $i2^{-n}$ for some $1\leq i\leq 2^{n}-1$ (i.e. there are
$|{\mathcal{X}}|=(2^{n}-1)^{d}$ points in ${\mathcal{X}}$). For each
$y\in{\mathcal{X}}$, let $L_{y}(x):=\nabla f(y)^{\top}(x-y)+f(y)$ be a
supporting hyperplane of the graph of $f$, where $\nabla f(y)$ is a
subgradient of $f$ evaluated at $y$. Then we have a convex approximation
$f_{\epsilon}(x):=\max_{y\in{\mathcal{X}}}L_{y}(x)$ which bounds $f$ from
below. Moreover, letting
$y_{|x}:=\operatorname*{arg\,min}_{y\in{\mathcal{X}}}||x-y||$, we have (for
$x\not\in{\mathcal{X}}$),

$\displaystyle f(x)-f_{\epsilon}(x)$
$\displaystyle=f(x)-\max_{y\in{\mathcal{X}}}L_{y}(x)$ $\displaystyle\leq
f(x)-L_{y_{|x}}(x)$ $\displaystyle=f(x)-f(y_{|x})-\sum_{i=1}^{d}\nabla
f(y_{|x})_{i}(x_{i}-y_{|x,i})$ $\displaystyle\leq
f(x)-f(y_{|x})+\sum_{i=1}^{d}|\nabla f(y_{|x})_{i}|\cdot|x_{i}-y_{|x,i}|$
$\displaystyle\leq
f(x)-f(y_{|x})+\sum_{i=1}^{d}\frac{\epsilon}{|x_{i}-y_{|x,i}|}\cdot|x_{i}-y_{|x,i}|$
$\displaystyle\leq(d+1)\epsilon$

Since $\epsilon$ is arbitrary, this construction forms a sequence of
approximations converging uniformly to $f$ from below. ∎

The following proposition shows that maxout units can be equivalently
represented by ICNN with the ReLU activation, and thus entails the density of
the latter (as well as ICNN with softplus activation).

######  Proposition 3.

ICNN with ReLU or softplus-type activation is dense in
${\mathcal{C}}_{\times}([0,1]^{d})$.

###### Proof.

Let $r(x)=\max(0,x)$ be the ReLU activation funciton. Any convex piecewise
linear function $f(x)$ can be represented by $f(x)=\max(L_{1},...,L_{k})$
where $L_{j}=a_{j}^{\top}x+b_{j}$, which can then be reduced to

$\displaystyle f(x)$
$\displaystyle=r(\max(L_{1}-L_{k},...,L_{k-1}-L_{k}))+L_{k}$
$\displaystyle=r(r(\max(L_{1}-L_{k-1},...,L_{k-2}-L_{k-1}))+L_{k-1}-L_{k})+L_{k}$
$\displaystyle=z_{k}$

where $z_{j}:=r(z_{j-1})+L_{j}^{\prime}$ for $2\leq j\leq k$,
$z_{1}=L_{1}-L_{2}$, $L_{j}^{\prime}:=L_{j}-L_{j+1}$ for $2\leq j\leq k-1$,
and $L_{k}^{\prime}:=L_{k}$.

Since by Proposition 2, pointwise maximum of affine functions is dense in
${\mathcal{C}}_{\times}$, so is ICNN with the ReLU activation function. The
same holds for softplus since softplus can be used to uniformly approximate
ReLU.

∎

See 4

###### Proof.

We let $x$ be a differentiable point of $G$. Since convergence of derivatives
wrt each coordinate can be dealt with independently, we assume $d=1$ without
loss of generality. We can write $f_{n}$ as

$f_{n}(x)=\lim_{m}f_{nm}(x)\quad\text{ where }\quad
f_{nm}(x)=\frac{F_{n}(x-1/m)-F_{n}(x)}{-1/m}$

The problem can be rephrased as proving 333Note that there is an implicit
dependency on $x$ since the result is pointwise.

$\displaystyle\lim_{n}\lim_{m}f_{nm}=\lim_{m}\lim_{n}f_{nm}$ (10)

Note that $f_{nm}$ is non-decreasing in $m$ since $F_{n}$ is convex, and thus
$f_{nm}\leq f_{n}$. Since $f_{nm}$ converges to $f_{n}$, for any $\epsilon>0$,
we can find an integer $\mu(\epsilon,n)$ such that for all
$m\geq\mu(\epsilon,n)$, $|f_{nm}-f_{n}|\leq\epsilon$. Let $m_{k}$ be a
subsequence of $\\{m\geq 1\\}$ defined as $m_{k}=\mu(2^{-k},n)$.

Then $|f_{nm_{k+1}}-f_{nm_{k}}|\leq 2^{-k}$, which is integrable wrt the
counting measure on positive integers $k$, since

$\int 2^{-k}=\lim_{K\rightarrow\infty}\sum_{k=1}^{K}2^{-k}=1$

Thus, letting $f_{nm_{0}}=0$, by the Dominated Convergence Theorem, we have

$\lim_{n}\lim_{K}f_{nm_{K}}=\lim_{n}\int
f_{nm_{k}}-f_{nm_{k-1}}=\int\lim_{n}f_{nm_{k}}-f_{nm_{k-1}}=\lim_{K}\lim_{n}f_{nm_{K}}$

Although we are only looking at the limit of the subsequence $m_{k}$, this is
sufficient for (10), since the LHS is equal to $\lim_{n}f_{n}$ (since each
$F_{n}$ is differentiable), and by linearity of the limit, the RHS is equal to
$\lim_{K}\frac{G(x-1/m_{K})-G(x)}{-1/m_{K}}=g(x)$ (since $G$ is differentiable
at $x$ by assumption).

Since the set of points over which $G$ is not differentiable is a set of
measure zero (Rockafellar, 1970, Thm. 25.5), the convergence holds almost
everywhere. ∎

See 4

###### Proof.

Assume $\mu$ and $\nu$ have finite second moments. Since $\mu$ is absolutely
continuous, by Brenier’s theorem, there exists a convex function
$G:\mathbb{R}^{d}\rightarrow\mathbb{R}$ such that $\nabla
G(X)\overset{d}{=}\nu$ (where the gradient is unique up to changes on a null
set). By Proposition 3, there exists a sequence of ICNN $F_{n}$ converging to
$G$ pointwise everywhere. Such a sequence can be found since we can let
$F_{n}$ approximate $G$ with a uniform error of $1/n$ on a compact domain
$[-n,n]^{d}$. Theorem 4 then implies the gradient map $f_{n}:=\nabla F_{n}$
converges to $\nabla G$ pointwise almost everywhere. This implies the weak
convergence of the pushforward measure of $f_{n}\circ X$.

Now remove the finite second moment assumption and let $X$ and $Y$ be random
variables distributed according to $\mu$ (with Lebesgue density $p$) and
$\nu$, respectively. Denote by $B_{k}$ a ball of radius $k>0$ centered at the
origin, i.e. $B_{k}:=\\{x:||x||\leq k\\}$. Let $X_{k}=X1_{X\in
B_{k}}+U_{k}1_{X\not\in B_{k}}$, where $U_{k}$ is an independent random
variable distributed uniformly on $B_{k}$, and let $\mu_{k}$ be the law of
$X_{k}$. Then $X_{k}\rightarrow X$ almost surely as $k\rightarrow\infty$, and
$\mu_{k}$ is still absolutely continuous wrt the Lebesgue measure, with its
density being $p(x)+\frac{1}{\operatorname{vol}(B_{k})}\mu(||X||>k)$ if
$||x||\leq k$ or $0$ otherwise. Let $Y_{k}$ and $\nu_{k}$ be defined similarly
(while $\nu_{k}$ may not be absolutely continuous wrt Lebesgue). From above,
since $X_{k}$ and $Y_{k}$ are bounded and admit a finite second moment, we
know fixing $k$, we can find a sequence of $f_{k,n}=\nabla F_{k,n}$ such that
$f_{k,n}\circ X_{k}\rightarrow Y_{k}$ in distribution as $n\rightarrow\infty$.
Now since weak convergence is metrizable, choose $n_{k}$ to be large enough
such that the distance between the pushforward of $f_{k,n_{k}}\circ X_{k}$ and
$\nu_{k}$ is at most $1/k$. An application of triangle inequality of the weak
metric implies $f_{k,n_{k}}\circ X_{k}\rightarrow Y$ in distribution as
$k\rightarrow\infty$. Finally, note that since $\\{||f_{k,n_{k}}\circ
X_{k}-f_{k,n_{k}}\circ X||>\epsilon\\}\subseteq\\{||X||>k\\}$, by
monotonicity, we have

$\displaystyle{\mathbb{P}}\left(\limsup_{k}||f_{k,n_{k}}\circ
X_{k}-f_{k,n_{k}}\circ
X||>\epsilon\right)\leq{\mathbb{P}}\left(\limsup_{k}||X||>k\right)=0$ (11)

and thus $||f_{k,n_{k}}\circ X_{k}-f_{k,n_{k}}\circ X||\rightarrow 0$ almost
surely (where ${\mathbb{P}}$ is the underlying probability measure of the
measure space). By Lemma 1 of Huang et al. (2020b) (or equivalently Lemma 2 of
Huang et al. (2020a)), $f_{k,n_{k}}\circ X$ converges in distribution to $Y$.

∎

##  Appendix C Optimality Proof

See 4

###### Proof.

We claim that if $F$ is a convex potential such that $Z=\nabla F(X)$ has $\nu$
as its law, then $\nabla F\equiv\nabla G$ almost surely. The proof of the
claim is originally due to Rüschendorf & Rachev (1990), but we present it here
for completeness. Let $Z^{\prime}$ be another random variable distributed by
$\nu$. Then by the Fenchel-Young inequality (applied to the convex potential
$F$),

$\mathbb{E}[X^{\top}Z^{\prime}]\leq\mathbb{E}[F(X)+F^{*}(Z^{\prime})]=\mathbb{E}[F(X)+F^{*}(Z)]=\mathbb{E}[F(X)+F^{*}(\nabla
F(X))]=\mathbb{E}[X^{\top}\nabla F(X)]$

This concludes the proof since $\nabla G$ uniquely solves the transportation
problem, which is equivalent to finding a transport map $\tilde{g}$ that
maximizes the covariance:

$\mathbb{E}[||X-\tilde{g}(X)||^{2}]=\mathbb{E}[||X||^{2}+||\tilde{g}(X)||^{2}]-2\mathbb{E}[X^{\top}\tilde{g}(X)]$

Let $F_{\infty}$ to be the pointwise limit of $F_{n}$. Then for any $x_{1}$,
$x_{2}$ and $t\in[0,1]$,

$\displaystyle F_{\infty}(tx_{1}+(1-t)x_{2})$
$\displaystyle=\lim_{n\rightarrow\infty}F_{n}(tx_{1}+(1-t)x_{2})$
$\displaystyle\leq\lim_{n\rightarrow\infty}tF_{n}(x_{1})+(1-t)F_{n}(x_{2})$
$\displaystyle=\lim_{n\rightarrow\infty}tF_{n}(x_{1})+\lim_{n\rightarrow\infty}(1-t)F_{n}(x_{2})=tF_{\infty}(x_{1})+(1-t)F_{\infty}(x_{2})$

That is, $F_{\infty}$ is convex. Now since $F_{n}$ is a convergent sequence of
convex functions, its gradient $\nabla F_{n}$ also converges pointwise almost
everywhere to $\nabla F_{\infty}$ by Theorem 4. Let $\rho$ denote the
Prokhorov metric, which metrizes the weak convergence, and by abuse of
notation, we write $\rho(X,Y)$ to denote the distance between the law of $X$
and $Y$. Then

$\rho(\nabla F_{\infty}(X),Y)\leq\rho(\nabla F_{\infty}(X),\nabla
F_{n}(X))+\rho(\nabla F_{n}(X),Y)$

which means $\nabla F_{\infty}(X)$ and $Y$ have the same law, $\nu$. Then by
the claim, $\nabla F_{\infty}\equiv\nabla G$, and thus $\nabla
F_{n}\rightarrow\nabla G$ a.s. as $n\rightarrow\infty$.

∎

##  Appendix D Experimental Details

###  D.1 Architecture details

#### Initialization

As ICNNs have positive weights, its initialization has a different dynamics
than a standard feed-forward network. If not stated otherwise, all parameters
are initialized using standard PyTorch modules (Paszke et al., 2019). To
parameterize positive weights, we modify the weight parameters of a standard
linear layer with the softplus activation. We then divide all the weights by
the total number of incoming units, so that the average magnitude of each
hidden unit will not grow as the dimensionality of the previous layer
increases.

In addition, we reparameterize the CP-Flow $F_{\alpha}$ as $F_{w_{0},w_{1}}$
defined as

$F_{w_{0},w_{1}}=s(w_{0}){||x||^{2}}/{2}+s(w_{1})F(x)$

where $s$ is the regular (logistic) softplus. $w_{0}$ is initialized to be
$s^{-1}(1)$, and $w_{1}$ is initialized to be $0$ so $F_{w_{0},w_{1}}$ is
closer to the identity map.

Finally, we insert the ActNorm layer (Kingma & Dhariwal, 2018) everywhere
before an activation function is applied with the data-dependent
initialization.

#### Activation function

We use the following convexity-preserving operators when designing an ICNN:

  1. 1.

invariance under affine maps: $g\circ f$ is convex if $f$ is linear and $g$ is
convex

  2. 2.

non-negative weighted sums: $\sum_{j}w_{j}f_{j}$ is convex if $w_{j}$s are
non-negative and $f_{j}$s are convex

  3. 3.

composing with non-decreasing convex functions: $g\circ f$ is convex if $f$
and $g$ are both convex and $g$ is non-decreasing

Notably, in 1., we do not require $g$ to be non-decreasing. We experiment with
a symmetrized version of softplus $g(x)=s(x)-0.5x$ where $s$ is a softplus-
type activation, whenever $g$ is used as the first activation. This way, the
derivative of $g$ is $s^{\prime}(x)-0.5$, which ranges between $\pm 0.5$ and
behaves more like tanh (than sigmoid). This can be used for the first hidden
layer of a regular ICNN, or the augmented layer of the input-augmented ICNN.

We also experiment with an offset version of softplus, which is defined as
$g(x)=s(x)-s(0)$. This way the output of the softplus is more symmetric since
it can be negative.

###  D.2 Toy examples

For toy examples, we compute the log-determinant of the Jacobian in a
bruteforce manner. We use the Adam optimizer with an initial learning rate of
0.005. We create a data set following the toy distribution of size 50000, and
train each model for 50 epochs with a minibatch size of 128. For MAF and NAF
we cap the gradient norm to be 10 for stability. For CP-Flow, we use the
Gaussian softplus as activation, and symmetrize it at the first layers.

Data | N. flows | N. hidden layers | N. hidden units
---|---|---|---
One moon | 5 | 3 | 32
Eight Gaussians | 5 | 3 | 32
Rings | 5 | 5 | 256
Table 5: Architectural details for toy density estimation.

###  D.3 Approximating Optimal Coupling

For the OT map approximation experiment, we simulate data (of size $50,000$)
from a Gaussian distribution ${\mathcal{N}}(\mu,\Sigma)$ with a prior
$\mu\sim{\mathcal{N}}(0,I)$ and $\Sigma\sim{\mathcal{W}}(I,d+1)$, where
${\mathcal{W}}$ is the Wishart distribution. For CP-Flow, we use a network of
5 hidden layers of $64$ hidden units, with the Gaussian-softplus and zero-
offset. We use the Adam optimizer with a minibatch size of $128$, trained for
two epochs to generate the figures.

###  D.4 Density Estimation

Model | Power | Gas | Hepmass | Miniboone | BSDS300
---|---|---|---|---|---
Real NVP | -0.17$\,\pm\,$0.01 | -8.33$\,\pm\,$0.14 | 18.71$\,\pm\,$0.02 | 13.55$\,\pm\,$0.49 | -153.28$\,\pm\,$1.78
Glow | -0.17$\,\pm\,$0.01 | -8.15$\,\pm\,$0.40 | 18.92$\,\pm\,$0.08 | 11.35$\,\pm\,$0.07 | -155.07$\,\pm\,$0.03
FFJORD | -0.46$\,\pm\,$0.01 | -8.59$\,\pm\,$0.12 | 14.92$\,\pm\,$0.08 | 10.43$\,\pm\,$0.04 | -157.40$\,\pm\,$0.19
MADE | 3.08$\,\pm\,$0.03 | -3.56$\,\pm\,$0.04 | 20.98$\,\pm\,$0.02 | 15.59$\,\pm\,$0.50 | -148.85$\,\pm\,$0.28
MAF | -0.24$\,\pm\,$0.01 | -10.08$\,\pm\,$0.02 | 17.70$\,\pm\,$0.02 | 11.75$\,\pm\,$0.44 | -155.69$\,\pm\,$0.28
TAN | -0.48$\,\pm\,$0.01 | -11.19$\,\pm\,$0.02 | 15.12$\,\pm\,$0.02 | 11.01$\,\pm\,$0.48 | -157.03$\,\pm\,$0.07
NAF | -0.62$\,\pm\,$0.01 | -11.96$\,\pm\,$0.33 | 15.09$\,\pm\,$0.40 | 8.86$\,\pm\,$0.15 | -157.73$\,\pm\,$0.04
CP-Flow | -0.52$\,\pm\,$0.01 | -10.36$\,\pm\,$0.03 | 16.93$\,\pm\,$0.08 | 10.58$\,\pm\,$0.07 | -154.99$\,\pm\,$0.08
Table 6: Test negative log-likelihood (in nats) of tabular datasets in Papamakarios et al. (2017) for density estimation models (lower is better). Results for compared models are taken from Grathwohl et al. (2018). Average and standard deviation report over 3 random seeds at the best hyperparameters found by grid search. For CP-Flow trained on Gas and BSDS300, only one seed converged at the time of submission, so we report N/A for the standard deviation.  Dataset | N. flows | N. hidden layers | N. hidden units | N. Parameters
---|---|---|---|---
Power | 10 | 5 | 512 | 5,463,272
Gas | 5 | 5 | 512 | 2,757,276
Hepmass | 5 | 5 | 512 | 2,923,897
Miniboone | 2 | 5 | 256 | 379,232
BSDS300 | 10 | 5 | 256 | 2,152,456
Table 7: Best hyperparameters found in our search and the consequent total
number of model parameters.

For each dataset, we search via grid search for the best hyperparameter
configuration by calculating the negative log-likelihood in the validation set
using an exact bruteforce computation of the log-determinant of the Jacobian.
Initially, we focus our search at basic hyperparameters that influence the
total number of parameters of the model, such as the number of flow blocks,
the number of hidden layers of each convex potential block, and the number of
hidden units per layer of each block. After a first selection of candidate
hyperparameters after a constant number of training steps, we instantiate
extra experiments considering variations in the softplus-type activations we
use. We find that Gaussian-softplus and symmetrizing it in the activations of
the first layer help to the performance, oftentimes accelerating the training.
In table 7, we report the final hyperparameter combinations we have used for
the results presented in table 6.

###  D.5 Generative Modeling with Convolutional ICNN

On MNIST and CIFAR-10, we used a multiscale architecture. For MNIST, we had 8
CP-Flow blocks, followed by an invertible downsampling (Dinh et al., 2017),
followed by 8 CP-Flow blocks, another downsampling, and final 8 CP-Flow
blocks. Before every block was an ActNorm (Kingma & Dhariwal, 2018) layer. For
CIFAR-10, we had 2 CP-Flow blocks, followed by an invertible downsampling
(Dinh et al., 2017), followed by 2 CP-Flow blocks, downsampling, 2 CP-Flow
blocks, downsampling, and final 2 CP-Flow blocks. Before every CP-Flow block
was an ActNorm (Kingma & Dhariwal, 2018) layer. All ICNN architectures had 4
hidden layers and 64 hidden units wide. We averaged across the final spatial
dimensions to obtain a scalar output for the convex potential.

We test the invertibility of a CP-Flow model trained on CIFAR-10 on a set of
out-of-distribution data sets constructed by Behrmann et al. (2020) in Table
8. Notably, we do not suffer from the exploding inverse problem and can
reliably invert all data sets, either in or out of distribution. Samples of
reconstructed images are shown in Figure 6 which show no visual difference
between the original images and their reconstructions.

Figure 6: (top) Data samples. (bottom) Reconstruction after passing it through a CP-Flow. Dataset | Glow | ResFlow | CP-Flow
---|---|---|---
CIFAR-10 (in-dist) | 1.12E-6 | 5.16E-4 | 1.96E-3
Uniform | Inf | 3.04E-4 | 3.62E-3
Gaussian | Inf | 1.31E-4 | 5.97E-3
Rademacher | Inf | 3.43E-5 | 5.27E-3
SVHN (Netzer et al., 2011) | 9.94E-7 | 1.31E-3 | 2.53E-3
Texture (Cimpoi et al., 2014) | Inf | 3.66E-4 | 1.69E-3
Places (Zhou et al., 2017) | Inf | 5.31E-4 | 1.76E-3
tinyImageNet | Inf | 6.26E-4 | 1.65E-3
Table 8: Reconstruction RMSE. Results for Glow and ResFlow are taken from
Behrmann et al. (2020).

###  D.6 Amortizing ICNN for Variational Infernece

We use the partially input convex neural network from Amos et al. (2017) with
multiplicative conditioning. There are some other options for conditioning,
such as with the hypernetwork (Ha et al., 2016) or feature-wise transformation
(Dumoulin et al., 2018). We remove the non-linear path of the conditioned
variable beyong the first layer, since we found that it hurts training of the
VAE. We use the Gaussian-softplus for all the experiments, symmetrizing it on
the first layers. For the Caltech experiment, we also use the offset version,
and the softplus is initialized with a multiplicative constant of 2, which we
found leads to faster convergence.

Data | N. flows | N. hidden layers | N. hidden units
---|---|---|---
FreyFaces | 4 | 4 | 256
Omniglot | 2 | 2 | 512
Caltech | 8 | 4 | 256
Table 9: Architectural details for the VAE experiment. Model | Freyfaces | Omniglot | Caltech
---|---|---|---
No flow (Kingma & Welling, 2013) | $4.53\pm 0.02$ | $104.28\pm 0.39$ | $110.80\pm 0.74$
Planar (Rezende & Mohamed, 2015) | $4.40\pm 0.06$ | $102.65\pm 0.42$ | $109.66\pm 0.42$
IAF (Kingma et al., 2016) | $4.47\pm 0.05$ | $102.41\pm 0.04$ | $111.58\pm 0.38$
Sylvester (Van Den Berg et al., 2018) | $4.45\pm 0.04$ | $99.00\pm 0.04$ | $104.62\pm 0.29$
CP-Flow (Ours) | $4.47\pm 0.02$ | $102.06\pm 0.03$ | $106.53\pm 0.55$
CP-Flow aug (Ours) | $4.45\pm 0.03$ | $100.82\pm 0.30$ | $105.17\pm 0.57$
Table 10: Negative ELBO of VAE (lower is better). For Freyfaces the results
are in bits per dim. The numbers are averaged over three runs of experiments.
Standard deviation is presented in the section D.6. For CP-Flow with input-
augmented ICNN trained on Omniglot, only one seed converged at the time of
submission, so we report N/A for the standard deviation.

##  Appendix E Additional Ablation

Figure 7: Simulated signal-to-noise ratio and bias of the gradient estimator
for different atol values

We perform an analysis on the effect of changing the absolute error tolerance
level (atol) on the gradient estimator’s bias and variance. In a synthetic
setting, we sample a $d\times d$ (where $d=10$) positive definite matrix $H$
from ${\mathcal{W}}(I,2d)$ and scale it down by $1/d$ (so that the diagonal
entries do not grow as $d$ increases). We linearly interpolate atol values
between $0$ and $1$, and compare the error statistics of the gradient
estimator against the ground truth gradient of $\log\det H$ w.r.t $H$. The
error statistics are (1) the signal-to-noise ratio (expected value of the
estimator, or the ground truth gradient, divided by the standard deviation),
and (2) the absolute bias (expected value of the absolute difference between
the estimate and the ground truth). Since these quantities are intractable, we
use Monte Carlo estimate of 100K samples (draw 100K i.i.d. samples of the
gradient estimator) to smooth out the curves. The estimates are averaged
across all entries of $H$ and $10$ different random $H$’s. Results are
presented in fig. 7.

On a more realistic setting, we monitor the average number of CG iterates (hvp
calls), per-iteration time, as well as validation loss on the Miniboone
dataset. Meanwhile, we compare the vanilla ICNN, input-augmented ICNN, as well
as a dense version of ICNN (Huang et al., 2017). Figure 8 shows that as
tolerance value becomes smaller, we indeed need more hvp calls, which
eventually saturates at the dimensionality of the data. However, there is not
much difference in terms of log-likelihood if the tolerance is sufficiently
small: the validation loss oscillates and diverges for
$\textnormal{atol}=0.1$, while the curves for different atol values smaller
than $0.001$ are almost indistinguishable and fairly stable. On the other
hand, there is a noticeable difference in performance if we simply replace the
vanilla ICNN with the input-augmented ICNN or the dense ICNN. This suggests
_tuning the architecture of the convex potential is more crucial for improving
the overall performance_. This set of ablation studies were all performed on
P100 NVIDIA GPUs and with a constant batch size of 1024.

#### Runtime of directly backpropogating through Lanczos

Furthermore, we include an analysis on the effect of the number of eigenvalues
used in the stochastic Lanczos quadrature method (i.e. size of the tridiagonal
matrix) on training time. In this experiment we directly backpropagated
through the Lanczos estimate to compute the stochastic gradient (instead of
using the proposed gradient estimator). Figure 9 shows that the runtime is
much higher than the runtime of the proposed gradient estimator using CG
(bottom left of Figure 8). We note that the experiments with $m=5$ diverged,
possibly due to the error in estimation. This complements the memory profile
shown in Figure 2, and accentuates the lower runtime and memory requirement of
the proposed method by contrast.

Figure 8: Ablation of different ICNN architectures and absolute error
tolerance for conjugate gradient. Left: the average number of CG iterates (hvp
calls) per flow layer (top row) and the corresponding average time (in
seconds) per iteration (bottom row). Top right: validation set negative log-
likelihood (exact estimate). Notice that, for $\text{atol}=$110-7$$, CG
iterations cap at $43$ per flow layer; this is the dimensionality of the input
data in the Miniboone dataset. Bottom right: per-iteration time (in seconds)
averaged over all training steps. Figure 9: Ablation of the effect of number
of eigenvalues used in SLQ on average time (in seconds) per iteration.