Optimal Neural Network Approximation of Wasserstein Gradient Direction via Convex Optimization
Authors: Yifei Wang
Peng Chen
Mert Pilanci
Wuchen Li
Venue: Submitted to ICLR 2023 (2023)
Keywords: Bayesian inference
convex optimization
semi-definite programming

Abstract: The computation of Wasserstein gradient direction is essential for posterior sampling problems and scientific computing. The approximation of the Wasserstein gradient with finite 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-definite 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. Numerical experiments including PDE-constrained Bayesian inference and parameter estimation in COVID-19 modeling demonstrate the effectiveness of the proposed method.

Under review as a conference paper at ICLR 2023
OPTIMAL NEURAL NETWORK APPROXIMATION OF WASSERSTEIN GRADIENT DIRECTION VIA CONVEX OPTIMIZATION
Anonymous authors Paper under double-blind review
ABSTRACT
The computation of Wasserstein gradient direction is essential for posterior sam- pling problems and scientific computing. The approximation of the Wasserstein gradient with finite samples requires solving variational problem. We study the variational problem the family of two-layer networks with squared-Rel UeLU activations, towards which we derive a semi-definite programming (SDP) relax- ation. This SDP can viewed as an approximation of the Wasserstein gradi- ent in 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 PDE- constrained Bayesian inference and parameter estimation in COVID-19 modeling demonstrate the effectiveness and efficiency of the proposed method.
INTRODUCTION
Bayesian inference plays an essential role in learning model parameters from the observational data with applications in inverse problems, scientific computing, information science, and machine learn- ing (Stuart) 2010). The central problem in Bayesian inference is to draw samples from posterior distribution. which characterizes the parameter distribution given data and prior distribution.
The Wasserstein gradient flow (Otto 2001 Ambrosio et al. 2005 Junge et al., 2017) has shown to be effective drawing samples from a posterior distribution, which attracts increasing attention in recent years. For instance, the Wasserstein gradient flow of Kulbback-Leibler (KL.) divergence con- nects 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 flow yields different viewpoint for sampling al- gorithms. In particular, the Wasserstein gradient direction also provides deterministic update of the particle system (Carrillo et al. 2021b). Based on the approximation or generalization of the Wasserstein gradient direction, many efficient sampling algorithms have been developed, including Wasserstein gradient descent (WGD) with kernel density estimation (KDE) (Liu et al. (201)), 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 leaminng complicated functions from data. They also have wide applications in Bayesian inverse problems (Rezende & Mohameed 2015; Onken al., 2020 Kruse et al., 2019; Lan et 202T). According to the universal approximation theorem of neural networks (Horiik et al. 1989 [u et al.) 2017), any arbitrarily complicated functions can learned by two-layer neural network with non- linear activations and sufficient 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 find the global optima of the training problem. Re- cently, based line of works (Pilanci & Ergen 2020 Sahiner al.) 2020 Bartan & Pilanci 2021a), the regularized training problem of two-layer neural networks with ReLUpplynomaal ac- tivation be formulated convex program. Indeed, by solving the convex program, we

---

Under review as a conference paper at ICLR 2023
construct the entire set of global optima of the nonconvex training problem (Wang al. 2020). The- oretical analysis (Wang et al. p022) shows that global optima of the training problem correspond to the simplest models with good generalization properties. Moreover, numerical results (Pilanci & Ergen, (020)) show that neural networks found by solving the convex program can achieve higher train accuracy and test accuracy compared neural networks trained by SGD with the same number of parameters.
In this paper, we study variational problem, whose optimal solution corresponds to the Wasser- stein gradient direction. Focusing on the family of two-layer neural networks with squared ReLU activation, we formulate the regularized variational problem terms of samples. Directly training the neural network to minimize the loss may get the neural network stuck at local minima or sad- dle 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-definite program (SDP) relax- ation by analyzing the geometry of dual constraints. The resulting SDP can be efficiently solved by convex optimization solvers such as CVXPY (Diamond & Boyd 2016). We then derive the corre- sponding 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 efficiency. Numerical results for experiments including PDE-constrained inference problems and Covid-19 parameter estimation problems illustrate the effectiveness and efficiency of our method.
1.1 RELATED WORKS
The time and spatial discretizations of Wasserstein gradient flows are extensively studied in liter-
ature (Jordan et al. 1998; Junge al.) 2017: Carrillo et al.] [20 b Bonet et al.) [202T] Liutkus
et al. 2019: Frogner & Poggio 2020). Recently, neural networks have been applied in solving or
approximating Wasserstein gradient flows (Mokrov et al. 2021 Lin et al.) 2021b a Alvarez-Melis ett al., [021;; Bunne et al.] 2021 Hwang et al., [021] 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 these studies, focus on a convex SDP relaxation of the varitional problem induced by the Wasserstein gradient direc- tion. Meanwhile, [Fngg et al. (202T) 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) 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.
2 BACKGROUND
In this section, we briefly review the Wasserstein gradient descent and present its variational for- mulation. In particular, we focus 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.1 WASSERSTEIN GRADIENT DESCENT
Consider an optimization problem in the probability space:
inf (pl(() PEP
= =
(1)g)() - log π (d)),
(1)
Here the integral is taken over R& and the objective functional DKL (pll is the KL divergence from P to π. The variable is the density function in the space P = (p E C (Re)) pdr = 1, p > 0). The function x E C (IR) is a known probability density function of the posterior distribution. By solving the optimization problem (T) can generate samples from the posterior distribution.
2

---

Under review as a conference paper at ICLR 2023
A known fact (Villani 2003 optimization problem
Chapter 8.3.1) is that the Wasserstein gradient descent flow for the satisfies
Pt V 8
Depe
V (pe(Vloge V log m))
8p:
@ Pt V
(pe V log
where p+(x)) p(z, (), apc the 1,2 first variation operator w.r.t. Pts V F denotes the divergence of vector valued function F: Rd + Rd and A is the Laplace operator. In step (a) we the fact that V log Pt = Vpt. This equation is also known as the gradient drift Fokker- Planckk equation. It corresponds to the following updates in terms of samples:
dr; - ( log (x(z)) V log (og)(xz))dds
(2)
where 2, follows the distribution of Pe. Clearly, when Pt π, the above dynamics reach the cquilibrium, which implies that the samples Vt are generated by the posterior distribution.
To solve the Wasserstein gradient flow (2). consider forward Eulerian discretization in time. In the 1-th iteration, suppose that {a:²) are samples drawn from pr. The update rule of Wasserstein gradient descent (WGD) on the particle system {af²} follows
x = 2 04 (xf),
where 4 Rd + R is a function which approximates log - loog π and aq > 0 the step size.
2.2 VARIATIONAL FORMULATION OF WGD
Given the particles ( J design the following variational problem to choose suitable func- tion & approximating the function log p log Consider
/ V
inf EC (R.)
log p(z) v log π (x)) (z)))
(4)
The objective functional evaluates the least-square discrepancy between V po V log π and 4 weighted by the density p. The optimal solution follows & log p log n, up constant shift. Let
) be a finite dimensional function space. The following proposition gives a formulation H C² (TRd) of 4 H.
Proposition 1 Let 7 C Ca (Rd) be function space. The variational problem in the domain H can be reformulated
[[V + / (AM())(a)) + (Vlog (T)))) p(x)dx.
inf 4EH
(5)
Remark A similar variational problem has been studied in (di Langosco et al. 2021). If we replace 4 for H by vector field in certain function family, then, the quantity in (5) is the negative regularized Stein discrepancy defined in (di Langosco et al. [021] between p and based on V. This problem also similar to the varitional problem for the score matching estimator in (Hyvärinen & Dayan 2005) by parameterizing & in given probabilistic model. In comparison, our method can be viewed as special case of score matching by using a two-layer neural network.
Therefore, by replacing the density p by finite samples {In finite samples forms
IN ~ p. the problem (5) in terms of
N Σ #=1
1 + AM(E.)
inf DEN
VV log (In), V)(xx))
(6)
N
==1
OPTIMAL NEURAL NETWORK APPROXIMATION OF WASSERSTEIN GRADIENT
In this section, we focus on functional space H. functions represented by two-layer neural net- works. We derive the primal and dual problem of the regularized Wasserstein variational problems.
3

