Functional Gradient Flows for Constrained Sampling

Summary

The paper proposes to adapt particle based variational infence to the context in which the distribution of interest is supported on a subset of R^d. It is supposed that the subset on which the distribution is supported is characterized as a lower level set of a function, g, and that this function has a gradient whose norm is bounded away from zero on the complement of the supporting set. Consequently, a gradient flow which descends g when outside the set and follows gradient dynamics for which (a mollified modification of) the original distribution of interest is a fixed point within that support might be expected to provide a good approach to sampling. The paper then concerns itself with constructing a particular (Wasserstein-type) gradient flow with such properties, establishing some theoretical properties of the proposed method and then exploring its performance empirically in a number of settings. The approach could be viewed as an extension of an existing method for dealing with distributions over manifolds in which the inequality constraint within the present work is replaced with an equality constraint and much of the novel work is concerned with addressing that difference in this particular framework.

Strengths

The problem of sampling from distributions constrained to subsets of R^d is an important one and extending the reach of modern sampling methods to this context is a valuable contribution. The general framework is an appealing one: it allows for fairly arbitary supports to be considered, provided only that a good choice of g can be obtained to characterize the support and the general approach is intuitively appealing; if within the support of the distribution of interest then pursue a standard gradient flow approach and outside that set modify the dynamics such that the flow moves towards the support. Reasonably competitive performance seems to be obtained in the numerical examples considered.

Weaknesses

I did find the paper a little difficult to read in places and the choice/definition of notations could be improved. For example, in the definition of the set $\Omega$ in line 97 it is not made explicit it is a subset of $\mathbb{R}^d$ and $g$ is never identified; something like $$\Omega = \{ x \in \mathbb{R}^d : g(x) \leq 0\} \text{ for some } g:\mathbb{R}^d \to \mathbb{R}$$ would be much clearer. And in line 243 (as an example of an odd choice which is made throughout the manuscript to write densities evaluated at a point with that point being entirely implicit): $$p(\beta) = \mathcal{N}(0,\sigma^2 I)$$ doesn't feature $\beta$ on the RHS and looking at (17) I at first thought you intended the $p(\beta)$ term to be uniform; the trivial edit to the more precise statement $p(\beta) = \mathcal{N}(\beta; 0,\sigma^2 I)$ would undoubtedly save many readers time. These are just some examples and while individually trivial my overall impression was that careful editing for presentation could significantly improve the manuscript. The examples all seem rather small by modern standards with "ground truth" estimates obtained by rejection sampling in one case (notwithstanding the fact that scale is not the only thing which makes inference challenging and there is always scope for methods which perform well in other important contexts). Is it feasible to scale the method up, for example to larger neural networks?

Questions

