Sampling with Trustworthy Constraints: A Variational Gradient Framework
Year: 2022

Abstract:
Sampling-based inference and learning techniques, especially Bayesian inference, provide an essential approach to handling uncertainty in machine learning (ML). As these techniques are increasingly used in daily life, it becomes essential to safeguard the ML systems with various trustworthy-related constraints, such as fairness, safety, interpretability. Mathematically, enforcing these constraints in probabilistic inference can be cast into sampling from intractable distributions subject to general nonlinear constraints, for which practical efficient algorithms are still largely missing. In this work, we propose a family of constrained sampling algorithms which generalize Langevin Dynamics (LD) and Stein Variational Gradient Descent (SVGD) to incorporate a moment constraint specified by a general nonlinear function. By exploiting the gradient flow structure of LD and SVGD, we derive two types of algorithms for handling constraints, including a primal-dual gradient approach and the constraint controlled gradient descent approach. We investigate the continuous-time mean-field limit of these algorithms and show that they have O(1/t) convergence under mild conditions. Moreover, the LD variant converges linearly assuming that a log Sobolev like inequality holds. Various numerical experiments are conducted to demonstrate the efficiency of our algorithms in trustworthy settings.

heading: 1 Introduction
text: 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.

heading: 2 Main Results
text: 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.

heading: 2.1 Review on Langevin Dynamics and SVGD: Sampling with Steepest Descent
text: 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.

