Functional Gradient Flows for Constrained Sampling

Thank you for your careful review and valuable questions! We address your comments and questions as below.

Weaknesses

W1: I did find (...) improve the manuscript.

A1: Thanks for your suggestion! We will modify the notations accordingly in our revision.

W2: The examples all (...) larger neural networks?

A2: Thanks for your question! Our method is feasible for learning larger neural networks and on higher dimensional real problems. Please refer to the third part of the global response.

Questions

Q1: Time complexity seems (...) to that question?

A1: Thanks for your question! Figure 8 in appendix D.2 should be able to answer your question with numerical results. Take the sample size $N=4000$ as an example. Given a quality of Wasserstein distance smaller than 0.6, the computation time of CFG is approximately 110s while MIED is approximately 320s. Changing the required quality to Wasserstein distance smaller than 0.4, the computation time of CFG change to approximately 160s while MIED change to approximately 400s. In general, CFG is more particle efficient as a functional gradient method.

Q2: The choice of (...) make that choice)?

A2: Thanks for your question! We would like to clarify that we focus on problems for which the constrained condition is provided in advance, so we choose the $g$ function to be the simplest way of representing the required constraint. For toy experiments, please refer to appendix D.1 for the analytical descriptions of the constrained domains $\Omega$. Corresponding to $\Omega$, we choose $g(x)=(||x||^2-1)(||x||^2-4)$ for the ring experiment, $g(x)=x_1^2+(\frac{6}{5}x_2-x_1^{\frac{2}{3}})^2-4$ for the cardioid experiment, $g(x)=e^{-2}-\frac{e^{-2(x_1-3)^2}+e^{-2(x_1+3)^2}}{e^{2(||x|| - 3)^2}}$ for the double moon experiment, and $g(x)=|x_1-x_2|+|x_1+x_2|-4$ for the block experiment. For Bayesian Lasso, $\Omega$ is the $l_q$ ball, and we choose $g(x)=||x||_q-r$ to match the problem scenario.

Q3: More generally, (...) $f_{net}$ and $z_{net}$, ...

A3: Thanks for the question! We experimented on choosing different number of particles in the Bayesian Lasso experiment, the results showed that increasing the number of particles can improve the sampling quality. Our experiment results are not very sensitive to the choice of architecture and the parameters of $f_{net}$ and $z_{net}$. Properly selecting the bandwidth $h$ according to the particle number can improve the results. One possible approach is to use the adaptive scheme to select the bandwidth as proposed in section 6.2.1 and appendix C. This scheme is backed by the error analysis in appendix C, and could be verified in the Bayesian Lasso experiment. Figure 1 in the attached pdf file showed that the adaptive scheme achieved slightly better energy distance results than the fixed bandwidth. More details on the parameter setting are provided in appendix D for each experiment.

Q4: I don't understand (...) over that support?

A4: Thanks for your question! This assumption is adopted from [1], which is to guarantee that the particles outside the domain can enter the constrained domain. In practice, we could take a sufficiently large domain that covers all the initial particles and the given constrained domain. $g$ could be bounded on this domain and $\nabla g$ is bounded away from zero, providing feasibility of training. This assumption does not serve for our main focus, which is the convergence analysis when the particles are inside the constrained domain.

[1] Zhang, R., Liu, Q., & Tong, X. T. Sampling in constrained domains with orthogonal-space variational gradient descent. Advances in Neural Information Processing Systems, 2022.

Q5: Do you have (...) the paper stronger.

A5: Thanks for your question! Firstly, from the best of our knowledge, there are not many methods that can directly deal with non-convex and complex hard constrained domains besides our method and the MIED method. As stated in the introduction, the mirror method is restricted to problems where mirror maps can be obtained, while projection methods suffers from high computation cost for complex domains because of the complexity in finding projection points. Additionally, compared to methods that can deal with this setting such as MIED, our method achieved better results on the cardioid, the double moon and the monotonic BNN experiment. We list Table 6 in appendix D.1 as below for the comparisons on the toy experiments. Secondly, using neural network to learn the optimal velocity field is more powerful way to implement ParVI, as it enhances the expressiveness and flexibility compared to classical kernel methods ([1][2][3]). This advantage can be verified by our better results compared to kernel based methods MSVGD, PD-SVGD and C-SVGD. Using the simple and flexible neural network framework, our method is also more scalable than MIED in terms of particle numbers N. Please refer to appendix D.2 for detailed discussion and verification.

Table: Wasserstein-2 distance and energy distance between the target distribution and the variational approximation on the three toy datasets.

[1] di Langosco, L. L., Fortuin, V., & Strathmann, H. (2021). Neural variational gradient descent. arXiv preprint arXiv:2107.10731.

[2] Dong, H., Wang, X., Lin, Y., & Zhang, T. (2022). Particle-based variational inference with preconditioned functional gradient flow. arXiv preprint arXiv:2211.13954.

[3] Cheng, Z., Zhang, S., Yu, L., & Zhang, C. (2024). Particle-based Variational Inference with Generalized Wasserstein Gradient Flow. Advances in Neural Information Processing Systems, 36.

