Functional Gradient Flows for Constrained Sampling

Authors: shiyue  zhang
longlin  yu
ziheng  cheng
cheng  zhang
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-13 09:08:30

Abstract: Recently, through a unified gradient flow perspective of Markov chain Monte Carlo (MCMC) and variational inference (VI), particle-based variational inference methods (ParVIs) have been proposed that tend to combine the best of both worlds. While typical ParVIs such as Stein Variational Gradient Descent (SVGD) approximate the gradient flow within a reproducing kernel Hilbert space (RKHS), many attempts have been made recently to replace RKHS with more expressive function spaces, such as neural networks. While successful, these methods are mainly designed for sampling from unconstrained domains. In this paper, we offer a general solution to constrained sampling by introducing a boundary condition for the gradient flow which would confine the particles within the specific domain. This allows us to propose a new functional gradient ParVI method for constrained sampling, called *constrained functional gradient flow* (CFG), with provable continuous-time convergence in total variation (TV). We also present novel numerical strategies to handle the boundary integral term arising from the domain constraints. Our theory and experiments demonstrate the effectiveness of the proposed framework.

Main paper: 
 introduction: 
 Efficiently approximating and sampling from unnormalized distributions is a fundamental and challenging task in probabilistic machine learning, especially in Bayesian inference. Various methods, including Markov Chain Monte Carlo (MCMC) and Variational Inference (VI), have been developed to address the intractability of the target distribution. In VI, the inference problem is reformulated as an optimization task that aims to find an approximation within a specific distribution family that minimizes the Kullback-Leibler (KL) divergence to the posterior (Jordan et al., 1999; Wainwright & Jordan, 2008; Blei et al., 2016) . Leveraging efficient optimization algorithms, VI is often fast during training and more scalable to large datasets. However, its approximation power may be limited depending on the chosen family of variational distributions. In contrast, MCMC methods generate One prominent particle-based variational inference technique is Stein Variational Gradient Descent (SVGD) (Liu & Wang, 2016) . It calculates the update directions for particles by approximating the gradient flows of the Kullback-Leibler (KL) divergence within a reproducing kernel Hilbert space (RKHS), where the approximation takes a tractable form (Liu, 2017; Chewi et al., 2020) . However, the performance of SVGD heavily depends on the choice of the kernel function and the quadratic computational complexity of the kernel matrix also makes it impractical to use a large number of particles. As kernel methods are known to have limited expressive power, many attempts have been made recently to expand the function class for gradient flow approximation (Hu et al., 2018; Grathwohl et al., 2020; di Langosco et al., 2021; Dong et al., 2023; Cheng et al., 2023) . By embracing a more expressive set of functions, such as neural networks, these functional gradient approaches have shown improved performance over vanilla SVGD while not requiring expensive kernel computation.While MCMC and VI methods have shown great success in sampling from unconstrained domains, they often struggle when dealing with target distributions supported on constrained domains. Sampling from constrained domains is an important and challenging problem that appears in various fields, such as topic modeling (Blei et al., 2003) , computational statistics and biology (Morris, 2002; Lewis et al., 2012; Thiele et al., 2013) . Recently, several attempts have been made to extend classical sampling methods like Hamiltonian Monte Carlo (HMC) or SVGD to constrained domains (Brubaker et al., 2012; Byrne & Girolami, 2013; Liu & Zhu, 2018; Shi et al., 2022) . However, these extensions either involve computationally expensive numerical subroutines such as solving nonlinear systems of equations, or rely on intricate implicit and symplectic schemes or mirror maps that require a case-by-case design effort tailored to specific constraint domains.In this paper, we propose a functional gradient ParVI method for sampling from probability distributions subject to constrained domains with general shapes. We demonstrate that functional gradient approaches for ParVIs can be seamlessly adapted to constrained domains by learning the gradient flows with a vector field that adheres to a boundary condition. Intuitively, this boundary condition would confine particles within the specified domain, and it is indeed a sufficient condition for gradient flows of probability measures confined to this domain. Following previous works (di Langosco et al., 2021; Dong et al., 2023; Cheng et al., 2023) , we employ the regularized Stein discrepancy objective for functional gradient estimation where we directly incorporate the boundary condition for the vector field into the design of its neural network approximation. Due to the domain constraints, integration by parts now would lead to an additional boundary integral term in the training objective. We, therefore, derived an effective approach for properly evaluating this term for general boundaries. Extensive numerical experiments across different constrained machine learning problems are conducted to demonstrate the effectiveness and efficiency of our method. 
related work: 
 To better capture the gradient flow, a number of functional gradient methods for ParVIs have been proposed recently that employ a larger and more expressive class of functions than reproducing kernel Hilbert spaces (RKHS). In particular, di Langosco et al. (2021) used neural networks to learn the Stein Discrepancy with an L 2 regularization term, and updated the particles based on the learned witness functions. Dong et al. (2023) and Cheng et al. (2023) provided extensions that accommodates a broader class of regularizers.Sampling from constrained domains is generally more challenging compared to the unconstrained ones. Brubaker et al. (2012) proposed a constrained version of HMC for sampling on implicit manifolds by applying Lagrangian mechanics to Hamiltonian dynamics, which requires solving a nonlinear system of equations for each numerical integration step. Lan et al. (2014) focused on constraints that can be transformed into hyper-spheres, which can be viewed a special case of Brubaker et al. (2012) that has close-form update formulas. Zhang et al. (2020) ; Ahn & Chewi (2021) ; Shi et al. (2022) extended Langevin algorithms and SVGD to constrained domains via mirror maps. However, these methods require explicit forms of transformations (e.g., spherical augmentation and mirror maps) that capture the constraints, which would limit their applications to simple domains. Bubeck et al. (2018) ; Brosse et al. (2017) ; Salim & Richtárik (2020) considered projected Langevin algorithms through projection oracle, which is of high computational cost for complex constrained domains. Zhang et al. (2022) proposed an orthogonal-space gradient flow approach for sampling in manifold domains with equality constraints, which employed a similar strategy to ours. Our method is different from theirs in that it is designed for domains with inequality constraints using ParVIs with neural network gradient flow approximations, and we also proposed novel numerical strategies to address boundary-related challenges.Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes] Justification: Please refer to the appendix E. Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