---

Under review as a conference paper at ICLR 2023
By amalyzing the dual constraints, convex SDP relaxation of the dual problem is obtained. We also present a practical implementation estimation V log p V log ग and discuss the choice of the regularization parameter.
Let di be an activation function. Consider the case where H. a class two-layer neural network with the activation function (xx):
H T E (R))) = a² (W* x)}
where O = (W, a:) is the parameter in the neural network with W E R and a E R Remark We can extend this model to handle the bias term by add an entry of in ,
For two-layer neural networks, we can compute the gradient and Laplacian of ④ E H. as follows:
m
Σ a, v ('
W WW a))
(8)
( = Σ Bill (w x).
(9)
Here represents the clement-wise multiplication. By adding regularization term to the variational problem 6
wC obtain
m a,w,t' an), V Vog (IX,)
min
a,w (wf xnn)
+ N
2N
(10)
m
a I B ( xn) + R()),
where 3 >>> is the regularization parameter. We focus on the squared ReLU activation v²(x) = (=) = (max(2. 0))² Note that a vanishnig second derivative is required for the Laplacian term (9. which makes the ReLU activation inadequate. For this activation function, we consider the regularization function R(0)) >x ([w:|]2 lax1").
Remark We note that (x) and (f) are all piece-wise degree-3 polynomials of the pa- rameters 0. Hence, we consider specific cubic regularization term above, analogous to (Bartan & Pilanci 2021a). By choosing this regularization term, we can derive simplified dual problem.
By utilizing the arithmetic and geometric mean (AM-GM) inequality,
rescale the first and second-layer parameters and formulate the regularized variational problem (10) as follows.
Proposition 2 (Primal problem) The regularized variational problem
110 can be reformulated to
m
min W.
Σ
awa (wf xn)
+
Σ n=1
as Iw 11 (w) xn)
#=1
N
(11)
+
Cw an). Vlog + B||a||1,
==1
s.f. ||w/|l2 $ 1 [m],
where 3 = 3-2 5/3 N/3 and we denote [m]] (L...., m].
In short, the optimal value of (10) and are the same. We can obtain the optimal solution of IT by rescaling the optimal solution of (1O) and vice versa. For simplicity, we write Y E N whose n-row is V log æ (I,) for E [N]. We introduce the slack variable on = - o,w w;) for
n e [N] and denote Z = [2]
2N]* E N Then, we can simplify the problem to
1
min W.aZ
-
2
+
ax a (wT an) + tr(Y 2) + $1|
2
1=
(12)
s.t. 2n =
Σ
aaw (2 w:) n E [N], w. 1|2 S 1, [m].
4

---

