Venue: NeurIPS 2022 Conference

Reviews:

final_score: rating: 6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
review_content: summary: This paper tackles the variational problem of Wasserstein gradient descent (a.k.a Wasserstein gradient flow with the KL discrepancy). The (variational) optimization problem is restricted to 2-layer neural networks with ReLU activation functions. The main contribution of this paper is to propose a convex (SDP) bi-dual problem that approximates the global solution of the original problem. 

strengths_and_weaknesses: # Strengths and weaknesses
1) Strengths:
- An interesting convexification of a non-convex variational problem, an idea that could be applied in a more general setting
- Even though strong duality does not hold, the authors provide a change of variable (primal -> bi-dual) such that the objective value of feasible solutions are equal

2) Weaknesses:
- The core idea could benefit from more intuitive explanations and comparisons with relatex work (see below)
- The experiments lack the important computational discussion

questions: # Questions

1. Theorem 1 shows that the bidual problem can still have feasible solutions with an objective equal to the primal problem: does the mapping go both ways ? 
2. Can we uncover additional properties about this mapping between the primal and the bidual ? At least at optimality ?
3. A few references about convexified 2-layer network problems are discussed in 38-41, how does the proposed method relate to these works ?
4. The obtained bidual is an SDP, is this equivalent to the idea of infinite width neural networks modeled as probability measures ? 
5. The result obtained in Equation (18) is a bit confusing, Since J corresponds to the bidual loss, I'm not sure what it means about the primal one (which is the one of interest), specially if eq (18) can be a strict inequality. How is that interpretable in terms of "approximating" the global minimum ? Can Z* and tilde(Z*) be analytically compared ? Since this justifies the claim "we solve the global minimization problem", I believe it deserves more clarification.
6. The authors briefly mention the computational burden of the proposed method in L232 being 50x slower than training a neural net. This is not surprising (SDPs are hard to solve) however, I believe this limitation is a very important aspect to investigate. Even though the experiments indicate that lower minima are obtained, the important question of quantifying the ratio:  obtained gain / large computational burden should be answered or at least discussed in more detail. Would such a method still work in large dimensions ? deeper neural networks ? large data sizes ?
7. I'm not sure I follow the point of the Covid-19 experiment, no comment is provided to interpret or analyze the obtained results. Moreover, why pSVGD is considered a reference ?

minor:
 when writing integrals, \int f dx does not make sense: either \int f or \int f(x) dx should be used


typos:
 L52.  we also present *a* practical ..
 L79. In *the* above third ..
 Eq (3): alpha_l is not defined

limitations: see above
review_title: null

final_score: rating: 6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
review_content: summary: The paper considers the problem of implementing Wasserstein gradient flow for the objective of sampling from a target distribution i.e. minimizing KL divergence. It is common to implement the Wasserstein gradient flow with a deterministic system of particles. For minimizing KL divergence, the vector-field  that governs the particles involves the interaction term $\nabla \log(p)$ where $p$ is the distribution of the particles. The most challenging step is to approximate $\nabla \log(p)$ in terms of particles, i.e. the direction of the gradient. This has been subject of numerous recent works in the literature. The authors consider the score-function minimization approach, which involves a stochastic optimization problem over space of functions. The novelty of the work is in the particular choice to representation the functions  and the optimization algorithm to approximate the optimal one. In particular, the authors propose to use a 2-layer neural network representation, consider the dual of the regularized version of the problem, and introduce a convex relation that produces an approximate solution to the original problem. The proposed approach is then numerically evaluated on several benchmark sampling examples.      
strengths_and_weaknesses: Strength: 
- The paper is well-written and clear
- Although there are numerous works on the topic, the topic is still interesting and there are open questions
- The idea of using this particular convexification approach in approximating the Wasserstein gradient is new and interesting

Weakness: 
- Computationally expensive approach (in its current form) 
- lack of theoretical discussion regarding the final proposed approach 
questions: Questions: 
1- Can you be more explicit and provide a discussion on the computational time of the proposed approach and how it scales with problem parameters? 
2- Are there any tuning happening for NN and CVX_NN approach?  
3- Does it help if you exclude estimating \nabla $\log (\pi)$ from the formulation because it is known? 
4- Can you say anything about the convergence of the gradient flow if the optimization problem to compute the gradient direction is solved up to a certain error? 
5- Can you be more explicit about the connection between the exact solution obtained from solving the relaxed convex dual problem and the exact solution to the original problem? 


limitations: Yes
review_title: null

final_score: rating: 5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
review_content: summary: This paper is to solve the Wasserstein gradient direction, which is the velocity in Wasserstein gradient flow, without training neural networks. The algorithm is based on a forward discretization of the Wasserstein gradient flow, so once the gradient direction is solved, you can calculate the positions of the particles in the next step. To solve the gradient direction, they firstly construct a least-square regression problem, and add a third-order polynomial regularization term. Then they derive the dual problem in several steps. Finally, they obtain a semi-definite programming relaxation problem and show the duality gap between the relaxed dual problem and the regularized dual problem is zero. The paper provides several illustrative examples.
strengths_and_weaknesses: **Strength:**

The method can obtain the velocity in Wasserstein gradient flow without training a neural network. They transform the primal problem based on a two-layer neural network into a SDP problem, which can enjoy a similar functional class richness. Unlike SGD in training neural networks, this algorithm can find the global minimum.

The paper gives rigorous derivation toward a relaxed SDP dual problem.

