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 extend gradient flow methods for variational inference to handle inequality constraints. These inequality constraints can capture fairness, safety, interpretability, and other such desiderata. They introduce two methods to solve this. First, they introduce a 'primal-dual gradient method': they formulate the equivalent minimax problem to the constrained optimization, and then they derive the dynamics for gradient descent ascent of the zero-sum game. Convergence is shown by explicitly providing a Lyapunov function. Next, they introduce a 'constraint controlled gradient descent', which first finds the feasible set, and then stays within the feasible set afterward (inspired by controlled barrier functions (CBFs) from control). Convergence is shown similarly to the work in CBFs.