C ONVEX P OTENTIAL F LOWS : U NIVERSAL P ROBABILITY D ISTRIBUTIONS WITH O PTIMAL T RANSPORT AND C ONVEX O PTIMIZATION
Anonymous authors
Paper under double-blind review
A BSTRACT
Flow-based models are powerful tools for designing probabilistic models with tractable density.
This paper introduces Convex Potential Flows (CP-Flow), a natural and efﬁcient 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 efﬁcient 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 constantmemory 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.
I NTRODUCTION
Normalizing ﬂows (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 efﬁcient 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 ﬂow?”
To gain a bit more intuition, we start from the one-dimension case.
If a 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 we are only allowed to move the probability mass continuously without ﬂipping 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 ﬂows from an optimal transport perspective, which allows us to deﬁne some notion of rearrangement cost (Villani, 2008).
It turns out, if we want the output of a ﬂow 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 ﬂows 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 ﬂow 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 efﬁciently estimate all quantities of interest.
In terms of the beneﬁts of parameterizing a ﬂow as a gradient ﬁeld, the convex potential is an R d → R function, which is different from most existing discrete-time ﬂows which are R d → 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 ﬂow’s parameterization, making CP-Flow a more ﬂexible general purpose density model.
B ACKGROUND : N ORMALIZING F LOWS AND O PTIMAL T RANSPORT
Normalizing ﬂows 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-ofvariable formula
[Equation]
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 ﬂexible 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 sufﬁces to show that the distribution induced by a normalizing ﬂow is universal.
Among many ways to establish distributional universality of ﬂow based methods ( e.g. Huang et al.
2018; 2020; 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 ) ∼ p Y if X ∼ p
X .
We seek to ﬁnd 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 ﬁnding the optimal transport map g that realizes the minimal expected cost
[Equation]
When the second moments of X and Y are both ﬁnite, 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):
Theorem 1 ( Brenier’s Theorem , Theorem 1.22 of Santambrogio (2015)) .
Let µ, ν be probability measures with a ﬁnite second moment, and assume µ has a Lebesgue density p X .
Then there exists a convex potential G such that the gradient map g :=
∇ G (deﬁned up to a null set) uniquely solves the Monge problem in eq.
(2) with the quadratic cost function c ( x, y ) =
|| x − y ||
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.
Our work sets itself apart from these prior works in that it is entirely likelihood-based, minimizing an empirical KL divergence as opposed to an approximate optimal transport cost.
C ONVEX P OTENTIAL F LOWS
Given a strictly convex potential F , we can deﬁne 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 deﬁnite.
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 ﬂows (3.3), and ﬁnally rigorously prove universality in the next section.
[Figure]
3. M ODELING
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.
to parameterize the convex potential, which has the following form
[Equation]
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 I NVERT ( F, y, CvxSolver ) 2:
Initialize x ← y 3: def closure (): 4:
Compute loss: l ← F ( x ) − y (cid:62) x 5: return l 6: x ← CvxSolver ( closure , x ) 7: return x
Invertibility and Inversion Procedure If the activation s is twice differentiable, then the Hessian
H F is positive semi-deﬁnite.
We can make it strongly convex by adding a quadratic term F α ( x ) =
α 2 || x ||
+ F ( x ) , such that H F α (cid:23) αI (cid:31) 0 .
This means the gradient map f α = ∇ F α is injective onto its image.
Furthermore, it is surjective since for any y ∈ R d , the potential x (cid:55)→ F α ( x ) − y (cid:62) x has a unique minimizer 1 satisfying the ﬁrst order condition ∇ F α ( x ) =
y , due to the strong convexity and differentiability.
We refer to this invertible mapping f
α as the convex potential ﬂow , or the CP-Flow.
The above discussion also implies we can plug in a black-box convex solver to invert the gradient map f α , 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 deﬁnite Jacobian matrix (as it is the
Hessian of the potential).
There exists numerous works on estimating spectral densities ( e.g. TalEzer & Kosloff, 1984; Silver & R¨oder, 1994; Han et al., 2018; Adams et al., 2018), of which this quantity is a special case.
See Lin et al.
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 α .
In particular, the log determinant can be rewritten in 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 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 ﬁrst constructing a symmetric tridiagonal matrix T ∈ R m × 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. O (1) -M EMORY U NBIASED ∇ log det H ESTIMATOR
[Figure]
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 ﬁg. 2).
Problem (i) is more speciﬁc to SLQ, because the gradient of an eigendecomposition is not deﬁned 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 efﬁcient 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 θ , we have the following identity:
[Equation]
Notably, in the last equality, we used the Hutchinson trace estimator (Hutchinson, 1989) with a
Rademacher random vector v , leading to a O (1) -memory, unbiased Monte Carlo gradient estimator.
Computing the quantity v (cid:62) 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
[Equation]
which has the unique minimizer z ∗
=
H − 1 v = v
[Equation]
H − 1 if H is symmetric positive deﬁnite.
Algorithm
S URROGATE O BJ ( F, x, CG )
Sample Rademacher random vector r 4: def hvp ( v ): 5: return v (cid:62) ∂∂x f ( x ) 6: z ← stop gradient ( CG ( hvp , r )) 7: return hvp ( z ) (cid:62) r
We use the conjugate gradient (CG) method, which is speciﬁcally designed for solving the unconstrained optimization problems in eq.
(4) with symmetric positive deﬁnite 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 ≤ 2 γ m ||
z 0 − z ∗
|| H , where z m is the estimate after m iterations.
The rate of convergence γ < 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 || ∞ <
τ is satisﬁed for some user-controlled tolerance.
Empirically, we ﬁnd 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: ddθ (cid:0) (
z m ) (cid:62) Hv (cid:1) , where only H is differentiated through (since z m is only used to approximate v (cid:62) H − 1 as a modiﬁer 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. C ONNECTION TO OTHER NORMALIZING FLOWS
Residual Flow For α = 1 , the gradient map f 1 resembles the residual ﬂow (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 sufﬁcient condition for invertibility.
In contrast, we enforce invertibility by using strongly convex potentials, which guarantees that the inverse of our ﬂow 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 beneﬁt 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 α = 1 , we have
F 1 =
|| x
||
+ L +2 ( s ( L 1 x ))
+ L 2 ( x ) .
Setting the weights of L 2 to zero, we have
[Equation]
We notice the above form bears a close resemblance to the Sylvester normalizing ﬂow (Van Den Berg et al., 2018) (with Q , R and (cid:101) R from Van Den Berg et al.
(2018) being equal to W (cid:62) 1 , diag ( w +2 ) and
I , respectively).
For the Sylvester ﬂow to be invertible, they require that R and (cid:101) 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 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 σ = s (cid:48) .
Then for the 1-dimensional case ( d = 1 ) and α = 0 (without the residual connection), we have
[Equation]
which is equivalent to the sigmoidal ﬂow 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 ﬂows (Kingma et al., 2016;
Huang et al., 2018; M ¨ uller 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`ere equation.
They do not require the potential to be convex since it is guaranteed to be invertible in the inﬁnitesimal setting of continuous normalizing ﬂows (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 ﬂows 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 ﬂow using scalar potentials.
T HEORETICAL A NALYSES
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 E
[ || x − f ( x ) || 2 ] .
Proofs of statements made in this section can be found in Appendices B and C.
To show (1), our ﬁrst step is to show that ICNNs can approximate arbitrary convex functions.
However, convergence of potential functions does not generally imply convergence of the gradient ﬁelds.
A classic example is the sequence F n = sin( nx ) / √ n and the corresponding derivatives f n = cos( nx ) √ n : F n → 0 as n → ∞ 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.
Theorem 2.
Let F n : R d
→ R be differentiable convex functions and G : R d → R be a proper convex function.
Assume F n → G .
Then for almost every x ∈ R d , G is differentiable and f n ( x ) :=
∇ F n ( x ) → ∇ 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.
Theorem 3 ( Universality ) .
Given random variables X ∼ µ and Y ∼ ν , with µ being absolutely continuous w.r.t.
the Lebesgue measure, there exists a sequence of ICNN F n with a softplus-type activation, such that ∇ F n ◦ X
→ Y in distribution.
N.B.
In the theorem we do not require the second moment to be ﬁnite, 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 ﬁnite 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.
Theorem 4 ( Optimality ) .
Let G be the Brenier potential of X ∼ µ and Y ∼ ν , and let F n be a convergent sequence of differentiable, convex potentials, such that ∇ F n ◦ X → Y in distribution.
Then ∇ F n converges almost surely to ∇ 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.
E XPERIMENT
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 difﬁculties in optimization and limit the effective expressivity of the model.
We propose an architectural enhancement of ICNN, deﬁned as follows (note the change in notation: instead of writing the pre-activations z , we use h to denote the activated units):
[Equation]
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. T OY EXAMPLES
Having distributional universality for a single ﬂow layer means that we can achieve high expressiveness without composing too many ﬂows.
We demonstrate this by ﬁtting the density on some
[Table]
data: x
[Equation]
( x ) z
[Equation]
) d =d
[Equation]
[Figure]
We compare with the masked autoregressive ﬂow (MAF, Papamakarios et al.
and the neural autoregressive ﬂow (NAF, (Huang et al., 2018)).
Results are presented in ﬁg. 3.
We try to match the network size for each data.
All models ﬁt the ﬁrst data well.
As afﬁne couplings cannot split probability mass, MAF fails to ﬁt to the second and third datasets 2 .
Although the last dataset is intrinsically harder to ﬁt (as NAF, another universal density model, also fails to ﬁt it well), the proposed method still manages to learn the correct density with high ﬁdelity.
5. A PPROXIMATING OPTIMAL COUPLING
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 ﬂow against J ||
x − y ||between the Gaussian data and the standard Gaussian prior (as there is a closed-form expression).
In ﬁg. 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 ﬂow (Kingma et al., 2016), which has the capacity to represent the multivariate Gaussian density, yet it does not learn the optimal coupling.
5. D ENSITY ESTIMATION
We demonstrate the efﬁcacy 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.
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 ﬂow blocks, the width and depth of the ICNN potentials.
See appendix 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 We also perform density estimation on standard image data sets, MNIST and CIFAR10.
For this, we use convolutional layers in place of fully connected layers.
Prior works have had
[Table] [Table]
to use large architectures, with many ﬂow blocks composed together, resulting in a large number of parameters to optimize.
While we also compose multiple blocks of CP-Flows
[Table]
we ﬁnd 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
[Table]
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 ﬂow blocks with the corresponding baseline.
We ﬁnd that on CIFAR-10, the baseline ﬂow models do not perform nearly as well as CP-Flow.
We believe this may be because CP-Flows are universal with just one ﬂow block
[Table]
whereas coupling and invertible residual blocks are limited in expressivity or Lipschitz-constrained.
5. A MORTIZING ICNN FOR V ARIATIONAL I NFERENCE
[Table]
Normalizing ﬂows also allow us to employ a larger, more ﬂexible 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.
and use the output of the encoder as the additional input for conditioning.
As table 3 shows, the performance of CPFlow 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.
C ONCLUSION
We propose a new parameterization of normalizing ﬂows using the gradient map of a convex potential.
We make connections to the optimal transport theory to show that the proposed ﬂow is a universal density model and leverage tools from convex optimization to enable efﬁcient training and model inversion.
Experimentally, we show that the proposed method works reasonably well when evaluated on standard benchmarks.
We also demonstrated that the performance can be improved by designing better ICNN architectures.
A S OFTPLUS T YPE A CTIVATION
In this section, we let r ( x ) = max(0 , x ) be the ReLU activation function.
Deﬁnition 1.
We say a function s is of the softplus type if the following holds
(a) s ≥ r (b) s is convex (c) | s ( x ) − r
( x ) | → 0 as | x | → ∞
Note that a softplus-type activation function is necessarily continuous, non-decreasing, and uniformly approximating ReLU in the following sense:
[Equation]
uniformly for all x ∈ R as a → ∞ .
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 ) =
(cid:82) x
−∞ F p ( y ) dy , where F p is the distribution function of p , and s is at least twice differentiable.
Proof.
We ﬁrst prove a claim (i): xF p ( x ) → 0 as x → −∞ .
First, for x ≤ 0 ,
[Equation]
Since 1 y ≤ x yp ( y ) → 0 as x → −∞ and | 1 y ≤ x yp ( y ) | ≤ | 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 deﬁnition, since (cid:82) yp ( y ) dy =
[Equation]
where we’ve used claim (i) to evaluate xF p ( x ) as x → −∞ .
On the other hand, integration by part implies
[Equation]
Twice differentiability follows from the differentiability of F p .
a
Now since r ( x ) is convex, Jensen’s inequality gives
[Equation]
(b) s is convex because s (cid:48)
=
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 | → ∞ (by the dominated or monotone convergence theorem).
It sufﬁces to show xF p ( x ) goes to 0 as x → −∞ , which is just claim (i), and x − xF p ( x ) > 0 goes to 0 as x → ∞ .
To show the latter, we can rewrite x − xF p ( x ) =
x (1 − F p ( x ))
= (cid:82) ∞ x 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 ﬁg. 5.
We experimented with the Gaussian-softplus and the logistic-softplus.
[Equation]
[Table]
[Equation]
Laplace
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
B U NIVERSALITY P ROOF
Notation: Given a convex set Ω ⊆ R d , we let C (Ω) denote the set of continuous functions on Ω , and C × (Ω) :=
{ f ∈ C (Ω) : f is convex } denote the set of convex, continuous functions.
We ﬁrst show that ICNNs with a suitable activation function are dense in C × .
A similar result can be found in Chen et al.
(2019b), where they use a different constructive proof: ﬁrst show that piecewise maximum of afﬁne 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 afﬁne functions, are a dense subset of C × .
Proposition 2.
Pointwise maximum of afﬁne functions is dense in C × ([0 , 1] d ) .
Proof.
Fix some (cid:15) > 0 .
Since f ∈ C × ([0 , 1] d ) is uniformly continuous on [0 , 1] d , there exists some δ > 0 such that | f ( x ) − f ( y ) | < (cid:15) provided that || x − y || <
δ .
Let n be big enough such that 2 − n < δ , and let X be the set of points whose coordinates sit on i 2 − n for some 1 ≤ i ≤ 2 n − 1 (i.e. there are |X| =
(2 n − 1) d points in X ).
For each y ∈ X , let L y ( x ) :=
∇ f ( y ) (cid:62) ( x − y )
+ f ( y ) be a supporting hyperplane of the graph of f , where ∇ f ( y ) is a subgradient of f evaluated at y .
Then we have a convex approximation f (cid:15) ( x ) := max y ∈X L y ( x ) which bounds f from below.
Moreover, letting y | x := arg min y ∈X || x − y
|| , we have (for x (cid:54)∈ X ),
[Equation]
Since (cid:15) 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 C × ([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 (cid:62) j x + b j , which can then be reduced to
[Equation]
where z j := r
( z j − 1 )
+ L (cid:48)
j for 2 ≤ j ≤ k , z
1 = L 1 − L 2 , L (cid:48)
j :=
L j − L j +1 for 2 ≤ j ≤ k − 1 , and
L (cid:48)
k :
= L k .
Since by Proposition 2, pointwise maximum of afﬁne functions is dense in C × , so is ICNN with the ReLU activation function.
The same holds for softplus since softplus can be used to uniformly approximate ReLU.
Theorem 2.
Let F n : R d
→ R be differentiable convex functions and G : R d → R be a proper convex function.
Assume F n → G .
Then for almost every x ∈ R d , G is differentiable and f n ( x ) :=
∇ F n ( x ) → ∇ G ( x ) =:
g ( x ) .
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
[Equation]
where
[Equation]
The problem can be rephrased as proving
[Equation]
Note that f nm is non-decreasing in m since F n is convex, and thus f nm ≤ f n .
Since f nm converges to f n , for any (cid:15) > 0 , we can ﬁnd an integer µ ( (cid:15), n ) such that for all m ≥ µ ( (cid:15), n ) , | f nm − f n | ≤ (cid:15) .
Let m k be a subsequence of { m ≥ 1 } deﬁned as m k = µ (2 − k , n ) .
Then | f nm k +1 − f nm k | ≤ 2 − k , which is integrable wrt the counting measure on positive integers k , since
[Equation]
Thus, letting f nm 0 = 0 , by the Dominated Convergence Theorem, we have
[Equation]
Although we are only looking at the limit of the subsequence m k , this is sufﬁcient 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 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.
Theorem 3 ( Universality ) .
Given random variables X ∼ µ and Y ∼ ν , with µ being absolutely continuous w.r.t.
the Lebesgue measure, there exists a sequence of ICNN F n with a softplus-type activation, such that ∇ F n ◦ X
→ Y in distribution.
Proof.
Assume µ
and ν have ﬁnite second moments.
Since µ is absolutely continuous, by Brenier’s theorem , there exists a convex function G : R d → R such that ∇ G ( X ) d
= ν (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 2 then implies the gradient map f n := ∇ F n converges to ∇ G pointwise almost everywhere.
This implies the weak convergence of the pushforward measure of f n ◦ X .
Now remove the ﬁnite second moment assumption and let X and Y be random variables distributed according to µ (with Lebesgue density p ) and ν , respectively.
Denote by B k a ball of radius k >centered at the origin, i.e. B n := { x : || x || ≤ k } .
Let X k
= X 1 X ∈ B k
+ U k 1 X (cid:54)∈ B k , where
U k is an independent random variable distributed uniformly on B k , and let µ k be the law of X k .
Then X k → X almost surely as k → ∞ , and µ k is still absolutely continuous wrt the Lebesgue measure, with its density being p ( x ) + 1 vol( B k ) µ ( || X
|| > k )
if || x || ≤ k or 0 otherwise.
Let
Y k and ν k be deﬁned similarly (while ν k may not be absolutely continuous wrt Lebesgue).
From above, since X k and Y k are bounded and admit a ﬁnite second moment, we know ﬁxing k , we can ﬁnd a sequence of f k,n = ∇ F k,n such that f k,n ◦ X k
→ Y k in distribution as n → ∞ .
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 ◦ X k and ν k is at most 1 /k .
An application of triangle inequality of the weak metric implies f k,n k ◦ X k
→ Y in distribution as k → ∞ .
Finally, note that since {|| f k,n k ◦ X k − f k,n k ◦ X || > (cid:15) } ⊆ {|| X
|| > k } , by monotonicity, we have
[Equation]
and thus || f k,n k ◦ X k − f k,n k ◦ X ||
→ 0 almost surely (where P is the underlying probability measure of the measure space).
By Lemma 1 of Huang et al.
(2020), f k,n k ◦ X converges in distribution to
Y .
C O PTIMALITY P ROOF
Theorem 4 ( Optimality ) .
Let G be the Brenier potential of X ∼ µ and Y ∼ ν , and let F n be a convergent sequence of differentiable, convex potentials, such that ∇ F n ◦ X → Y in distribution.
Then ∇ F n converges almost surely to ∇ G .
Proof.
We claim that if F is a convex potential such that Z = ∇ F ( X ) has ν as its law, then ∇ F ≡ ∇ G almost surely.
The proof of the claim is originally due to R¨uschendorf & Rachev (1990), but we present it here for completeness.
Let Z (cid:48) be another random variable distributed by ν .
Then by the Fenchel-Young inequality (applied to the convex potential F ),
[Equation]
This concludes the proof since ∇ G uniquely solves the transportation problem, which is equivalent to ﬁnding a transport map ˜ g that maximizes the covariance:
[Equation]
Let F ∞ to be the pointwise limit of F n .
Then for any x 1 , x 2 and t ∈ [0 , 1] ,
[Equation]
That is, F ∞ is convex.
Now since F n is a convergent sequence of convex functions, its gradient ∇ F n also converges pointwise almost everywhere to ∇ F ∞ by Theorem 2.
Let ρ denote the Prokhorov metric, which metrizes the weak convergence, and by abuse of notation, we write ρ ( X, Y ) to denote the distance between the law of X and Y .
Then
[Equation]
which means ∇ F ∞ ( X ) and Y have the same law, ν .
Then by the claim, ∇ F ∞ ≡ ∇ G , and thus ∇ F n → ∇ G a.s. as n → ∞ .
D E XPERIMENTAL D ETAILS D. A RCHITECTURE 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 α as F w 0 ,w 1 deﬁned as
[Equation]
where s is the regular (logistic) softplus.
w 0 is initialized to be s − 1 (1) , and w 1 is initialized to beso 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. invariance under afﬁne maps: g ◦ f is convex if f is linear and g is convex 2.
non-negative weighted sums: (cid:80) j w j f j is convex if w j s are non-negative and f j s are convex 3.
composing with non-decreasing convex functions: g ◦ 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 .
5 x where s is a softplus-type activation, whenever g is used as the ﬁrst activation.
This way, the derivative of g is s (cid:48)
( x ) − 0 . 5 , which ranges between ± 0 .
5 and behaves more like tanh (than sigmoid).
This can be used for the ﬁrst 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 deﬁned as g ( x ) =
s ( x ) − s (0) .
This way the output of the softplus is more symmetric since it can be negative.
D. T OY 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 ﬁrst layers.
[Table]
D. A PPROXIMATING O PTIMAL C OUPLING
For the OT map approximation experiment, we simulate data (of size 50 , 000 ) from a Gaussian distribution N ( µ, Σ) with a prior µ ∼ N (0 , I ) and Σ ∼ W ( I, d + 1) , where W is the Wishart distribution.
For CP-Flow, we use a network of 5 hidden layers of 64 hidden units, with the Gaussiansoftplus and zero-offset.
We use the Adam optimizer with a minibatch size of 128 , trained for two epochs to generate the ﬁgures.
D. D ENSITY E STIMATION
[Table] [Table]
For each dataset, we search via grid search for the best hyperparameter conﬁguration by calculating the negative log-likelihood in the validation set using an exact bruteforce computation of the logdeterminant of the Jacobian.
Initially, we focus our search at basic hyperparameters that inﬂuence the total number of parameters of the model, such as the number of ﬂow blocks, the number of hidden layers of each convex potential block, and the number of hidden units per layer of each block.
After a ﬁrst 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 ﬁnd that
Gaussian-softplus and symmetrizing it in the activations of the ﬁrst layer help to the performance, oftentimes accelerating the training.
In table 7, we report the ﬁnal hyperparameter combinations we have used for the results presented in table 6.
[Table]
Netzer
[Table] [Table]
Texture (
[Table] [Table]
Places (Zhou et al
[Table] [Table]
[Table]
D.
G ENERATIVE M ODELING WITH C ONVOLUTIONAL 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 ﬁnal 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 ﬁnal 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 ﬁnal 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.
D. A MORTIZING ICNN FOR V ARIATIONAL I NFERNECE
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 ﬁrst layer, since we found that it hurts training of the VAE.
We use the Gaussian-softplus for all the experiments, symmetrizing it on the ﬁrst 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.
[Table] [Table]
E A DDITIONAL A BLATION
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 × d (where d = 10 ) positive deﬁnite matrix H from W ( I, 2 d ) 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 ﬁg. 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 sufﬁciently small: the validation loss oscillates and diverges for 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 proﬁle shown in Figure 2, and accentuates the lower runtime and memory requirement of the proposed method by contrast.
[Figure]