O PTIMAL N EURAL N ETWORK A PPROXIMATION OF W ASSERSTEIN G RADIENT D IRECTION VIA C ONVEX O PTIMIZATION
Anonymous authors
Paper under double-blind review
A BSTRACT
The computation of Wasserstein gradient direction is essential for posterior sampling problems and scientiﬁc computing.
The approximation of the Wasserstein gradient with ﬁnite samples requires solving a variational problem.
We study the variational problem in the family of two-layer networks with squared-ReLU activations, towards which we derive a semi-deﬁnite programming (SDP) relaxation.
This SDP can be viewed as an approximation of the Wasserstein gradient in a broader function family including two-layer networks.
By solving the convex SDP, we obtain the optimal approximation of the Wasserstein gradient direction in this class of functions.
We also propose practical algorithms using subsampling and dimension reduction.
Numerical experiments including PDEconstrained Bayesian inference and parameter estimation in COVID-19 modeling demonstrate the effectiveness and efﬁciency of the proposed method.
I NTRODUCTION
Bayesian inference plays an essential role in learning model parameters from the observational data with applications in inverse problems, scientiﬁc computing, information science, and machine learning (Stuart, 2010).
The central problem in Bayesian inference is to draw samples from a posterior distribution, which characterizes the parameter distribution given data and a prior distribution.
The Wasserstein gradient ﬂow (Otto, 2001; Ambrosio et al., 2005; Junge et al., 2017) has shown to be effective in drawing samples from a posterior distribution, which attracts increasing attention in recent years.
For instance, the Wasserstein gradient ﬂow of Kullback-Leibler (KL) divergence connects to the overdampled Langevin dynamics.
The time-discretization of the overdamped Langevin dynamics renders the classical Langevin Monte Carlo Markov Chain (MCMC) algorithm.
In this sense, the computation of Wasserstein gradient ﬂow yields a different viewpoint for sampling algorithms.
In particular, the Wasserstein gradient direction also provides a deterministic update of the particle system (Carrillo et al., 2021b).
Based on the approximation or generalization of the
Wasserstein gradient direction, many efﬁcient sampling algorithms have been developed, including
Wasserstein gradient descent (WGD) with kernel density estimation (KDE) (Liu et al., 2019), Stein variational gradient descent (SVGD) (Liu & Wang, 2016), and neural variational gradient descent (di Langosco et al., 2021), etc.
Meanwhile, neural networks exhibit tremendous optimization and generalization performance in learning complicated functions from data.
They also have wide applications in Bayesian inverse problems (Rezende & Mohamed, 2015; Onken et al., 2020; Kruse et al., 2019; Lan et al., 2021).
According to the universal approximation theorem of neural networks (Hornik et al., 1989; Lu et al., 2017), any arbitrarily complicated functions can be learned by a two-layer neural network with nonlinear activations and a sufﬁcient number of neurons.
Functions represented by neural networks naturally provide an approximation towards the Wasserstein gradient direction.
However, due to the nonlinear and nonconvex structure of neural networks, optimization algorithms including stochastic gradient descent may not ﬁnd the global optima of the training problem.
Recently, based on a line of works (Pilanci & Ergen, 2020; Sahiner et al., 2020; Bartan & Pilanci, 2021a), the regularized training problem of two-layer neural networks with ReLU/polynomial activation can be formulated as a convex program.
Indeed, by solving the convex program, we can construct the entire set of global optima of the nonconvex training problem (Wang et al., 2020). Theoretical analysis (Wang et al., 2022) shows that global optima of the training problem correspond to the simplest models with good generalization properties.
Moreover, numerical results (Pilanci & Ergen, 2020) show that neural networks found by solving the convex program can achieve higher train accuracy and test accuracy compared to neural networks trained by SGD with the same number of parameters.
In this paper, we study a variational problem, whose optimal solution corresponds to the Wasserstein gradient direction.
Focusing on the family of two-layer neural networks with squared ReLU activation, we formulate the regularized variational problem in terms of samples.
Directly training the neural network to minimize the loss may get the neural network stuck at local minima or saddle points and it often leads to biased sample distribution from the posterior.
Instead, we analyze the convex dual problem of the training problem and study its semi-deﬁnite program (SDP) relaxation by analyzing the geometry of dual constraints.
The resulting SDP can be efﬁciently solved by convex optimization solvers such as CVXPY (Diamond & Boyd, 2016).
We then derive the corresponding relaxed bidual problem (dual of the relaxed dual problem).
Thus, the optimal solution to the dual problem yields an optimal approximation of the Wasserstein gradient direction in a broader function family.
We also analyze the choice of the regularization parameter and present a practical implementation using subsampling and parameter dimension reduction to improve computational efﬁciency.
Numerical results for experiments including PDE-constrained inference problems and
Covid-19 parameter estimation problems illustrate the effectiveness and efﬁciency of our method.
1. R ELATED WORKS
The time and spatial discretizations of Wasserstein gradient ﬂows are extensively studied in literature (Jordan et al., 1998;
Junge et al., 2017; Carrillo et al., 2021a;b; Bonet et al., 2021; Liutkus et al., 2019; Frogner & Poggio, 2020).
Recently, neural networks have been applied in solving or approximating Wasserstein gradient ﬂows (Mokrov et al., 2021; Lin et al.
, 2021b;a; Alvarez-Melis et al., 2021;
Bunne et al., 2021; Hwang et al., 2021; Fan et al., 2021).
For sampling algorithms, di Langosco et al.
(2021) learns the transportation function by solving an unregularized variational problem in the family of vector-output deep neural networks.
Compared to these studies, we focus on a convex SDP relaxation of the varitional problem induced by the Wasserstein gradient direction.
Meanwhile, Feng et al.
(2021) form the Wasserstein gradient direction as the mininimizer the
Bregman score and they apply deep neural networks to solve the induced variational problem.
In comparison to previous works on the convex optimization formulations of neural networks using
SDP (Bartan & Pilanci, 2021a;b), they focus on the polynomial activation and give the exact convex optimization formulation (instead of convex relaxation).
In comparison, we focus on the neural networks with the squared ReLU activation, which has not been considered before.
Our method can also apply to the analysis of supervised learning problem using squared ReLU activated neural networks.
B ACKGROUND
In this section, we brieﬂy review the Wasserstein gradient descent and present its variational formulation.
In particular, we focus on the Wasserstein gradient descent direction of KL divergence functional.
Later on, we design a neural network convex optimization problem to approximate the
Wasserstein gradient in samples.
2. W ASSERSTEIN GRADIENT DESCENT
Consider an optimization problem in the probability space:
[Equation]
Here the integral is taken over R d and the objective functional D KL ( ρ (cid:107) π ) is the KL divergence from ρ to π .
The variable is the density function ρ in the space P = { ρ ∈ C ∞ ( R d ) | (cid:82) ρdx
= 1 , ρ > 0 } .
The function π ∈ C ∞ ( R d ) is a known probability density function of the posterior distribution.
By solving the optimization problem (1), we can generate samples from the posterior distribution.
A known fact (Villani, 2003, Chapter 8.3.1) is that the Wasserstein gradient descent ﬂow for the optimization problem (1) satisﬁes
[Equation]
where ρ t ( x ) =
ρ ( x, t ) , δδρ t is the L 2 ﬁrst variation operator w.r.t.
ρ t , ∇ · F denotes the divergence of a vector valued function F : R d → R d and ∆ is the Laplace operator.
In step (a) we uses the fact that ρ t ∇ log ρ t = ∇ ρ t .
This equation is also known as the gradient drift Fokker-Planck equation.
It corresponds to the following updates in terms of samples:
[Equation]
where x t follows the distribution of ρ t .
Clearly, when ρ t = π , the above dynamics reach the equilibrium, which implies that the samples x t are generated by the posterior distribution.
To solve the Wasserstein gradient ﬂow (2), we consider a forward Eulerian discretization in time.
In the l -th iteration, suppose that { x nl } are samples drawn from ρ l .
The update rule of Wasserstein gradient descent (WGD) on the particle system { x nl } follows
[Equation]
where Φ l : R d → R is a function which approximates log ρ l − log π and α l > 0 is the step size.
V ARIATIONAL FORMULATION OF WGD
Given the particles { x n } Nn =1 , we design the following variational problem to choose a suitable function Φ approximating the function log ρ − log π .
Consider
[Equation]
The objective functional evaluates the least-square discrepancy between ∇ log ρ − ∇ log π and ∇ Φ weighted by the density ρ .
The optimal solution follows Φ = log ρ − log π , up to a constant shift.
Let
H ⊆ C 1 ( R d ) be a ﬁnite dimensional function space.
The following proposition gives a formulation of (4) in H .
Proposition 1 Let H ⊆
C 1 ( R d ) be a function space.
The variational problem (4) in the domain H can be reformulated to
[Equation]
Remark 1 A similar variational problem has been studied in (di Langosco et al., 2021).
If we replace ∇ Φ for Φ ∈ H by a vector ﬁeld Ψ in certain function family, then, the quantity in (5) is the negative regularized Stein discrepancy deﬁned in (di Langosco et al., 2021) between ρ and π based on Ψ .
This problem is also similar to the varitional problem for the score matching estimator in (Hyv¨arinen & Dayan, 2005) by parameterizing Φ in a given probabilistic model.
In comparison, our method can be viewed as a special case of score matching by using a two-layer neural network.
Therefore, by replacing the density ρ by ﬁnite samples { x n } Nn =1 ∼ ρ , the problem (5) in terms of ﬁnite samples forms
[Equation]
O PTIMAL NEURAL NETWORK APPROXIMATION OF W ASSERSTEIN GRADIENT
In this section, we focus on functional space H of functions represented by two-layer neural networks.
We derive the primal and dual problem of the regularized Wasserstein variational problems.
By analyzing the dual constraints, a convex SDP relaxation of the dual problem is obtained.
We also present a practical implementation estimation of ∇ log ρ − ∇ log π and discuss the choice of the regularization parameter.
Let ψ be an activation function.
Consider the case where H is a class of two-layer neural network with the activation function ψ ( x ) :
[Equation]
where θ =
( W, α ) is the parameter in the neural network with W ∈ R d × m and α ∈ R m .
Remark 2 We can extend this model to handle the bias term by add an entry of 1 in x 1 , . . .
, x n .
For two-layer neural networks, we can compute the gradient and Laplacian of Φ ∈ H as follows:
[Equation]
Here ◦ represents the element-wise multiplication.
By adding a regularization term to the variational problem (6), we obtain
[Equation]
where β > 0 is the regularization parameter.
We focus on the squared ReLU activation ψ ( z ) =
( z ) 2+ = (max { z, 0 } ) 2 .
Note that a non-vanishing second derivative is required for the Laplacian term in (9), which makes the ReLU activation inadequate.
For this activation function, we consider the regularization function R ( θ ) =
[Equation] mi =
[Equation]
32 + | α
i | 3 ) .
Remark
We note that ∇ Φ θ ( x ) and ∆Φ θ ( x ) are all piece-wise degree- 3 polynomials of the parameters θ .
Hence, we consider a speciﬁc cubic regularization term above, analogous to (Bartan &
Pilanci, 2021a).
By choosing this regularization term, we can derive a simpliﬁed dual problem.
By utilizing the arithmetic and geometric mean (AM-GM) inequality, we can rescale the ﬁrst and second-layer parameters and formulate the regularized variational problem (10) as follows.
Proposition 2 (Primal problem)
The regularized variational problem (10) can be reformulated to
[Equation]
where ˜ β
= 3 · 2 − 5 / 3 Nβ
[-Remark 
-]and we denote [ m ] = { 1 , . . . , m } .
In short, the optimal value of (10) and (11) are the same.
We can obtain the optimal solution of (11) by rescaling the optimal solution of (10) and vice versa.
For simplicity, we write Y ∈ R N × d whose n -row is ∇ log π ( x n ) for n ∈ [ N ] .
We introduce the slack variable z n = (cid:80) mi =
α i w i ψ (cid:48) ( x Tn w i ) for n ∈ [ N ] and denote Z =
[ z. . .
z N ] T ∈ R N × d .
Then, we can simplify the problem (11) to
[Equation]
To derive the convex relaxtion of neural network training problem, the dual problem plays an import role.
By applying the Lagrangian duality, we can derive the dual problem of (12) as follows.
Proposition 3 (Dual problem)
The dual problem of the regularized variational problem (12) is
[Equation]
which provides a lower-bound on (12) .
We note that the dual problem can be infeasible if the regularization parameter ˜ β is below certain threshold.
In other words, if the regularization term is missing or the regularization parameter is not large enough, the optimal value of the dual problem is −∞ and the primal problem is not lower bounded.
3. A NALYSIS OF DUAL CONSTRAINTS AND THE RELAXED DUAL PROBLEM
Now, we analyze the constraint in the dual problem.
We note that it is closely related to the regularization parameter, which we will discuss later.
For simplicity, we take ψ (cid:48)(cid:48) (0) =
as the subgradient of ψ (cid:48) ( z ) at z = 0 , i.e., taking the left derivative of ψ (cid:48) ( z ) at z = 0 .
Let X =
[ x 1 , . . .
, x N ] T ∈ R N × d .
Denote the set of all possible hyper-plane arrangements corresponding to the rows of X as
[Equation]
Here I ( s ) = 1 if the statement s is correct and I ( s ) = 0 otherwise.
Let p = |S| be the cardinality of S , and write S
=
{ D 1 , . . .
, D p } .
According to (Cover, 1965), we have the upper bound p ≤ 2 r (cid:16) e ( N − 1) r (cid:17) r , where r = rank ( X ) .
Based on the analysis of the dual constraints, we can derive a convex SDP as a relaxed dual problem.
Proposition 4 (Relaxed dual problem) The relaxed dual problem is the following SDP:
[Equation]
) where we denote [ p ] = { 1 , . . .
, p } .
For j ∈ [ p ] , we denote A j (Λ) = − Λ T D
j X − X T D j Λ , B j = 2 tr( D j ) I d , ˜
A j (Λ) =
(cid:20)
A j (Λ)  (cid:21) , ˜ B
j = (cid:20) B
j   (cid:21) , H ( j )
= (cid:20) I
d  −(cid:21) and H ( j ) n = [Equation] [Equation] n
[Equation] nn ) x Tn (cid:21) , n ∈ [ N ] The vector e d +1 ∈ R d +1 satisﬁes that ( e d +1 ) i =
for i ∈ [ d ] and ( e d +1 ) d +1 = 1 .
The optimal value of (15) gives a lower bound on the dual problem (13) , and hence on the primal problem (12) .
The relaxed bi-dual problem provides insights on approximating the primal problem via convex optimization, which is derived as follows.
As an equivalent formulation of the convex dual problem (15), it can be viewed as a convex relaxation of the primal problem (12).
Proposition 5 (Relaxed bi-dual problem)
The dual of the relaxed dual problem (15) is as follows
[Equation]
Here A ∗ j is the adjoint operator of the linear operator A j .
As (15) is a convex problem and the Slater’s condition is satisﬁed, the optimal values of (15) and (16) are same.
The bi-dual problem (16) is closely related to the primal problem (12).
Indeed, any feasible solutions of the primal problem (11) can be mapped to feasible solutions of (16).
We note that the mapping from the primal solution to the bi-dual solution cannot go both ways, unless these two problems are equivalent.
Theorem 1 Suppose that ( Z, W, α ) is feasible to the primal problem (12) .
Then, there exist matrices { S ( j, +) , S ( j, − ) } pj =1 constructed from ( W, α ) such that ( Z, { S ( j, +) , S ( j, − ) } pj =1 ) is feasible to the relaxed bi-dual problem (16) .
Moreover, the objective value of the relaxed bi-dual problem (16) at ( Z, { S ( j, +) , S ( j, − ) } pj =1 ) is the same as objective value of the primal problem (12) at ( Z, W, α ) .
Let J ( Z, { S ( j, +) , S ( j, − ) } pj =1 ) denote the objective value of the relaxed bi-dual problem (16) at a feasible solution ( Z, { S ( j, +) , S ( j, − ) } pj =1 ) .
Let ( Z ∗ , W ∗ , α ∗ ) denote a globally optimal solution of the primal problem (12).
By Theorem 1, there exist matrices { S ( j, +) , S ( j, − ) } pj =1 such that ( Z ∗ , { S ( j, +) , S ( j, − ) } pj =1 ) is a feasible solution of the relaxed bi-dual problem (16) and
J ( Z ∗ , { S ( j, +) , S ( j, − ) } pj =1 ) is the same as the objective value of (12) at its global minimum ( Z ∗ , W ∗ , α ∗ ) .
On the other hand, let ( ˜ Z ∗ , { ˜ S ( j, +) ,
˜ S
( j, − ) }
pj =1 ) denote an optimal solution of the relaxed bi-dual problem (16).
From the optimality of ( ˜ Z ∗ , { ˜ S ( j, +) , ˜ S ( j, − ) } pj =1 ) , we have
[Equation]
Note that at ( Z ∗ , W ∗ , α ∗ ) we obtain the optimal approximation of ∇ log ρ −∇ log π at x 1 , . . .
, x N in the family of two-layer squared-ReLU networks (7).
Smaller or equal objective value of the relaxed bi-dual problem (16) can be achieved at ( ˜ Z ∗ , { ˜ S ( j, +) , ˜ S
( j, − ) }
pj =1 ) than at ( Z ∗ , { S ( j, +) , S ( j, − ) } pj =1 ) .
Therefore, we can view ˜
Z ∗ gives an optimal approximation of ∇ log ρ − ∇ log π evaluated on x 1 , . . .
, x N in a broader function family including the two-layer squared ReLU neural networks.
From the derivation of the relaxed bi-dual problem, we have the relation ˜
Z ∗ = − Λ ∗ − Y , where (Λ ∗ , { r ( j, +) , r ( j, − ) ) is optimal to the relaxed dual problem (15) and ( ˜ Z ∗ , { ˜ S ( j, +) , ˜ S ( j, − ) } pj =1 ) is optimal to the relaxed bi-dual problem (16).
Therefore, by solving Λ ∗ from the relaxed dual problem (15), we can use − Λ ∗ − Y as the approximation of ∇ log ρ − ∇ log π evaluated on x 1 , . . .
, x N .
Remark
We note that solving the proposed convex optimization problem 15 renders the approximation of the Wasserstein gradient direction.
Compared to the two-layer ReLU networks, it induces a broader class of functions represented by { S ( j, +) , S ( j, − ) } pj =1 .
This contains more variables than the neural network function.
3. P RACTICAL IMPLEMENTATION
Although the number p of all possible hyper-plane arrangements is upper bounded by 2 r (( N − 1) e/r ) r with r = rank( X ) , it is computationally costly to enumerate all possible p matrices D 1 , . . .
, D p to represent the constraints in the relaxed dual problem (4).
In practice, we ﬁrst randomly sample M i.i.d.
random vectors u 1 , . .
, u M ∼ N (0 , I d ) and generate a subset
ˆ S = { diag ( I ( Xu j ≥ 0)
| j ∈ [ M ] } . of S .
Then, we optimize the randomly sub-sampled version of the relaxed dual problem based on the subset ˆ S and obtain the solution Λ .
Here − Λ − Y is used as the direction to update the particle system X .
If the regularization parameter is too large, then we will have − Λ − Y = 0 , which makes the particle system unchanged.
Therefore, to ensure that ˜ β is not too large, we decay ˜ β by a factor γ 1 ∈ (0 , 1) .
This also appears in (Ergen et al., 2021).
On the other hand, if ˜ β is too small resulting the relaxed dual problem (4) infeasible, we increase ˜ β by multiplying γ − 1 2 , where γ 2 ∈ (0 , 1) .
Detailed explanation of the adjustment of the regularization parameter can be found in Appendix D.
The overall algorithm is summarized in Algorithm 1.
Require: initial positions { x n 0 } Nn =1 , step size α l , initial regularization parameter ˜ β 0 , γ 1 , γ 2 ∈ (0 , 1) .
1: while not converge do 2:
Form X l
Y l
Solve Λ l from the relaxed dual problem (15) with ˜ β
= ˜ β l .
4: if the relaxed dual problem with ˜ β
= ˜ β l is infeasible then 5:
Set X
X n ∈ [ N ] and set β + γ − 1 β
6: else 7:
Update X α
[-l +
-](Λ
+ Y l ) n [ N ] set β l + γ 1 ˜ β l .
8: end if 9: end while
Applying the standard interior point method (Boyd et al., 2004) leads to the computational time
O ((max { N, d 2 } ˆ p ) 6 ) .
For high-dimensional problems, i.e., d is large, the computational cost of solving (15) can be large.
In this case, we apply the dimension-reduction techniques (Zahm et al., 2018; Chen & Ghattas, 2020;
Wang et al., 2021a) to reduce the parameter dimension d to a data-informed intrinsic dimension ˆ d , which is often very low, i.e., ˆ d (cid:28) d , which can dramatically decrease the computational time (18).
N UMERICAL EXPERIMENTS
In this section, we present numerical results to compare WGD approximated by neural networks (WGD-NN) and WGD approximated using convex optimization formulation of neural networks (WGD-cvxNN).
The performance of compared methods is assessed by the sample goodness-of-ﬁt of the posterior.
For WGD-NN, in each iteration, it updates the particle system using (3) with a function Φ represented by a two-layer squared ReLU neural network.
The parameters of the neural network is obtained by directly solving the nonconvex optimization problem (10).
For high-dimensional problems, we apply the dimension reduction technique and compare the projected versions (pWGDNN and pWGD-cvxNN).
We note that although the cost for solving the relaxed dual problem (15) using standard convex optimization solvers in WGD-cvxNN can be higher compared to that by a direct neural network training in WGD-NN, this cost difference is negligible in the entire optimization dominated by the likelihood evaluation when the model (e.g., PDE) is expensive to solve.
In such cases WGDcvxNN and WGD-NN have similar computational complexity but WGD-cvxNN achieves better performance.
We use the standard convex optimization solver CVXPY (Diamond & Boyd, 2016) with MOSEK(ApS, 2019) inner solver.
Applying randomized SDP solvers (Yurtsever et al., 2021), randomized second-order methods (Pilanci & Wainwright, 2017; Lacotte et al., 2021) or advanced
SDP solvers (Zhao et al., 2010; Yang et al., 2015; Wang et al., 2021b) for large-scale problem can improve the computation time.
Moreover, the induced SDPs have speciﬁc structures of many similar constraints.
Solving the SDP (15) can be accelerated by designing a specialized convex optimization solver, which is left for future work.
4. A TOY EXAMPLE
We test the performance of WGD on a bimodal 2-dimensional double-banana posterior distribution introduced in (Detommaso et al., 2018).
We ﬁrst generate 300 posterior samples by a Stein variational Newton (SVN) method (Detommaso et al., 2018) as the reference, as shown in Figure 1.
We evaluate the performance of WGD-NN and WGD-cvxNN by calculating the maximum mean discrepancy (MMD) between their samples in each iteration and the reference samples.
In the comparison, we use N = 50 samples and run for 100 iterations with step sizes α l = 10 − 3 .
For
WGD-cvxNN, we set β = 1 , γ 1 = 0 .
95 and γ 2 = 0 .
For WGD-NN, we use m = 200 neurons and optimize the regularized training problem (10) using all samples with the Adam optimizer (Kingma & Ba, 2014) with learning rate 10 − 3 for 200 sub-iterations.
We also set the regularization parameter β
= 1 and decrease it by a factor of 0 .
95 in each iteration.
We ﬁnd that this setup of parameters is more suitable.
The posterior density and the sample distributions by WGD-cvxNN and WGD-NN at the ﬁnal step of 100 iterations are shown in Figure 1.
It can be observed that WGD-cvxNN provides more representative samples than WGD-NN for the posterior density.
In Figure 2, we plot the MMD of the samples by WGD-cvxNN and WGD-NN compared to the reference SVN samples at each iteration.
We observe that the samples by WGD-cvxNN achieves much smaller MMD than those of WGD-NN compared to the reference SVN samples, which is consistent with the results shown in Figure 1.
4. PDE- CONSTRAINED NONLINEAR B AYESIAN INFERENCE
In this experiment, we consider a nonlinear Bayesian inference problem constrained by the following partial differential equation (PDE) (Chen & Ghattas, 2020) with application to subsurface (Darcy) ﬂow in a physical domain D
= (0 , 1) 2 ,
[Equation]
where u is pressure, v is velocity, h is force, e x is a random (permeability) ﬁeld equipped with a Gaussian prior x ∼ N ( x 0 , C ) with covariance operator C = ( − δ ∆ + γI ) − α where we set δ = 0 .
1 , γ = 1 , α = 2 and x 0 = 0 .
This problem is widely used in many areas, for instance, estimating permeability in groundwater ﬂow, thermal conductivity in material science or electrical impedance in medical imaging, We impose Dirichlet boundary conditions u = 1 on the top boundary and u = 0 on the bottom boundary, and homogeneous Neumann boundary conditions on the left and right
We use a ﬁnite element method with piecewise linear elements for the discretization of the problem, resulting in 81 dimensions for the discrete parameter.
The data is generated as pointwise observation of the pressure ﬁeld at 49 points equidistantly distributed in (0 , 1) 2 , corrupted with additive 5% Gaussian noise.
We use a DILI-MCMC algorithm Cui et al.
(2016) with 10000 effective samples to compute the sample mean and sample variance, which are used as the reference values to assess the goodness of the samples.
[Figure]
We run pWGD-cvxNN and pWGD-NN with 64 samples for ten trials with step size α l = 10 − 3 , where we set β = 10 , γ 1 = 0 .
95 , and γ 2 = 0 .
95 10 for both methods.
The RMSE of the sample mean and sample variance are shown in Figure 3 for the two methods at each of the iterations.
We can observe that pWGD-cvxNN achieves smaller errors for both the sample mean and the sample variance compared to pWGD-NN at each iteration.
Moreover, pWGD-cvxNN provides much smaller variation of the sample mean and sample variance for the ten trials compared to pWGD-NN.
Furthermore, by an effective reduction of the parameter dimension from 81 to data-informed 20 in our pWGD-cvxNN, as used and analyzed in (Zahm et al., 2018; Chen & Ghattas, 2020; Wang et al., 2021a), the time for solving the SDP is signiﬁcantly reduced from about 800 seconds in average to less than 1 second (about 0.7 in average), making our pWGD-cvxNN computationally efﬁcient.
4. B AYESIAN INFERENCE FOR COVID-
In this experiment, we use Bayesian inference to learn the dynamics of the transmission and severity of COVID-19 from the recorded data for New York state.
We use the model, parameter, and data as in Chen & Ghattas (2020).
More speciﬁcally, we use a compartmental model for the modeling of the transmission and outcome of COVID-19.
We take the number of hospitalized cases as the observation data to infer a social distancing parameter, a time-dependent stochastic process that is equipped with a Tanh–Gaussian prior to model the transmission reduction effect of social distancing, which becomes 96 dimensions after discretization.
We use the projected Stein variational gradient descent (pSVGD) method Chen & Ghattas (2020) as the reference to evaluate the goodness of samples.
We run pWGD-cvxNN and pWGD-NN using 64 samples for 100 iterations with step size α l = 10 − 3 , where we set β = 10 , γ 1 = 0 .
95 , and γ 2 = 0 .
95 10 for both methods as in the last example.
From Figure 4 we can observe that pWGDcvxNN produces more consistent results than pWGD-NN compared to the reference pSVGD results, for both the sample mean and 90% credible interval, both in the inference of the social distancing parameter and in the prediction of the hospitalized cases.
C ONCLUSION
In the context of Bayesian inference, we approximate Wasserstein gradient direction by the gradient of functions in the family of two-layer neural networks.
We propose a convex SDP relaxation of the dual of the variational primal problem, which can be solved efﬁciently using convex optimization methods instead of directly training the neural network as a nonconvex optimization problem.
In particular, we established that the gradient obtained by the new formulation and convex optimization is at least as good as the one approximated by functions in the family of two-layer neural networks, which is demonstrated by various numerical experiments. By stacking the two-layer neu- [Figure]
ral networks in each step together, our proposed method formulate a deep neural network to learn the transportation map from prior to posterior.
In future studies, specialized solvers for structured SDPs, including the relaxed dual problem, can lead to drastic accelerations of our proposed method and it is of central importance for the practical applications of our algorithms to real-world problems.
We also expect to extend our convex optimization formulation of neural networks to the calculation/approximation of generalized Wasserstein ﬂows.
We also expect to apply deep neural networks for the approximation of Wasserstein gradient ﬂows based on recent works on convex optimization formulations of deep neural networks (Wang et al., 2021c; Ergen & Pilanci, 2021a;b).
A C ODES FOR NUMERICAL EXPERIMENT
All codes for the numerical experiment can be found in https://github.com/ai-submit/
OptimalWasserstein .
B C OMPARISON WITH PREVIOUS WORKS ON CONVEX OPTIMIZATION FORMULATION OF NEURAL NETWORKS
Previous works on convex optimization formulation of neural networks mainly focus on the supervised learning problem of two-layer neural networks using convex loss functions (e.g., squared loss, logistic loss).
Our work utilizes a similar convex analytic framework to solve the variational problem of approximating the Wasserstein gradient direction, which is different from supervised learning.
The convex optimization approach is related to the idea of inﬁnite width neural networks modeled as probability measures.
The dual problem itself is equivalent to the convex dual problem when the neural network in the primal problem has inﬁnitely many neurons.
However, the convex optimization approach tackles networks of arbitrary width that are able to learn useful representations, while the inﬁnite width limit is quite limited (limited to basically kernel methods).
C A DDITIONAL NUMERICAL EXPERIMENT C. PDE- CONSTRAINED LINEAR B AYESIAN INFERENCE
In this experiment, we consider a linear Bayesian inference problem constrained by a partial differential equation (PDE) model for contaminant diffusion in environmental engineering in domain
D = (0 , 1) ,
[Equation]
where x is a contaminant source ﬁeld parameter in domain D , u is the contaminant concentration which we can observe at some locations, κ and ν are diffusion and reaction coefﬁcients.
For simplicity, we set κ, ν = 1 , u (0) =
u (1) = 0 , and consider 15 pointwise observations of u with 1% noise, equidistantly distributed in D .
We consider a Gaussian prior distribution x ∼ N (0 , C ) with covariance given by a differential operator C =
( − δ ∆+ γI ) − α with δ, γ, α > 0 representing the correlation length and variance, which is commonly used in geoscience.
We set δ = 0 .
1 , γ = 1 , α = 1 .
In this linear setting, the posterior is Gaussian with the mean and covariance given analytically, which are used as reference to assess the sample goodness.
We solve this forward model by a ﬁnite element method with piece-wise linear elements on a uniform mesh of size 2 k , k ≥ 1 .
We project this high-dimensional parameter to the data-informed low dimensions as in Wang et al.
(2021a)
to alleviate the curse of dimensionality when applying WGD-cvxNN and WGD-NN, which we call pWGD-cvxNN and pWGD-NN, respectively.
For k = 4 we have 17 dimensions for the discrete parameter and 4 dimensions after projection.
We run pWGD-cvxNN and pWGD-NN using 16 samples for 200 iterations with α l = 10 − 3 , β = 5 , γ 1 = 0 .
95 , and γ 2 = 0 .
95 10 for both methods.
We use m = 200 neurons for pWGD-NN and train it by the Adam optimizer for 200 sub-iterations as in the ﬁrst example.
From Figure 5, we observe that pWGD-cvxNN achieves better root mean squared error (RMSE) than pWGD-NN for both the sample mean and the sample variance compared to the reference.
[Figure]
D C HOICE OF THE REGULARIZATION PARAMETER
As the constraints in the relaxed dual problem (15) depends on the regularization parameter ˜ β , it is possible that for small ˜ β , the relaxed dual problem (15) is infeasible.
Consider the following SDP
[Equation]
[Equation]
Here the variables are ˜ β, Λ and { r ( j, +) , r ( j, − ) } pj =1 .
Let ˜
β
1 be the optimal value of the above problem.
Then, only for ˜ β ≥ ˜ β 1 , there exists Λ ∈ R N × d satisfying the constraints in (15).
In other words, the relaxed dual problem (15) is feasible.
We also note that ˜ β 1 only depends on the samples
X
and it does not depend on the value of ∇ log π evaluated on x 1 , . . .
, x N .
On the other hand, consider the following SDP
[Equation]
where the variables are ˜ β and { r ( j, +) , r ( j, − ) } pj =1 .
Let ˜
β
2 be the optimal value of the above problem.
For ˜ β ≥ ˜ β 2 , as Y is feasible for the constraints in (15), the optimal value of the relaxed dual problem (15) is 0 .
In short, only when ˜ β ∈ [˜ β 1 , ˜ β 2 ] , the variational problem (15) is non-trivial.
To ensure that solving the relaxed dual problem (15) gives a good approximation of the Wasserstein gradient direction, we shall avoid choosing ˜ β either too small or too large.
E P ROOFS E. P ROOF OF P ROPOSITION
P ROOF We
ﬁrst note that
[Equation]
We notice that the term 12 (cid:82) (cid:107)∇ log ρ −∇ log π (cid:107) 22 ρdx does not depend on Φ .
Utilizing the integration by parts, we can compute that
[Equation]
Therefore, the variational problem (4) is equivalent to
[Equation]
[Equation]
By restricting the domain C ∞ ( R d ) to H , we complete the proof.
E. P ROOF OF P ROPOSITION
P ROOF Suppose that ˆ w i = β − 1 i w i
and ˆ α i =
β 2 i α i , where β i > 0 is a scale parameter for i ∈ [ m ] .
Let θ (cid:48)
= { ( ˆ w i , ˆ α i ) } mi =1 .
We note that
[Equation] [Equation]
25) and ˆ α
[Equation]
This implies that Φ θ ( x ) =
Φ θ (cid:48) ( x ) and ∇ · Φ θ ( x ) = ∇ · Φ θ (cid:48) ( x ) .
For the regularization term R ( θ ) , we note that
[Equation]
The optimal scaling parameter is given by α i = 2 − 1 / 9 (cid:107)
[Equation]
i
[Equation]
| α i |
[Equation]
As the scaling operation does not change (cid:107) w i (cid:107) 22 | α i | , we can simply let (cid:107) w i (cid:107) 2 = 1 .
Thus, the regularization term β 2 R ( θ ) becomes ˜ β N (cid:80) mi =1 (cid:107) u
i (cid:107) 1 .
This completes the proof.
E. P ROOF OF P ROPOSITION
P ROOF Consider the Lagrangian function
[Equation]
For ﬁxed W , the constraints on Z and α are linear and the strong duality holds.
Thus, we can exchange the order of min Z,α and max Λ .
Thus, we can compute that
[Equation]
By exchanging the order of min and max
[Equation]
we can derive the dual problem:
[Equation]
This completes the proof.
E. P ROOF OF P ROPOSITION
P ROOF Based on the hyper-plane arrangements D 1 , . .
, D p , the dual constraint is equivalent to that for all j ∈ [ p ] , [Equation] ˜
[Equation]
(31) holds for all w ∈ R d satisfying [Equation] w [Equation]
2 ≤ 1 , (2 D j − I )
Xw ≥ 0 .
This is equivalent to say that for all j ∈ [ p ]
[Equation]
From a convex optimization perspective, the natural idea to interpret the constraint (32) is to transform the minimization problem into a maximization problem.
We can rewrite the minimization problem in (32) as a trust region problem with inequality constraints:
[Equation]
As the problem (33) is a convex problem, by taking the dual of (33) w.r.t.
w , we can transform (33) into a maximization problem.
However, as (33) is a trust region problem with inequality constraints, the dual problem of (33) can be very complicated.
According to (Jeyakumar & Li, 2014), the optimal value of the problem (33) is bounded by the optimal value of the following SDP
[Equation]
34) from below.
Lemma
The dual problem of SDP (34) takes the form
[Equation]
in variables r = 
[Equation]
 ∈ R N +1 and γ ∈ R .