Under review as a conference paper at ICLR 2023
To derive the convex relaxtion of neural network training problem, the dual problem plays 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
is
N
2 AA +1
s.t.
max will
Σ
12ll2 (x", w) W/ w) 3 (13)
which provides lower-bound on
We note that the dual problem can be infeasible if the regularization parameter 3 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 -o and the primal problem not lower bounded.
3.1 ANALYSIS OF DUAL CONSTRAINTS AND THE RELAXED DUAL PROBLEM
Now, we analyze the constraint in the dual problem. We note that it is closely related the regular- ization parameter, which we will discuss later. For simplicity, we take v² = 0 as the subgradient of (z) at = 0, i.e., taking the left derivative of (=) at Denote the set of all possible hyper-plane arrangements corresponding to the rows of X
Let X = ,..,.... NN)" Nxd
= (diag((T( (T( 2 ()) E Rd w #0..
(14)
Here T(s) if the statement & correct and T(s) 0 otherwise. Let [S] be the cardinality of S, and write S {DI..... Dp}. According to (Cover 1965), have the upper bound p s 2r (c(N-1) where rank(X) Based on the analysis of the dual constraints, we derive r
convex SDP as a relaxed dual problem.
Proposition 4 (Relaxed dual problem) The relaxed dual problem is the following SDP:
max A,(rU:),,p 7, Dg 1
2 A
s.1. A,(A) + Bj + Σ
U. H() + Beat 2 (,-)) 2 20, (pl. (15)
- -A((A) - B3 + Σ 1
(i++) HD) + Bea e 2 0,r(itt) 20, E [p).
n=0
where we denote [p] fl,...,pp. For j
[p].
we denote A,(A)
-AT XTD,A,
(A)(A)
B
0 (1
Ta
o and H!!
By 2tr(D,) Aj(A)
0 B,
0
0
0
(1 1--(((D)))))) mm )kn)
n [N] The vector ed41 E IR++1 satisfies that (ea+1)s
2l1-2(D))nnaal 0.for i [d] and (ca+1)at1 = 1.
0
The optimal value of [5) gives a lower bound on the dual problem problem
13 and hence on the primal
The relaxed bi-dual problem provides insights approximating the primal problem via convex optimization, which is derived follows. As an equivalent formulation of the convex dual problem 15 it can be viewed convex relaxation of the primal problem (12
5

---

Under review as a conference paper at ICLR 2023
Proposition 5 (Relaxed bi-dual problem) The dual of the relaxed dual problem 13
is as follows
min 2,((8 80.
2 1 +
YI
2 1
Σ
tr (rB -s(;:)))
-=1
3 Σ
(Istit) + s SU. Jeattehii))
(16)
s.t. Z
Σ
s(itt),
j=1
(r(S H < O, (r(((i+)) )2)) 0, 0,...,N,55 lp.
Here A3 is the adjoint operator of the linear operator A₁-
As (15) convex problem and the Slater's condition is satisfied, the optimal values of (15) and (16) are same. The bi-dual problem (T6) is closely related to the primal problem (12). Indeed, any feasible solutions of the primal problem (T]) 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 Suppose that (Z, W,W) is feasible to the primal problem 12 Then, there exist matrices ( -( ===1 constructed from (W, a) such that (Z, (St.tt 5 feasible to the relaxed bi-dual problem (16). Moreover, the objective value of the relaxed bi-dual problem (16) at
(Z, (
is the same as objective value of the primal problem (12 at (Z, W a)).
Let J(Z, (St++), (U-) denote the objective value of the relaxed bi-dual problem [6 at
feasible solution (Z, (SU++), 35 Let (Z", "y*" denote a globally optimal solu- tion of the primal problem 11). By Theorem there exist matrices (05,+),5655-133
such
that "Z"{{5(i+) (15(++),,(ū-))) )==) is feasible solution of the relaxed bi-dual problem (66) and
the same as the objective value of [12) at its global minimum (Z*, W*, a²). On the other hand, let (55(5 (6.-)) denote an optimal solution of the relaxed bi-dual problem (16). From the optimality of (2", (5(++), (00-))))00))00 -( we have s
J(Z", (SU++), ( ((=-));=)))
J j(2", ( 55. S J(Z", , ( 1z
(17)
Note that at (Z*, W* aa) we obtain the optimal approximation of V log p- log ा at I...... NN in the family of two-layer squared-Rede ReeL networks (7). Smaller or equal objective value of the relaxed bi-dual problem [16] can be achieved at (Z*, (55.)))+) S
) than at (Z*,
Therefore,
view 2* gives optimal approximation of V log p V Vlog evaluated on .I,...,7N N in a broader function family including the two-layer squared ReI.U neural networks.
From the derivation of the relaxed bi-dual problem, we have the relation 2 = - Y, where (A*, (rlist), ((ār)) is optimal to the relaxed dual problem (15) and (ž", ( (000-+)),0-,))0 S optimal to the relaxed bi-dual problem (6).. Therefore, by solving A from the relaxed dual problem (15). we can use -A* Y as the approximation of V log V log evaluated on æ1, ........ NN.
/==1 is
Remark 4 We note that solving the proposed convex optimization problem 15 renders the approxi- mation of the Wasserstein gradient direction. Compared to the two-layer ReLU networks, it induces broader class of functions represented by (, } This contains more variables than the neural network function.
3.2 PRACTICAL IMPLEMENTATION
Although the number p of all possible hyper-plane arrangements is upper bounded by 2r((N 1)e/r)" with rank(X), it is computationally costly enumerate all possible p matri- ces Di...., Dp to represent the constraints in the relaxed dual problem (4) In practice, we first randomly sample M i.i.d. random vectors u1,..., UM NO, Ta) and generate subset

---

Under review as a conference paper at ICLR 2023
S = [diag([(Xx) 0)1i E [M]]. of S. Then, optimize the randomly sub-sampled version of the relaxed dual problem based on the subset 8 and obtain the solution A. Here -A Y is used as the direction to update the particle system X. If the regularization parameter is too large, then we will have- -A - Y = O, which makes the particle system unchanged. Therefore, to ensure that | is not too large, we decay 3 by factor "/1 E (0, 1). This also appears in (Ergen et al. 2021). On the other hand, if 3 is too small resulting the relaxed dual problem 4 infeasible, we increase 3 by multiplying 7/2 where 2/2 E (0, Detailed explanation of the adjustment of the regularization The overall algorithm is summarized in Algorithm I
parameter can be found in Appendix D
Algorithm Convex neural Wasserstein descent
Require: initial positions {a8
N n= step size on, initial regularization parameter fo. 71 7/2 (0, 1).
1: while not converge do
2: Form X; and Yi based on {>] J and {V log (x)) [
3: Solve A₁ from the relaxed dual problem 15 with 3 = 3/.
4: if' the relaxed dual problem with β 3 is infeasible then
5: Set X1+1 X1 for [N] and set B1+1 - 7/2 B..
6: else
7: Update X X2 + ar(Ar Yi) for n e [AV] and set B1+1 =113.
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 (p)). ))
(18)
For high- high--dimensionl problems, i.e., d is large, the computational cost of solving In this case, we apply the dimension reduction techniques (Zahm et al.] 20118 Chen & Ghattas 2020 -eneeniioorreucttioon Wang et al. 2021a) to reduce the parameter dimension to a data-informed intrinsic dimension d, which is often very low, i.e., < d, which can dramatically decrease the computational time (18)
(15) can be large.
4 NUMERICAL 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-evXNN). The performance of compared methods is assessed by the sample goodness-of of the posterior. For WGD-NN, in each iteration, it updates the particle system using (3) with a function + represented by two-layer squared ReLU neural network. The parameters of the neural network obtained by directly solving the nonconvex optimization problem (10). For high-dimensional problems. we apply the dimension reduction technique and compare the projected versions (pWGD- NN 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 direct neural network training in WGD-NN, this cost difference is negligible the entire optimization dominated by the likelihood evaluation when the model (c.g.. PDE) is expensive to solve. In such cases WGD- cvxNN and WGD-NN have similar computational complexity but WGD-evxNN achieves better performance. We use the standard convex optimization solver CVXPY (Diamond & Boyd, 2016) with MOSEK 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. (021b) for large-scale problem can improve the computation time. Moreover, the induced SDPs have specific 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.

---

Under review as a conference paper at ICLR 2023
4.1 A TOY EXAMPLE
We test the performance of WGD on a bimodal 2-dimensional double-banana posterior distribu- tion introduced in (Detommaso et al. 2018). We first generate 300 posterior samples by a Stein variational Newton (SVN) method (Detommaso et al. 2018) as the reference, as shown in Fig- ure ! 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 og 10-3, For 0.95 For WGD-NN, we use rons and optimize the regularized training problem (10) using all samples with the Adam optimizer (Kingma & Ba 2014) with learning rate 10 10-3 for 200 sub-iterations. We also set the regularization parameter 3 and decrease it by a factor of 0.95 in each iteration. We find that this setup of parameters is more suitable.
WGD-cvXNN, we set 3 = 1, "1 0.95 and
200 neu-
The posterior density and the sample distributions by WGD-evxNN and WGD-NN at the final step of 100 iterations shown in Figure [ It can be observed that WGD-evxNN provides more rep- resentative samples than WGD-NN for the posterior density. In Figure 2 we plot the MMD of the samples by WGD-evxNN 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
WGD CvxNN
WGD NN
3
3
1
-1
-1
11
-1
Figure 1: Posterior density and sample distributions by WGD-evxNN and WGD-NN the final step of 100 iterations, compared to the reference SVN samples (right).
0.003
WGD+cvaNN WGD-NN
MND
0.002
20
100
iteratioo
Figure 2: MMD of WGID-vxNNN and WGD-NN samples compared to the reference SVN samples.
4.2 PDE-CONSTRAINED NONLINEAR BAYESIAN INFERENCE
In this experiment, we consider nonlinear Bayesian inference problem constrained by the following partial differential equation (PDE) (Chen & Ghattas 2020) with application to subsurface (Darcy) flow in physical domain D (0, 1)²
v + Vu 0 in D. V vvv= in D.
(19)
where u pressure, v is velocity, h is force, e" random (permeability) field equipped with Gaussian prior
(2o, C) with covariance operator C (- > where we set 5 = 0.1.7 = a, = 2 and Øy = 0. This problem widely used in many areas, for instance, estimating permeability in groundwater flow, thermal conductivity in material science electrical impedance in medical imaging, We impose Dirichlet boundary conditions u = on the top boundary and u 0 the bottom boundary, and homogeneous Neumann boundary conditions the left
8

---

Under review as a conference paper at ICLR 2023
and right boundaries for We finite element method with piecewise linear elements for the discretization of the problem, resulting in 81 dimensions for the discrete parameter. The data is generated pointwise observation of the pressure field at 49 points equidistantly distributed in 1)², corrupted with additive 5% Gaussian noise. We use a DILI-MCMC algorithm Cui et al. (206)) with 10000 effective samples compute the sample mean and sample variance, which used as the reference values to assess the goodness of the samples.
pAF Kh
SIND
(ARSD cx/M
SINGD eVINN
Log 10S10009EE Sareinea
-0.4-
0.
-05-
-08
-14
20
Figure 3: Ten trials and the RMSE of the sample mean (top) and sample variance (bottom) by pWGD-NN and PWGD-cvXNN at different iterations. Nonlinear inference with PDE constraint.
We run PWGID-evXNN and PWGD-NN with 64 samples for ten trials with step size ay = 10 where we set B = 10, 71 = 0.95, and 72 = 0.9510 for both methods. The RMSE of the sam- ple mean and sample variance are shown in Figure 3 for the two methods at each of the iterations. We can observe that DWGD-evxNN achieves smaller errors for both the sample mean and the sam- ple 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 data-informed 20 in our PWGD-cvxNN, as used and analyzed in (7ahm et al.) [0188] Chen & Ghattas 2020) Wang et al.) 2021a), the time for solving the SDP is significantly reduced from about 800 seconds in average to less than second (about 0.7 in average), making our PWGD-cvxNN computationally efficient.
4.3 BAYESIAN INFERENCE FOR COVID-19
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 specifically, we use 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, time-dependent stochastic process that is equipped with a Fonh-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 DWGD-evxND and PWGD-NN using 64 samples for 100 iterations with step size ai = 10-3 where we set / = 10, 7i 0.95, and 722 = 0.95 for both methods in the last example. From Figure we can observe that pWGD- cVxNN produces more consistent results than PWGD-NN compared to the reference pS VSGD 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.
5 CONCLUSION
In the context of Bayesian inference, we approximate Wasserstein gradient direction by the gradient of functions 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 efficiently using convex optimiza- tion methods instead of directly training the neural network as a nonconvex optimization problem. In particular,
established that the gradient obtained by the new formulation and convex opti- mization 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-

---

Under review as a conference paper at ICLR 2023
NN vs pSVGD social distancing
NN pSVGD # hospitalized
0.0
m0..1 +R JOYCD
JN
23000
mont
main LOVCD
23000
32000
02
5000
10
THE coch
poocn
22000
0.4
Aa
Age CVXNN pSVGD social distancing
CVXNN 5SVGD hospitalized
Figure 4: Comparison of pWGD-cvxNN and PWGD-NN to the reference by PSVGD for Bayesian inference the social distancing parameter (left) from the data of the hospitalized cases (right) with sample mean and 90% credible interval.
ral networks in each step together, our proposed method formulate 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, is of central importance for the practical applications of our algorithms to real-world problems. We also expect extend our convex optimization formulation of neural networks to the calcula- tion/approximation of generalized Wasserstein flows. We also expect to apply deep neural networks for the approximation of Wasserstein gradient flows based on recent works on convex optimization formulations of deep neural networks (Wang et al. 2021c
lead to drastic accelerations of our proposed method and
Ergen & Pilanci 2021a b).
REFERENCES
David Alvarez-Melis, Yair Schiff, and Youssef Mroueh. Optimizing funetionals the space of probabilities with input convex neural networks. arXiv preprint arXiv:2106.00774 2021.
Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media, 2005.
MOSEK ApS. Mosek optimization suite, 2019.
Burak Bartan and Mert Pilanci. Neural spectrahedra and semidefinite lifts: Global convex op- timization of polynomial activation neural networks in fully polynomial-ime. arXiv preprint arXiv:2101.02443 2021a.
Burak Bartan and Mert Pilanci. Training quantized neural networks to global optimality via semidef- inite programming. In International Conference on Machine Learning, pp. 694-704. PMLR, 2021b.
Clément Bonet, Nicolas Courty, François Septier, and Lucas Drumetz. Sliced-wasserstein gradient flows. arXiv preprint 10X 2021.
Stephen Boyd, Stephen P Boyd, and Lieven Vandenberghe. Convex optimization. Cambridge uni- versity press, 2004.
Charlotte Bunne, Laetitia optimal transport modeling of population dynamics. arXiv preprint arXiv:2106.06355 2021.
Andreas Krause, and Marco Cuturi. Jkonet: Proximal
Jose A Carrillo, Katy Craig, Li Wang. and Chaozhen Wei. Primal dual methods for wasserstein gradient flows. Foundations of Computational Mathematics, pp. 1-55, 2021a.
Jose A Carrillo, Daniel Matthes, and Marie-Therese Wolfram. Lagrangian schemes for wasserstein gradient flows. Handbook of Numerical Analysis, 22:271-311 2021b.
Peng Chen and Omar Ghattas. Projected stein variational gradient descent. Advances in Neural Information Processing Systems, 13:194711958 2020.
10

---

Under review as a conference paper at ICLR 2023
Thomas M Cover. Geometrical and statistical properties of systems of linear inequalities with ap- plications in pattern recognition. IEEE transactions on electronic computers, (3):326-334, 1965.
Tiangang Cui, Kody JH Law, and Youssef M Marzouk. Dimension-iindeeenden likelihood-informed
memc. Journal of Computational Physics, 304: 10t:1009-177, 2016.
Gianluca Detommaso, Tiangang Cui, Alessio Spantini, Youssef Marzouk, and Robert Scheichl. A stein variational newton method. arXiv preprint arXiv:/806.0:000 1806.03885, 2018.
Lauro Langosco di Langosco, Vincent Fortuin, and Heiko Strathmann. Neural variational gradient descent. arXiv preprint arXiv:2107. 10731, 2021.
Steven Diamond and Stephen Boyd. CVXPY: A Python-embedded modeling language for convex optimization. Journal of Machine Learning Research, 17(83):1-5, 2016.
Tolga Ergen and Mert Pilanci. Global optimality beyond two layers: Training deep relu networks via convex programs. In International Conference on Machine Learning, pp. 2993-3003, PMLR, 2021a.
Tolga Ergen and Mert Pilanci. Path regularization: A convexity and sparsity inducing regularization parallel relu networks. arXiv preprint urXiv:2110.095:5 0954488 2021b.
Tolga Ergen, Arda Sahiner, Batu Ozturkler, John Pauly, Morteza Mardani, and Mert Pilanci. Demys- tifying batch normalization in relu networks: Equivalent convex optimization models and implicit regularization. arXiv preprint 203i:2/03300199 2021.
Jiaojiao Fan, Amirhossein Taghvaei, and Yongxin Chen. Variational wasserstein gradient flow. arXiv preprint arXiv:21 2021.
Xingdong Feng, Yuan Gao, Jian Huang, Yuling Jiao, and Xu Liu. Relative entropy gradient sampler for unnormalized distributions. arXiv preprint 1rX00:2110.02787, 2021.
Charlic Frogner and Tomaso Poggio. Approximate inference with wasserstein gradient flows. In International Conference Artificial Intelligence and Statistics, pp. 2581-2590. PMLR, 2020.
Kurt Homik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are uni- versal approximators. Neural networks, 2(5):359-366, 1989.
Hyung Ju Hwang, Cheolhyeong Kim, Min Sue Park, and Hwijae Son. The deep minimizing move- ment scheme. arXiv preprint arXiv:2109. 14851, 2021.
Aapo Hyvärinen and Peter Dayan. Estimation of non-normalized statistical models by score match- ing. Journal of Machine Learning Research, 6(4), 2005.
Vaithilingam Jeyakumar and GY Li. Trust-region problems with linear inequality constraints: exact sdp relaxation. global optimality and robust optimization. Mathematical Programming, 147(1): 171-206, 2014.
Richard Jordan, David Kinderlehrer, and Felix Otto. The variational formulation of the fokker- planck equation. SIAM journal on mathematical analysis. 29(1):1-17, 1998.
Oliver Junge, Daniel Matthes, and Horst Osberger. A fully discrete variational scheme for solving nonlinear fokker-planck equations in multiple space dimensions. SIAM Journal on Numerical Analysis, 55(1):419-44 2017.
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv:1412.6980. 2014.
Jakob Kruse, Gianluca Detommaso, Robert Scheichl, and Ulirich Köthe. Hint: Hierarchi- cal invertible neural transport for density estimation and bayesian inference. arXiv preprint arXiv:1905.10687 2019.
Jonathan Lacotte, Yifei Wang, and Mert Pilanci. Adaptive newton sketch: Linear-time optimization with quadratic convergence and effective hessian dimensionality. In International Conference on Machine Learning, pp. 5926-5936. PMLR, 2021.

---

Under review as a conference paper at ICLR 2023
Shiwei Lan, Shuyi Li, and Babak Shahbaba. Scaling up bayesian uncertainty quantification for inverse problems using deep neural networks. arXiv preprint arXiv:2101.03900 2021.
Alex Tong Lin, Samy Wu Fung, Wuchen Li, Levon Nurbekyan, and Stanley Osher. Alternating the population and control neural networks to solve high-dimensional stochastic mean-field games. Proceedings of the National Academy of Sciences, 118(31), 202
Alex Tong Lin, Wuchen Li, Stanley Osher, and Guido Montúfar. Wasserstein proximal of gans. arXiv preprint arXiv:2102.06862. 2021b.
Chang Liu, Jingwei Zhuo, Pengyu Cheng, Ruiyi Zhang, and Jun Zhu. Understanding and acceler- ating particle-based variational inference. In International Conference on Machine Learning, pp- 4082-4092. PMLR, 2019.
Qiang Liu and Dilin Wang. Stein variational gradient descent: general purpose bayesian inference algorithm. In Advances in neural information processing systems, pp. 2378-22-8 2386, 2016.
Antoine Liutkus, Umut Simsekli, Szymon Majewski, Alain Durmus, and Fabean-Robert Stöter. Bliced-wasserstein flows: Nonparametric generative modeling via optimal transport and diffu- sions. In International Conference Machine Leurning, pp. 4104-4113. PMLR, 2019.
Zhou Lu, Hongming Pu, Feicheng Wang, Zhiqiang Hu, and Liwei Wang. The expressive power of neural networks: A view from the width. In Proceedings of the 31st International Conference on Neural Information Processing Systems, pp. 6232-6240, 2017.
Petr Mokrov, Alexander Korotin, Lingxiao Li, Aude Genevay, Justin Solomon, and Evgeny Burnaev. Large-scale wasserstein gradient flows. arXiv preprint arXiv:2106.00736. 2021.
Derek Onken, Samy Wu Fung, Xingjian Li, and Lars Ruthotto. Ot-flow: Fast and accurate continu- ous normalizing flows via optimal transport. arXiv preprint wrXiv:2006000104 2020.
Felix Otto. The geometry of dissipative evolution equations: the porous medium equation. Commu- nications Partial Differential Equations, 161--441 2001.
Mert Pilanci and Tolga Ergen. Neural networks are convex regularizers: Exact polynomial-time con-
vex optimization formulations for two-layer networks. In International Conference Machine Learning, pp. 7695-7705. PMLR, 2020.
Mert Pilanci and Martin Wainwright. Newton sketch: A near linear-time optimization algorithm with linear-quadratic convergence. SIAM Journal on Optimization, 27(1):205-245, 2017.
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. Interna- tional conference on machine learning, pp. 1530-1538. PMLR, 2015.
Arda Sahiner, Tolga Ergen, John Pauly, and Mert Pilanci. Vector-output relu neural network prob- lems are copositive programs: Convex analysis of two layer networks and polynomial algo- rithms. arXiv preprini urXiv:2012.133: 2020.
Andrew M Stuart. Inverse problems: Bayesian perspective. Acta numerica, 19:451-559, 2010.
Cédric Villani. Topics in optimal transportation. American Mathematical Soc., 2003.
Yifei Wang, Jonathan Lacotte, and Mert Pilanci. The hidden convex optimization landscape of two- layer relu neural networks: an exact characterization of the optimal solutions. arXiv preprint arXiv:2006.05900 2020.
Yifei Wang, Peng Chen, and Wuchen Li. Projected wasserstein gradient descent for high- dimensional bayesian inference. arXiv preprint arXiv:2102.06330, 2021a.
Yifei Wang, Kangkang Deng. Haoyang Liu, and Zaiwen Wen. A decomposition augmented la- grangian method for low-rank semidefinite programming. arXiv preprint arXiv:2109.11707. 2021b.
Yifei Wang, Tolga Ergen, and Mert Pilanci. Parallel deep neural networks have zero duality gap. arXiv preprint arXiv:21/0006483 2021c.
12

---

Under review as a conference paper at ICLR 2023
Yifei Wang, Yixuan Hua, Emmanuel Candés, and Mert Pilanci. Overparameterizee relu neu- ral networks learn the simplest models: Neural isometry and exact recovery. arXiv preprint NrXiv:2209.15265, 2022.
Liuqin Yang, Defeng Sun, and Kim-Chuan Toh. Sdpnal
+: majorized semismooth newton-cg augmented lagrangian method for semidefinite program- ming with nonnegative constraints. Mathematical Programming Computation, 7(3):331-385 2015.
Alp Yurtsever, Joel A Tropp, Olivier Fercoq, Madeleine Udell, and Volkan Cevher. Scalable semidefinite programming. SIAM Journal on Mathematics of Data Science, 3(1):171-200, 2021.
Olivier Zahm, Tiangang Cui, Kody Law, Alessio Spantini, and Youssef Marzouk. Certified dimen- sion reduction in nonlinear bayesian inverse problems. arXiv preprint arXiv:1807.037122 2018.
Xin- Yuan Zhao, Defeng Sun, and Can-Chuan Toh. newton-cg augmented lagrangian method for 2010.
semidefinite programming. SIAM Journal on Optimization,
13

---

Under review as a conference paper at ICLR 2023
A CODES FOR NUMERICAL EXPERIMENT
All codes for the numerical experiment can be found in / www.gg.....ooaaaaeeeeeeomoo wwt.m//mmaaauuiitttaa.a.- submit/
Opt
B COMPARISON WITH PREVIOUS WORKS ON CONVEX OPTIMIZATION FORMULATION OF NEURAL NETWORKS
Previous works on convex optimization formulation of neural networks mainly focus on the su- pervised learning problem of two-layer neural networks using convex loss functions (e.g., squared loss. logistic loss). Our work utilizes 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 infinite 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 infinitely many neurons. However, the convex optimization approach tackles networks arbitrary width that are able to learn useful representa- tions, while the infinite width limit is quite limited (limited to basically kernel methods).
C ADDITIONAL NUMERICAL EXPERIMENT
C.1 PDE-CONSTRAINED LINEAR BAYESIAN INFERENCE
In this experiment, we consider a linear Bayesian inference problem constrained by a partial dif- ferential equation (PDE) model for contaminant diffusion in environmental engineering in domain D = (0, 1).
-kAu vu= = in D,
where is a contaminant source field parameter in domain D, u is the contaminant concentration which we can observe at some locations, K and y are diffusion and reaction coefficients. For sim- plicity, set K, 1, u(0) u(1) O, and consider 15 pointwise observations of with 1% noise, equidistantly distributed in D. We consider a Gaussian prior distribution x ~ WN C') with covariance given by a differential operator C (-6A+72) with d, 7, 0 representing the cor- relation length and variance, which is commonly used in geoscience. We set ó = 0.1, 7 = 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 finite element method with piece-wise linear elements on a uniform mesh of size 2$ & 2 1. We project this high-dimensional parameter to the data-informed low dimensions in Wang et al. (2021a) to alleviate the curse of dimensionality when applying WGD-cvxNN and WGD-NN, which we call PWGD-evxNN and PWGD-NN, respectively. For k 4 we have 17 dimensions for the discrete parameter and dimensions after projection.
We run PWGD-cvXNN p PWGD-NN using 16 samples for 200 iterations with 10-3,-03 5, Y1 = 0.95, and f/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 in the first example. From Figure 5 observe that PWGD-cvxNV achieves better root mean squared error (RMSE) than PWGD-NN for both the sample mean and the sample variance compared the reference.
14

---

Under review as a conference paper at ICLR 2023
PW/D MN
PINGD NN
PWGG--VXANN
-0.2
Loged/ensE of meant
0.0
-0.4
Loglo/RASE
0.44
-0.8
-0.6
-1.0
Rerations
3 Iterations
Figure Ten trials and the RMSE of the sample mean (top) and sample variance (bottom) by PWGD-NN and pWGD-cvxNN at different iterations. Linear inference problem.
D CHOICE OF THE REGULARIZATION PARAMETER
As the constraints in the relaxed dual problem 15 depends the regularization parameter B, is possible that for small B, the relaxed dual problem 15 infeasible. Consider the following SDP
min B, s.t. A,(A) By + Σ
( + Beatie ł 0,
n==
(20)
- AAç(A) B, + Σ
(1+) H + Best1 20,
n
n=
(1-) 2 0,r( 0, E (p).
Here the variables are 3, A and (r(itt,,,ttt ,(b-))33 Let 31 be the optimal value of the above prob- lem. Then, only for 3 2 31, there exists E Nxd satisfying the constraints in (15). In other words, the relaxed dual problem (15) is feasible. We also note that 31 only depends on the samples X and it does not depend on the value of V log π evaluated on ;1,...,NN: On the other hand, consider the following SDP
min 8, s.t. A,(Y) By
Σ
H()
Best1c441
==0
(21)
- -A,(Y) B3 + Σ rn
(1++) H Bea+1 2
(. 2 0, ((++) 2 0.j E lp).
where the variables are 3 and (,(itt),(h- ,()-))pp pp 1- Let A2 be the optimal value of the above problem. For 3 2 32, as feasible for the constraints in (I5). the optimal value of the relaxed dual problem (15) is 0. In short, only when 3 E [,3,, 32], the variational problem (15) is non-trivial. To ensure that solving the relaxed dual problem (15) gives good approximation of the Wasserstein gradient direction, shall avoid choosing 3 either too small too large.
15

---

Under review as a conference paper at ICLR 2023
E PROOFS
E. PROOF OF PROPOSITION] 1
PROOF We first note that
1
| Vlogg + V log π p122dd
| +
V log ग V log VΦ) pdr
(22)
2
/ log 2
1
V log क plppe...
We notice that the term f | log p- V log क ppd does not depend t. Utilizing the integration by parts, we can compute that
f V log VΦ) pda
/
Vp V
pdr
/ Pp, V
dx
(23)
/ Adpd..
Therefore, the variational problem
is equivalent to
1
/
V | [pd + / (V (Vlog,, v pdx +
/ A
inf t (Rd) 2
(24)
By restricting the domain C (RR) to H. we complete the proof.
E.2 PROOF OF PROPOSITION
PROOF Suppose that d; 3, w, and â₂ 83 a4, where 3r > 0 is scale parameter for E [m].
Let 0' {(w,, dm)) We note that
â anway (1o xnd = Broawi (3,
xn) (f awway (w xn);
(25)
and
ā 6.110:,1120 (10 x.,) = aullw; ( m.)) = asllwa113v (w xn).
(26)
This implies that (o(x)) t (x) and V + V (x)).. For the regularization term R(0), we note that
| + 1, 3 + 3, ||
1
1 +
= 3 + 3
w
11
31 2
(27)
=3 - 2 2/3
The optimal scaling parameter is given by a, 2-1 //
As the scaling operation does not change [wu:1B/os]] simply let ||æe|| Thus, the regularization term & R(0) becomes
la. 1/3
i Dm [| 1- This completes the proof.
N
16

---

Under review as a conference paper at ICLR 2023
E.3 PROOF OF PROPOSITION
PROOF Consider the Lagrangian function
m
L(Z, W. a, A)
1211
+
a alllsll22 1 xn) + +t((Y?2 22 2)+ + B||a||1
N
+
n
Zn
a v (I wi)
==1
(28)
=B1|a|11
(1 12 (w: (W.TTa))))))))))()) xn) Wn w)))))
2=1
1
+ 2
1217 + tr((Y + Z)) 2
For fixed W, the constraints on 7 and are linear and the strong duality holds. Thus, we can exchange the order of minz,a and max,. Thus, we can compute that
min max L(Z, W, a. A) Z.W.a A
min max min L(Z, W. ar A) W
A a,2
m
min max min
B||a|| a.
(1wwllBW" (wf (wfam) xn) X u.abb (oeeb((aTTu) ws)) + 11112 + tr((Y + A)
W
A
N
min max
A +
+
max 1ec-llon
re 12 (w: an)) yn warl (2 (x] us)
W
2
I
1=1
(29)
By exchanging the order of min and max, we can derive the dual problem:
m
max min II
+
Y
max 2
Σ Hew 11200 an) - v way aa (aHws)
5
A
W
2
mell
= max A
I|A 112
s.t.
max
| (uF xn) - yn (aTTw))
s 3, [m]
wxlla
1
max A
V s.t.
max w: will
1|ll 2 (wTxx)= xn) - wib
will wnb (x w)
- 8 [m]
(30)
This completes the proof.
E.4 PROOF OF PROPOSITION
PROOF Based on the hyper-plane arrangements DI,..., Dp. the dual constraint is equivalent to that for all j [p].
[tr(D)) - 2w TT A Xw 3 holds for all w e Rd satisfying |/|| < (2D, - I)X 2 j E [p]
(31)
This is equivalent to say that for all
-B 2 min 2tr(D;)) 2w A
(32)
s.t. ||æ]| < 1, 1,2(D; 1)Xw 0, 3 < max m 2n A D'D,WW,
s.t. ||l|| 1,2(D, 1)Xw 2 0.
17