heading: 2.2 Primal-Dual Gradient Method
text: Introducing a Lagrange multipler λ, Problem (2) is equivalent to the following minimax problem: min q∈P max λ≥0 {L(q, λ) = KL(q || p∗0) + λEq[g(θ)]} . (13)
We solve the minimax problem by alternatively updating λ and q with gradient descent. Note that we can rewrite L(q, λ) into L(q, λ) = KL(q || p∗λ)+Φ(λ), with p∗λ(θ) = p ∗ 0(θ) exp(−λg(θ) + Φ(λ)), and Φ(λ) = min
q∈P L(q, λ) = − logEθ∼p∗0 [exp(−λg(θ))],
and p∗λ = arg minq∈P L(q, λ). Therefore, if we want to minimize L(q, λt) with a fixed λt, we should update qt to move it towards p∗λt , which can be done by following the functional gradient descent with either SVGD or Langevin dynamics: dθt/dt = rqt,p∗λt (θt). With a fixed q, we perform standard projected gradient ascent on λ shown in (6):
λt = max(λt−1 + h̃Eqt [g(θ)], 0), (14) where h̃ is a step size. In the continuous time limit (h̃ → 0), we find the λt follows a projected ordinary differential equation. It can be written jointly with the density qt of θt as a primal-dual gradient flow (PDGF):
dqt dt = −∇ · (φtqt) = −∇ · (rqt,p∗λt qt), dλt dt = [ηEqt [g]]λt,+, (15)
with [a]λ,+ = I(λ ≤ 0) max(a, 0) + I(λ > 0)a and η = h̃/h characterizes the update speed of λt relative to qt. Again, rqt,p∗λt is defined by (11) for the Langevin dynamics, and (12) for the SVGD, as long as p∗0 is replaced by p ∗ λt . Numerical Implementation To implement (15) through a particle-based algorithm, we use Euler discretization and approximate qt with empirical measure ∑n i=1 δθi,t/n, where {θi,t}ni=1 are the particles at time t. This yields the updates of SVGD and Langevin shown in (4) and (5). Note that in (5), we run n parallel chains of Langevin dynamics, which are coupled since they share the common λ. The empirical distribution of {θi,t} is used to approximate Eqt [g(θ)] in (14). See Algorithm 1 and more details in Algorithm 3.
Convergence Analysis
One problem of primal-dual gradient method is that the algorithm does not try to meet the constraint directly. Instead, it tries to minimize KL(q || p∗λ) which involves the constraint. So in order to prove the convergence, we need to assume that KL(q || p∗λ) and its maximum descending rate DF (q, p∗λ) dominate the constraint function g. In particular, we consider the following assumptions:
Assumption 2.1. There exist positive constants c1, c2 <∞, such that for any t ≥ 0, (Eqt [g]− Ep∗λt [g]) 2 ≤ c1DFt(qt, p∗λt) 2 (16)
(Eqt [g]− Ep∗λt [g]) 2 ≤ c2KL(qt || p∗λt). (17)
Algorithm 1 Primal-Dual Method Initialize the particles {θi,0}ni=1 and λ0. Let qt = ∑ i δθi,t/n.
for iteration t do Update the particles {θi,t}ni=1 by Eq. (4) (SVGD), or Eq. (5) (Langevin). Update λt by Eq. (14). end for
Algorithm 2 Constraint Controlled Method Initialize the particles {θi,0}ni=1. Let qt =∑ i δθi,t/n.
for iteration t do If SVGD, update {θi,t}ni=1 by Eq. (4), and then update λt by Eq. (8) If Langevin, update {θi,t}ni=1 Eq. (5), and then Update λt by Eq. (7). end for
Condition (17) always holds when g is a bounded function, due to the Pinsker’s inequality (See e.g., Lemma A.1 of Cui & Tong (2021)). Condition (16) holds if g can be written into a form of g(θ) = a +∇ log p∗λ(θ)>ψ(θ) +∇>ψ(θ) for some a ∈ R, ψ ∈ Ft and ‖ψ‖ 2 Ft ≤ c1 (see Lemma B.1 in Appendix). It can also be obtained by (17) and a log-Sobolev like inequality (18), which we will discuss later.
Assumption 2.2. There exists 0 < v0 ≤ v1 <∞ such that v0 ≤ varθ∼p∗λt [g(θ)] ≤ v1 for ∀t ≥ 0.
This is a mild assumption, and holds when g is bounded and is not a constant on the support of p∗λt .
Our analysis is based on the following Lyapunov function for the minimax problem on L(q, λ):
E(q, λ) = L(q, λ)− 2Φ(λ) + max λ≥0 Φ(λ).
Lemma 2.3. We have E(q, λ) ≥ 0, and E(q, λ) = 0 iff (q, λ) is a saddle point of L(q, λ).
Under mild conditions, we can show E(qt, λt) is decaying along PDGF. Moreover, we can find a solution that meets the KKT requirement approximately.
Theorem 2.4. If (16) holds and 0 < η < 1/(2c1), then E(qt, λt) decreases monotonically following the primal-dual gradient flow in (15),
− d dt E(qt, λt) ≥ ∆(qt, λt) := (1− 2c1η)DFt(qt, p∗λt) 2 + 1 2 η(Eqt [g])2 × I(λ > 0 or Eqt [g] > 0).
Also (15) finds an approximate solution with an O(1/T ) rate: mint∈[0,T ] ∆(qt, λt) ≤ 1 T E(q0, λ0),∀T ≥ 0.
Note that ∆(q, λ) is a measure of optimality. This is because ∆(q, λ) = 0 implies the KKT condition of the constrained optimization (1) holds, which is q = p∗λ, Eq[g] ≤ 0 and λEq[g] = 0. Linear Convergence with Log-Sobolev like Condition We can further show a linear convergence of KL divergence to the optimal solution by assuming a Logarithmic Sobolev like inequality.
Assumption 2.5. There exists a positive constant c3 <∞, such that for any t ≥ 0,
KL(qt || p∗λt) ≤ c3DFt(qt, p ∗ λt) 2. (18)
Theorem 2.6. Suppose problem (13) has a solution (q∗, λ∗), and Assumptions 2.1,2.2 and 2.5 hold. The PDGF following (15) with 0 < η < 1/(2c1) will converge to q∗ linearly in KL divergence:
KL(qt||q∗) ≤ α1E(qt, λt) ≤ α1 exp(−α2t)E(q0, λ0),
where α1 = max{1 + 12c2, v1 2v0 + 12v0 } and α2 = min{ 1 c3 (1− c1η), 12ηv0}.
Assumption 2.5 is equivalent to the log-Sobolev inequality in the Langevin case (when Ft = Lqt,2 and DF (·, ·)2 is Fisher divergence), which holds when log p∗λ is a bounded perturbation of a strongly concave function (Holley & Stroock, 1986). In particular, it holds when p∗0 is strongly log-concave and g is convex. So it can be seen as a strong convexity (concavity) assumption. In the optimization literature, e.g. Nesterov et al. (2018), it is well known such assumption is necessary for linear convergence; without it, we can usually only show O(1/t) convergence similar to Theorem 2.4.
For SVGD, Assumption 2.5 is less well understood. If one replaces qt with any probability measure q, then (18) would fail to hold for kernel Stein discrepancy (KSD) because if qt is a discrete particle measure, we would have KL(qt || p∗λ) = ∞ but DFt(qt, p∗λ) < ∞ under some mild conditions on the kernel kt; see also Lemma 36 of Duncan et al. (2019). On the other hand, it is unclear if (17) will hold for the smaller class of densities where {qt : t ∈ [0,∞)} takes place in SVGD given that the initialization q0 is sufficiently regular (and in what sense). Therefore, Theorem 2.6 is not readily applicable to SVGD. We hope future works can draw more understandings on this issue.