Time complexity seems important in comparing methods. Other than Appendix D.2 this isn't much discussed in this manuscript and that appendix isn't particularly enlightening. I think most readers would like an answer to the question "How much computation is required to obtain an answer of a given quality (however that is assessed) using this and competing methods and how does that vary as the required quality varies?" Are you able to give at least some heuristic numerical answer to that question? The choice of $g$ function only seems to be made explicit for the monotonic neural network example where it is not a simple function. What were the choices used in the other numerical examples and can you offer users any guidance on how they should go about specifying the particular choice of $g$ for other problems (while it may be straightforward to specify a $g$ such that it is zero on the boundary of the domain there are clearly a great many possible choices and it isn't obvious a priori how sensitive performance will be to the choice or how to make that choice)? More generally, how should users specify the various tuning parameters other than $g$ itself -- number of particles, architecture and parameters for the neural networks f_net and z_net, ... I don't understand the second part of Assumption 5.1: in line 139 it is stated that $p_0$ is generally assumed to be supported on all of $\mathbb{R}^d$; in the first part of the assume the gradient of $g$ is assumed to be bounded away from zero outside the support set. How can $g$ have gradient bounded away from zero across the support of $p_0$ (less $\Omega$) if that support is all of $R^d$ and $g$ be bounded over that support? Do you have in mind particular settings where the greater flexibility enjoyed by the CFG method in specifying constrained domains is important? I was disappointed that the numerical study didn't show such a setting. This flexibility seems like it should be a major advantage of the method and one thing that was absent from the numerical study was evidence that the proposed method either solves problems which existing methods cannot or performs substantially better in those settings. Such evidence would, I think, make the paper stronger.

Rating

5

Confidence

3

Soundness

3

Presentation

2

Contribution

3

Limitations

I think this is done reasonably well. I would have liked to see more discussion of the tuning of the algorithmic parameters.

Summary

The paper proposed a particle based method (using neural networks) to sample probability densities which are supported on subdomain of $\mathbb{R}^d$. This is done in the spirit of Stein variational gradient descent but using neural network instead of kernels and by incorporating the constraint into the functional which leads to a boundary integral which is handled via a band approximation. A complete convergence to equilbrium analysis is given, under a suitable assumption on the neural network approximation. The algorithm is tested on a variety of synthetic and real data sets and the results are competitive compared to other methods on the market (e.g MIED).

Strengths

This a very well written paper which propose a new approach to constrained sampling. The approach is very well motivated and is a natural extension of SVGD. The reviewer enjoyed very much the clarity of the exposition and the conceptual clarity of the approach. A nice convergence analysis in continuous time is provided which is relatively straightforward and does not provide strong guarantees.There is lots more to be done on the theoretical side but this is a nice first step. The model is tested on a variety of examples with satisfactory results and the model compares quite well (similar accuracy) with other algorithms on the market.

Weaknesses

Relative weakness is that the proposed method does not seem to outperform existing algorithms. Maybe a more detailed comparison with existing method interms of costs would be useful.

Questions

1) Can the authors accommodate multiple constraints and more complicated geometries? Are those generalizations possible? A discussion of these issues would be welcome. 2) How could the theoretical analysis be strengthened? The rate of convergence TV norm is slow and one would expect (?!) much faster convergence in suitable norm.

Rating

7

Confidence

4

Soundness

4

Presentation

4

Contribution

3

Limitations

The limitations are addressed adequately.

Summary

The authors develop a solution to constrained sampling by introducing a boundary condition for the gradient flow which would confine the particles within the specific domain. This gives a new functional gradient ParVI method for constrained sampling, called constrained functional gradient flow (CFG), with provable continuous-time convergence in total variation (TV). They also present numerical strategies to handle the boundary integral term arising from the domain constraints. They provide theoretical and experimental study to show the effectiveness of the proposed framework.

Strengths

The authors provide extensive theoretical analysis to support the proposed method. The authors provide solid theoretical results along with the proposed method. The authors also provide experiments on different datasets to validate the approach. The paper is well written.

Weaknesses

Experiments on a real-world application could strengthen the paper.

Questions

There is no question from the reviewer.

Rating

7

Confidence

3

Soundness

3

Presentation

4

Contribution

3

Limitations

N/A

Summary

This paper addresses the challenging problem of constrained sampling from an unnormalized distribution. Building on the principles of Stein variational gradient descent (SVGD), which merges the strengths of variational inference and Markov Chain Monte Carlo (MCMC), this work innovates by applying SVGD in a constrained setting. Specifically, it introduces boundary conditions for the gradient flow to confine particles within a designated domain, demonstrating convergence in total variation. Numerical experiments highlight the method's effectiveness.

Strengths

1. The paper is well-written. 2. The proposed method is technically sound and demonstrates convergence in total variation. 3. Experiments on both synthetic and real-world datasets showcase the method's potential.

Weaknesses

In Bayesian neural network experiments, the paper overlooks numerous baselines of existing Bayesian Neural Networks (BNNs) using variational inference, MCMC, and Laplace approximation.

Questions

If we have additional equality constraints, does the method still work?

Rating

7

Confidence

2

Soundness

3

Presentation

3

Contribution

3

Limitations

The paper has discussed the limitations.

Summary

The authors study functional Wasserstein gradient descent methods with constraints in support of target distribution. They study the variational problem of vector fields using the least square formula between the current step and the gradient of score functions. To maintain the constrained set, the authors restrict the choices of vector fields to particular linear combinations of vector fields and use the neural network functions to approximate these vector fields. The authors prove the convergence of the proposed method in total variation. Numerical examples demonstrate the effectiveness of the proposed method.

Strengths

1. The authors carefully analyze the algorithm's importance in handling the sampling domain's boundary. In particular, example 4.2 is a good example. 2. The algorithm variational formulation is well-designed with nice numerical examples.

Weaknesses