P ROOF Consider the Lagrangian
[Equation]
36) where r ∈ R N +1+ and γ ∈ R .
By minimizing L ( Z, r, γ ) w.r.t.
Z ∈ S d +
+ , we derive the dual problem (35).
The constraints on Λ in the dual problem (13) include that the optimal value of (34) is bounded from below by − ˜ β .
According to Lemma 1, this constraint is equivalent to that there exist r ∈ R N +1 and γ such that
[Equation]
As e d +1 e Td
+1 is positive semi-deﬁnite, the above condition on Λ is also equivalent to that there exist r ∈ R N +1 such that
[Equation]
Therefore, the following convex set of Λ
[Equation]
is a subset of the set of Λ satisfying the dual constraints
[Equation]
On the other hand, the constraint on Λ
[Equation] [Equation]
is equivalent to
[Equation]
By applying the previous analysis on the above trust region problem
[Equation]
the following convex set of Λ
[Equation]
is a subset of the set of Λ satisfying the dual constraints
[Equation]
Therefore, replacing the dual constraint
[Equation]
˜ β by
[Equation]
we obtain the relaxed dual problem.
As its feasible domain is a subset of the feasible domain of the dual problem, the optimal value of the relaxed dual problem gives a lower bound for the optimal value of the dual problem.
E. P ROOF OF P ROPOSITION
P ROOF Consider the Lagrangian function
[Equation] [Equation]
where we write
[Equation]
Here we write S d +1+
[Equation]
By maximizing w.r.t.
Λ and r , we derive the bi-dual problem (16).
E. P ROOF OF T HEOREM
Suppose that ( Z, W, α ) is a feasible solution to (11).
Let D j 1 , . . .
, D j k be the enumeration of { diag ( I ( Xw i ≥ 0))
| i ∈ [ m ] } .
For i ∈ [ k ] , we let
[Equation]
and
[Equation]
For j / ∈ { j 1 , . . .
, j k } , we simply set S ( j, +) = 0 , S ( j, − )
= 0 .
As (cid:107) w i (cid:107) 2 ≤ 1 and D j i = I
( Xw i ≥ 0) , we can verify that tr( S ( j, − ) H ( j ) n ) ≤ 0 , tr( S ( j, +) H ( j ) n ) ≤ 0 are satisﬁed for j = j 1 , . . .
, j m and n = 0 , 1 , . .
N .
This is because for n = 0 , as H ( j i )
[Equation]
, it follows that
[Equation]
For n = 1 , . . .
, N , we have
[Equation]
Based on the above transformation, we can rewrite the bidual problem in the form of the primal problem (12).
For S ∈ S d +1 , we note that
[Equation]
where S 1: d, 1: d denotes the d × d block of S consisting the ﬁrst d rows and columns.
This implies that ˜
A ∗ j ( S ) = − 2 D j XS 1: d, 1: d .
Hence, we have
[Equation]
Therefore, we have
[Equation]
As n -th row of Z satisﬁes that z n = 2 (cid:80) m i
=
α i w i ( x Tn w i )
+ , this implies that
[Equation]
Hence ( Z, { ( S ( j, − ) , ( S ( j, − ) } pj =1 ) is feasible to the relaxed bi-dual problem (16).
We can also compute that
[Equation]
and
[Equation]
Thus, the primal problem (12) with ( Z, W, α ) and the relaxed bi-dual problem (16) with ( Z, { ( S ( j, − ) , ( S ( j, − ) } p j =1 ) have the same objective value.