heading: 2.3 Constraint Controlled Gradient Descent
text: The results of the primal-dual gradient method can be sensitive to the initialization λ0 and the learning rate η of λ. We now propose a “constraint controlled” method, which finds a constrained variant of steepest descent direction φt that yields the steepest descent on the KL objective like Section 2.1, while ensuring that the solution converges to the feasible set rapidly when it is violated.
Following Section 2.1, assume we update qt by dθt = φt(θt)dt, and want to decide the optimal φt. To solve the constrained optimization (2), in addition to maximizing the descending rate of KL
divergence as Section 2.1, we also want to ensure that the constraint Eqt [g] is properly controlled, such that 1) if the constraint is not met (i.e., Eqt [g] ≥ 0), we should monotonically decrease Eqt [g], and 2) after the constraint is met (i.e., Eqt [g] ≤ 0), we should monotonically descend the loss KL(qt, p ∗ 0), while ensuring that the solution stays within the constraint set. Note that
d dt Eqt [g] = Eqt [∇g(θ)>φt(θ)] := Sqt,g[φt] = 〈sqt,g, φt〉Ft ,
where Sqt,g is a linear operator onFt, whose Riesz representation is assumed to be sqt,g.Generalizing the functional steepest descent idea in (10), we propose to set φt to be the solution of
φt = arg min φ∈Ft
{ −〈rqt,p∗0 , φ〉Ft + 1
2 ‖φ‖2Ft s.t. 〈sqt,g, φ〉Ft ≤ −αEqt [g]
} , (19)
where α ≥ 0 is a control coefficient. This ensures that, if the constraint is not met (i.e. Eqt [g] > 0), the constraint is descending ( ddtEqt [g] ≤ −αEqt [g] < 0). If the constraint is met (Eqt [g] ≤ 0), then it allows Eqt [g] to increase, but with a rate smaller than −αEqt [g], which decreases towards zero when the solution approaches the constraint boundary {q : Eq[g] = 0}; this on one hand allows us to have the flexibility to choose φ to decrease KL(q || p∗0)), and on the other hand ensures the solution is confined inside the constraint set. The high-level idea here is similar to the control barrier functions in control theory (e.g., Ames et al., 2019), which works in completely different settings. Note that the constraint in (19) can be viewed as a linearization of the constant Eq[g] ≤ 0 around qt, as αEqt+1/α [g] = αEqt [g] + d dtEqt [g] = αEqt [g] + 〈sqt,g, φ〉Ft , while the loss in (19) can be viewed as a simple quadratic approximation of KL(q || p∗0). Therefore, (19) can be viewed as a functional variant of sequential quadratic programming (Nocedal & Wright, 2006).
Using the Lagrange duality of (19), it is easy to derive that φt = rqt,p∗0 − λtsqt,g = rqt,p∗λt , with
λt = max
( αEqt [g] + 〈rqt,p∗0 , sqt,g〉Ft
‖sqt,g‖ 2 Ft
, 0 ) . (20)
See Lemma B.2 in Appendix. Therefore, the density following the constraint controlled gradient flow (CCGF) should satisfy:
dqt dt = −∇ · (φtqt) = −∇ · (rqt,p∗λt qt). (21)
This is similar to PDGF (15) as qt is driven towards p∗λt . But the λt in PDGF is recursively updated by the algorithm, while CCGF has the explicit formula for λt. Intuitively, CCGF will be more efficient.
Next, we obtain Langevin and SVGD variants by taking Ft to be Lqt,2 and RKHS, respectively. Langevin Dynamics and its implementation Assume Ft = Lqt,2. We have rqt,p∗0 = ∇ log p ∗ 0 − ∇ log qt and sqt,g = ∇g. Hence φt = ∇ log p∗0 − λt∇g − ∇ log qt, with λt defined in (7). The evolution of qt associated with dθ = φt(θ)dt is standard Fokker Planck equation associated with p∗λt , which can be realized by the same parallel Langevin particle dynamics in (5), when setting qt = ∑n i=1 δi,t/n. Note that λt depends on qt only through the expectation Eqt [·], which can be replaced by the empirical average of the particles. See Algorithm 2 and more details in Algorithm 4.
The SVGD Case Let Ft be the RKHS of kernel kt(θ, θ′). Then it is easy to show that sqt,g(·) = Eθ∼qt [∇θg(θ)kt(θ, ·)]. Therefore, φt(·) = Eθ∼qt [(∇ log p∗0(θ) − λt∇g(θ))kt(θ, ·) + ∇θkt(θ, ·)], with λt defined in (8). Similar to regular SVGD, we set qt = ∑n i=1 δθi,t/n, which is iteratively updated with SVGD updates with p∗λt as the target distribution as (4). Also, λt again only depends on Eqt [·], which should be evaluated with the empirical mean of the particles. See Algorithm 2 and the detailed version in Algorithm 4.
Convergence Analysis
Since CCGF tries to meet the constraint explicitly, we have the following attractive properties:
Theorem 2.7. Suppose ‖sqt,g‖Ft 6= 0 in the CCGF (21). Then with φt defined in (20), we have Eqt [g] ≤ exp(−αt)Eq0 [g], ∀t ≥ 0.
If Eqt0 [g] ≤ 0 for some t0 ≥ 0, then Eqt [g] ≤ 0 and d dtKL(qt || p ∗ 0) ≤ 0 for all t ≥ s.
Note that the requirement ‖sqt,g‖Ft 6= 0 is necessary for the λt in CCGF to be well defined. If fact, if ‖sqt,g‖Ft = 0, then ddtEqt [g] = 0 along any update direction φ. And if Ft = Lqt,2, one can show g is a constant function qt-a.s.. In other words, this is an ill-posed scenario where local methods like CCGF have no chance to solve. An assumption slightly stronger than λt 6=∞ would be assuming it is bounded from above:
Assumption 2.8. There exists an upper bound λmax,+ < ∞, such that for all time t when the constraint is not satisfied (i.e., Eqt [g] > 0), we have λt ≤ λmax,+ in (20).
This is a mild regularity condition, which holds, for example, if Eqt [g] and ‖rqt‖ are upper bound, and ‖sqt,g‖ > c− > 0. In the following, we show that, under Assumption 2.8, our algorithm meets the KKT condition with a O(1/t) rate. Further, if Assumption 2.5 holds, then KL(qt || q∗) converges to zero exponentially fast. Note that the conditions here are milder than that of the primal-dual method (which needs Assumption 2.1 and 2.2).
Theorem 2.9. Let {qt} be the density function of dθt = φt(θ)dt with φt defined in (20). Under Assumption 2.8, we have
min t∈[0,T ]
∆2(qt, λt) := DFt(qt, p∗λt) 2 + α(−λtEqt [g])+ ≤
1 T (KL(q0, p ∗ 0) + λmax,+(Eq0 [g])+) .
Note that, given that Eqt [g] ≤ 0 (see Theorem 2.7), ∆2(qt, λt) = 0 implies the rest of KKT condition (i.e., qt = p∗λt and (−λtEqt [g])+ = 0).
We establish below the linear convergence of KL divergence under the log-Sobolev condition (18).
Theorem 2.10. Suppose problem (13) has a solution (q∗, λ∗), ‖sqt,g‖Ft 6= 0, and Assumption 2.5 also holds. Then the CCGF following (21) with α ≤ 1/c3 will converge to q∗ linearly in KL divergence:
KL(qt || q∗) ≤ e−αtKL(q0 || q∗), ∀t ∈ [0,+∞).
Similar to Theorem 2.6, this result is not readily applicable to the SVGD case because (18) can not be verified for SVGD.