---

Under review as a conference paper at ICLR 2023
From convex optimization perspective, the natural idea to interpret the constraint (32) is to trans- form the minimization problem into a maximization problem. We can rewrite the minimization problem in (32
trust region problem with inequality constraints:
min w" (B, + AA(A) w, WERd
(33)
s.t. || - 1,D, 1)x 20. 0.
As the problem 33 is a convex problem, by taking the dual of 333 w.r.t. we can transform (33
into maximization problem. However, (33) is region problem with inequality constraints, the dual problem of (33) can be very complicated. According to (Jeyakumar & Li (014)) the optimal value the problem 33 is bounded by the optimal value of the following SDP
tr((A,(A) + B,)2), ZE8
min
s.t. tr 2) - 0, O...., N, 14+1 1, 1, - 0.
(34)
from below.
Lemma
The dual problem of SDP
34 takes the form
max -, s.t. S = (A) B3
rall 71s41ee441 2 0,52 2
(35)
in variables r
E N+ and y E R.
INN
PROOF Consider the Lagrangian
L(Z,r,7) (r)(((())))) B,)2) + Σ Tn tr(l(()Z)) + 1(t((Z 1),
(36)
n=
where r E N+++ and E IR. By minimizing L(Z,r,1 7) w.r.t E $ we derive the dual problem 35
The constraints on A in the dual problem (13) include that the optimal value of (34 bounded from
below by -3. According to Lemma this constraint equivalent to that there exist r E R+1+ and such that
77 -35 (A(A) + B3 Mr H() tyeat1eh S, 0.
(37)
As 6d416/41 is positive semi-definite, the above condition on also equivalent to that there exist r E N+ such that
A,(A) + By
Tw H() + Bea4te ł 0,r 2
(38)
Therefore, the following convex of
N
{A A,(A) By + Σ a subset of the set satisfying the dual constraints
-.. + Beat 2 0, p(d--) r(0--)20} 0)
(39)
==0
" (B, (B,+ (A()) w 2 25
min
(40)
[|w]] (20, 1w20
18

---

Under review as a conference paper at ICLR 2023
On the other hand, the constraint on A
max [| -0
w T (BB++A(())) w S
(41)
is equivalent to
min | =2
=" (B, A2(A)) -B.
(42)
By applying the previous analysis the above trust region problem, the following convex set of
N
{ (A:-A,(A) - By + Σ H( + B = U, rlit) 2 20}
(43)
a subset of the of A satisfying the dual constraints
{ {A: A:
max T + A,(A)) w S 8 (200:
(44)
Therefore, replacing the dual constraint ma/u:l/lllas1 En (w xn) - w/// (srTw))
s
by
A,(A) + B + Σ
G0) H
Be841c = 0. lp).
n
(45)
AA(A) B3 + Σ
( 1 + Beat1e = 0,J lpl.
==
,..-) 2 0, (..t)) 2 0, E [p).
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.5 PROOF PROPOSITION
PROOF Consider the Lagrangian function
L(A,r,) = - | +
P Σtn
(SU.-) (A(A) +By + Σ
HU4 + zeatierr))
tr
n
1=1
n=0
Σ tr g(++t)
-AA(A)) - B, +
Σ
N
(( + gesteedt))
(1+t)
n
j=1
n=0
(46)
where we write
r= p (1,-) ((P.-), (1.t) ..... rlatt) e (I-
(47)
S (s0,-) s(P.-) ,((,,)) s(ett) E 1s4+))2P.
Here we write s++1
{S gd+1ls 2 0}. By maximizing w.r.t. and r, we derive the bi-dual
problem (6).
E.6 PROOF THEOREM
Suppose that (Z, W.,) is feasible solution to
11
Let Dj....,, Din be the enumeration of
(diag(I(Xw, 2 0)) e [m]]. For i [k], we let
wi]
sU++)
Σ
ay
s(,.)) 0,
(48)
1
Io di.g (x)
19