background: 
 Notations We use x to denote particles in R d and Ω = {x|g(x) ≤ 0} to denote the constrained domain. Notation 1 Ω is the indicator function of Ω. Let ∥ • ∥ denote the standard Euclidean norm of a vector. Let P(R d ) denote the set of probability distributions on R d that are absolute continuous with respect to the Lebesgue measure, and we do not distinguish the probability measure with its density function. Let D KL be the Kullback-Leibler divergence and TV be the total variation distance. We use C(X , Y) to denote the space of continuous mappings from X to Y, and use the shorthand C(X ) for C(X , X ).Let p * ∈ P(R d ) be the target distribution we wish to sample from. We can frame the problem of sampling into a KL divergence minimization problemEQUATIONwhere Q ⊂ P(R d ) is the space of probability measures. Particle-based variational inference methods (ParVIs) can be described within this framework where Q is represented as a set of particles. Starting from an initial distribution p 0 and an initial particle x 0 ∼ p 0 , we update the particle x t following dx t = v t (x t )dt where v t : R d → R d is the velocity field at time t. The density p t of x t follows the continuity equation dp t /dt = -∇ • (v t p t ), and the KL divergence decreases with the following rate:EQUATIONParVIs aim to find the optimal velocity field v t that minimizes (2) in a Hilbert space H with the squared norm regularizer as followsmin vt∈H -E pt ⟨∇ log p * p t , v t ⟩ + 1 2 ∥v t ∥ 2 H .(3)Different choices of the space H in (3) lead to different gradient flow algorithms. SVGD Liu & Wang (2016) chooses H to be a Reproducing Kernel Hilbert Space with kernel k(•, •). This way, the optimal velocity field has a close form solutionEQUATIONalbeit the design of kernels can be restrictive for the flexibility of the method.When H = L 2 (p t ), then v * t = ∇ log p * pt is also known as the Wasserstein gradient of KL divergence (Jordan et al., 1998) . Since the score function of particle distribution ∇ log p t is generally inaccessible, recent works propose to parameterize v t as a neural network and train it through (3) di Langosco et al. (2021) ; Dong et al. (2023) ; Cheng et al. (2023) . Due to Stein's identity, (3) has a tractable formEQUATION)where F is the neural network family. This method is called functional Wasserstein gradient in contrast to the kernelized version.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Answer: [Yes]Answer: [Yes] Justification: Please refer to the codes and toy datasets in the supplemental materials. Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