**Weakness:**

The primal problem is regularized, which introduces some bias to the gradient direction. 

Several limitations: the target function is only for KL divergence. The function class is only two-layer neural networks. 

There is not enough discussion on the computational complexity. I'm rather interested in the complexity dependency on the dimension d, neural network width m, and the number of particles. 

The main advantage of this method is no need to train the neural network, which can be very time-consuming. However, the method even takes a longer time than training nn, which makes me a bit disappointed. Can you add more explanations about why WGD-cvxNN takes much longer time? 
questions: **Question/suggestion**

I suggest moving some intermediate propositions to the supplementary materials and adding more transitional contexts, for example, between theorem 1 and the paragraph afterward. The authors can also add more explanations to the notation, e.g. what's the difference between $r^{(j,-)}$ and $r^{(j,+)}$? And I suggest presenting the optimization variables right below min or max in all the optimization problems, like in eq 16 and eq 17.

The dimension reduction technique (row 200) is quite important to accelerate the algorithm. I suggest the authors add more details there.

What does p mean in pWGD?

What does the y-axis mean in the social distancing plot in Figure 4? The social distance in meters?

I suggest adding a reference in section 4.2 about the equation of PDE.

In experiments, different examples use different baseline methods to compare, which is confusing. Why do you use so many different baselines?

In Figure 1, how do you calculate the posterior density? Do you use kernel density estimation or you can precisely calculate it?

**Typo:**

In Figure 4 captions, it should be left / right instead of top/bottom.

In row 210 bracket, compared to WGD-NN only

Overall, I think the writing of this paper can be improved. For example, the equations can be presented in a more concise manner. And the computational complexity needs to be more carefully analyzed.
limitations: discussed in weakness
review_title: null

final_score: rating: 5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
review_content: summary: Implementability of Wasserstein gradient flows, which find many applications in sampling and numerical PDEs, hinges on the computation of the Wasserstein gradient. This paper proposes to find the best two-layer squared ReLU activation neural network approximation for the Wasserstein gradient direction. It is shown that this approximation can be formulated as a convex program, albeit with too many constraints to be tractable. The program is then solved approximately by randomly sampling a subset of the constraints. Experiments show that the resulting algorithm has reasonable performance.
strengths_and_weaknesses: Seeking new implementations of Wasserstein gradient flows is an important practical problem, and in that respect the approach of this paper is novel and interesting. However, it also seems to fall short when evaluated via either theoretical interest or practical utility.

Regarding practical utility: although the experiments are reasonable, they are largely toy experiments. There is no discussion of the runtime of the method, which leads me to believe that the method is not scalable (since each step requires solving a convex program, it seems computationally heavy). In applications, it appears that this approach yields no benefits over simply parameterizing the Wasserstein gradient via neural networks and training the neural network directly, rather than solving a convex program.

Therefore, it seems more appropriate to evaluate the paper on its theoretical merits, but unfortunately there is not much to speak of. One of the main claims of the paper is that one can use principled convex optimization solvers to implement the method, but in the end the convex program cannot be implemented directly and an approximation must be used. Moreover, with the use of a convex program, one hopes that it comes with some theoretical guarantees, but there are none to be found here.

In short, this paper does not have enough substance. Although the idea is promising, in its current stage the submission is premature and needs more work.

I have also found numerous grammatical errors in the writing, so it would benefit from another round of proofreading.

questions: As mentioned above, can you provide runtime comparisons for the experiments?

Also, the formulation of the optimization problem in (5) is similar to work on score matching, see, e.g., [H05], which is used extensively in generative modeling. Can you cite this literature in your work and include a comparison and discussion in the main text?

[H05] Aapo Hyvärinen, Estimation of non-normalized statistical models by score matching.

limitations: No, the authors do not thoroughly discuss the limitations of their approach.
review_title: null

Meta-review: metareview: This work proposes a SDP approach to computing Wasserstein gradient direction for 2-layers NNs, without the need of training the underlying NN. To compute the gradient direction, the authors construct a least-square regression problem, and add a polynomial regularization term. Then, they show that the (relaxed) dual is an SDP problem.

_Pros_
- The idea of casting the Wasserstein gradient direction as an SDP is novel, and interesting. It also paves the way to more general formulations
- The obtained optima is *global*

_Cons_
- The exposition is lacking some motivation at some points. I think the authors could have move some technical discussion (for instance after Thm 1) to give a better insight on the motivations. Some previous works is also sometimes not put in context (for instance regarding the dimensionality reduction).
- For a (mostly) theoretical paper, the statements are sometimes not precise enough, e.g., what is an "equivalent" problem: having the same minimum? the same argmin? In Prop 1, what properties are required on the function space? Don't you need hypothesis on $\psi$? etc.
- From a pratical point of view -- but I don't think that practicality is the core aspect of the paper -- the computational cost is totally prohibitive as discussed by all referees.
- It does not seems that there is a strong practical improvements with respect to training directly the NN after the parameterization of the Wasserstein gradient.

I believe the idea of casting the Wasserstein gradient direction as an SDP problem is interesting, but with respect to the ratio of pros/cons above, and the lack of a strong positive opinion on this work, I recommend to reject this submission in its current state.
I encourage the authors to revise the manuscript in the light of the comments by all reviewers and my own comments for a future submission. In particular, the revision should include the discussion with reviewer RwmS which better highlights your work.
confidence: Certain
Decision: Reject