Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: In this work the authors introduce an innovative discretisation of the Jordan Kinderleher Otto scheme (which describes the evolution of the gradient flow of relative entropy in the Wasserstein metric).  To do so the authors iteratively construct a transformation map as a composition of gradients of input-convex neural networks, which satisfy a universal approximation theory among convex functions.  Using Brenier's theorem, an iterative variational formulation for the density at discrete time-steps is obtained, from which one can generate samples and or evaluate an unnormalised density.

A variety of applications are presented, including finding the stationary distribution, as well as nonlinear filtering.   
main_review: This submission brings two important ideas together -- the idea of using input-convex neural networks with Brenier's theorem to formulate a variational problem for the density.   While straightforward, it is a nice idea which provides a means of evaluating the density of an overdamped Langevin dynamics which overcomes the curse of dimensionality associated with spatial discretisation approaches. 

While I find the idea interesting and the submission is clearly written, I would argue that it provides a very elegant solution to a quite narrow problem.   Evaluating the density of an SDE beyond the usual overdamped langevin SDE is not supported (granted one can adjust the metric tensor to accomodate some cases), in particular, non-ergodic diffusions cannot be approached via this formulation.

Secondly, for problems which involve stationary sampling as presented in the examples, we are severely limited by the various costs, mainly the quadratic cost in time-step which forbids long time -simulations.

Finally, it's worth remembering that this is still a discretisation in time, and it is not clear from this manuscript whether discretisation error induced by finite h has really been considered or studied.  I feel that this should have been addressed more carefully in the examples.
limitations_and_societal_impact: this has been adequately addressed
review_title: null

final_score: rating: 5: Marginally below the acceptance threshold
review_content: summary: The Fokker-Planck equation can be treated as gradient descent over entropy functionals in Wasserstein space. To solve this problem, Jordan Kinderlehrer and Otto proposed the JKO scheme. Later, through the connection between the Brenier theorem and this scheme, [10] treat it as an optimization problem which aims at finding the convex Brenier functional $\psi$ in equation (7). However, the solution given by [10] is usually intractable in high dimensional space. Thus, this paper proposes to use the input convex neural networks (ICNNs) proposed by [5] to replace $\psi$. And then use SGD to optimize the model parameter. This forms the core contribution of this work. Additionally, the experiments are interesting.
main_review: - My main concern about this paper is its lack of novelty. As summarized above, the core contribution is to replace the convex functional $\psi$ in equation (7) proposed by [10] with the input convex neural networks (ICNNs) proposed by [5]. It is a common trick to replace the convex Brenier potential with the ICNN network in optimal transport community. The only difference is the background of the applications.
- In theorem 1, how to compute the high dimensional determinant $det \nabla T(x_n)$? 
- In Sec. 4.1, since $W_2(\rho_{k-1}, \rho_k)$ for different $k$ has been computed during the optimization, it should be reported. The plot can help illustrate the performance of the proposed method.
- The comparison between the proposed method and [10] in low dimensional space in terms of memory usage/training time/accuray is necessary, because it is a direct generalization of [10].
limitations_and_societal_impact: N/A
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: This paper is on the computation of Wasserstein gradient flow, a type of gradient flow based on the optimal transport theory. The proposed algorithm is built on the famous JKO scheme. To implement the JKO scheme, input-convex neural networks (ICNNs) are used to model the optimal transport map between the distributions at consecutive steps. 
main_review: Originality:
Computing Wasserstein gradient flow is a classical problem in optimal transport. Most existing algorithms requires discretization of the underlying space and are thus not scalable. The present work starts from the JKO scheme, an implicit scheme, for Wasserstein gradient flow, and utilizes ICNNs to realize a JKO step. The objective free energy is also estimated efficiently using samples as in Theorem 1. The combination of these two techniques is new. A related work on particle-based Fokker-planck is [1].

Quality:
The main results on the Wasserstein gradient flow of the Fokker-planck free energy are technically sound. 
1. Since this gradient flow is mainly for sampling purpose, the authors can include more comprehensive comparison with more advanced sampling algorithms other than overdamped Langevin (EM). 
2. There is only one example (4.2) to illustrate the temporal behavior of the gradient flow and only two time points (t=0.5, 0.9) are evaluated. More examples are needed.
3. The claim ``our method admits straightforward generalization to any F that ….,’’ should either be removed or more properly explained. The proposed algorithm does use some special property of the Fokker-planck free energy that cannot be generalized.
4. The results of dual JKO in figure 4 seems to be in consistent with [2].
5. The algorithm seems to have larger variance than EM, especially for low dimensional problems (Figure 1, 3, 4).

Clarity:
The paper is over-all well-written. 
1. The experiments can be improved by adding more details. For instance, how is Accuracy in Table 1 defined?
2. The notation T is used in Theorem 1 and Section 5 for different purposes.
3. It should be argmax in (11) and the following equations

Significance:
The computation of Wasserstein gradient flow is an important problem. The proposed algorithm provides a possible way to calculate Wasserstein gradient flow in high dimensional setting. One limitation of this algorithm is that it is only applicable to the Fokker-planck free energy.

[1] K.F. Caluya et al, Proximal recursion for solving the Fokker-Planck equation
[2] C. Frogner et al, Approximate inference with Wasserstein gradient flows
limitations_and_societal_impact: Limitations and societal impact are properly discussed.
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: This paper proposes an input-convex neural networks (ICNNs) based method to approximate the JKO scheme for discretizing Wasserstein gradient flows. Following a previous work [1], the key idea is to reformulate the distribution optimization problem in each JKO iteration as optimization over convex functions due to the well-known Brenner’s polar factorization theorem [2]. ICNNs are used to parameterize these convex functions from computational considerations. The proposed method is applied to simulate Wasserstein gradient flows for sampling from unnormalized density and nonlinear filtering.
[1] Discretization of functionals involving the Monge-Ampere operator.
[2] Polar factorization and monotone rearrangement of vector-valued functions.

main_review: It is a novel idea to introduce ICNNs for addressing the computational challenges of the JKO scheme. ICNNs may be a promising tool to make the JKO scheme feasible for large-scale computations.
Quality: The work is technically sound.
Pros: The methodology part is well-organized, and the advance of the proposed method is clear.
The potential applications are interesting, such as sampling from unnormalized density, nonlinear filtering.
Cons: (1) Maybe the title is a little ambiguous. It would be better to include input-convex neural networks (ICNNs) in the title since the proposed method is entirely based on ICNNs.
(2) As the title indicates, the focus of this work is a large-scale setting. However, it is hard to conclude that the large-scale setting is well considered in the experimental part. For example, the considered data dimension is at most 32 in Section 4.1, 12 in Section 4.2, 60 in Section 4.3, 1 in Section 4.4.


limitations_and_societal_impact: No.
review_title: null

Meta-review: This work focuses on a new method to approximate Wasserstein gradient flows, allowing to sample from and compute the density at each time step. The use of ICNN to play the role of the convex potential allows to discretize the time dimension efficiently. Based on the overall evaluation by the reviewers - which is mostly positive - and my own, I recommend to accept this work.

There was a serious and fruitful discussion between reviewers and authors, the latter providing interesting perspective on the motivation of their work, that should be included in a camera-ready version. 
Decision: Accept (Poster)