main method: 
 Consider sampling from a target distribution p * (x) supported on Ω = {g(x) ≤ 0}, where g : R d → R is a continuously differentiable function. Throughout the work, we assume that p * admits differentiable positive density function on Ω.We first reformulate the problem into a constrained optimization in the space of probability measures:EQUATIONHowever, note that p * is only supported on Ω, the problem is ill-posed if p(Ω) < 1, since in this case p will not be absolute continuous with respect to p * and D KL (p∥p * ) cannot be defined. Besides, note that the initial particle distribution p 0 is generally assumed to be fully supported on R d and thus does not meet the constraint.To fix this issue, a natural idea is to drive the particle distribution p t towards Ω to satisfy p t (Ω) = 1 and train the functional gradient flow on constrained domains by minimizing the regularized Stein discrepancy (RSD)EQUATIONThese stringent requirements make the construction of velocity field for constrained sampling far more volatile than conventional functional gradient for unconstrained domains.In the rest of this section, we first present the idea of velocity design in Section 4.1 and then derive the tractable training objective in Section 4.2. The practical algorithm is proposed in Section 4.3.We first show that a globally continuous velocity field may fail to achieve the exact minimum of RSD under the constraint. To ensure that particles inside Ω will never escape, the velocity should satisfy the boundary conditionEQUATIONIn fact, if (8) holds, by Stoke's formula and the continuity equation∂ ∂t p t (x) = -∇ • (p t (x)v t (x)), we have d dt p t (Ω) = d dt E pt 1 Ω = Ω ∂ ∂t p t (x)dx = - Ω ∇ • (p t (x)v t (x))dx = - ∂Ω p t (x)v t (x) • ⃗ ndS.(9) Thus we can concludeProposition 4.1. If v t • ⃗ n ≤ 0 on ∂Ω, then p t (Ω) will not decrease.However, it is possible that there does not exist an optimal solution of (7) that meets the boundary condition ( 8) and is continuous on ∂Ω, as illustrated in the following 1D-example.Example 4.2. Let p * ∝ exp(-x 2 2 ) • 1 {x 2 ≤1} and p ∝ exp(-(x-1 2 ) 2 2). We have∇ log p * p = -1 2 in Ω and L RSD = 1 2 1 -1 p|v(x) + 1 2 | 2 dx -1 8 . The constraint on boundary is v(1) ≤ 0, v(-1) ≥ 0. Hence there is no optimal v in C(R).Therefore, it is necessary to design velocities for particles inside and outside separately. For particles located outside Ω (contributing to the probability p(g(x) > 0)), it is prudent to make the velocity drive the particles into the domain. On the other hand, for those already inside Ω, the velocity field should be able to fit the target distribution. Overall, it is reasonable to use the following piece-wise construction:EQUATION)Here λ > 0 is a constant. The gradient direction ∇g will push the particles into Ω. And h net is a continuous neural net to learn the optimal direction in Ω. We observed that the magnitudes of the gradients of different particles vary significantly in the numerical experiments and thus use the normalized gradient ∇g ∥∇g∥ instead.Based on the velocity design in (10), we are now ready to derive a tractable training objective for h net to learn the optimal direction in Ω through (7).Note that in principle, h net should be trained after p t (Ω) = 1. However, in order to enhance efficiency in practice, we use inner particles to minimize RSD before all the particles are driven in the constrained domain. At this phase, we replace p t in (7) with the conditional measure pt := p t (•|Ω).For any v ∈ C(Ω, R d ), expand (7) and we getEQUATIONSubstitute v with the continuous extension of h net on Ω and we obtain the training objective for h net .Note that the first two terms in (11) are aligned with (5) by substituting in the integration domain, while the last term is a boundary integral specific to inequality constraints. This is one of the essential differences between training functional gradients on constrained and unconstrained domains.The boundary integral ∂Ω pv • ⃗ ndS is computationally intractable in general. Inspired by the co-area formula (Federer, 2014) , we can estimate it using a heuristic band-wise approximation as followsEQUATION) where p(x) ∝ p(x) such that Sh p(x)dx = 1, and Sh := {x ∈ Ω : d(x, ∂Ω) ≤ h} is a band-like area near the boundary with a small bandwidth h > 0, which further has a proxy asEQUATION)Please refer to Appendix A for more explanations. This way, we can use the particles in the band-like area to obtain a Monte Carlo estimate of the boundary integral. Suppose there are m particles in Ω, of which {x j } n j=1 are in the band-like area. Then p(x ∈ Sh ) ≈ n m , and thusEQUATION)We summarize all the techniques above and propose the Constrained Functional Gradient (CFG) in Algorithm 1. We use the neural network structure designed as in Appendix D for h net .Algorithm 1 CFG: Constrained Functional Gradient Require: Unnormalized target distribution p * , initial particles {x i 0 } N i=1 , initial neural network parameters w 0 = {η 0 , ξ 0 }, weight parameter λ, iteration number L, L ′ , particle step size α, parameter step size η, bandwidth hfor k = 0, • • • , L -1 do Assign w 0 k = w k Identify {x ij k } n j=1 := {x i k |x i k ∈ Sh }, {x lr k } m r=1 := {x l k |x l k ∈ Ω} for t = 0, • • • , L ′ -1 do Set h w t k = f η t k -z 2 ξ t k • ∇g Compute L RSD (w) = 1 m m r=1 [-∇ log p * (x lr k ) T h w (x lr k ) -∇ • h w (x lr k ) + 1 2 ∥h w (x lr k )∥ 2 ] + 1 mh n j=1 h w (x ij k ) T ∇g(x ij k ) ∥∇g(x ij k )∥ Update w t+1 k = w t k + η∇ w L RSD (w t k ) end for Assign w k+1 = w L ′ k Compute v w k+1 = (f η k+1 -z 2 ξ k+1 • ∇g) • 1 {g<0} -λ ∇g ∥∇g∥ • 1 {g≥0} Update particles x i k+1 = x i k + αv w k+1 (x i k ) for i = 1, • • • , N end for return Particles {x i L } N i=15 Theoretical AnalysisIn this section, we present the convergence guarantee of CFG in terms of TV distance. Consider the continuous dynamic dx t = v t (x t )dt where v t is defined in (10). Let p t be the law of x t . We first investigate when particles can get into the domain.Assumption 5.1. ∥∇g(x)∥ ≥ C > 0 for x ∈ Ω c and M 0 = max{g(x 0 ) :x 0 ∈ supp(p 0 )} < ∞.For simplicity, we do not delve into the case that the particles outside Ω may get stuck at local extremal point of g before getting in Ω. Based on this assumption, the particles will enter the constrained domain in finite time:Theorem 5.2. There exists a finite t 0 ≤ M0 λC , such that g(x t ) ≤ 0, ∀t ≥ t 0 .The proof is deferred to Appendix B.1. Since domains with inequality constraints do not necessarily require the velocity field to gradually decay near the boundary as in Zhang et al. (2022) , directly using λ ∇g ∥∇g∥ • 1 {g≥0} allows for a swift penetration of the particles into the constrained domain. Once the particles are in the constrained domain, an important technique is to extend the target distribution to be non-trivially supported on the whole space by convolution with a Gaussian kernel and use it as a proxy. Denote the convolution of the target distribution and the Gaussian distribution as p * = p * * N (0, σ 2 I), where σ > 0 can be arbitrarily small. Now the KL divergence D KL (p t ∥p * ) is well defined on R d . Moreover, by the convolution properties, TV(p * ∥p * ) → 0 as σ → 0 (Stein & Shakarchi, 2005) .Assumption 5.3 (Poincaré inequality). The target distribution p * satisfies var p * [f ] ≤ κE p * [∥∇f ∥ 2 ]for any smooth function f . This is a common assumption in convergence analysis that describes the degree of convexity of the target distribution (Chewi et al., 2022) . In contrast to Zhang et al. (2022) , our approach relies solely on the Poincaré property of the original target distribution, rather than requiring it for the conditional measure on a zero measure set. Following Cheng et al. (2023) , we make the following assumption on the approximation error of neural networks.Assumption 5.4. For any t > t 0 , Ω p t ∥h net (t) -∇ log p * pt ∥ 2 dx ≤ ϵ.Now, we are ready to present the main result. Proposition 5.5. Suppose that p * satisfies κ-PI and the data distribution p t is smooth. DenoteF (p t , p * ) := Ω p t ∥∇ log p * p t ∥ 2 dx,then the following bound holds for any t > t 0EQUATIONRelating the gradient of KL divergence and the Fisher information and letting σ → 0, we have Theorem 5.6. Under assumption 5.4 and the assumptions in Proposition 5.5, suppose also that the KL divergence D KL (p t0 ∥p * ) < ∞. The following bound holdsEQUATION)Please refer to Appendix B.2 and B.3 for detailed proofs. 
