Optimal Neural Network Approximation of Wasserstein Gradient Direction via Convex Optimization
Abstract
The computation of Wasserstein gradient direction is essential for posterior sampling problems and scientiﬁc computing.
The approximation of the Wassersteingradient with ﬁnite samples requires solving a variational problem.
We study thevariational 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 broaderfunction family including two-layer networks.
By solving the convex SDP, we obtain the optimal approximation of the Wasserstein gradient direction in this class offunctions.
Numerical experiments including PDE-constrained Bayesian inferenceand parameter estimation in COVID-19 modeling demonstrate the effectiveness ofthe proposed method.
Introduction
Bayesian inference plays an essential role in learning model parameters from the observationaldata with applications in inverse problems, scientiﬁc computing, information science, and machinelearning (Stuart, 2010).
The central problem in Bayesian inference is to draw samples from a posteriordistribution, 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 beeffective in drawing samples from a posterior distribution, which attracts increasing attention in recentyears.
For instance, the Wasserstein gradient ﬂow of Kullback-Leibler (KL) divergence connects tothe overdampled Langevin dynamics.
The time-discretization of the overdamped Langevin dynamicsrenders the classical Langevin Monte Carlo Markov Chain (MCMC) algorithm.
In this sense, thecomputation of Wasserstein gradient ﬂow yields a different viewpoint for sampling algorithms.
Inparticular, the Wasserstein gradient direction also provides a deterministic update of the particlesystem (Carrillo et al., 2021b).
Based on the approximation or generalization of the Wassersteingradient direction, many efﬁcient sampling algorithms have been developed, including Wassersteingradient descent (WGD) with kernel density estimation (KDE)
(Liu et al., 2019), Stein variationalgradient descent (SVGD) (Liu & Wang, 2016), and neural variational gradient descent (di Langoscoet al., 2021), etc.
Meanwhile, neural networks exhibit tremendous optimization and generalization performance inlearning complicated functions from data.
They also have wide applications in Bayesian inverseproblems (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 non-linear activations and a sufﬁcient number of neurons.
Functions represented by neural networksnaturally provide an approximation towards the Wasserstein gradient direction.
However, due to the nonlinear and nonconvex structure of neural networks, optimization algorithmsincluding 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, 2021), theregularized training problem of two-layer neural networks with ReLU/polynomial activation canbe formulated as a convex program.
The optimal solution of the convex program renders a globaloptimum of the nonconvex training problem.
In this paper, we study a variational problem, whose optimal solution corresponds to the Wassersteingradient 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 neuralnetwork to minimize the loss may get the neural network stuck at local minima or saddle points andit often leads to biased sample distribution from the posterior.
Instead, we analyze the convex dualproblem of the training problem and study its semi-deﬁnite program (SDP) relaxation by analyzingthe geometry of dual constraints.
The resulting SDP is practically solvable and it can be efﬁcientlyoptimized by convex optimization solvers such as CVXPY (Diamond & Boyd, 2016).
We then derivethe corresponding relaxed bidual problem (dual of the relaxed dual problem).
Thus, the optimalsolution to the dual problem yields an optimal approximation of the Wasserstein gradient directionin a broader function family.
We also present a practical implementation and analyze the choice ofthe regularization parameter.
Numerical results including PDE-constrained inference problems and
Covid-19 parameter estimation problems illustrate the effectiveness of our proposed method.
1.1 Related 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 Langoscoet al.
(2021) learns the transportation function by solving an unregularized variational problem in thefamily of vector-output deep neural networks.
Compared to these studies, we focus on a convex SDPrelaxation of the varitional problem induced by the Wasserstein gradient direction.
Meanwhile, Fenget al.
(2021) form the Wasserstein gradient direction as the mininimizer the Bregman score and theyapply deep neural networks to solve the induced variational problem.
Background
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 divergencefunctional.
Later on, we design a neural network convex optimization problems to approximate the
Wasserstein gradient in samples.
2.1 Wasserstein 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 ( ρ k π ) is the KL divergence fromρ to π .
The variable is the density function ρ in the space P = { ρ ∈ C ∞ ( R d ) |
R ρdx = 1 , ρ > 0 } .
The function π ∈ C ∞ ( R d ) is a known probability density function of the posterior distribution.
Bysolving 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 theoptimization problem (1) satisﬁes
[Equation]
where ρ t ( x ) =
ρ ( x, t ) and δδρ t is the L
ﬁrst variation operator w.r.t.
ρ t .
In the above third equality,a fact ρ t ∇ log ρ t = ∇ ρ t is used.
Here ∇ · F denotes the divergence of a vector valued function 80 F :
R d → R d and ∆ is the Laplace operator.
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 theequilibrium, 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 Wassersteingradient descent (WGD) on the particle system { x nl } follows
[Equation]
(3) where Φ l : R d → R is a function which approximates log ρ l − log π and α l > 0 is the step size.
2.2 Variational formulation of WGD
Given the particles { x n } Nn =1 , we design the following variational problem to choose a suitable 90 function Φ approximating the function log ρ − log π .
Consider
[Equation]
The objective functional evaluates the least-square discrepancy between ∇ log ρ − ∇ log π and ∇ Φ 92 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 formulationof (4) in H .
Proposition 1 Let H ⊆
C 1 ( R d ) be a function space.
The variational problem (4) in the domain H 96 is equivalent 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 99 regularized Stein discrepancy deﬁned in (di Langosco et al., 2021) between ρ and π based on Ψ .
Thisproblem is also similar to the varitional problem for the score matching estimator in (Hyvärinen &
Dayan, 2005) by parameterizing Φ in a given probabilistic model.
In comparison, our method can beviewed as a special case of score matching by using a two-layer neural network model.
Therefore, by replacing the density ρ by ﬁnite samples { x n } Nn =1 ∼ ρ , the problem (5) in terms of 104 ﬁnite samples forms
[Equation]
Optimal neural network approximation of Wasserstein 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 alsopresent a practical implementation estimation of ∇ log ρ − ∇ log π and discuss the choice of the 110 regularization parameter.
Let ψ be an activation function.
Consider the case where H is a class of two-layer neural networkwith 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] [Equation]
Here ◦ represents the element-wise multiplication.
By adding a regularization term to the variationalproblem (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 termin (9), which makes the ReLU activation inadequate.
For this activation function, we consider theregularization function R ( θ ) =
{+P mi =1 ( k w i k 
 + | α i | 3 ) .+}
Remark
We note that ∇ Φ θ ( x ) and ∆Φ θ ( x ) are all piece-wise degree- 3 polynomials of theparameters θ .
Hence, we consider a speciﬁc cubic regularization term above, analogous to (Bartan &
Pilanci, 2021).
By choosing this regularization term, we can derive a simpliﬁed convex dual problem.
By rescaling the ﬁrst and second-layer parameters, the regularized variational problem (10) can beformulated as follows.
Proposition 2 (Primal problem)
The regularized variational problem (10) is equivalent to
[Equation]
11) where ˜ β
= 3 · 2 − 5 / 3 Nβ .
For simplicity, we write Y =  
[Equation]
π
[Equation]
x 1 ) T .
[Equation]
∇
[Equation]

 ∈ R N × d .