Limitations

I would (...) parameters.

A1: Please refer to A3 of Q3 of our response for more details.

Thank you for your careful review and valuable questions! We address your comments and questions as below.

Weaknesses

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

A1: We achieved better results than MSVGD in the block experiment, and achieved better results mostly than MIED in the ring, cardioid and the double moon experiment. The latter results are in appendix D.1, which are listed below as Table 1. For large-scale experiments monotonic BNN, our method also outperforms other methods, including MIED. We list the monotonic BNN results in our main text (Table 1 in section 6.3) below as Table 2.

Table 1: Wasserstein-2 distance and energy distance between the target distribution and the variational approximation on the three toy datasets.

Table 2: Results of monotonic Bayesian neural network under different monotonicity threshold. The results are averaged from the last 10 checkpoints for robustness.

Questions

Q1: Can the authors accommodate multiple constraints and more complicated geometries? Are those generalizations possible? A discussion of these issues would be welcome.

A1: Thanks for your question! The generalizations are possible. Please refer to the first part of the global response.

Q2: 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.

A2: Thanks for the question! Our TV convergence rate is comparable to [1]. TV metric is broadly used in many machine learning problems because it is well defined for any two distributions. KL metric needs the initial distribution being absolute continuous with respect to the target distribution, while TV is not restricted by this requirement. Considering other suitable norm that applies to the constrained domain problem is very interesting, and we leave it for future work.

[1] Zhang, R., Liu, Q., & Tong, X. T. Sampling in constrained domains with orthogonal-space variational gradient descent. Advances in Neural Information Processing Systems, 2022.

Thank you for your careful review! We address your comments as below.

Weaknesses

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

A1: Thanks for the suggestion! In our experiments, the monotonic BNN is applied to real-world problems. The setting is originally derived from [1] about modelling in law. Additionally, it is also widely applied in social problems, including predicting admission decisions ([2]). For higher dimensional real-world applications, please refer to the third part of the global response.

[1] Karpf, J. (1991, May). Inductive modelling in law: example based expert systems in administrative law. In Proceedings of the 3rd international conference on Artificial intelligence and law (pp. 297-306).

[2] Liu, X., Tong, X., & Liu, Q. (2021). Sampling with trusthworthy constraints: A variational gradient framework. Advances in Neural Information Processing Systems, 34, 23557-23568.

Thank you for your careful review and valuable questions! We address your comments and questions as below.

Weaknesses

W1: The authors construct (10) as the projected vector field. The choices of lambda seem to be lacking.

A1: Thanks for the suggestion! We choose lambda to be 1 in toy experiments and Bayesian Lasso experiments. We choose lambda to be 100 in monotonic BNN experiments. Lambda is used to control how fast the particles can be pushed inside the constrained domain. According to the problem scale and dimensionality, we select lambda in a reasonable range. We will make this point clearer in our revision.

W2: The authors miss essential literature in the field. One needs to discuss the related references carefully.

Wang, Y., Li, W. Accelerated Information Gradient Flow. J Sci Comput 90, 11 (2022).

A. Lin et.al. Wasserstein Proximal of GANs. GSI, 2022.

Wang, Y. et al. Optimal Neural Network Approximation of Wasserstein Gradient Direction via Convex Optimization. 2022.

A2: Thanks for the suggestion! These references implement Wasserstein gradient flows in multiple regions. Wang & Li (2022) considered the Accelerated Information Gradient flow, which is the generalization of the Hamiltonian flow integrating multiple metrics. They provided the particle version for the flow with detailed theoretical analysis, and they also experimented on common sampling tasks including BLR and BNN. Lin et.al. (2022) applied the Wasserstein proximal method to GAN, the underlying theoretical guarantee is about Wasserstein natural gradient flow. Wang et.al. (2022) considered a similar problem of using neural networks to approximate the Wasserstein gradient flow as [1], the main difference is that Wang et.al. (2022) considered parameterizing the velocity field using the gradient of a neural network, rather than directly using a neural network. All the above references do not consider the problem of constrained domain sampling, for which we designed functional gradient flows. We will add these references in our revision.

[1] Zhang, R., Liu, Q., & Tong, X. T. Sampling in constrained domains with orthogonal-space variational gradient descent. Advances in Neural Information Processing Systems, 2022.

Questions

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

A1: Thanks for the question! Empirically, we pick lambda to be 1 for toy experiments and Bayesian Lasso experiments, and we pick lambda to be 100 for monotonic BNN experiments. Lambda is used to control how fast the particles can be pushed inside the constrained domain. According to scale of the problem, relatively large lambda may provide swift entering for some constrained domains, but may incur some instability if lambda is too large that particles overshoot the constrained domains. In practice, for small scale experiments where the step size can be large, one may choose a small lambda. For large scale problems (e.g., BNN) where the step size is small, one may want to use a large lambda to drive the particles to the constrained domain faster.

Thank you for your careful review and valuable questions! We address your comments and questions as below.

Weaknesses

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