experiments: 
 In this section, we compare CFG to other baseline methods for constrained domain sampling, including MSVGD (Shi et al., 2022) , MIED (Li et al., 2022) , Spherical HMC (Lan et al., 2014) , PD-SVGD and Control SVGD (Liu et al., 2021) . Throughout the section, we use neural networks with 2 hidden layers and initialize our particles with Gaussian distributions unless otherwise stated. We use the neural network structure designed as in Appendix D. All the experiments were implemented with Pytorch. More details can be found in Appendix D. The code is available at https://github. com/ShiyueZhang66/Constrained-Functional-Gradient-Flow.We first conduct 2-D experiments to test the effectiveness of CFG on sampling from truncated Gaussian distributions within various constrained domains, including ring-shaped, cardioid-shaped and double-moon-shaped domains. We covered non-convex domains, including unconnected domains. We use 1000 particles for all domains. From the left of Figure 1 we can see that the particles quickly converge to the target distribution without escaping the constrained domain. Comparison to MIED (via various distributional metrics) on the first three domains are deferred to Appendix D.1, where CFG provides better approximation in most cases.To compare with MSVGD and MIED, we additionally conduct the experiment of sampling from truncated gaussian mixture distribution within the block-shaped domain. The initial distribution is the uniform distribution. The right of Figure 1 shows that our method outperformed MSVGD and achieved comparable results to MIED in terms of Wasserstein-2 distance and energy distance.The Lasso method is broadly used in model selection to avoid overfitting by imposing a penalty term on model parameters. Park & Casella (2008) introduced a Bayesian alternative, named Bayesian Lasso, that replaces the plenty term with the double exponential prior distribution p prior (β) ∝ exp(-λ∥β∥ 1 ). Lan et al. (2014) then proposed Spherical HMC method, which introduced more flexibility in choosing priors with explicit ℓ q -norm constraints: p prior (β) ∝ p(β)1 {∥β∥q≤r} . For q = 1, it corresponds to the Lasso method. For more general case when q > 1, it is called Bayesian Bridge regression.Following Lan et al. (2014) , we choose the prior to be the truncated Gaussian distribution, where p(β) = N (0, σ 2 I). This leads to the Bayesian regularized linear regression model (Lan et al., 2014) : The posterior distribution for β is β|y, X, σ 2 ∝ N (β * , σ 2 (X T X + I) -1 )1(∥β∥ q ≤ r), where β * = (X T X + I) -1 X T y. This posterior has an inequality constraint and our method can apply.EQUATIONWe first generate a 20-dimensional dataset (X, y), where X ∈ R 1000×20 , y ∈ R 1000 and y = Xβ true + ϵ, where ϵ ∼ N (0, 25I). We set the true regression coefficients to be β true = (10, . . . , 10, 0, . . . , 0), where the first half of components is tens and the second half is zeros. Let βOLS denote the estimates obtained by ordinary least squares (OLS) regression. We follow Park & Casella (2008) and set r = ∥ βOLS ∥ 1 .We compare our method on Wasserstein-2 distance and energy distance with Spherical HMC and MIED using different number of particles. The ground truth is obtained via rejection sampling using 100,000 posterior samples. We use the h 0 (dN ) -1 3 scheme to adapt the bandwidth to the number of particles N (see Appendix C), where we choose h 0 d -1 3 = 0.1. From Figure 2 , we see that CFG performs the best in both metrics when N is small. This indicates the sample efficiency of particle-based variational inference for constrained sampling, which aligns with the findings in Liu & Wang (2016) for unconstrained sampling. All methods provide similar results when N is large.It is worth noting that as a kernel-based method, the time complexity of MIED scales quadratically as N increases, while functional gradient methods like CFG scale linearly. Meanwhile, we observed in our experiments that MIED tends to require more iterations to converge when the number of particles increases. These issues add to the overall time cost for MIED to achieve accurate approximation, especially when a large N is used (see Figure 9 in Appendix D.2). Similar issues of SVGD are also stated in Dong et al. (2023) . Please refer to Appendix D.2 for detailed information.Following Lan et al. (2014) , we also evaluate our method using the diabetes dataset discussed in Park & Casella (2008) . We compare the posterior median estimates given by Spherical HMC, CFG and MIED for the Lasso regression model (q = 1) and the Bridge regression model (q = 1.2), as the shrinkage factor s := r/∥ βOLS ∥ 1 varies from 0 to 1. We use 5000 particles for CFG and MIED. Figure 3 : Left: Bayesian Lasso (q = 1) using Spherical HMC (upper left), CFG (upper middle) and MIED (upper right). Bayesian Bridge Regression (q = 1.2) using Spherical HMC (lower left) CFG (upper middle) and MIED (upper right). Right: Results of monotonic Bayesian neural network with ϵ = 0.01. Only the portion below 0.02 is shown on the y-axis to better display the performance of models satisfying constraint.From Figure 3 we can see that our method aligns well with Spherical HMC and MIED, which shows the effectiveness of our method.We use the COMPAS dataset following the setting in Liu et al. (2021) . Given a Bayesian neural network ŷ(•, θ), define the constraint function g(θ) = ℓ mono (θ) -ϵ with ℓ mono (θ) = E x∼D [∥(-∂ xmono ŷ(x; θ)) + ∥ 1 ].Here, x mono denotes the subset of features to which the output should be monotonic. Note that in Liu et al. (2021) , the constraint is only in the sense of expectation, i.e., E p [g(θ)] ≤ 0. In the context of monotonic neural network (Karpf, 1991; Liu et al., 2020) , however, the monotonicity constraint is g(θ) ≤ 0, which defines a domain with inequality constraints and thus our method can be applied.We use two-layer ReLU neural network with 50 hidden units. With different threshold ϵ ∈ {0.005, 0.01, 0.05}, we compare our CFG with MIED and two SVGD variants in Liu et al. (2021) : PD SVGD and Control SVGD (C-SVGD). The number of particles is 200 and the results are shown in Table 1 . We observe that in all the tasks, our method outperforms the other two baselines in terms of both test accuracy and test likelihood. "Ratio Out" denotes the proportion of particles outside Ω = {g(θ) ≤ 0} to the total number of particles. All methods can successfully force almost all the particles into the domain. We further plot in the right of Figure 3 the averaged test monotonic loss ℓ mono against test accuracy during the training process for ϵ = 0.01. For MIED, the test accuracy is higher when more particles are outside the domain, while lower when forcing the particles into the domain during training. Notably, Vanilla SVGD exhibits lower accuracy and higher monotonic loss, indicating the importance of constraint-sampling methods. Furthermore, our method can scale up to even higher dimensional real problems. We experimented on the COMPAS dataset using larger monotonic BNNs and on the 276-dimensional larger dataset Blog Feedback Liu et al. (2020) . Please refer to Appendix D.3 for detailed information. Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: Please refer to appendix D and the codes and toy datasets in the supplementary materials.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
