Particle-based Variational Inference with Generalized Wasserstein Gradient Flow

Summary

This paper proposes to compute gradient flows of the KL divergence in the space of probability measures endowed by a generalization of the Wasserstein distance, which uses more general cost based on Young functions, and called Generalized Wasserstein gradient flows (GWGF). Authors provide the forward Euler scheme of such gradient flow and observe that the choice of the Young function has an impact on the convergence rate of the flow. Thus, besides proposing an algorithm to solve the GWGF with neural networks and analyzing its convergence rate, they also propose an algorithm to adapt the Young function in order to have accelerated convergence rate. Finally, they propose several applications and compare with other gradient flows based methods such as the SVGD or the preconditioned functional gradient flow. They show that they obtain consistently better results and also illustrate the benefits of the adaptative algorithm.

Strengths

The paper is overall pretty clear and the idea to perform gradient flows in this class of OT problems with more general costs is very appealing. - Interesting idea well motivated - Good Theoretical analysis - Many applications with comparisons with other methods

Weaknesses

- It is not a big weakness, but the class of Young functions investigated seems not very big (as it only experimentally considers $c'x,y)=\|x-y\|_p^p$ with $p\in [1.1,4]$). A test with a less common cost could have been interesting.

Questions

The existence of the OT map for the Wasserstein cost is not discussed and is assumed. In [1] (Section 1.3), the result is stated when the cost is of the form $c(x,y)=g(x-y)$ with $g$ strictly convex. In the Young function class, $g$ is only assumed to be convex. So I think I am missing something here. In equation (7), I believe that a $\frac{1}{2h}$ is missing under the Wasserstein cost. In Figure 1, the SVGD is not reported. Is it because it is hard to find good hyperparameters which converge for multimodal distributions? I think it would have been nice to add a plot of the evolution of p in the adaptative version of GWG. Equation (13) is referred to as the Euler-Maruyama scheme. In my opinion, this is the forward-euler discretization and not the Euler-Maruyama scheme which refers to the discretization of SDEs. [1] Santambrogio, F. (2015). Optimal transport for applied mathematicians. Birkäuser, NY, 55(58-63), 94.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

3 good

Contribution

4 excellent

Limitations

The limitations are addressed in appendix.

Summary

The paper's primary focus is the use of a generalized Wasserstein gradient flow of the KL divergence for solving the sampling problem in a particle-based variational inference framework, which they name Generalized Wasserstein Gradient Descent (GWG). Unlike the usual Wasserstein gradient flow of the KL, which corresponds to the Langevin diffusion, the implementation of this generalized gradient flow is more complex, hence the authors propose learning the gradient vector field via a neural network. They further back their proposal with a theoretical analysis which indicates that if the neural network learns the vector fields sufficiently well, guarantees can be obtained in a $L^q$ variant of the Fisher information.

Strengths

The paper presents a novel sampling algorithm using generalized Wasserstein gradient flows, which provides an interesting avenue for designing alternative approaches to SVGD. The algorithm is backed by theory. The authors present robust guarantees about the convergence of the proposed method given certain conditions on the learning ability of the neural network.

Weaknesses

While the approach presented in the paper is interesting, the technical novelty appears to be limited, as the underlying analysis follows the framework presented by Balasubramanian et al. Assumptions 1 and 2, which concern the neural network's capability to accurately learn vector fields, form a significant portion of the work's theoretical underpinning. However, the realism of these assumptions in practical scenarios is unclear.

Questions

In the context of existing literature, the reference to Bernton (2018) on line 88 seems to be out of place. Would it not be more accurate to reference the Jordan–Kinderlehrer--Otto (JKO) paper? Lines 170-171 contain an inaccurate claim regarding the dissipativity assumption in Balasubramanian et al. (the dissipativity assumption they use is much weaker than the log-Sobolev inequality) – could this be corrected? The paper heavily relies on Assumptions 1 and 2 regarding the accuracy of the neural network's learning of the vector fields. Could the authors shed more light on how realistic these assumptions are and discuss any empirical evidence that might support these assumptions?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The fact that the analysis hinges on an assumption for which it is unclear if it holds in practice, namely accurate estimation of the vector field via neural networks, should be addressed more clearly.

Summary

This paper proposes a general version to solve Wasserstein Gradient Flow for Particle-based VI. The authors show that this approach offers strong convergence guarantees and better performance over SVGD and other methods, as evidenced by extensive experiments on both simulated and real data sets. The paper also introduces a novel theoretical convergence guarantee of ParVIs with neural-net-estimated vector fields for generalized Wasserstein gradient flow.

Strengths

GWG is an interesting extension of ParVI, which extend the quadratic distance to more general metrics. Both theoretical and empirical justification suggest the improvement of the proposed algorithm.

Weaknesses

The title Particle-based Variational Inference with Generalized Wasserstein Gradient Flow is too close to (Dong et al. 2023). I would suggest to use a new title to highlight the improvements. More intuitive illustration of the method would be more helpful. The implementation details and codes are missing.

Questions

What is a proper heuristic strategy to choose $p$? Is there any geometric interpretation of the proposed $g$?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

The paper lacks some intuitive demonstration of the propose $g$, which would be very helpful for readers.

Summary

This paper proposes a novel particle-based variational inference framework based on generalized Wasserstein gradient flow of KL divergence, named generalized Wasserstein gradient descent (GWG). The strong convergence guarantee of the proposed algorithm is provided. The authors also provide an adaptive version based on Wasserstein metric to accelerate convergence.

Strengths

- The paper is well-written and clearly-organized. - It is quite novel to replace the standard Wasserstein-2 metric by using a Wasserstein metric with cost function defined by Young function. - Numerical results support effectiveness of proposed method.

Weaknesses

- The description of main results (Theorem 2 and Theorem 3) are unclear. It is not specified how the vector-valued function vk is learned. Is vk the maximizer of equ (14) with finite samples? - The convergence result only holds for \bar \mu_{Nh}, which is the average of path instead of the end-iteration solution \mu_{Nh}. - The numerical results do not show a quanitative convergence improvement of the proposed Ada-GWG with other methods like SVGD, L2-GF or PFG. For instance, It would be better to compare these method in the form of Figure 2, while you plot KL divergence or even maximum mean discrepancy between current iteration and posterior as a function of iteration number or cpu time. Typos: - In Algorithm 1, it should be \hat A(p_k) instead of \hat A(p).

Questions

- How Assumption 3 on the smoothness of neural networks are ensured in practice? How to estimate the constants G_p and M_p given a neural network?

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

Yes.