We introduce the slack variable z n =
P mi[-=
-]
[Equation] 1 α i w i ψ 0 ( x Tn w i ) for n ∈ [ N ] and denote Z = [ z. . .
z N ] T ∈ R N × d .
Then, we can simplifythe problem (11) to
[Equation]
[Equation] Based on the above reformulation, 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) .
3.1 Analysis of dual constraints and the relaxed dual problem
Now, we analyze the constraint
[Equation]
in the dual problem.
We note that this constraint is closely related to the regularization parameter,which we will discuss later.
For simplicity, we take ψ 00 (0) =
as the subgradient of ψ 0 ( z ) at z = 0 ,i.e., taking the left derivative of ψ 0 ( z ) at z = 0 .
Let X =
[ x 1 , . . .
, x N ] T ∈ R N × d .
Denote the set ofall 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 cardinalityof 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.
It gives a lower bound for the optimal value of the dual problem (13).
Proposition 4 (Relaxed Dual problem) Consider the following SDP:
[Equation]
The variables are Λ ∈ R N × d and r ( j, − ) , r ( j, +) ∈ R n +1 for j ∈ [ 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 (Λ) 0 (cid:21) , ˜ B
j = (cid:20)
B j  0 (cid:21) ,
H ( j )=
(cid:20)
I d  − 1 (cid:21) and H ( j ) n = (cid:20) (1 − 2( D j ) nn ) x n (1 − 2( D j ) nn ) x
Tn (cid:21) , n {+∈ [ N ] The vectore d +1 ∈ R d +1 satisﬁes that ( e d +1 ) i = 
 for i ∈ [+}
The optimal value of (15) gives a lower bound on the dual problem (13) , and hence on the primalproblem (12) .
In the following proposition, we derive the relaxed bi-dual problem.
It can be viewed as a convexrelaxation of the primal problem (12).
Proposition 5 (Relaxed bi-dual problem)
The dual of the relaxed dual problem (15) is as follows
[Equation]
in variables Z ∈ R N × d , S ( j, +) , S ( j, − ) ∈ S
d +1+ for j ∈ [ p ] .
Here A ∗ j is the adjoint operator of the 157 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.
We can show that any feasible solutions of the primal problem (11) can be mapped tofeasible solutions of (16).
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 therelaxed 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) ata 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 suchthat
( 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 therelaxed 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 173 in the family of two-layer squared-ReLU networks (7).
Smaller or equal objective value of the relaxedbi-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 176 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 ) isoptimal 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 abroader class of functions represented by { S ( j, +) , S ( j, − ) } pj =1 .
This contains more variables than theneural network function.
3.2 Practical implementation
Although the number p of all possible hyper-plane arrangements is upper bounded by 2 r (( N − 1) e/r ) rwith r = rank( X ) , it is computationally costly to enumerate all possible p matrices D 1 , . . .
, D p torepresent the constraints in the relaxed dual problem (4).
In practice, we ﬁrst randomly sample Mi.i.d.
random vectors u 1 , . .
, u M ∼ N (0 , I d ) and generate a subset ˆ S of S as follows:
[Equation]
Then, we optimize the randomly sub-sampled version of the relaxed dual problem based on the subsetˆ S and obtain the solution Λ .
We then use − Λ − Y 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 particlesystem 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 dualproblem (4) infeasible, we increase ˜ β by multiplying γ − 1 2 , where γ 2 ∈ (0 , 1) .
Detailed explanationof the adjustment of the regularization parameter can be found in Appendix C.
The overall algorithmis summarized in Algorithm 1.
We note that the randomly subsampled version of the relaxed dual problem (15) involves 2 N ˆ p inequality constraints and 2ˆ p linear matrix inequality constraints with size ( d + 1) ×
( d + 1) .
Applying the standard interior point method (Boyd et al., 2004) leads to the computational time up to
O ((max { N, d 2 } ˆ p ) 6 ) . 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 and Y l based on { = and { log } Nn = Solve Λ l from the relaxed dual problem (15) with ˜ β
= ˜ β 4: if the relaxed dual problem with ˜ β = ˜ β l is infeasible then 5: Set
X l +1 = X l for n ∈ [ N ] and set ˜ β l +1 = γ − 1 2 ˜ β l . 6: else 7: Update X
{+l +
 = X l+} + α l (Λ l + Y l ) for n ∈ [ N ] and set ˜ β l +1 = γ 1 ˜ β l .
8: end if 9: end while
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., 2021) to reduce the parameter dimension d to a data-informed intrinsic dimension ˆ d ,which is often very low, i.e., ˆ d (cid:28) d .
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-cvxNN).
The performance of the two methods is assessed by the sample goodness-of-ﬁtof the posterior.
For WGD-NN, in each iteration, it updates the particle system using (3) with afunction Φ represented by a two-layer squared ReLU neural network.
The parameters of the neuralnetwork is obtained by directly solving the nonconvex optimization problem (10).
We note thatit takes longer time by WGD-cvxNN (compared to WGD-NN) to solve the convex optimizationproblem.
However, this optimization time is often dominated by the time in likelihood evaluation ifthe model is expensive to solve.
Moreover, the induced SDPs have speciﬁc structures of many similarconstraints, whose solution can be accelerated by designing a specialized convex optimization solver.
This is left for future work.
4.1 A toy example
We test the performance of WGD on a bimodal 2-dimensional double-banana posterior distributionintroduced 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 evaluatethe 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, weuse N = 50 samples and run for 100 iterations with step sizes α l = 10 − 3 .
For WGD-cvxNN, weset β = 1 , γ 1 = 0 .
95 and γ 2 = 0 .
For WGD-NN, we use m = 200 neurons and optimize theregularized 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 anddecrease 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 ﬁnalstep of 100 iterations are shown in Figure 1.
It can be observed that WGD-cvxNN provides morerepresentative 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 thereference SVN samples at each iteration.
We observe that the samples by WGD-cvxNN achievesmuch smaller MMD than those of WGD-NN compared to the reference SVN samples, which isconsistent with the results shown in Figure 1.
For WGD-cvxNN, it takes 572s in total, while for WGDNN, it takes 16s in total.
WGD-cvxNN takes much longer time than WGD-NN as WGD-cvxNNaims to solve for the global minimum of the relaxed convex dual
4.2 PDE-constrained nonlinear Bayesian inference
In this experiment, we consider a nonlinear Bayesian inference problem constrained by the followingpartial differential equation (PDE) (Chen & Ghattas, 2020) with application to subsurface (Darcy)ﬂow in a physical domain D
= (0 , 1) 2 ,
[Equation]
D
[Equation]
(19) where u is pressure, v is velocity, h is force, e x is a random (permeability) ﬁeld equipped witha Gaussian prior x ∼ N ( x 0 , C ) with covariance operator C = ( − δ ∆ + γI ) − α where we set 240 δ = 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 electricalimpedance in medical imaging, We impose Dirichlet boundary conditions u = 1 on the top boundaryand u = 0 on the bottom boundary, and homogeneous Neumann boundary conditions on the leftand right boundaries for u .
We use a ﬁnite element method with piecewise linear elements for thediscretization of the problem, resulting in 81 dimensions for the discrete parameter.
The data isgenerated 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 areused as the reference values to assess the goodness of the samples by pWGD-cvxNN and pWGD-NN.
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 samplemean and sample variance are shown in Figure 3 for the two methods at each of the iterations.
Wecan observe that pWGD-cvxNN achieves smaller errors for both the sample mean and the samplevariance compared to pWGD-NN at each iteration.
Moreover, pWGD-cvxNN provides much smallervariation of the sample mean and sample variance for the ten trials compared to pWGD-NN.
80 # iterations 1.
L o g 10 ( R
M S E o f m e a n ) pWGD
[Figure]
NN pWGD
[Figure]
cvxNN    80 # iterations 0.
L o g 10 ( R
M S E o f v a r i a n c
[Figure] [Figure]
4.3 Bayesian inference for COVID-
In this experiment, we use Bayesian inference to learn the dynamics of the transmission and severityof COVID-19 from the recorded data for New York state, as studied in Chen & Ghattas (2020).
We use the model, parameter, and data as in Chen & Ghattas (2020).
More speciﬁcally, we use acompartmental model for the modeling of the transmission and outcome of COVID-19.
We take thenumber of hospitalized cases as the observation data to infer a social distancing parameter, a timedependent stochastic process that is equipped with a Tanh–Gaussian prior to model the transmissionreduction effect of social distancing, which becomes 96 dimensions after discretization.
We run a projected Stein variational gradient descent (pSVGD) method Chen & Ghattas (2020) asthe reference, and run pWGD-cvxNN and pWGD-NN using 64 samples for 100 iterations with stepsize α l = 10 − 3 , where we set β = 10 , γ 1 = 0 .
95 , and γ 2 = 0 .
95 10 for both methods as in the lastexample.
From Figure 4 we can observe that pWGD-cvxNN produces more consistent results withpSVGD than pWGD-NN for both the sample mean and 90% credible interval, both in the inferenceof the social distancing parameter and in the prediction of the hospitalized cases.
[Figure]
Conclusion
In the context of variational Wasserstein gradient descent methods for Bayesian inference, we considerthe approximation of the Wasserstein gradient direction by the gradient of functions in the family oftwo-layer neural networks.
We propose a convex SDP relaxation of the dual of the variational primalproblem, which can be solved efﬁciently using convex optimization methods instead of directlytraining the neural network as a nonconvex optimization problem.
In particular, we established thatthe gradient obtained by the new formulation and convex optimization is at least as good as theoptimal approximation of the Wasserstein gradient direction by functions in the family of two-layerneural networks, which is demonstrated by various numerical experiments.
In future works, weexpect to extend our convex neural network approximations to generalized Wasserstein ﬂows.
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for information onhow to answer these questions.
For each question, change the default [TODO] to [Yes] , [No] , or[N/A] .
You are strongly encouraged to include a justiﬁcation to your answer , either by referencingthe appropriate section of your paper or providing a brief inline description.
For example:
• Did you include the license to the code and datasets?
[Yes] See Section ?? .
• Did you include the license to the code and datasets?
[No] The code and the data areproprietary.•
Did you include the license to the code and datasets?
[N/A]
Please do not modify the questions and only use the provided macros for your answers.
Note that the
Checklist section does not count towards the page limit.
In your paper, please delete this instructionsblock and only keep the Checklist section heading above along with the questions/answers below.
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’scontributions and scope?
[Yes](b) Did you describe the limitations of your work?
[Yes](c)
Did you discuss any potential negative societal impacts of your work?
[N/A](d)
Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes]2.
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes](b) Did you include complete proofs of all theoretical results?
[Yes]3.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes](b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
[Yes](c)
Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[No]
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[No]4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[N/A]
(b) Did you mention the license of the assets?
[N/A]
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you’reusing/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identiﬁable information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, ifapplicable?
[N/A](b)
Did you describe any potential participant risks, with links to Institutional ReviewBoard (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation?
[N/A]