conclusion: 
 We proposed a new functional gradient ParVI method for constrained domain sampling, named CFG, which uses neural networks to design a piece-wise velocity field that satisfies a non-escaping boundary condition. We presented novel numerical strategies to deal with boundary integrals arising from the domain constraints. We showed that our method has TV distance convergence guarantee. Empirically, we demonstrated the effectiveness of our method on various machine learning tasks with inequality constraints.Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes] Justification: Please refer to the right of figure 1 , figure 2 and figure 8 .Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [Yes] Justification: Please refer to appendix F. As far as we are concerned, our paper has no potential negative societal impacts. We do not engage in negative fairness, privacy or security issues. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: Our paper poses no such risks for misuse.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: We cite the original papers in the experiments. Please refer to Section 6.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
 
 Appendix: The boundary integral ∂Ω pv • ⃗ ndS is a crucial factor that distinguishes ParVI on constrained domain from unconstrained case, due to its computational intractability in general. We derive a heuristic band-wise approximation (12) and thus the remaining issue is to identify which particle is in the band-wise area {d(x, ∂Ω) ≤ h}. Proposition A.1. For any h > 0, the sufficient condition for {x ∈ Ω :d(x, ∂Ω) ≤ h} is g(x) ≤ 0, g x + h ∇g(x) ∥∇g(x)∥ ≥ 0. (18) Proof. For any x ∈ Ω, define f (r) = g x + r ∇g(x) ∥∇g(x)∥ . Then f (•) is continuous with f (0) = g(x) ≤ 0. If f (h) ≥ 0, by intermediate value theorem, there exists h 0 ∈ [0, h] such that f (h 0 ) = 0.And we have either g(x) = 0 or g(x 0 ) = 0 wherex 0 = x + h 0 ∇g(x) ∥∇g(x)∥ . Therefore x ∈ ∂Ω or x 0 ∈ ∂Ω, implying d(x, ∂Ω) ≤ h since ∥x -x 0 ∥ = h 0 ≤ h.In fact, if g(x) ≤ 0, then d(x, ∂Ω) ≤ h is equivalent to max ∥y-x∥≤h g(y) ≥ 0. When h is small and ∇g(x) ̸ = 0, the maximum point of LHS is approximately y = x + h ∇g(x) ∥∇g(x)∥ since ∇g(x) is the steepest ascent direction. Hence ( 18) is a reasonable proxy when h is small.B.1 Proof of Theorem 5.2 Theorem B.1. Let M 0 = max{g(x 0 )} < ∞, there exists a finite t 0 ≤ M0 λC , such that g(x t ) ≤ 0, ∀t ≥ t 0 .Proof. It suffices to prove that for any z 0 , there exists a finite t 0 ≤ g(z0)λC , such that g(z t ) ≤ 0, ∀t ≥ t 0 . Sinced dt g(z t ) = (∇g) T v = (∇g) T h net • 1 {g<0} -λ∥∇g∥ • 1 {g≥0}Using assumption ∥∇g∥ > C, we can state that there exists a finite t 0 ≤ g(z0) λC , such that g(z t ) ≤ 0, ∀t ≥ t 0 .Firstly prove that there is a t 0 ≤ g (z0) λC such that g(z t0 ) ≤ 0. By contradiction, if such t 0 does not exist, we have g(z t ) > 0 for all t ∈ [0, g(z0) λC ], then d dt g(z t ) = -λ∥∇g∥, integrate both sides yields g(z t ) -g(z 0 ) = -g(z 0 ) λC 0 λ∥∇g∥dt ≤ -g(z 0 ), then g(z t ) ≤ 0, contradiction.Then prove that g(z t ) ≤ 0, ∀t ≥ t 0 .If not, there exists a minimum t > t 0 , such that g(z t ) > 0, then d dt g(z t ) = -λ∥∇g∥ < 0, which means for sufficiently small δ > 0, g(z t-δ ) > 0, contradiction.B.2 Proof of Proposition 5.5 Proposition B.2. Suppose that p * satisfies κ-PI and the data distribution p t is smooth. Denote F (p t , p * ) := Ω p t ∥s p * -s pt ∥ 2 dx, then the following bound holds for t > t 0EQUATIONProof. By equivalent expression of TV distance, it suffices to prove that for any function f satisfying |f | ≤ 1,|E p * [f ] -E pt [f ]| ≤ 8(κ + σ 2 )F (p t , p * ) + O(σ)Note that for t > t 0 , p t (Ω) = 1|E p * [f ] -E pt [f ]| = 1 {g≤0} p * (x)f (x)dx -1 {g≤0} p t (x)f (x)dx + 1 {g>0} p * (x)f (x)dx ≤ 1 {g≤0} (p * (x) -p t (x))f (x)dx + 1 {g>0} p * (x)f (x)dx ≤ 1 {g≤0} (p * (x) -p t (x))f (x)dx + ϵ 2 .whereϵ 2 = 1 {g>0} (p * (x) -p * (x))f (x)dx = O(σ) by the L 1 convergence of Gaussian convo- lution Stein & Shakarchi (2005).The first term can be bounded by1 {g≤0} (p * (x) -p t (x))f (x)dx 2 = 1 {g≤0} p * (x)( p t (x) p * (x) -1)f (x)dx 2 ≤ 1 {g≤0} p * (x)( p t (x) p * (x) -1) 2 dx • 1 {g≤0} p * (x)( p t (x) p * (x) + 1) 2 f (y) 2 dx ≤ 2 1 {g≤0} p * (x)( p t (x) p * (x) -1) 2 dx • 1 {g≤0} p * (x)( p t (x) p * (x) + 1)dx ≤ 4 1 {g≤0} p * (x)( p t (x) p * (x) -1) 2 dx ≤ 8(1 -1 {g≤0} p t (x)p * (x)dx) = 8(1 - p t (x)p * (x)dx) = 8(1 -b) with b = p * (x) p t (x) p * (x) dx ≤ 1.Then note thatvar p * p t (x) p * (x) = E p * [ p t (x) p * (x) ] -b 2 = 1 -b 2 .From Courtade. (2020), we have the Poincaré coefficient of N (0, σ 2 I) is σ 2 , and the Poincaré coefficient of p * = p * * N (0, σ 2 I) is no larger than κ + σ 2 .Therefore, by the Poincaré inequality, we have1 -b ≤ (1 -b 2 ) = var p * p t (x) p * (x) ≤ (κ + σ 2 ) p * (x) ∇ p t (x) p * (x) 2 dx = (κ + σ 2 ) 1 {g≤0} p t (x)∥s pt (x) -s p * (x)∥ 2 dx = (κ + σ 2 )F (p t , p * ) And |E p * [f ] -E pt [f ]| ≤ 8(κ + σ 2 )F (p t , p * ) + ϵ 2 with ϵ 2 = O(σ).B.3 Proof of Theorem 5.6Theorem B.3. Following the same assumptions in Proposition 5.5. Suppose also that the KL divergence D KL (p t0 , p * ) < ∞, the following bound holds as σ → 0EQUATIONProof. By equivalent expression of TV distance, it suffices to prove that for any functionf satisfying |f | ≤ 1, min t≤T |E p * [f ] -E pt [f ]| ≤ O(T -1 2 + ϵ 1 2 )We now only consider the time after t 0 , which means that p t (Ω) = 1.First we state the relation between p * = p * * N (0, σ 2 I) and p * (x). When σ → 0, by the pointwise convergence of the convolution Stein & Shakarchi (2005) , we have ∇p * (x) → ∇p * (x) and p * (x) → p * (x). Note that for g(x) ≤ 0, p * (x) and p * (x) are strictly positive, thus s p * (x) → s p * (x), as σ → 0 for any x ∈ Ω.Note that p * (x) is supported on R d , we haved dt D KL (p t (x)∥p * (x)) = -p t v T t (s p * -s pt )dx = -p t (h net • 1 {g<0} -λ ∇g ∥∇g∥ • 1 {g≥0} ) T (s p * -s pt )dx = -p t (h net • 1 {g≤0} ) T (s p * -s pt )dx + p t ((h net + λ ∇g ∥∇g∥ ) • 1 {g=0} ) T (s p * -s pt )dx = -p t • 1 {g≤0} ∥s p * -s pt ∥ 2 dx -p t • 1 {g≤0} (h net -(s p * -s pt )) T (s p * -s pt )dx ≤ -p t • 1 {g≤0} ∥s p * -s pt ∥ 2 dx + ϵ 0 p t • 1 {g≤0} ∥s p * -s pt ∥ 2 dx + 1 4ϵ 0 p t • 1 {g≤0} ∥h net -(s p * -s pt )∥ 2 dx ≤ -(1 -ϵ 0 ) p t • 1 {g≤0} ∥s p * -s pt ∥ 2 dx + 1 2ϵ 0 p t • 1 {g≤0} ∥h net -(s p * -s pt )∥ 2 dx + 1 2ϵ 0 p t • 1 {g≤0} ∥s p * -s p * ∥ 2 dxFor ϵ 0 < 1, the second term is O(ϵ), and the third term is O(σ).Using the notation F (p t , p * ) := p t • 1 {g≤0} ∥s p * -s pt ∥ 2 dx, integrating both sides yields the following:T t0 F (p t , p * )dt ≤ T 0 F (p t , p * )dt ≤ 1 1 -ϵ 0 D KL (p 0 , p * ) + T (O(ϵ) + O(σ)). So min t0≤t≤T F (p t , p * ) ≤ 2 (1 -ϵ 0 )T D KL (p 0 , p * ) + O(ϵ) + O(σ) = O(T -1 ) + O(ϵ) + O(σ) (21)Plugging the result of F (p t , p * ) (21) in Proposition 5.5 and let σ → 0, we havemin t0≤t≤T |E p * [f ] -E pt [f ]| ≤ O(T -1 2 + ϵ 1 2 ).Finally, by triangular inequality:|E p * [f ] -E pt [f ]| ≤ |E p * [f ] -E pt [f ]| + |E p * [f ] -E p * [f ]|Using the result that Gaussian convolution converge to original L 1 function in L 1 (Stein & Shakarchi, 2005) , we have|E p * [f ] -E p * [f ]| → 0 when σ → 0 for any p * . This means min t0≤t≤T |E p * [f ] -E pt [f ]| ≤ O(T -1 2 + ϵ 1 2), the proof is complete.Error Analysis of Boundary Integral Estimation We first present the error analysis of band-wise approximation, based on which we derive a heuristic O((dN ) -1 3 ) scheme for bandwidth selection when adjusting the number of particles.e = 1 mh n j=1 v(x j ) T ∇g(x j )/∥∇g(x j )∥ -∂Ω pv • ⃗ ndS. We claim that the mean square error E[e 2 ] = O(h 2 + 1 n ).The idea is to leverage bias-variance decomposition.We estimate the bias of the band-wise approximation, using Lagrange's Mean Value TheoremEQUATIONThen we turn to bound variance. From (12) we haveEQUATION)The variance of Monte Carlo estimation isE[A 2 3 ] = O( 1 n ). By Central Limit Theorem, E[( n m -p(x ∈ Sh )) 2 ] = O( 1 m ), and thus E[A 2 4 ] = O( 1 m ).Finally, through bias-variance decomposition, we have However, n depends on h implicitly. We use a simple example to determine this relation. Consider the target distribution being the uniform distribution, and the constrained domain being the ball inEQUATION)(MSE) v = v1 v = v2 v = v3R d with radius R, then p(x ∈ Sh ) ≈ R d -(R-h) d R d= O(dh). In this case, the total mean square error is of order O(h 2 + 1 dN h ), which implies the optimal h is O((dN ) -1 3 ), in this case the total mean square error is of order O((dN ) -2 3 ).In practice, one can set h = h 0 (dN ) -1 3 and tune h 0 through sparse grid search.To verify the above analysis, we conduct a toy 2D simulation study on estimating boundary integral of 3 different velocities at the boundary using 5 different numbers of samples from 3 different distributions. The constrained domain is the block areaΩ = {x | |x 1 | ≤ 2, |x 2 | ≤ 2}.The number of particles are 10 2 , 10 3 , 10 4 , 10 5 , 10 6 . The types of velocity and distribution are listed in Table 2 . The corresponding true values of boundary integral ∂Ω pv • ⃗ ndS are listed in the Table 3 .Table 2 : The types of velocities and distributions.Types 1 2 3v(x) ⃗ n (x 2 , x 1 ) (x 2 2 , x 2 1 ) p(x)Unif(Ω) N ((0, 0), I) on Ω N ((0, -2), I) on Ω Table 3 : The true values of boundary integral in the verification experiment.Velocitiesv 1 v 2 v 3 p 1 1 0 0 p 2 0.226259 0 0 p 3 0.911333 0 -0.617187We estimate E[e 2 ] by 10 trials for different number of particles and plot the curves as in Figure 4 . We set h = h 0 (dN ) -1 3 and h 0 d -1 3 = 0.5. The slope of the log-log plot is approximately -2 3 for all types of velocities and distributions, which supports the previous analysis.Additionally, if we do not follow the scheme of h = h 0 (dN ) -1 3 , and for example, fixing h = 0.5 • (10 2 ) -1 3 or h = 0.5 • (10 6 ) -1 3 (the starting and ending edgewidth of the adaptive edgewidth scheme h = 0.5 • N -1 3 ). Figure 5 shows that the estimated E[e 2 ] will tend to go up in certain cases, which supports the rationality of the h = h 0 (dN ) -1 3 scheme. Structure of Neural Networks We parameterize h net in (10) using two neural networks f net and z net as followsEQUATIONwheref net : R d → R d , z net : R d → R.Here, the -z 2 net • ∇g term relates to the reflection process in reflected stochastic differential equation Pilipenko (2014) , and always points inward Ω. It can be regarded as a reflection term that forces the particle to stay inside Ω as it neutralizes the outward normal-pointing component. Heuristically, -∇g serves as the reflection direction and z 2 net contributes in learning the magnitude. Empirically, we observe that incorporating the gradient of the boundary ∇g substantially alleviates the issue of particles adhering to the boundary during training, and may slightly improve the sampling quality. Please refer to D.1 below.Setting Details On toy experiments, we conduct CFG on four 2-D constrained distributions. The experiments are implemented on Intel 2.30GHz CPU with RAM 16384MB and NVIDIA GeForce RTX 3060 Laptop GPU with total memory 14066MB. The density and constrained domain are in Table 4 . The truncated gaussian mixture is 9 equally weighted gaussian distribution with all standard variance equal to 0.2, the centers are (-1.7, -1.7), (-1.7, 0), (-1.7, 1.7), (0, -1.7), (0, 0), (0, 1.7), (1.7, -1.7), (1.7, 0), (1.7, 1.7). p * (x) ∝ N (0, I) on Ω Ω = {x | 1 ≤ ∥x∥ 2 ≤ 4} Cardioid p * (x) ∝ N (0, I) on Ω Ω = {x | x 2 1 + ( 6 5 x 2 -x 2/3 1 ) 2 ≤ 4} Double-moon p * (x) ∝ q on Ω Ω = {x| -log q(x) ≤ 2, where q(x) = e -2(x 1 -3) 2 +e -2(x 1 +3) 2 e 2(∥x∥-3) 2 } Block Truncated gaussian mixture on Ω Ω = {x | |x 1 | ≤ 2 and |x 2 | ≤ 2}For CFG, f net and z net are three-layer neural networks with LeakyReLU activation (negative slope=0.1). The number of hidden units is 128, except for the RING on which 256 is used. Both neural nets are trained by Adam optimizer with learning rate 0.002, except for RING on which 0.005 is used. The number of inner-loop of gradient updates for f net and z net is set to 10, except for the RING on which 3 is used. The total number of iterations is 2000 and the step size of particle is 0.005 except for the RING on which 0.01 is used. The band width is set to 0.05 except for BLOCK on which 0.001 is used. λ in the piece-wise construction of the velocity field is chosen to be 1.On the BLOCK experiment, the total number of iterations is 2000. The ground truth is obtained by rejection sampling using 10 4 posterior samples. We run the accuracy experiments on 3 random seeds. The average results and the standard errors of the means are represented in the figure using lines and shades. For MSVGD, we follow Shi et al. (2022) and set the learning rate to 0.05 (selected from {0.01, 0.05, 0.1}). The kernel is the IMQ kernel and the kernel width is 0.1 (selected from {0.05, 0.1, 0.2}). For MIED, we use the default hyperparameters in Li et al. (2022) and set the particle stepsize to be 0.01 (selected from {0.005, 0.01, 0.02}).Additional Comparison of MIED Table 5 shows the additional comparison of CFG and MIED on the first three constrained domains. CFG achieved better results in most of the cases. For additional equality constraints, our proposed framework can still apply by adopting the idea of Zhang et al. (2022) . We could use velocity fields like v ♯ (x) to guided the particles to satisfy the equality constraints, and use our method to propose velocity fields substituting v ⊥ (x) to satisfy the inequality constraints. Adding these two types of velocity presents the desired velocity field. To illustrate the idea, we demonstrate the training process of a 3D toy ring distribution example. Suppose the coordinate of the particle is x = (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 ≤ x 2 + y 2 ≤ 4. The initial distribution is the 3D standard gaussian distribution. We choose v ⊥ (x) = -sign(z)|z| 1.5 . From figure 7 we can see that the particles collapse to the xOy plane and converge inside the ring domain.For additional inequality constraints, we need multiple velocity field outside the constrained domain for particles to enter the constrained domain, and the boundary integral term can be similarly estimated using band-wise approximation.Setting Details The experiments are implemented on Intel 2.30GHz CPU with RAM 16384MB and NVIDIA GeForce RTX 3060 Laptop GPU with total memory 14066MB. For CFG, f net and z net are three-layer neural networks with LeakyReLU activation (negative slope=0.1). The number of hidden units is 256. Both neural nets are trained by Adam optimizer with learning rate 0.0005 for 10 iterations in the inner loop. The total number of iterations is 1000 and the stepsize of particle is 0.004. λ in the piece-wise construction of the velocity field is chosen to be 1.For Spherical HMC, we follow the same main setting as Lan et al. (2014) . The number of burn-in epochs is 4000, and the number of leap-frog is selected from {50, 100}.For MIED, we use the default hyperparameters in Li et al. (2022) and set the particle stepsize to gradually decay when the particle number N goes up. We set particle stepsize to 0.5 for N ∈ {5, 15, 30, 70}, 0.35 for N ∈ {180, 400}, 0.3 for N = 900, and 0.2 for N = 2000. The number of iterations is 4000 for convergence except for N = 2000, which needs 10000 iterations to converge.3 2 1 0 1 2 3 3 2 1 0 1 2 3 3 2 1 0 1 2 3 Iter: 0 3 2 1 0 1 2 3 3 2 1 0 1 2 3 3 2 1 0 1 2 3Iter: 1003 2 1 0 1 2 3 3 2 1 0 1 2 3 3 2 1 0 1 2 3Iter: 10003 2 1 0 1 2 3 3 2 1 0 1 2 3 3 2 1 0 1 2 3Figure 7 : Illustration of generalizing our method to accommodating equality and inequality constraints. On real diabetes dataset, for CFG, f net and z net are three-layer neural networks with LeakyReLU activation (negative slope=0.1). The number of hidden units is 50. Both neural nets are trained by Adam optimizer with learning rate 0.005 for 10 iterations in the inner loop. The total number of iterations is 300 and the bandwidth h = 1. The number of particles is 5000 and the stepsize of particle is selected from {0.9, 1.05, 1.2}.For Spherical HMC, we follow the same setting as Lan et al. (2014) . The number of total epochs is 11000 and the number of burn-in epochs is 1000.For MIED, we use the default hyperparameters in Li et al. (2022) and the particle stepsize is selected from {0.1, 0.2, 0.5}. The number of iterations is 2000.Additional Comparisons on the Synthetic Dataset We further compare the sampling results of SPH, CFG and MIED on different dimensions. We choose dimensions 5, 10, 15 and 20. From the design of β true , the samples of first two dimensions are close to 10, and the last two dimensions are close to 0. Figure 8 shows that CFG achieves comparable results to MIED on Wasserstein-2 distances, while slightly better result on energy distance in most of the dimensions. SPH eventually catches up when the particle number grows up. This alignes with the results stated in Section 6.2.1.The Nonlinear Time Complexity of MIED Larger number of particles is important for higher sample qualities. Similar to Dong et al. (2023) , CFG is more scalable than MIED in terms of particle numbers N . This is because CFG is obtained with iterative approximation, the complexity is O(N ). On the other hand, MIED incurs complexity O(N 2 ) due to the calculation of the weight denominator i,j e Ii,j . Figure 9 plot the Wasserstein-2 Distance versus the computation time using at most 1000 iterations for CFG and 10000 iterations for MIED. The number of particles are N ∈ {900, 2000, 4000}. It is clear that MIED entails more computation cost with the increase of particle numbers for larger N . From the ending points of each curves (denotes the computation time for 1000 iterations of CFG and 10000 iterations of MIED), the time cost of CFG only grows linearly (from 130s to 280s to 600s approximately), while the time cost of MIED grows quadratically (from 100s to 220s to 830s approximately), which supports the above claim of time complexity.Additionally, figure 9 indicates that for large N , the number of iterations for MIED should also increase. 4000 iterations is sufficient for N = 900, while 10000 iterations is needed for N = 2000, and more than 10000 iterations is needed for MIED to completely converge to achieve better sampling results.The Effect of Choosing the Adaptive Bandwidth Figure 10 compared the energy distance between the adaptive bandwidth scheme h = 0.1N -1 3 and fixed bandwidth h = 0.1, 0.01 of CFG. We can see that the adaptive scheme achieved slightly better energy distance results than the fixed bandwidth.Setting Details The experiments are implemented on Nvidia GeForce RTX 4080 Laptop GPU with memory 12 GB. We use 200 particles and 200 data batch size for stochastic gradient in all methods.For CFG, both f net and z net are three layers with LeakyReLU activation (negative slope = 0.1). f net have 300 hidden units while z net have 200. Both neural nets are trained by Adam optimizer with learning rate 0.001 for 10 iterations in the inner loop. The bandwidth h = 0.02 and the stepsize of particle is 5e-5. The total number of iterations is 1200. λ in the piece-wise construction of the velocity field is chosen to be 100.For PD SVGD and Control SVGD, we run 3000 iterations with particle stepsize selected from {0.001, 0.002, 0.005} to ensure convergence.For MIED, we run 1500 iterations with default hyperparameters in Li et al. (2022) and particle stepsize selected from {2, 3, 5} × 10 -4 .Additional Experiments on Scaling Up to Higher Dimensions To test our method on higher dimension, we first experiment on the COMPAS dataset using larger monotonic BNNs. The number of neurons in the layer of BNN increased from 50 to 100, making the particle dimension increase from 903 to 1502. From Table 7 , our method still achieved the best accuracy results compared to other methods, while achieving competitive log-likelihood results.For even higher dimension, we additionally experiment on the 276-dimensional larger dataset Blog Feedback Liu et al. (2020) using monotonic BNN. The particle dimension is 13903. From Table 8 , our method still achieved the best result. Extensions to SVGD. The main idea of our work is enforcing the particles into the constrained domain, and then using neural networks to capture the target distribution. The latter part can also be accomplished by using kernels like in SVGD. The boundary integral may be treated in the similar way.We believe that our framework can be adapted to a larger family of variational sampling methods.More Efficiently Choosing Bandwidth h. In this paper, we only consider the simple case of uniform target distribution in choosing a better adaptive scheme for bandwidth h. Proposing an adaptive scheme based on the current particles is an interesting topic, which will be left for future work.Extensions to GWG. Cheng et al. (2023) proposed a variational framework including general geometries by using l p norm regularization. In this case, the treatment of boundary integral is the same as in the RSD loss expansion. The effect of using other regularization terms is left for future study.Relations with the Reflected Stochastic Differential Equation. Contrast to SDE-based sampling methods, our work is based on ODE. Both SDE and ODE can simulate trajectories including reflection. We proposed a heuristic understanding of the relation between SDE and ODE reflection in this work, a more thorough theoretical analysis may be presented. We leave this to future research.This paper presents work whose goal is to advance the field of machine learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here. As far as we are concerned, our paper has no potential negative societal impacts.1. Claims Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? Answer: [Yes] Justification: The last four sentences of the abstract and the last four sentences of the introduction reflect the paper's contributions and scope. Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 


