Sampling in Constrained Domains with Orthogonal-Space Variational Gradient Descent
Anonymous Author(s)
AfﬁliationAddress email
Abstract
Sampling methods, as important inference and learning techniques, are typically designed for unconstrained domains.
However, constraints are ubiquitous in machinelearning problems, such as those on safety, fairness, robustness, and many otherproperties that must be satisﬁed to apply sampling results in real-life applications.
Enforcing these constraints often leads to implicitly-deﬁned manifolds, makingefﬁcient sampling with constraints very challenging.
In this paper, we proposea new variational framework with a designed orthogonal-space gradient ﬂow (OGradient) for sampling on a manifold G 0 deﬁned by general equality constraints.
O-Gradient decomposes the gradient into two parts: one decreases the distanceto G 0 and the other decreases the KL divergence in the orthogonal space.
Whilemost existing manifold sampling methods require initialization on G 0 , O-Gradientdoes not require such prior knowledge.
We prove that O-Gradient converges to thetarget constrained distribution with rate e
O (1 / the number of iterations ) under mildconditions.
Our proof relies on a new Stein characterization of conditional measurewhich could be of independent interest.
We implement O-Gradient through both
Langevin dynamics and Stein variational gradient descent and demonstrate itseffectiveness in various experiments, including Bayesian deep neural networks.
Introduction
Sampling methods, such as Markov chain Monte Carlo (MCMC) [1] and Stein variational gradient descent (SVGD) [22, 21] have been widely used for getting samples from or approximatingintractable distributions in machine learning (ML) problems, such as estimating Bayesian neuralnetwork posteriors [35], generating new images [30] and training energy-based models [14].
Whilebeing powerful, most sampling methods usually can only be used in unconstrained domains or somespecial geometric spaces.
This greatly limits the application of sampling to many real-life tasks.
We consider sampling from a distribution ⇡ with an equality constraint g ( x ) = 0 where g : R d ! R isa general differentiable function.
The domain in this case is the level set G
= { x 2 R d : g ( x ) = 0 }which is a submanifold in R d .
We do not require additional information about G 0 , such as explicitparameterization or known in-domain points, which is in contrast, often demanded by previousmethods [3, 32, 2, 18].
The problem deﬁned above includes many ML applications, such as diseasediagnosis with logic rules constraint, policymaking with fairness constraint for different demographicsubgroups, and autonomous driving with robustness constraint to unseen scenarios.
In this paper, we propose a new variational framework which transforms the above constrainedsampling problem into a constrained functional minimization problem.
A special gradient ﬂow,denoted orthogonal-space gradient ﬂow (O-Gradient), is developed to minimize the objective.
Asillustrated in Figure 1a, the direction of O-Gradient v can be decomposed into two parts: the ﬁrst
[Figure]
part v ] drives the sampler towards the manifold G 0 following r g and keeps it on G 0 once arrived; the 36 second part v ? makes the sampler explore G 0 following the density ⇡ ( x ) .
We prove the convergenceof O-Gradient in the continuous-time mean-ﬁeld limit.
O-Gradient can be applied to both Langevindynamics and SVGD, resulting in O-Langevin and O-SVGD respectively.
As shown in Figure 1b&c,both methods can converge to the target distribution on the manifold.
In particular, O-Langevinconverges following a noisy trajectory while O-SVGD converges smoothly, similar to their standardunconstrained counterparts.
We empirically demonstrate the sampling performance of O-Langevinand O-SVGD across different constrained ML problems.
We summarize our contributions as follows:
• We reformulate the hard-constrained sampling problem into a functional optimizationproblem and derive a special gradient ﬂow, O-Gradient, to obtain the solution.
• We prove that O-Gradient converges to the target constrained distribution with ratee O (1 / the number of iterations ) under mild conditions.
Our proof technique includes a new
Stein characterization of conditional measure which could be of independent interest.•
We implement O-Gradient through both Langevin dynamics and SVGD and demonstrate
its effectiveness in various experiments, including a constrained synthetic distribution,
income classiﬁcation with fairness constraint, loan classiﬁcation with logic rules and image
classiﬁcation with robust Bayesian deep neural networks.
Related Work
Sampling on Explicitly Deﬁned Manifolds Manifolds with special shapes, such as geometricor physics structures, can sometimes be explicitly parameterized in lower dimension spaces.
Forexample, a torus embedded in R 3 can be explicitly deﬁned in two dimensions using polar coordinates.
Variants of classical methods have been developed to sample on such manifolds, including rejectionsampling [6], Langevin dynamics [32], Hamiltonian Monte Carlo (HMC) [3] and Riemannianmanifold HMC [15, 27].
However, explicit parameterization is only applicable to a few special casesand cannot be used for general machine learning problems.
In contrast to this line of work, ourmethod is able to work with more general manifolds deﬁned in the original domain R d .
Sampling on Implicitly Deﬁned Manifolds Many common applications are not endowed withsimple manifolds, such as molecular dynamics [16], matrix factorization [26] and free energycalculations [31].
Motivated by these applications, sampling methods on implicitly deﬁned manifoldshave been developed.
Brubaker et al.
[2] has proposed a family of constrained MCMC methodsby applying Lagrangian mechanics to Hamiltonian dynamics.
Zappa et al.
[34] has introduced aconstrained Metropolis-Hastings (MH) with a reverse projection check to ensure the reversibility.
Later, this method has been extended to HMC [18] and multiple projections [19].
However, theimplementation and analysis of these methods often assume the algorithm starts on the manifold andnever leaves it, requiring prior known points on the manifold and expensive projection subroutines,such as Newton’s method [2, 34, 17–19] or a long time ordinary differential equation (ODE) [36, 28,13].
In contrast, our method works with distributions supported on the ambient space and thus getsrid of the above strong assumptions, leading to a much faster update per iteration.
This makes ourmethod especially suitable for complex ML models such as deep neural networks.
Sampling with a Moment Constraint Recently, sampling with a general moment constraint, suchas E q [ g ]  0 where q is the approximated distribution, has been studied [24].
However, this typeof constraint can not guarantee every sample to satisfy g ( x ) = 0 .
From a technical view, the targetdistribution with a moment constraint is usually not singular w.r.t.
⇡ , so the problem is conceptuallyless challenging compared to the problem considered in this work.
Preliminaries
Variational Framework We review the derivation of Langevin dynamics and SVGD from a uniﬁedvariational framework.
The variational approach frames the problem of sampling into a KL divergenceminimization problem: min q 2 P KL( q || ⇡ ) where P is the space of probability measures.
Suppose 83 at time t , a velocity ﬁeld v t ( x ) is applied to q t to decrease the KL divergence,
[Equation]
where A ⇡ v ( x ) = s ⇡ ( x ) > v ( x )+ r · v ( x ) is the Stein operator, and s p = r log p is the score functionof distribution p .
The optimal v t is obtained by solving an optimization in a Hilbert space H ,
[Equation]
The above objective makes sure that v t decreases the KL divergence as fast as possible.
Langevin Dynamics and SVGD Algorithms Both Langevin dynamics and SVGD can be derivedfrom this variational framework by taking H to be different spaces.
Taking H to be L 2 q , the velocityﬁeld becomes v t ( · ) =
r s ⇡ ( · ) +  q t ( · ) which can be simulated by Langevin dynamics dx t = 90 s ⇡ ( x t )
dt + dW t with W t being a standard Browninan motion.
After discretization, the update stepof Langevin dynamics is x t +1 = x t + Langevin( x t ) , where
Langevin
[Equation]
x [Equation]
=
[Equation]
and ⌘ is a step size.
Taking H to be the reproducing kernel Hilbert space (RKHS) of a continuously differentiable kernelk , the velocity ﬁeld becomes v t ( · ) =
E x ⇠ q t [ k t ( · , x ) s ⇡ ( x ) + r x k t ( · , x )] .
After discretization, the 95 update step of SVGD for particles { x i } ni =1 is x i,t +
= x
i,t + ⌘ · SVGD k ( x i,t ) , for i = 1 , . . . , n , 96 where ⌘ is a step size and
[Equation]
Main Method
In this section, we formulate the constrained sampling problem into a constrained optimizationthrough the variational lens in Section 4.1, and introduce a new gradient ﬂow to solve the problemin Section 4.2.
We apply this general framework to Langevin dynamics and SVGD, leading to twopractical algorithms in Section 4.3.
4.1 Constrained Variational Optimization
Recall that our goal is to draw samples according to the probability of ⇡ , but restricted to a lowdimensional manifold speciﬁed by an equality: G 0 :=
{ x 2 R d : g ( x ) = 0 } .
Similar to the standardvariational framework in Section 3, we can formulate the problem into a constrained optimization inthe space of probability measures: min q 2 P KL( q || ⇡ ) , s.t. q ( g ( x ) = 0)
= 1 .
However, this problemis in general ill-posed, because when q satisﬁes the constraint, both dqd ⇡ and KL( q || ⇡ ) are ill-deﬁned.
Although the problem is ill-posed, we are actually still able to derive a KL -gradient ﬂow to solvethe problem by considering q supported on R d .
The intuition of the derivation is that, in addition tominimizing the objective as in Eq.
(1), the velocity ﬁled v t should also push q towards G 0 to satisfythe constraint.
Surprisingly, the distribution q t following such a gradient ﬂow indeed converges tothe target distribution on the manifold.
We will focus on the derivation of the gradient ﬂow here andleave its rigorous justiﬁcation in Section 5.
4.2 Orthogonal-Space Gradient Flow (O-Gradient)
As mentioned above, besides maximizing the decay of KL( q || ⇡ ) , the velocity ﬁeld v t also needs todrive q towards the manifold satisfying g ( x ) = 0 . We formulate this as solving
[Equation]
where ( x ) is an increasing odd function.
We use ( x ) =
↵ sign ( x ) | x | 1+  with ↵ > 0 and  2 (0 , 1]in this paper.
The objective is the same as Eq.
(2) in the standard variational framework while theconstraint ensures that v t pushes q t towards the manifold and keeps it stay.
It is easy to see that thesolution of the above problem can be decomposed as v t = v ] + v ? where
[Equation]
Note that v ] is parallel to r g and the remaining is to deﬁne v ? , which we now write as
[Equation]
The projection operator D makes sure that v ? ? r g holds for any u , of which the optimal value wecan get by maximizing the unconstrained objective in Eq.
[Equation]
While the solution u always exists, its expression depends on H , which we will discuss in Section 4.3.
Overall, we obtain the velocity ﬁeld v t by ﬁrst formulating a constrained optimization and then transforming it into an unconstrained optimization via orthogonal decomposition.
We call v t OrthogonalSpace Gradient Flow (O-Gradient) and it drives q t to the target distribution without requiring priorknowledge of the manifold G 0 .
4.3 Practical Algorithms
After deriving O-Gradient for general Hilbert spaces H , we explain how to implement it using SVGDand Langevin dynamics.
The resulting O-SVGD and O-Langevin are outlined in Algorithm 1.
At ahigh level, our algorithms keep the original SVGD or Langevin dynamics movement in the directionsperpendicular to r g , while pushing the density towards G 0 along the r g direction.
O-SVGD
We apply O-Gradient to SVGD ﬁrst since it is fairly straightforward.
Recall that v ] canbe obtained using Eq.
(6), we solve Eq.
(8) to get v ? through the following lemma.
Lemma 4.1.
When H is the RKHS, a solution to Eq.
(8) is v ? =
Du ( x ) =
E y ⇠ q t ( k ? ( x, y ) s ⇡ ( y ) +r y k ? ( x, y )) with the orthogonal-space kernel k ?
( x, y ) =
k ( x, y ) D ( x ) D ( y ) .
Then the combined velocity is obtained using the original SVGD with the kernel k ? ,
[Equation]
Numerically, we iteratively update a set of n particles { x i,t } ni =1 ⇢ R d , such that its empirical 140 distribution, P ni =1  ✓ i,t /n is an approximation of q t in a proper sense when step size ⌘ ! 0 and 141 particle size n ! + 1 .
Similar to the update of standard SVGD in Eq.
(4), the update of O-SVGD isx i,t +1 = x
i,t + ⌘ · ( v ] ( x i,t ) + SVGD k ? ( x i,t ))
where
[Equation]
[Equation]
It is worth noting that SVGD k ? is identical to Eq.
4 but with kernel k ? rather than k .
given: Initialization { x i, 0 } ni =1 for O-SVGD and x 0 for O-Langevin, step size ⌘ . loop compute v ] ( x ) =
( g ( x ))
r g ( x ) kr g ( x ) kand the projection operator D ( x ) = I  r g ( x ) r g ( x ) T kr g ( x ) k 2 .
if O-SVGD , update { x i,t } ni =1 by Eq.
if O-Langevin , update x t by Eq.
end loop
O-Langevin
The Langevin implementation requires some additional derivation.
First of all, with
H =
L 2 q , the target velocity ﬁeld is v t ( x ) =
( x )  D ( x ) s q t ( x ) where  ( x ) =
v ] ( x ) + D ( x ) s ⇡ ( x ) .
This leads to a density ﬂow
[Equation]
Next, we try to design a stochastic differential equation (SDE) of which the Fokker–Plank equation(FPE) is identical to Eq.
Following the projection intuition, one natural guess would be [Equation] However, we ﬁnd that the FPE of this SDE does not follow Eq. We then ﬁx it by adding acorrection drift term r and obtain the following theorem.
Theorem 4.2.
Consider a vector ﬁeld r ( x ) = r · D ( x ) , or its component-wise formulation r i ( x ) =
P dj =1 @ x j D i,j ( x ) , where x j denotes the j th dimension of x . Consider the SDE
[Equation]
11) with  ( x ) = v ] ( x ) + D ( x ) s ⇡ ( x ) , then its FPE is identical to Eq.
Moreover, i) the value g ( x t )has deterministic decay ddt g ( x t )
=  ( x t ) ; ii) for any f such that r f ? r g = 0 , the generator ofx t satisﬁes the Langevin equation: L f ( x ) =
r f T ( x ) s ⇡ ( x ) +  f ( x ) .
It is worth pointing out that if g ( x 0 ) = 0 , then g ( x t ) ⌘ 0 , that is, x t always stays on G 0 .
In this case,
Eq.
(11) degenerates to manifold Langevin dynamics studied in previous work [9, 32].
However, our
SDE does not have this requirement since it is still well-deﬁned off G 0 .
This is especially useful fornumerical implementations, leading to a fast algorithm without expensive projection steps.
Similar to the standard Langevin dynamics update in Eq.
(3), the update rule of O-Langevin isx t {++
+}
= x t + ⌘ · v ] ( x t ) + Langevin ?
( x t ) where
[Equation]
Theoretical Analysis
We theoretically justify the convergence of O-Gradient in this section.
To do so, we ﬁrst describethe target measure as a conditioned measure ⇧ 0 , then derive its associated orthogonal-space Fisherdivergence, and ﬁnally prove that O-Gradient converges to ⇧ 0 .
5.1 Conditioned measure and its Stein equation
First of all, conditioning on a zero-measure set is a challenging concept.
Assume we have a distribution⇧ with density ⇡ .
Let A be a set with ⇧ ( A ) 6 = 0 , then ⇧ ( B | A )
=
⇧ ( B \ A ) ⇧ ( A ) .
However, if ⇧ ( A ) = 0 ,this deﬁnition is ill-posed.
Instead, a standard way to deﬁne a conditioned measure on a zero measureset is using disintegration theorem [5].
Let ⇧ z ( · ) =
⇧ ( · | g ( x ) =
z ) be the measure such that
[Equation]
(13) where ⇧ g is the pushforward measure of ⇧ under g .
It is worth noting that ⇧ 0 is not the measurewith Hausdorff density ⇡ on G 0 .
Instead, under regularity conditions, ⇧ 0 can also be deﬁned through
Hausdorff density ⇡ ( x ) / | r g ( x ) | on G 0 , see Lemma 6.4.1 of [29].
One can also avoid this rather abstract deﬁnition, and consider an alternative deﬁnition which is moreintuitive and Bayesian.
Suppose we observe z = g ( x ) + ⇠ , where ⇠ ⇠ N (0 , 1) .
Under prior x ⇠ ⇡ ,the posterior ⇡  2 ,z ( x ) / ⇡ ( x ) exp(  | z  g ( x ) | 2 / 2  2 ) .
Then we can show that ⇧ z is the (weak)limit of ⇡  2 ,z as  ! 0 , moreover, it satisﬁes a Stein characterization of conditional measure .
Weneed some regularity assumptions to state our results.
Deﬁnition 5.1.
A density q is g -regular
if there is a L so that for small enough ⌘ >
[Equation]
Note that Eq.
(14) holds if q g is a Lipschitz function.
Moreover, while Deﬁnition 5.1 requires Eq.
(14)to hold for all z , in practice we only concern ⇧ 0 , thus only need the equation to hold for z near 0 .
Proposition 5.2.
Suppose ⇡ is g -regular.
Then the weak limit of ⇡ ⌘ ,z ( x ) / ⇡ ( x ) exp(  12 {+⌘ ( g ( x )  z ) 2 )as ⌘ ! 
+}
concentrates on G z and satisﬁes Eq.
Moreover, for each z ,
[Equation]
The ﬁrst part of Proposition 5.2 also provides us with a way to sample ⇧ z , as we can sample asequence of distribution ⇡ ⌘ k ,z with ⌘ k ! 0 .
However, this method involves double loops and isusually much more expensive than single-loop algorithms.
The second part of Proposition 5.2 can be used to formulate as a orthogonal-space Stein equation.
Inspeciﬁc, given a q concentrated on G z , we consider checking whether E q A ⇡  is zero for  2 H ?
={  :  ?
r g } , which is a necessary condition for q = ⇧ z .
Notably, this equation is well deﬁned in R d without using any parameterization of G z .
Compared with the standard Stein equation, we add an additional restriction that  2 H ? .
We argue 192 such restriction is actually very natural.
To see that, we note ⇧ z is concentrated on G z , so we should 193 check the Stein equation with vector ﬁeld that “live” only on G z , namely the tangent bundle T G z .
Notably, if  ( x ) 2 T G z , then  ( x ) ? r g ( x ) , therefore it is natural to consider H ? .
5.2 Orthogonal-Space Fisher divergence
Next we consider how to generalize Fisher divergence for constrained sampling.
Since the orthogonalspace Stein equation can only discern discrepancy in H ? , it is natural to consider an orthogonal-space
Fisher divergence
[Equation]
where D is the H ? -projection operator deﬁned pointwisely using Eq.
Its easy to see F ? is welldeﬁned for a density q on R d .
But it remains unclear whether it can be used to measure the distancebetween q and ⇧ z .
To answer this question, we need some form of strong convexity conditions.
Recall that in theunconstrained case, such conditions can be characterized by Poincaré inequalities (PI) when H is L 2 q ,so we will focus on this setting.
PI counterparts of SVGD remain an open question [10, 7], which wehope to be addressed by future works.
When constrained on G z , PI’s formulation involves gradientsso it depends on the Riemannian metric used in G z .
One natural choice of Riemannian metric is themetric inherited from R d which equals k D ( x ) r f ( x ) k , since D ( x ) is also the projection onto thetangent bundle of G z .
This leads to the following deﬁnition.
Deﬁnition 5.3.
⇧ z follows a  -PI, if for any smooth f on R d , var ⇧ z
[ f ]   E ⇧ z [ k D r f k 2 ] .
PI on Riemannian manifold has been studied in previous work [11, 20].
In particular, Theorem 2.10of [11] shows that PI holds for the Riemannian measure if the manifold G z is compact, and hence italso holds for any equivalent distributions.
Proposition 5.4.
Suppose that ⇧ z satisﬁes  -PI for | z |   , and q is supported on { x : | g ( x ) |   } .
Then for any function f such that | f |  1 , the following holds
[Equation]
In particular, Proposition 5.4 shows that if we have a q such that it is 1) supported close to G 0 and 2)
F ? ( q, ⇡ ) is small, then q is close to ⇧ 0 in total variation.
We can also interpret the bound alternativelyas a decomposition of the difference between q and ⇧ 0 : the ﬁrst part measures the difference alongwith H ? directions and it is controlled by F ? ; the second part measures the difference along r g 219 direction, and it is controlled by the distance between q ’s support to G 0 .
As a ﬁnal remark, the orthogonal-space Stein equation (15) and orthogonal-space Fisher divergence(16) can be applied to densities supported on R d , and their formulations do not require informationof G 0 such as parameterization and geodesic.
So they can be used as very computationally friendlystatistical divergence in practice.
Proposition 5.5.
Suppose we apply v t = v ] + Du with (6) and (8) to the density ﬁeld q t .1.
Let M t = max {| g ( x ) | , x 2 supp ( q t ) } .[-Let S t be the solution of an ordinary differentialequation ˙ S t = − ψ ( S t ) . Suppose S 
-]
Suppose M 0 < 1 , then ddt M t =  ( M t ) .2.
Suppose g is also smooth enough so that
[Equation]
Then the KL-divergence follows ddt KL
[Equation]
q
[Equation]
The ﬁrst result shows that the support of q t will shrink towards G 0 .
But in order to avoid possiblestagnation case, e.g. k g k is too small, it is necessary to impose (17).
The second result provides ussome hint on the choice of .
Note that under q t for large enough t , g ( x ) will take values close
to0 , so we need (0) = ˙
= 0 .
Since ( x ) should have the same sign with x , one natural choicewould be ( y ) =
↵ sign ( y ) | y | 1+  .
Theorem 5.6.
Suppose we choose ( y ) =
↵ sign ( y ) | y | 1+  for an ↵ > 0 ,  2 (0 , 1] and theconditions of Proposition 5.5 hold, suppose also that KL( q 0 , ⇡ ) < 1 , then M T
= O ( T  1  ) ,min t  T F ? ( q t , ⇡ ) =
O (log T/T ) .
The orthogonal-space Stein equality holds approximately, in thesense that for any  ( x ) ? r g ( x ) and k  k H = 1 , | E q t A ⇡  |  p
F ? ( q t , ⇡ ) =
O ( p log T/T ) .
Theorem 5.6 shows that we can obtain a q t that is supported on G [   ,  ] =
{ x : | g ( x ) |   } , while it 238 has close to zero orthogonal-space Fisher divergence and the Stein identity holds approximately on
H ? .
Combining it with Proposition 5.4, we have
Corollary 5.7.
Under the conditions of Theorem 5.6, for a bounded f with H =
L 2 q , suppose h f ( z ) =
E ⇡ [ f | g ( X ) =
z ] =
E ⇧ z [ f ] is Lipschitz with near z = 0 , then | E q t [ f ]  E {+⇧ 
+}
[ f ] |  O ( p log T/T ) .
Experiments
We demonstrate the effectiveness of our methods on various tasks, including a constrained syntheticdistribution, an income classiﬁcation with fairness constraint, a loan application with logic constraintand Bayesian deep neural networks with robustness constraint.
For all MCMC methods, we run nparallel chains and collect the ﬁnal samples.
For all SVGD, we use n particles with an RBF kernel.
Synthetic Distribution We ﬁrst demonstrate our methods on a two-dimensional synthetic distribution where the ground truth samples are available.
Let y ⇠ N (0 , I ) , and we transform y tox =
[ x 1 , x 2 ] by x =   1 ( y ) where  ( x ) =
[ x 1 + x 32 , x 2 ] > .
We aim to sample from ⇡ ( x ) withconstraint g ( x ) = x 1 + x 32 = 0 .
We are able to generate ground truth samples by ﬁrst generatingy {+
 ⇠+} N (0 , 1) and then setting x gt =
[  y 30 , y 0 ] .
We report the energy distance between the collected samples and the ground truth samples (measuring how well samples approximate the targetdistribution) and the mean absolute error (MAE) | g ( x ) | (measuring how well the samples satisfy theconstraint).
We set n = 50 for this task.
In order to compare with previous manifold sampling methods, which all require initializations onthe manifold, we ﬁrst consider the case when the sampler starts on the manifold.
We compare withconstrained Langevin dynamics (CLangevin) and constrained Hamiltonian Monte Carlo (CHMC),which are two advanced manifold samplers [18].
From Figure 2a, we observe that O-SVGD and
[Figure]
[Figure] O-Langevin converge after 2000 and 4000 steps respectively whereas CLangevin and CHMC havenot fully converged even after 5000 steps.
These results indicate that projection-based methods couldbe inefﬁcient due to ﬁnding sub-optimal solutions.
In terms of runtime, we observe that O-Langevinconverges the fastest (see Appendix for comparisons).
Figure 2b shows that previous methods havealmost zero MAE all the time since they are required to stay on the manifold.
Our methods havereasonably small MAE (but not zero due to operating in the ambient space) even without expensiveprojection steps.
O-Langevin has more wiggly trajectories than O-SVGD due to Gaussian noise.
Next, we test our methods starting outside the manifold, where previous manifold sampling methodsare not applicable.
From Figure 2d&e, we observe again that O-SVGD converges the fastest in termsof iterations and both methods converge very quickly to the manifold and are able to stay close to it.
Income Classiﬁcation with Fairness Constraint ML Fairness has gained increasing attentionin recent years, which aims to guarantee unbiased treatments for individuals w.r.t.
gender, race,disabilities, etc.
Following previous work [25, 23], we predict whether an individual’s annual incomeis  50 , 000 unfavorably in terms of the gender.
We use Adult Income dataset [12] and train a
Bayesian neural network with two-layer multilayer perceptron (MLP).
During testing, we use CVscore [4] as the measure of disparate impact.
Training and testing performance of our methods andstandard algorithms are shown in Figure 3.
We ﬁnd that during training and testing, both O-Langevinand O-SVGD satisfy the constraint very well while maintaining high training LL and test accuracy.
In contrast, standard Langevin and SVGD violate the constraint signiﬁcantly, indicating that they failto make fair predictions.
These results demonstrate the power of O-Gradient which keeps the sampleron the manifold while still correctly sampling from ⇡ ( x ) .
Loan Classiﬁcation with Logic Rules Following previous work [24], we predict whether to lendloans to applicants and encode two logic rules: (1) an applicant must be rejected if not employedand having the lowest credit rank; (2) an applicant must be approved if employed over 15 yearsand having the highest credit rank.
We use a Bayesian logistic regression model and set g ( x ) =
l logic ( x ) =
E ( µ, ⌫ ) ⇠ D logic [ Loss ( ⌫ , ˆ ⌫ ( µ ; x ))] where D logic is a uniform distribution over all datapoints( µ, ⌫ ) that satisfy the constraint, Loss refers to the classiﬁcation loss and ˆ ⌫ is the prediction made bythe model.
In Figure 4, we see that both of our methods satisfy the logic rules well whereas standardalgorithms violate the constraint severely.
A drop in the accuracy of our methods might be causedby conﬂicts between the data and the logic rules, as indicated in Liu et al.
Nevertheless, ourmethods can obtain models satisfying the constraint, thus guarantee their safety and interpretability.
[Figure] [Table] Prior-Agnostic Bayesian Neural Networks Bayesian neural network (BNN) has been widelyused in deep learning (DL) for quantifying uncertainty.
The posterior of BNNs is determined by theprior which reﬂects our prior knowledge, and the likelihood which quantiﬁes the data ﬁtness.
Due tocomplex NN architectures, specifying an appropriate prior for BNNs has been shown to be difﬁcult [8].
A poor prior often leads to a bad posterior and thus bad inference results.
To automatically control theinﬂuence of the prior, one approach is to sample from the posterior with the constraint of a reasonablyhigh data ﬁtness.
We apply our methods to image classiﬁcation on CIFAR10 with ResNet-18.
Sincethe training loss is typically very small in DL, we set the constraint to be g ( x ) =
Loss ( x ) .
Similarto stochastic Langevin dynamics (SGLD)
[33] and SVGD, our methods are easy to be combinedwith stochastic gradients which are used on this task.
We also ignore the second-order derivativeterms in our methods for speedup.
Besides standard algorithms, we also compare with a tempered
SGLD which has been used in the literature for good performance [35].
We report test error, expectedcalibration error (ECE) and AUROC on the SVHN dataset (measuring out-of-distribution detection).
From Table 2, we see that O-Langevin and O-SVGD improve signiﬁcantly over their unconstrainedcounterparts on all three metrics.
O-Langevin has the lowest test error and the highest AUROC while
O-SVGD has the lowest ECE.
These results indicate that the constraint can automatically limit theeffect of the prior, leading to a better posterior and thus better generalization and calibration.
Conclusion and Limitations
We propose a new variational framework with a designed gradient ﬂow (O-Gradient) for sampling ingeneral, implicitly deﬁned constrained domains.
O-Gradient is formed by two orthogonal directionswhere one drives the sampler towards the domain and the other explores the domain by decreasing a
KL divergence.
We prove the convergence of O-Gradient and apply the framework to both Langevindynamics and SVGD.
We empirically demonstrate the power of our methods on a variety of ML tasks.
While ML achieves impressive performance, we must take realistic constraints into considerationwhen deploying it in daily life.
This work provides a principled framework for solving constrainedsampling problem with theoretical guarantees and practical algorithms.
One possible limitation of our methods is that the obtained samples are not exactly on the manifolddue to working in an ambient space, though can be made arbitrarily close to the manifold by settinghyperapameters.
This might be solved by including a projection step.
Another limitation is that weonly consider an equality constraint.
It will be interesting to incorporating inequality constraintsinto the framework.
Besides, our theory only considers the convergence in the continuous-timelimit.
Future work could provide convergence analysis in discrete time, which may characterize theproperties of practical algorithms more accurately.
[-In other words, it will notmatch (10) unless ( II ) = 0 , which happens when (cid:80) i,j ∂ 
-]
[-( z ) z 1+ β , we find that for some c 0 that depends on M 
-]
[-Moreover, we have that for some constant C 
-]
[-Inpractice, we recommend to set β = 
-]