heading: 3 Experiments
text: Algorithms and Settings We summarize the tested methods and the unconstrained baselines here: Langevin: Vanilla Langevin dynamics. We run n parallel chains of Langevin dynamics for sampling; SVGD: Vanilla SVGD with codes from Liu & Wang (2016); Primal-Dual+Langevin: Langevin dynamics with Primal-dual gradient descent; Primal-Dual+SVGD: SVGD with Primaldual gradient descent; Control+Langevin: Langevin dynamics with constraint controlled gradient descent; Control+SVGD: SVGD with constraint controlled gradient descent. For the SVGD methods, we follow the configuration in Liu & Wang (2016). We use the RBF kernel with bandwidth chosen by the standard median trick, that is, we use kt(θ, θ′) = exp(−‖θ − θ′‖2 /w2t )) where the bandwidth wt is set by wt = Median{‖θi,t − θj,t‖ : i 6= j} based on the particles {θi,t}ni=1 at the t-th iteration. As a remark regarding the choice of kernel, Gorham & Mackey (2017) provided a counter-example that suggests that Stein discrepancy with Gaussian RBF kernel may fail to metrize the weak convergence and suggested to use inverse multi-quadratic (IMQ) kernel which does not suffer from the problem. However, this counter-example assumes to use a fixed bandwidth and does not hold for Gaussian RBF kernel equipped with the median trick which can adapt to the scale of the data better. We leave the study of better and adaptive choice of kernels to future works. In our implementation, we adopt the same decaying step size as suggested in Welling & Teh (2011) for both SVGD and Langevin dynamics, where ht = h0(1.0 + t)−0.55 and h0 is a hyper-parameter. The step size for the Lagrangian multiplier in primal-dual methods is chosen from {0.001, 0.01, 0.1, 1, 10, 100, 1000}. The hyper-parameters are determined by grid-search to reach the smallest constraint loss in each experiment.
PrLPDl-DuDl+LDngevLn
PrLPDl-DuDl+LDngevLn
PrLPDl-DuDl+SVGD
Embedding Logic Rules into Black-box Models ML models can provide accurate predictions but are difficult to interpret and control explicitly. Our method can be used to enforce the ML model to be consistent with a set of logic rules to improve the interpretability and safety. We illustrate this with a loan classification problem of predicting whether to lend loans to a specific applicant. We impose the two logic rules: (1) an applicant must be denied if she has the lowest credit rank and not employed; (2) an applicant must be approved if she has the highest credit rank and has been employed over 15 years. To apply our method, we set p∗0 to be the typical posterior distribution of Bayesian logistic regression, and define the constraint to g(θ) = `logic(θ)− , with `logic(θ) = E(x,y)∼Dlogic [Loss(y, ŷ(x; θ))], where Dlogic is the uniform distribution on the (x, y) that satisfy the logic constraints, Loss is the classification loss and ŷ(x; θ) is the prediction by the Bayesian logistic regression model. We vary the threshold in {0.001, 0.01, 0.02, 0.03, 0.04, 0.05} and find the corresponding training log-likelihood (LL). In combination, they are plotted as a trade-off curve shown in in Figure 1(a, left). In Figure 1(a, right), we plot the trade-off curve of the testing accuracy vs. the constraint violation ((Eθ∼q[`logic(θ)]− )+) on the testing data. We can see that the controlled-based methods tend to enforce the constraints better (reaching lower values in logic loss) in both training and testing set, and the SVGD methods tend to achieve higher test accuracy.
Training Monotonic Bayesian Neural Networks In some applications, it can be desirable to enforce the ML prediction to be monotonic w.r.t certain attributes (Karpf, 1991; Sill, 1998). For example, when predicting admission decisions, a fair ML system must admit a student with higher GPA over the students with lower GPA, given that they are identical in the other features. We apply our method to enforce monotonicity in Bayesian neural networks. We use the COMPAS dataset following the setting in Liu et al. (2020). In this case, p∗0 is the posterior of a Bayesian neural network on the data, and g is defined as g(θ) = `mono(θ)− with `mono(θ) = Ex∼D [‖(−∂xmono ŷ(x; θ))+‖1]. Here, xmono denotes the subset of features to which the output should be monotonic. By varying in {0.0001, 0.01, 0.1, 0.5}, we plot in Figure 1(b) the trade-off curve of LL vs. the monotonic loss on training/testing set, which follow similar trends as Figure 1(a).
Training Fair Bayesian Neural Networks We use the constraint defined in (3), and use the Adult Income dataset (Kohavi, 1996), which is a classification problem of predicting whether the annual income of a person is≥ $50,000, with gender as the protected attribute. For the experiment, we follow the setting in Martinez et al. (2020); Liu & Vicente (2020). With = {0.0001, 0.001, 0.005, 0.01}, Figure 2(a, left) shows the trade-off curve of training LL vs. fairness loss, and Figure 2(a, right) shows the testing accuracy vs. CV score (Calders & Verwer, 2010), a standard measure of disparate impact. Figure 2(b) shows an example of training curves. We can again see that the control based methods enforce the constraints more strictly and in a faster speed.
Constrained Sampling on Generative Models We apply our method on Noise Conditional Score Networks (NCSN) (Song & Ermon, 2019). NCSN models the distribution p∗0 of images by training a neural network to estimate its score function ∇ log p∗0, then generates images by running Langevin dynamics with the learned score function. Here, we add constraints to draw rare samples from the generative model. In this experiment, we constrain the center of the generated images to be a black square, by defining the constraint function to be g(θ) = ‖θΩ‖ − , where θ denotes the whole image, θΩ is a image patch on region Ω, which is the square at the center of image. We use a threshold of = 0.0001. We use Control+Langevin to draw samples under the constraint, as the original NCSN model is based on Langevin dynamics. As shown in Fig. 3, our method successfully generates diverse images without violating the constraints.

heading: 4 Related Works, Limitations, Conclusion
text: Primal-Dual vs. Constraint-Control This work has proposed two methods, primal-dual and CCGF. These two methods are derived from two different perspectives. CCGF tries to fit the constraints directly, while the primal-dual operates in a more indirect manner. This is reflected in the proof techniques, as we can directly show the constraint decays linearly for the CCGF method, but we can only control the constrain through a Lyapunov function with the primal dual method. We believe that the conditions required by CCGF is intrinsically weaker than that of primal-dual. This can also be told from the fact that primal-dual method’s performance depends on the initialization of λt and the step size η. But the CCGF method does not have this issue.
Related Works Although sampling on unconstrained domain has been extensively studied, efficient algorithms for constrained settings are largely lacking. A body of works have developed to extend Langevin dynamics and SVGD on Riemannian manifolds (Girolami & Calderhead, 2011; Patterson & Teh, 2013; Liu & Zhu, 2018); these methods, however, rely on computationally efficient characterization of the manifolds and as a result does not work efficiently with constraints specified by general nonlinear inequalities. Projection-based methods (e.g, Sen et al., 2018) can also be developed for sampling on constrained spaces, but again only work for simple constraints with efficient projection maps. Note that our method considers a different setting when the constraint is applied on the population (rather than individual particles) and works for general nonlinear inequality constraints. In terms of problem formulation, our constrained optimization in (2) coincides with the posterior regularization framework of Zhu et al. (2014); but we focus on particle-based inference based on SVGD and Langevin, which was not studied before.
Limits, Impacts, Future Directions Our new methods cast a full spectrum of theoretical questions that can not be completed with a single work. In the current work, we mainly focus on the continuous time and the infinite particle limit n→ +∞ when qt is a smooth probability density. A natural future direction is to study the guarantees with finite particle size n and discrete time. Another valuable direction is to study whether and how linear convergence can be established for SVGD; see Duncan et al. (2019) for an in-depth discussion of this issue for unconstrained SVGD.
On the practical side, since our methods leverage SVGD and Langevin dynamics, they necessarily inherent their limits, such as the sensitivity on the choice of kernel and step sizes. Another practical issue is that we can only enforce constraints on the training set, and hence generalization error need to be considered if we need to strictly enforce constraints on the testing set. Also, this work focuses on a single constraint for simplicity, which is sufficient for most practical applications since different constraints can be combined into a single constraint easily. Extensions to multiple constraints are straightforward but will be studied separately.
In terms of social impact, our method aims to impose trustworthy constraints in Bayesian sampling, and can help increase the reliability, fairness, and interpretability of ML systems applied to daily life. We will release code freely online to promote the applications of our approach, although it does open the possibility of malicious use of our techniques for adversarial purposes.

heading: Acknowledgements
text: The work is conducted in part in the statistical learning and AI group in computer science at UT Austin, which is supported in part by CAREER-1846421, SenSE-2037267, EAGER-2041327, and Office of Navy Research, and NSF AI Institute for Foundations of Machine Learning (IFML). Xingchao Liu is supported in part by a funding from BP. X. T. Tong’s research is supported by Singapore MOE Academic Research Funds R-146-000-292-114. The authors thank the reviewers for all the suggestions made in the reviewing process.

Review summaries:
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.

The paper presents SVGD variants that can handle constraints on the optimized distribution in the form of an expectation under it. The methods are based on Lagrangian multiplier (primal-dual form) and constraint controlled optimization, respectively. Their convergence analyses in the context of a constraint are developed. Experiment verifies a good fit to the target distribution with the satisfaction of the constraint.

This paper proposes to modify existing approximation posterior inference algorithms to satisfy an additional moment constraint that might be related to trustworthiness, fairness, or something else. Two algorithms are proposed, which can be used with either a discretized Langevin algorithm or Stein variational gradient descent. Convergence properties of the underlying gradient flows are also characterized. 

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.