Decision: This paper investigates the constrained sampling problem in Particle-based Variational Inference methods and proposes a framework called Constrained Functional Gradient Flow (CFG) to restrict samples within a specific domain. This is achieved by introducing a boundary condition for the gradient flow. The paper also provides a convergence guarantee in total variation for the proposed method.

Almost all reviewers agree that the paper merits acceptance, and after carefully reading the rebuttal and discussion, I tend to concur. Please incorporate the reviewers’ feedback into the final version of the paper, with particular attention to the following concerns:

1- Include the geometric generalization of your method and clarify the assumptions about entering the boundary and how to set $\lambda$ and other hyper-parameters. (Response to Rev s6Xw, and YhS1)

2- Include the experimental results of real-world applications. 

3- Include the related baseline and references and their complexity comparison. (Response to Rev YhS1, PQmA, and VuAH)

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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?
limitations: limitations: Limitations have been adequately addressed.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: Can authors illustrate how to pick lambda? How does lambda affect the algorithm, at least for numerical reasons?
limitations: limitations: There are no limitations.
paper_summary: paper_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.
strength_weakness: strength_weakness: In Bayesian neural network experiments, the paper overlooks numerous baselines of existing Bayesian Neural Networks (BNNs) using variational inference, MCMC, and Laplace approximation.
questions: questions: If we have additional equality constraints, does the method still work?
limitations: limitations: The paper has discussed the limitations.
paper_summary: paper_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.
strength_weakness: strength_weakness: Experiments on a real-world application could strengthen the paper.
questions: questions: There is no question from the reviewer.
limitations: limitations: N/A
paper_summary: paper_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).
strength_weakness: strength_weakness: 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: 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.
limitations: limitations: The limitations are addressed adequately.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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.
limitations: limitations: I think this is done reasonably well.

I would have liked to see more discussion of the tuning of the algorithmic parameters.

Comments: title: Thanks for your response
Thanks for your response
Thanks for your response
Response to authors
Thanks for your response
Thanks for the detailed response
Reply to authors
comment: Thanks for raising the score! Please feel free to let us know if you have further questions!
Thanks for your careful review! Please feel free to let us know if you have further questions!
Thanks for raising the score! We will modify our paper in our revision according to the responses made in the rebuttal.
Thank you for the comparison and the detailed answer. I think that this is a solid and innovative paper and I will raise my score a bit to a 7
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.
Thanks for your suggestion! We will modify our paper accordingly in our revision.
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.
The authors have addressed my questions. I will keep my score.