Sampling  with Trusthworthy Constraints:  A Variational Gradient Framework   

1 Introduction . Efficient approximation and sampling methods of intractable distributions plays a key role in probabilistic machine learning , especially Bayesian inference . Traditionally , this problem has been framed as Monte Carlo sampling . Recently , variational optimization ideas , which frame the sampling problem into functional minimization problem of KL divergence in the space of distributions , have been popularized , as exemplified by Langevin dynamics ( through the lens of Wasserstein gradient flow ( Jordan et al. , 1998 ) ) , and Stein variational gradient descent ( Liu & Wang , 2016 ; Liu , 2017 ; Liu & Wang , 2018 ) . Let p∗0 be a probability density function supported on Rd from which it is intractable to draw samples exactly ; p∗0 can be the posterior distribution in Bayesian inference , or a probabilistic graphical model constructed based on data and expert knowledge . The variational methods consider ∗ : Equal contribution . Code is available at https : //github.com/gnobitab/ConstrainedSampling . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . to minimize the following KL divergence objective in the space of distributions , denoted by P , min q∈P KL ( q || p∗0 ) , ( 1 ) which , if solved exactly , recovers the target distribution q = p∗0 . Langevin dynamics and SVGD can be viewed as solving ( 1 ) following two different types of gradient flow under different metrization of P ; in practice , both Langevin dynamics and SVGD yield a particle approximation , a set of points ( a.k.a . particles ) { θi } ni=1 ⊂ Rd , whose empirical measure q = ∑n i=1 δθi/n approximates p ∗ 0 in weak convergence ; the difference is that Langevin dynamics obtains the particles with a random diffusion process , while SVGD evolves the particles with a deterministic repulsive interacting particle system . However , most existing works focus on sampling from unconstrained domains , reflecting the fact that Problem ( 1 ) is an unconstrainted optimization on P . In many cases , especially trustworthy machine learning , we would like to impose a constraint of high practical importance in addition to approximating p∗0 , yielding a constrained variant of ( 1 ) min q∈P KL ( q || p∗0 ) , s.t . Eq [ g ( θ ) ] ≤0 , ( 2 ) where we find the best approximation of p∗0 subject to a moment constraint Eq [ g ( θ ) ] ≤ 0 , with g a general nonlinear function specified by the users . If p∗0 is the posterior distribution in Bayesian inference , Problem ( 2 ) can be viewed as a form of posterior regularization ( Zhu et al. , 2014 ) . Our framework provides a powerful tool for enabling reliable trustworthy machine learning with fairness , safety , interpretability , and other constraints , which we illustrate as follows . Fairness Constraints A motivating example is Bayesian fairness ( Chakraborty , 2020 ; Ji et al. , 2020 ; Dimitrakakis et al. , 2019 ) , in which , in addition to approximating the posterior distribution , we want to enforce equalized odds or other fairness constraints on different demographic subgroups . Specifically , suppose we are given a dataset D = { x ( i ) , y ( i ) , z ( i ) } Ni=1 consisting of the feature vector x ( i ) , a label y ( i ) , and a protected attribute z ( i ) ( e.g. , male vs. female ) , and we want to fit it with a prediction model ŷ ( x ; θ ) described by a parameter θ . In typical Bayesian inference , we are interested in finding a distribution q on θ to approximate the posterior distribution p∗0 ( θ ) = p ( θ | D ) . In fair Bayesian inference , we also hope q to satisfy certain fairness constraints . For example , to control the disparate impact , one can introduce a constraint of Eq [ g ( θ ) ] ≤ 0 with g ( θ ) = ` fair ( θ ) − , ` fair ( θ ) = ( covD [ z , ŷ ( x ; θ ) ] ) 2 , ( 3 ) where we enforce the prediction ŷ ( x ; θ ) to be uncorrelated with the protected attribute z . Therefore , solving ( 2 ) allows us to obtain the best approximation of p∗0 subject to the fairness constraint . Safety and Interpretability Constraints It is difficult to control and make sense of the behavior of large AI models such as deep neural networks . In many applications , such as healthcare , robotics , and AI-based systems , it is essential to add guardrails to ensure that the AI models stay within a pre-specified safety region . For example , we may want to ensure that the prediction ŷ ( x , θ ) must fall inside a pre-specified interval [ ŷ0 , − ( x ) , ŷ0 , + ( x ) ] , which can be ensured by Eq [ g ( θ ) ] ≤ 0 with g ( θ ) = Dist ( ŷ ( x , θ ) , [ ŷ0 , − ( x ) , ŷ0 , + ( x ) ] ) , where Dist ( · , · ) is a point-to-set distance . Similar approach can also be used to control and increase the interpretability of deep neural networks ( DNNs ) , by enforcing that the prediction of DNNs is close to simple interpretable models such as linear classifiers and logic rules . Prior-Agnostic Bayesian Inference The posterior of typical Bayesian inference depends on both the likelihood ` ( θ ) which represents the fitness on data and a prior distribution ( let it be p∗0 ) . A poor choice of prior may heavily influence the posterior and yield poor fitness . One approach to automatically limiting the influence of prior is to find a q that minimizes KL ( q || p∗0 ) , subject to the moment constraint with g ( θ ) = − ` ( θ ) + , so that we always have a guaranteed high data fitness ( i.e. , high ` ( θ ) ) regardless of the choice of the prior . Note that p∗0 here is the prior , while it is the posterior in the two settings above . Our Contribution We advance the frontier of sampling with constraints in the following ways : i ) We propose two general approaches to extend SVGD and Langevin dynamics to the constrained setting ( 1 ) , including a primal-dual gradient method ( see Algorithm 1 ) , and a novel constraint controlled gradient descent method ( see Algorithm 2 ) . ii ) We develop novel theoretical analysis on both methods for both convex and non-convex settings . iii ) We demonstrate the power of our approaches on a variety of tasks related to trustworthy machine learning , including fair Bayesian classification , incorporating logic rules into black-box models , and training monotonic neural networks . Overview of the Main Algorithms We provide a quick summary of the practical algorithms that we develop for solving ( 1 ) , so that the readers with a practical interest can skip the main theoretical derivation , which is in Section 2.1 . Our algorithms iteratively update a set of n particles { θi , t } ni=1 ⊂ Rd through iteration t , such that its empirical distribution , denoted as qt = ∑n i=1 δθi , t/n , approximately solves the constrained Problem 2 in a proper sense when t → +∞ and n → +∞ . The updates of our algorithm for the SVGD and Langevin cases are : SVGD : θi , t+1 = θi , t + hEθ∼qt [ ( ∇ log p∗0 ( θ ) − λt∇g ( θ ) ) k ( θ , θi , t ) +∇θkt ( θ , θi , t ) ] , ( 4 ) Langevin : θi , t+1 = θi , t + h ( ∇ log p∗0 ( θi , t ) − λt∇g ( θi , t ) ) + √ 2hξi , t , ( 5 ) where { ξt , i } are i.i.d . standard Gaussian noise and h is a step size . The updates modify the standard SVGD and Langevin update rules by introducing an extra λt∇g ( θ ) term to account the constraint in ( 2 ) , where λt ≥ 0 serves as a Lagrange multipler , whose update rule is described below . Primal-dual Gradient Method A straightforward approach is to update λt by performing projected gradient descent on the dual problem of ( 2 ) . As we show in Section 2.2 , for both the SVGD and Langevin case , this can be achieved by , λt = max ( λt−1 + h̃Eqt [ g ( θ ) ] , 0 ) , ( 6 ) where h̃ is a step size that can different from h. This is similar to the typical primal-dual gradient method on finite dimensional optimization : we increase λt−1 if Eqt [ g ( θ ) ] ≥ 0 ( constraint violated ) and decrease it if the constraint is met . See Algorithm 1 , and the detailed version in Algorithm 3 in the appendix . Constraint Controlled Gradient Method Because λt is iteratively updated in ( 14 ) , the result can be sensitive to the initialization λ0 and the learning rate η of λ . In Section 2.3 , we propose an alternative “ constraint controlled ” method so ( 4 ) and ( 5 ) yield the steepest descent on the KL objective while ensuring that the solution converges to the feasible set rapidly when it is violated . The derived update λt is different for the Langevin and SVGD cases and is shown below : Langevin : λt = max ( αEqt [ g ] + Eqt [ ( ∇ log p∗0 ) > ∇g +∇ > ∇g ] Eqt [ ‖∇g‖ 2 ] , 0 ) , ( 7 ) SVGD : λt = max ( αEθ∼qt [ g ( θ ) ] + Eθ , θ′∼qt [ ∇θ′g ( θ′ ) > ( ∇ log p∗0 ( θ ) +∇θ ) kt ( θ , θ′ ) ] Eθ , θ′∼qt [ ∇g ( θ ) > ∇g ( θ′ ) kt ( θ , θ′ ) ] , 0 ) , ( 8 ) where α > 0 is a coefficient . Note that the λt here is completely decided by the information at the t-th iteration and does not need to be iteratively updated from λt−1 like ( 6 ) . See Algorithm 2 and more details in Algorithm 4 in Appendix . 2 Main Results . We first introduce the background on SVGD and Langevin through the view of minimizing ( 1 ) via functional steepest descent in Section 2.1 . We then introduce our two approaches to extending SVGD and Langevin dynamics for solving the constraint optimization in ( 2 ) in Section 2.2 and Section 2.3 . We provide theoretical analysis of both approaches with and without convexity assumptions . 2.1 Review on Langevin Dynamics and SVGD : Sampling with Steepest Descent . We provide a unified introduction to Langevin dynamics and Stein variational gradient descent ( SVGD ) , which can be viewed as minimizing KL ( q || p∗0 ) with two different types of gradient flow on the space of distributions P . Our results apply to both algorithms . For audience with special interest in one of the two algorithms , we recommend ignoring the other one in the first read . Assume we start from an initial density q0 and the associated random variable θ0 ∼ q0 . We consider moving the random variable θt along a time-dependent vector field φt : Rd → Rd . In other words , θt is driven by an ordinary differential equation ( ODE ) : dθt/dt = φt ( xt ) . Let qt be the distribution of θt at time t , which is known to follow the Fokker-Planck equation dqt/dt = −∇ > ( φtqt ) , where ∇ > f ( θ ) : = ∑d ` =1 ∂θ ` f ` ( θ ) denotes the divergence operator of a vector-valued function f : Rd → Rd with θ ` and f ` the ` -th element of θ and f , respectively ; here we use the “ > ” notation because we formally view∇ as a d-dimensional column vector . To decrease KL ( qt || p∗0 ) as fast as possible , we want to choose φt , from a pre-decided candidate function set Ft , so that the decreasing rate − ddtKL ( qt || p ∗ 0 ) is maximized . One can show that − d dt KL ( qt || p∗0 ) = Eθ∼qt [ ( ∇ log p∗0 ( θ ) −∇ log qt ( θ ) ) > φt ( θ ) ] : = Rqt , p∗0 [ φt ] , ( 9 ) which is a linear functional Rqt , p∗0 acting on Ft . Assume Ft is a Hilbert space , equipped with an inner product 〈· , ·〉Ft and a norm ‖·‖Ft . By Riesz representation theorem , assuming Rqt , p∗0 is a continuous linear functional , there exists an element rqt , p∗0 ∈ Ft , such that Rqt , p∗0 [ φ ] = 〈rqt , p∗0 , φ〉Ft . We call rqt , p∗0 the Riesz representation of Rqt , p∗0 in Ft . The optimal choice of φt in Ft can be framed into an optimization : φt = arg min φ∈Ft { −〈rqt , p∗0 , φ〉Ft + 1 2 ‖φ‖2Ft } = rqt , p∗0 , DFt ( qt , p ∗ 0 ) 2 = ∥∥rqt , p∗0∥∥2 , ( 10 ) where we add a regularization ‖φ‖2Ft to constrain the scale of φ , yielding the optimal solution φt = rqt , p∗0 ; here we defined DFt ( qt , p ∗ 0 ) 2 to be the corresponding maximum descending rate of KL divergence , which can be viewed as a discrepancy measure between qt and p∗0 . In this work , we always assume that Ft is sufficiently large , so that DFt ( qt , p∗0 ) = 0 implies qt = p∗0 . Langevin dynamics and SVGD can be viewed as using different Hilbert spaces Ft , hence yielding different Riesz representation rqt , p∗0 for the linear operator Rqt , p∗0 . Langevin Dynamics Taking Ft to be Lqt,2 , the Hilbert space of Rd → Rd maps equipped with inner product 〈φ , φ′〉Lqt,2 : = Eθ∼qt [ φ ( θ ) > φ′ ( θ ) ] . Then it is immediate to see from ( 9 ) that φt ( θ ) = rqt , p∗0 ( θ ) = ∇ log p ∗ 0 ( θ ) −∇ log qt ( θ ) . ( 11 ) In this case , DFt ( qt , p∗0 ) 2 is the Fisher divergence Eθ∼qt [ ‖∇ log p∗0 ( θ ) −∇ log qt ( θ ) ‖ 2 ] . Note that the Fokker Planck equation , dqt/dt = −∇ > ( ( ∇ log p∗0 −∇ log qt ) qt ) = −∇ > ( ∇ log p∗0 ( θ ) qt ( θ ) ) + ∇ > ∇qt ( θ ) , coincides with the density functions of the Langevin diffusion process dθt = ∇ log p∗0 ( θt ) dt+ √ 2dWt , whose time-discretization yields the ( unadjusted ) Langevin Monte Carlo method ; here Wt is the standard Wiener process . Meanwhile , other discretization methods also exist , see e.g . Wibisono ( 2018 ) ; Salim et al . ( 2020 ) . Stein Variational Gradient Descent ( SVGD ) In SVGD , we take Ft to the reproducing kernel Hilbert space of a positive definite and continuously differentiable kernel kt ( θ , θ′ ) . By the reproducing property of RKHS and Stein identity ( see Liu & Wang ( 2016 ) ) , one can show that the optimal φt is φt ( · ) = rqt , p∗0 ( · ) = Eθ∼qt [ ∇ log p ∗ 0 ( θ ) kt ( θ , · ) +∇θkt ( θ , · ) ] , ( 12 ) Related , DFt ( qt , p∗0 ) reduces to kernel Stein discrepancy ( Liu et al. , 2016 ; Chwialkowski et al. , 2016 ) : DF ( qt , p∗0 ) 2 = ∥∥rqt , p∗0∥∥2Ft = Eθ∼qt [ ( ∇θ log p∗0 ( θ ) +∇θ ) > ( ∇θ′ log p∗0 ( θ′ ) +∇θ′ ) kt ( θ , θ′ ) ] . A nice property of the φt in ( 12 ) , compared with the φt in ( 11 ) , is that it depends on qt only through the expectation operator Eqt , which enables a direct particle implementation of dθt = φtdt , instead of resolving to a diffusion process . Specifically , if we initialize q0 = ∑n i=1 δθi,0/n to be the empirical measure of a set of particles { θi,0 } i , then qt remains to be an empirical measure , that is , qt = ∑n i=1 δθi , t/n , where the particles are evolved by d dt θi , t = Eθ∼qt [ ∇ log p∗0 ( θ ) kt ( θ , xi , t ) +∇θkt ( θ , θi , t ) ] , ∀i = 1 , . . . , n , This is a set of differential equations coupled by the mean field qt . Following Liu & Wang ( 2016 ) , the bandwidth parameter of the kernel kt ( θ , θ′ ) at time t can depend on the current particles { θi , t } i via the median trick .

Summary: The authors propose two ways to extend Unadjusted Langevin Algorithm (ULA) and Stein Variational Gradient Descent (SVGD) for sampling from a distribution known up to its normalization constant, under moment constraints. The first one is a primal-dual gradient descent scheme, while the second one is a controlled gradient descent one.  They provide convergence guarantees in continuous time for both schemes: a 1/T rate for a given optimality criterion Delta in the more generic case (Th 3.4 and 3.8), and linear convergence rate in KL if some form of log-concavity is satisfied for the current target p_{lamba_t}* (Th 3.5 and 3.9).  These methods perform reasonably well on synthetic and real world experiments, with a better performance for controlled-based methods.