A1: Thanks for the advice! Existing methods are useful in common unconstrained problems, but these methods generally do not apply to problems with constrained domains.

Questions

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

A1: Thanks for the question! Our proposed framework can still apply to these types of problems by adopting the idea of [1]. We could use velocity fields like $v_{\sharp}(x)$ to guided the particles to satisfy the equality constraints, and use our method to propose velocity fields substituting $v_{\perp}(x)$ to satisfy the inequality constraints. We add these two types of velocity to obtain the desired velocity field. To illustrate the idea, we demonstrate the training process of a 3D toy example in Figure 2 of the attached pdf file. Suppose the coordinate of the particle is $(x,y,z)$. The target distribution is a truncated standard gaussian located in the ring shaped domain in xOy plane. This corresponds to the equality constraint $z=0$ and the inequality constraint $1\le x^2+y^2\le 4$. The initial distribution is the 3D standard gaussian distribution. We can see that the particles collapse to the xOy plane and converge inside the ring domain.

[1] Zhang, R., Liu, Q., & Tong, X. T. Sampling in constrained domains with orthogonal-space variational gradient descent. Advances in Neural Information Processing Systems, 2022.

Thank you for your insightful review and valuable questions! We address your comments and questions as below.

Weaknesses

W1: 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 (...) non-trivial.

A1: Thanks for the question! First, we want to emphasize that in practice the step-size needs not to be very small (the algorithm would still convergence inside the constrained domain as the functional gradient approaches zero). Assumption 5.1 provides positive velocity field on the boundary. Also, when particles get stuck near the boundary, the boundary integral part would provide extra "force" to push these particles towards the inner side of the domain if the local density of the particles is higher than the target. With proper step size, empirically the discontinuity on the particle movements does not affect the overall convergence. Please refer to Figure 6 and Table 6 (listed below) in appendix D.1 for empirical results.

W2: The handling of the boundary integral is challenging, (...) This isn't really addressed in the text.

A2: Thanks for the question! The bandwidth $h$ needs to be adjusted according to the number of particles $N$ to get good performance. To put it simple, $h$ should be large if $N$ is small and vice versa. The choice of $h$ needs to strike a good balance between the band-wise numerical approximation error of the integral and the variance due to Monte Carlo estimation. Please refer to appendix C for detailed error analysis and simulation verification. We will make this reference clearer in our revision.

W3: I believe (though happy to be corrected on this), (...) simply impossible.

A3: Thanks for the question! For non-convex and complex boundaries where the mirror map is infeasible, it is hardly possible to use the mirror method. We have experimented on these cases in our paper, including the ring, cardioid, double-moon experiments in section 6.1 and the monotonic BNN experiments. Additionally, when mirror method is applicable such as the block experiment, we have shown in the right of Figure 1 that we still achieved better results than the mirror method MSVGD.

Questions

Q1: How limiting is the assumption that $||\nabla g||\ge 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?

A1: Thanks for the question! Firstly, we would like to be clear that the main focus of our theoretical analysis is on the convergence behavior when the particles are inside the constrained domain. The assumption that $||\nabla g||\ge C$ is mainly to guarantee that the particles would enter the constrained domain. In fact, any assumption would work here as long as it is sufficient enough that the particles will not get stuck at local modes outside the domain. Secondly, this assumption is not very limiting, which is adopted from [1]. In general, if $g$ has local extrema outside the constrained domain which hinder the entering, the particles may stuck or collapse at these local modes and affect the training. This would be left for future work.

[1] Zhang, R., Liu, Q., & Tong, X. T. Sampling in constrained domains with orthogonal-space variational gradient descent. Advances in Neural Information Processing Systems, 2022.

Q2: Can the authors clearly explain the bias-variance trade-off in choosing $h$, and the implications of that for the performance of the method.

A2: Please refer to appendix C for detailed error analysis and simulation verification. We will refer to this discussion earlier in our revision.

Q3: Can the authors explain (...) large discontinuity.

A3: Thanks for the question! Generally speaking, the additional boundary integral term in RSD will endow the trained velocity field with a repulsive factor that prevents clustering near the boundary unless the target distribution encourages the particles to do so. This is different from the unconstrained cases and demonstrates the importance of properly estimating the boundary integral. Empirically, for positive step size, we can see from Figure 6 in appendix D that the 'sticking' phenomenon does not occur.

Q4: 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?

A4: Thanks for the questions! For your first question, for complex real problems such as the monotonic BNN, it may be sample inefficient to use rejection sampling for a uniform distribution simply to guarantee that the initial particles are inside the constrained domain. The constrained domain of some problems may be difficult to locate and may have a low acceptance rate for rejection sampling. For your second question, even if the initial particles are all inside the domain, the boundary condition is still essential for correct sampling of the constrained distribution. Please refer to appendix D.1 for detailed ablation study on with and without estimating the boundary integral (Table 6 in appendix D.1 is listed below). The latter suffers from the 'sticking' to the boundary problem as you stated before while the former does not.

Table: Ablation results of not estimating boundary integral, with and without $z_{net}$.