1. The authors construct (10) as the projected vector field. The choices of lambda seem to be lacking. 2. The authors miss essential literature in the field. One needs to discuss the related references carefully. 1. Wang, Y., Li, W. Accelerated Information Gradient Flow. J Sci Comput 90, 11 (2022). 2. A. Lin et.al. Wasserstein Proximal of GANs. GSI, 2022. 3. Wang, Y. et al. Optimal Neural Network Approximation of Wasserstein Gradient Direction via Convex Optimization. 2022.

Questions

Can authors illustrate how to pick lambda? How does lambda affect the algorithm, at least for numerical reasons?

Rating

6

Confidence

5

Soundness

4

Presentation

3

Contribution

3

Limitations

There are no limitations.

Summary

The paper introduces an extension of the ParVI method for sampling on measures supported on constrained subset $\Omega$ of $\mathbb{R}^d$, where $\Omega$ is the sublevel set of a phase-field function $g$. The paper presents a construction in terms of a discontinuous particle velocity field, with the discontinuity occurring on the boundary of the domain. Outside the domain, the velocity field pulls the particles towards the boundary of the domain, while within the domain the normal parVI approach is used. To leverage the usual Stein formulation through integration by parts, a new boundary term occurs which must be calculated. This is approximated by a spatial integral along a thin band within the domain. The authors prove that (i) particles eventually arrive at the domain boundary (using the assumption that there are no stationary points in the level set function $g$ outside the domain); and (ii) under a universal approximation assumption, the empirical distribution converges in TV, assuming a Poincare inequality holds. The author demonstrate the method on a number of synthetic and real examples.

Strengths

The paper presents an alternative approach to the mirror descent strategies to handling sampling from constrained domains. This is potentially very useful in settings where there is not an obvious parametrisation of the constrained measure which lends itself to mirror-descent approaches -- to my knowledge this is novel. The methodology is generally applicable, and the authors have been very clear about their assumptions and provided some good theoretical results. The experiments are sufficiently challenging.

Weaknesses

The main weakness to me lies in how the boundary term is handled, for two reasons. My main concern is what happens when the step-size is very small. In this case, I can easily see scenarios where: (1) The particles reach the boundary and remain stuck there. What is pushing the particles across the boundary. When the step-size is not small this is ok, because particles will 'overshoot' and then get picked up by the other side of the velocity field. This isn't really addressed in the text, nor is the effect of the discontinuity on the particle movements, which I expect to be non-trivial. (2) The handling of the boundary integral is challenging, because there is a clear bias-variance trade-off. The smaller $h$ is, the less likely it is to find particles in the domain, and thus we would expect the variance of the integral estimator to be huge. This isn't really addressed in the text. (3) I believe (though happy to be corrected on this), that most of the examples could be handled with the mirror approach? It would have been nice to see examples where the alternative is simply impossible.

Questions

(1) How limiting is the assumption that $\lVert \nabla g\rVert \geq C$? I can see this being a reasonable assumption for simple constraints, but what about others where the boundary is complex, e.g. corrugated, or is non-convex? (2) Can the authors clearly explain the bias-variance trade-off in choosing $h$, and the implications of that for the performance of the method. (3) Can the authors explain why the particles do not 'stick' at the boundary when the step-size gets very small? Similarly, can they explore the effects of the large discontinuity. (4) For the problems considered in this paper, could the authors not simply instead use rejection sampling for a uniform distribution within the constrained domain and then use that as an initial condition for the constrained parvi? Potentially then you wouldn't even need to consider the boundary condition?

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

Limitations have been adequately addressed.

Reply to authors

The authors have addressed my questions. I will keep my score.

Thanks for the detailed response

Thanks for the very detailed response. Having read the collected reviews and rebuttals, I feel more positive about the manuscript and have increased my score from 4 to 5. I would still have liked to see a more comprehensive numerical validation of the approach and I hope that the example described in the global rebuttal at least can find its way into any subsequent versions of the paper.

Response to authors

I thank the authors for the detailed rebuttal. I am still not entirely convinced by the argument around the small step-size limit of the algorithm, but I admit I have not derived the continuum limit and studied the effect of the boundary term. The empirical experiments suggest it works sufficiently well. Based on the responses + discussion I will increase my score. I hope the responses made in the rebuttal make their way into the paper.