---

Under review as a conference paper at ICLR 2023
and
gla+t) = =0, S(0
Σ
al
w w/
Wi
(49)
For j $ f we simply set Sli+) = 0, S(,-) = 0. As l|wlll S and Dj. ((ws; 2 0),
can verify that (r($6:))11)) I 0, 0,tr(((i+)) S 0 are satisfied for j fI, ....... . Im and
0.1,...... N. This is because for O, as H())
Ia
it follows that
0 -1
(r(S( (())))
=
Σ
on 20rl21 1) S 0,
1a:20,dlag (1 (Xxr20))=Di
(50)
S H
Σ
ar(1/ur12 - 1) s 0.
For n = ,,...,,,, we have
tr(S(h++))H) (())))=
Σ
201 - 2(D) 31 dun) wy 50
ti 20,dnni (m)))))))
(51)
tr(s(h--)
Σ
an ((D)) wi S 0.
Lor (0
Based on the above transformation, we can rewrite the bidual problem in the form of the primal problem (2... For S E gd+1, note that
tr(SAA(A))
+ X" --2tr(A D
where S1:d,1:d denotes the x d block S consisting the first rows and columns. This implies that A3(S) = -2D, Hence, we have
(Stiit -S(h-))= [ 2w)
Σ
20, X Xwmwx =-
Σ
201(X (Xw)))
[dlla([(Xe:00)
Therefore, we have
a; (Xw:)++ Σ Az(5(h-) - = _o 1=1
1=1
As n-th row of satisfies that on = 2 2m (.W;(IT w;)+ this implies that
m
p
Z= 2 _ a (Xw:)
-
S(i+t)..
1=1
Hence (Z, (SSU)-) (SU.-)) ((5)))))) feasible to the relaxed bi-dual problem (16).
We can also compute that
tr(B, - - 2
I [s w. > 0) || w: w
1=1 n=1
and
Σ ((tiit) -li. ((i T
Σ las--
Thus, the primal problem 12) with W.) and the relaxed bi-dual problem 16 with (Z, (SU.)) (SS0- }) have the same objective value.
20


Reviews:
[{"review_id": "zUIsKd-n9C9", "reviewer": "Reviewer_efky", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The manuscript proposes a convex relaxation for the approximation of sample based Wasserstein gradients with a KL divergence as an objective. The motivation for this is the application to Bayesian inference problems. \nTo derive the relaxation, the drift in the ODE describing the evolution of the particles can be characterized as the minimizer of an energy functional, which can be used as a training objective for a shallow neural network. A convex relaxation as a semidefinite program of the dual of this optimization problem is presented. Further, the dual of this  relaxation is given and used in order to show that the relaxation  provides a lower bound to the primal problem. \nThe benefits of the proposed method are demonstrated on three different problems, on a synthetic toy problem and on two Bayesian inference problems, which also include real world data.  ", "summary_of_the_review": "The paper studies an important and timely problem and provides a nice extension of recent works on SDP relaxations to neural network approximation of Wasserstein gradients. Overall, I enjoyed reading the manuscript as the overall idea is clearly laid out. For more detailed comments on the different aspects, I would like to refer to my comments above.\nCurrently, I feel that the manuscript requires improvements in different aspects, which I summarize below:\n\nThings that I find necessary to see addressed:\n\n* *Relation to prior works on convex relaxations:* I think it is important to highlight the difference of the proposed relaxation to existing works on convex relaxations of neural network training problems. \n* *Clarity:* Particularly Section 3 is not easy to digest in its current form. A lot of things can be improved by proof reading the section and fixing formatting (e.g. in Remark 4 and Lemma 4, Proposition 4 and the paragraph below Theorem 1). Concrete comments for Proposition 4: There is currently no statement made in the proposition, so it causes some confusion in its current form; the formatting should definitely be improved; ideally, the symbols should be introduced before the equation; $[p]$ notation is not introduced as a notation; there is a comma at the end of (15) and it continues with upper case; there is a fullstop missing after $[N]$; $A_j(\\Lambda)$ is introduced but doesn’t appear in (15). \n* *Optimality:* In its current form the title suggests evidence on the optimality of the proposed method for the approximation of Wasserstein gradients. Further, you state in the conclusion that „the gradient … is at least as good“. Currently I do not see results supporting this in the manuscript. Do you see this as a consequence of Theorem 1 or of your experiments? Theorem 1 implies that the regularized dual gives a lower bound on the primal and your experiments show that the resulting WGD converges faster. This does however not imply that the approximation of the Wasserstein gradient is more exact as far as I understand it. I think the claim of optimality needs to be supported better or needs to be adjusted accordingly.\n\nThings I recommend to address:\n* *Experiments:* It would be good to have experiments comparing the primal and dual problem in the following sense: What is the deviation of the two update directions of WGD-NN and WGD-cvxNN away from the true Wasserstein gradient. In particular, this could be an empirical proof that the approximation of the Wasserstein gradient is indeed better with the proposed method. \n* *Consequences of Theorem 1:* The manuscript would benefit from an elaboration of the implications of Theorem 1. In particular, what does it say about the relation between the proposed relaxation and the primal problem?\n* *More general setup:* The proposed method is only presented for the KL divergence as an objective. Although this is very well motivated by the application in Bayesian inference tasks, the contributions of the manuscript would be stronger, if if was formulated and analyzed in a more general setting.\n\nSpecific questions:\n* Why do you call the proposed approximation optimal, i.e., does it satisfy any optimality criteria?\n* What are the implications of Theorem 1?\n* Proposition 1:\n\t* What happens with the boundary terms in the partial integration?\n\t* Why do you formulate Proposition 1 for a function space? I think it is just an alternative expression of the energy due to partial integration. Being nit-picky, later you consider the problem on a function class (so no linear structure), which might confuse people if they read „function space“ is something linear. I would rather simply say that the energy in (4) takes this form. \n\nFurther thoughts (not so important): \n* The experiments could be stronger if not only compared against other methods relying on Wasserstein gradient approximations, but rather on  Bayesian inference techniques.\n* *Variational formulation:* I am not entirely sure, whether I appreciate the wording here. The reason for this is that I am unsure whether (4) is the variational formulation of (2) in the sense that I am unsure, whether its Euler-Lagrange equations characterize the Wasserstein gradient. Think for example of a linear PDE $-\\Delta u = f$ (lets say with zero boundary values), which is both the unique minimizer of $\\frac12\\int \\lvert\\nabla u\\rvert^2dx-\\int fu dx$ and of $\\int \\lvert \\Delta u+f\\rvert dx$, where the first one is called the variational formulation and the second one a residual minimization. \n* *Formatting (pretty nit-picky):* Line breaks in section titles are not good; a single subsection in a section is not too nice, maybe a \\paragraph does the job; ", "strength_and_weaknesses": "**Strengths:**\n* The Introduction is very well written and the overall idea of the paper presented clearly, which leads to a smooth overall read. \n* I appreciate the idea of the convex SDP relaxation of the NN training problem encountered in the variational formulation of the Wasserstein gradient approximation with shallow networks.\n* The empirical results show that the approach can be made effective even for larger problems and include problems with real world data.\n\n**Weaknesses:**\n* Where the introduction is written very nicely, the later sections are not always a very smooth read. This is mainly due to a) linguistic errors, which can easily be fixed b) rather heavy notation and suboptimal formatting, which can also be addressed.\n* The discussion of related works in the introduction does a good job giving an overview over approximations of WGD. However, the relation to prior works on SDP relaxations of neural network training is not sufficiently well described. Some of these works are mentioned in the introduction, but their precise relation and in particular the difference of the manuscript to prior works here is not described. This would be important to add in order to make the contributions of the manuscript clear.\n*  The experiments would be stronger if not only compared against other methods relying on Wasserstein gradient approximations, but rather on  Bayesian inference techniques.\n* The title of the paper implies that using the proposed relaxation yields a better approximation of the Wasserstein gradient compared to solving the primal problem. This is however neither supported sufficiently well in an experiment nor a theoretical result as far as I can see.  \n* The theoretical aspects of the proposed method are not very well elaborated. In particular, the implications of Theorem 1 for the method are not laid out very transparently. \n* The proposed method is only presented for the KL divergence as an objective. Although this is very well motivated by the application in Bayesian inference tasks, the contributions of the manuscript would be stronger, if if was formulated and analyzed in a more general setting.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "gUePIiDGWpZ", "reviewer": "Reviewer_4MPo", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper uses an SDP relaxation for the problem of computing the Wasserstein gradient. They use numerical algorithms that make their proposed method suitable for practical scenarios involving Bayesian inference.", "summary_of_the_review": "Overall the paper uses a straightforward technique to a problem that has limited practical applications (learning two-layers neural nets).", "strength_and_weaknesses": "Strength\n======\n\n- The idea of using a convex SDP relaxation of the dual of the variational primal problem is interesting to understand the behaviour of a particular family of two-layer neural nets.\n\n\nWeaknesses\n=========\n\n- Using convex SDP relaxations is not a new idea per se: see, e.g., \"Semidefinite Relaxation of Quadratic Optimization Problems\" by Tom Luo *et al*.\n\nFeedback\n=======\n\n- This sentence: \"However, due to the nonlinear and nonconvex structure of neural networks, optimization algorithms\nincluding **stochastic gradient descent may not find the global optima** of the training problem\" is presented as if, in general, not finding the global optima is an issue in SGD. In training neural nets you precisely do not want to train till the global optima as that would likely mean **overfitting**.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "4TD9OPa52_P", "reviewer": "Reviewer_mWPZ", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper tackles the issue of computing Wasserstein gradient direction for 2-layers NNs thanks to a SDP approach with main properties that there is no need to train the underlying NN. The gradient direction is computed with the dual formulation of a least-square + a polynomial regularization term.", "summary_of_the_review": "I believe the paper is based on an interesting idea but lacks clarity in its exposition, and fails to convince on the practical aspect.", "strength_and_weaknesses": "**Strenghts**\n- This is probably the first work to cast this problem as a SDP problem\n- The optimal is global (despite being a non convex problem initially)\n- The paper is well-written\n\n**Weaknesses**\n- The exposition is lacking an I believe some part of the paper could be moved to the appendix to free some space to better put in perspective the related work on that matter\n- Theoretical statements lacks precision.\n- There is no clear advantage both theoretically and practically to not train directly the NN on the Wasserstein gradient.\n- Finally, and probably the most important issue, is that the computational cost is prohibitive. It is not clear why at least there is no discussion on how to solve the SDP problem when d is large.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject