Title: Constrained Sampling with Primal-Dual Langevin Monte Carlo
Venue: NeurIPS 2024
Accepted: yes
Keywords: Sampling

optimization

constrained optimization

duality
Abstract: This work considers the problem of sampling from a probability distribution known up to a normalization constant while satisfying a set of statistical constraints specified by the expected values of general nonlinear functions. This problem finds applications in, e.g., Bayesian inference, where it can constrain moments to evaluate counterfactual scenarios or enforce desiderata such as prediction fairness. Methods developed to handle support constraints, such as those based on mirror maps, barriers, and penalties, are not suited for this task. This work therefore relies on gradient descent-ascent dynamics in Wasserstein space to put forward a discrete-time primal-dual Langevin Monte Carlo algorithm (PD-LMC) that simultaneously constrains the target distribution and samples from it. We analyze the convergence of PD-LMC under standard assumptions on the target distribution and constraints, namely (strong) convexity and log-Sobolev inequalities. To do so, we bring classical optimization arguments for saddle-point algorithms to the geometry of Wasserstein space. We illustrate the relevance and effectiveness of PD-LMC in several applications.

Constrained Sampling with Primal-Dual Langevin Monte Carlo

Abstract

This work considers the problem of sampling from a probability distribution known up to a normalization constant while satisfying a set of statistical constraints specified by the expected values of general nonlinear functions. This problem finds applications in, e.g., Bayesian inference, where it can constrain moments to evaluate counterfactual scenarios or enforce desiderata such as prediction fairness. Methods developed to handle support constraints, such as those based on mirror maps, barriers, and penalties, are not suited for this task. This work therefore relies on gradient descent-ascent dynamics in Wasserstein space to put forward a discretetime primal-dual Langevin Monte Carlo algorithm (PD-LMC) that simultaneously constrains the target distribution and samples from it. We analyze the convergence of PD-LMC under standard assumptions on the target distribution and constraints, namely (strong) convexity and log-Sobolev inequalities. To do so, we bring classical optimization arguments for saddle-point algorithms to the geometry of Wasserstein space. We illustrate the relevance and effectiveness of PD-LMC in several applications.

1 Introduction

Sampling is a fundamental task in statistics, with applications to estimation and decision making, and of growing interest in machine learning (ML), motivated by the need for uncertainty quantification and its success in generative tasks [1-3]. In these settings, the distribution we wish to sample from ( target distribution ) is often known only up to its normalization constant. This is the case, for instance, of score functions learned from data or posterior distributions of complex Bayesian models. Markov Chain Monte Carlo (MCMC) algorithms can be used to tackle this problem [4, 5] and Langevin Monte Carlo (LMC), in particular, has attracted considerable attention due to its simplicity, theoretical grounding, and effectiveness in practice [3, 6-9]. These sampling algorithms, however, do not naturally incorporate requirements on the samples they generate. Specifically, standard MCMC methods do not enforce restrictions on the target distributions, such as support (e.g., truncated Gaussian), conditional probabilities (e.g., fairness), or moments (e.g., portfolio return) constraints. This limitation is often addressed by post-processing, transforming variables, or by introducing penalties in the target distribution. Though successful in specific settings, these approaches have considerable downsides. Post-processing techniques such as rejection sampling (see, e.g., [10, 11]) may substantially reduce the effective number of samples (number of samples generated per iteration of the algorithm). Variable transformations based on link functions, projections, or mirror/proximal maps (see, e.g., [12-18]) only accommodate (deterministic) support constraints and are not suited

Mohammad Reza Karimi

for statistical requirements such as robustness or fairness [19-22]. Though modifying the target distribution directly offers more flexibility (see, e.g., [23]), it does not guarantee that constraints are satisfied (Table 1). We refer the reader to Appendix A for a more detailed literature review.

This paper overcomes these issues by directly tackling the constrained sampling problem. Explicitly, it seeks to sample not from a target distribution π on R d , but from the distribution µ ⋆ that solves

where P 2 ( R d ) denotes the set of probability measures on R d with bounded second moments and the functions g i , h j represent the requirements. Note that (PI) only considers measures µ against which g i , h j are integrable. Otherwise, the expectations are taken to be + ∞ , making the corresponding measure infeasible. Observe that (PI) is more general than the support-constrained sampling problem considered in, e.g., [12-18]. Indeed, it constrains the distribution µ rather than its samples x (Table 1). Algorithms based on projections, barrier, or mirror maps are not suited for this type of constraints (see Section 2.2 for more details). To tackle (PI), this paper instead derives and analyzes a primal-dual LMC algorithm (PD-LMC) that is the sampling counterpart of gradient descent-ascent (GDA) methods from (Euclidean) optimization. A dual ascent algorithm was previously proposed to tackle (PI), but it requires the exact computation of expectations with respect to intractable distributions [24]. This paper not only overcomes this limitation, but also provides convergence guarantees for a broader class of constraint functions.

The main contributions of this work include:

a discrete-time constrained sampling algorithm (PD-LMC, Algorithm 1) for solving (PI) that precludes any explicit integration (Section 3);

an analysis of PD-LMC proving that it converges sublinearly (in expectation) with respect to the Kullback-Leibler divergence (convex case) or Wasserstein distance (strongly convex case). The analysis is performed directly on the discrete-time iterations and requires only local Lipschitz continuity and bounded variance assumptions (Section 3.1);

an extension of these result for target distributions satisfying a log-Sobolev inequality (LSI) for a variant of PD-LMC (Algorithm 2, Section 3.2);

numerical examples illustrating applications of (PI) and the effectiveness of PD-LMC (Section 4).

2 Problem formulation

2.1 Background on Langevin Monte Carlo

Consider a target distribution π ∈ P 2 ( R d ) absolutely continuous with respect to Lebesgue measure whose density (also denoted π ) can be expressed as π ( x ) = e -f ( x ) /Z for some normalization constant Z . Define the Kullback-Leibler (KL) divergence of µ with respect to π as

where dµ dπ is the Radon-Nikodym derivative, V is the potential energy , and H is the negative entropy if µ is absolutely continuous with respect to π ; and + ∞ otherwise. For a wide class of functions f (e.g.,

smooth and strongly convex), samples from π can be obtained from the path of the Langevin diffusion process, whose instantaneous values x ( t ) have distributions µ ( t ) evolving according to the Fokker-Planck equation [25]. Explicitly,

for a d -dimensional Brownian motion W ( t ) , where ∇ · q denotes the divergence of q and ∇ W 2 KL( µ ∥ π ) denotes the Wasserstein-2 gradient of KL( ·| π ) at µ [26, Theorem 10.4.17] (see Appendix B for more details). Indeed, the Langevin diffusion (2) brings the distribution µ ( t ) of x ( t ) progressively closer to the target π . In fact, the Fokker-Planck equation can be interpreted as a gradient flow of the KL divergence with respect to the Wasserstein-2 distance [7, 25].

However, computing the path of the stochastic differential equation in (2) is not practical and discretizations are used instead. Chief among them is the (forward) Euler-Maruyama scheme, which leads to the celebrated Langevin Monte Carlo (LMC) algorithm [6]

for a step size γ k > 0 , where I d denotes the d -dimensional identity matrix. Notice that it is not necessary to know Z in order to evaluate (3). This has made LMC and its variants widely popular in practice and the subject of extensive research. Despite (3) being a biased time-discretizations of the Langevin diffusion in (2) [7], rates of convergence of LMC have been obtained for smooth and strongly convex [8, 27] or convex [28] potentials or when the target distribution π verifies an LSI [29].

2.2 Constrained sampling

Our goal, however, is not to sample from π itself, but from a distribution close to π that satisfies a set of statistical requirements. Explicitly, we wish to sample from a distribution µ ⋆ that solves (PI). Since (PI) constrains the distribution µ rather than its samples x , it can accommodate more general requirements than the support constraints typically considered in constrained sampling (e.g., [12-18]). Next, we illustrate the wide range of practical problems that can be formulated as (PI). These examples are further explored in Section 4 and more details on their formulations are provided in Appendix E.

Sampling from convex sets : though we have stressed that (PI) accommodates other types of requirements, it can also be used to constrain the support of π , i.e., to sample from

for a closed convex set C ⊂ R d . Indeed, let C be the intersection of the 0 -sublevel sets of convex functions { s i } i =1 ,...,I . Such a description always exists (see Appendix E). Then, (PII) can be cast as (PI) using g i ( x ) = [ s i ( x )] + for [ z ] + = max(0 , z ) . Notice that the g i are convex and that although they are not everywhere differentiable, I ( s i ( x ) > 0) ∇ s i ( x ) is a subgradient of g i , where I ( E ) = 1 on the event E and 0 otherwise. Observe that support constraints can also be imposed using projections, mirror/proximal maps, and barriers as in [12-18]. These methods, however, constrain the samples x rather than their distribution µ as in (PII).

Rate-constrained Bayesian models : rate constraints have garnered attention in ML due to their central role in fairness [20, 21]. Consider data pairs ( x, y ) , where x ∈ X are features and y ∈ { 0 , 1 } labels, and a protected (measurable) subgroup G ⊂ X . Let π be a Bayesian posterior of the parameters θ of a model q ( · ; θ ) denoting the probability of a positive outcome (based, e.g., on a binomial model). We wish to enforce statistical parity, i.e., we wish the prevalence of positive outcomes within the protected group G to be close to or higher than in the whole population. We cast this problem as (PI) by constraining the average probability of positive outcome as in

where δ > 0 denotes our tolerance [22]. Naturally, multiple protected groups can be accommodated by incorporating additional constraints. Hence, constrained sampling provides a natural way to encode fairness in Bayesian inference.

Counterfactual sampling : rather than imposing requirements on probabilistic models, constrained sampling can also be used to probe them by evaluating counterfactual statements. Indeed, let π denote a reference probabilistic model such that sampling from π yields realizations of the 'real world.' Consider the counterfactual statement 'how would the world have been if E [ g ( x )] ≤ 0 ?' Constrained sampling not only gives realizations of this alternative world, but it also indicates its 'compatibility' with the reference model, namely the value P ⋆ of (PI).

More concretely, consider a Bayesian stock market model. Here, π is a posterior model for the (log-)returns of I assets, e.g., distributed as Gaussians N ( ρ, Σ) . Here, the vector ρ describes the mean return of each stock and Σ their covariance. We can investigate what the market would look like if, e.g., the mean and variance of each stocks were to change by solving

Due to correlations in the market, certain choices of ¯ ρ i or ¯ σ 2 i may be more 'unrealistic' than others. Additionally, it could be that some of these conditions are vacuous conditioned on the others. As we show next, our approach to tackling (PI) effectively isolates the contribution of each requirement in the solution µ ⋆ , thus enabling us to identify which are (conditionally) vacuous and which are most at odds with the reference model π .

2.3 Lagrangian duality and dual ascent algorithms

Although directly sampling from µ ⋆ does not appear straightforward, it admits a convenient characterization based on convex duality that is amenable to be sampled using the LMC algorithm (3). Indeed, let g : R d → R I and h : R d → R J be vector-valued functions collecting the constraint functions g i and h j respectively. The Lagrangian of (PI) is then defined as

for λ ∈ R I + and ν ∈ R J , where

and a normalization constant Z λν . Notice that P ⋆ = min µ max λ ≥ 0 , ν L ( µ, λ, ν ) , which is why (PI) is referred to as the primal problem .

To obtain the dual problem of (PI), define the dual function

Notice from (4) that the minimum in (6) is achieved for µ λν from (5), the Lagrangian minimizer , so that d ( λ, ν ) = log( Z/Z λν ) . The solution of (6) is therefore a tilted version of π , whose tilt is controlled by the dual variables ( λ, ν ) . Since (6) is a relaxation of (PI), it yields a lower bound on the primal value, i.e., d ( λ, ν ) ≤ P ⋆ for all ( λ, ν ) ∈ R I + × R J . The dual problem seeks the tilts ( λ ⋆ , ν ⋆ ) that yield the best lower bound, i.e.,

The set Φ ⋆ = argmax λ ≥ 0 , ν d ( λ, ν ) of solutions of (DI) is called the set of Lagrange multipliers . Note from (6) that (DI) depends on the distributions µ and π through its objective d .

The dual problem (DI) has several advantageous properties. Indeed, while the primal problem (PI) is an infinite dimensional, smooth optimization problem in probability space, the dual problem (DI) is a finite dimensional, non-smooth optimization problem in Euclidean space. What is more, it is a concave problem regardless of the functions f, g, h , since the dual function (6) is the minimum of a set of affine functions in ( λ, ν ) [30, Prop. 4.1.1]. These properties are all the more attractive given that, under mild conditions stated below, (DI) can be used to solve (PI).

Assumption 2.1. There exists µ † ∈ P 2 ( R d ) with KL( µ † ∥ π ) ≤ C < ∞ such that E µ † [ g i ] ≤ -δ < 0 and E µ † [ h j ] = 0 for all i, j .

Proposition 2.2. Under Assumption 2.1, the following holds:

(i) P ⋆ = D ⋆ ;

(ii) there exists a finite pair ( λ ⋆ , ν ⋆ ) ∈ Φ ⋆ ;

(iii) for any solution µ ⋆ of (PI) and ( λ ⋆ , ν ⋆ ) of (DI) , it holds that

(iv) the solution of (PI) is µ ⋆ = µ λ ⋆ ν ⋆ for ( λ ⋆ , ν ⋆ ) ∈ Φ ⋆ ;

(v) consider the perturbation of (PI)

Then, ( λ ⋆ , ν ⋆ ) are subgradients of P ⋆ (0 , 0) = P ⋆ , i.e., P ⋆ ( u, v ) ≥ P ⋆ -λ ⋆ ⊤ u -ν ⋆ ⊤ v , and if P ⋆ ( u, v ) is differentiable at (0 , 0) , then ∇ u P ⋆ ( u, v ) = -λ ⋆ and ∇ v P ⋆ ( u, v ) = -ν ⋆ at (0 , 0) .

Proof. In finite dimensional settings, (i)-(v) are well-known duality results (see, e.g., [30]). While they also hold for infinite dimensional optimization problems, their proofs are slightly more 'scattered.' We collect their reference below. The objective of (PI) is a convex function and its constraints are linear functions of µ . Hence, (PI) is a convex program. Under Slater's condition (Assumption 2.1), it is (i) strongly dual ( P ⋆ = D ⋆ ) and (ii) there exists at least one solution ( λ ⋆ , ν ⋆ ) of (DI) (see [31, Sec. 8.6, Thm. 1] or [32, Cor. 4.1]). This implies (iii) the existence of the saddle-point (7) [33, Prop. 2.156], (iv) that µ ⋆ ∈ argmin µ L ( µ, λ ⋆ , ν ⋆ ) = { µ λ ⋆ ν ⋆ } , since the KL divergence is strongly convex and its minimizer is unique [34, Thm. 7.3.7], and (v) that ( λ ⋆ , ν ⋆ ) are subgradients of the perturbation function P ⋆ ( u, v ) [33, Prop. 4.27]. ■

Proposition 2.2 shows that given solutions ( λ ⋆ , ν ⋆ ) of (DI), the constrained sampling problem (PI) reduces to sampling from µ λ ⋆ ν ⋆ ∝ e -U ( · ,λ ⋆ ,ν ⋆ ) (see Appendix F for an explicit example of this result). It is important to note that this results only relies on the KL divergence being (strongly) convex in the standard L 2 geometry, i.e., along mixtures of the form tµ 0 +(1 -t ) µ 1 for t ∈ [0 , 1] . This does not imply that it is (geodesically) convex in the Wasserstein sense [35, Section 9.1.2]. This would require U in (5) to be convex for all λ ≥ 0 and ν ∈ R J , i.e., for f, g to be convex and h to be linear.

Hence, Proposition 2.2 reduces the constrained sampling problem (PI) to that of finding the Lagrange multipliers ( λ ⋆ , ν ⋆ ) . Despite their finite dimensionality, however, computing these parameters is intricate. Indeed, since (DI) is a concave program, we could obtain ( λ ⋆ , ν ⋆ ) ∈ Φ ⋆ using

for η k > 0 , where we used the fact that E µ λν [ g ] and E µ λν [ h ] are (sub)gradients of the dual function (6) at ( λ, ν ) [36, Thm. 2.87]. This procedure is known in optimization and game theory as dual ascent or best response [37, 38]. Notice, however, that (8) is not a practical algorithm as it requires explicit integration with respect to the intractable distribution µ λν from (5).

This issue was partially addressed in [24] (in continuous time and without equality constraints, i.e., J = 0 ) by replacing the Lagrangian minimizer µ λ k ν k in (8) by the distribution of LMC samples, as in

Note that since J = 0 , we omit the argument ν of U for clarity. Nevertheless, the updates in (9) still require an explicit integration. While it is now possible to sample from µ k (namely, using the x k ), empirical approximations of E µ k [ g ] may not only require an exponential (in the dimension d ) number of samples (e.g., [39, Thm. 1.2]), but it introduces errors that are not taken into account in the analysis of [24]. In the sequel, we address these drawbacks by replacing these dual ascent algorithms by a saddle-point one.

Algorithm 1 Primal-dual LMC

1: Inputs: η k > 0 (step size), x 0 ∼ µ 0 , and ( λ 0 , ν 0 ) = (0 , 0) .

2: for k = 0 , . . . , K -1

3: x k +1 = x k -η k ∇ x U ( x k , λ k , ν k ) + √ 2 η k β k , for β k ∼ N (0 , I d )

4: λ k +1 = [ λ k + η k g ( x k ) ] +

6: end

5: ν k +1 = ν k + η k h ( x k )

3 Primal-dual Langevin Monte Carlo

Consider the GDA dynamics for the saddle-point problem (DI) in Wasserstein space. Explicitly,

for the Lagrangian L defined in (4), where [ z ] λ, + = z for λ > 0 and [ z ] λ, + = max( a, 0) otherwise [40, Sec. 2.2]. Observe that ∇ λ i L ( µ, λ, ν ) = E µ [ g i ] and ∇ ν j L ( µ, λ, ν ) = E µ [ h j ] . Hence, the algorithm from [24] described in (9) involves a deterministic implementation of (10b) that fully integrates over µ ( t ) . In contrast, we consider a stochastic , single-particle implementation of (10) that leads to the practical procedure in Algorithm 1.

Explicitly, we also use an Euler-Maruyama time-discretization of the Langevin diffusion associated to (10a) (step 3), but replace the expectations in (10b)-(10c) by single-sample approximations (steps 45). Algorithm 1 can therefore be seen as a particle implementation of the deterministic Wasserstein GDA algorithm (10). As such, it resembles a primal-dual counterpart of the LMC algorithm in (3), which is why we dub it primal-dual LMC (PD-LMC). Alternatively, Algorithm 1 can be interpreted as a stochastic approximation of the dual ascent method in (9). This suggests that the gradient approximations in steps 4-5 could be improved using mini-batches, which is in fact how [24] approximates the expectation in (9). Our theoretical analysis and experiments show that these mini-batches are neither necessary nor always worth the additional computational cost (see Section 3.1 and Section 4). Note that the 'stochastic approximations' in Algorithm 1 refer to the dual updates (steps 4-5) rather than the LMC update (step 3) as in stochastic gradient Langevin [41]. Though these methods could be combined, it is beyond the scope of this work.

The remainder of this section is dedicated to analyzing the convergence properties of PD-LMC for both stochastic dual gradients (as in Algorithm 1) and exact dual gradient (as in (9)). For the latter, we obtain guarantees for the discrete implementation (9) under weaker assumptions than the continuous-time analysis of [24]. We consider strongly log-concave target distributions in Section 3.1 and those satisfying an LSI in Section 3.2.

3.1 PD-LMC with (strongly) convex potentials

As opposed to the traditional LMC algorithm (3) or the deterministic updates in (9), Algorithm 1 involves three coupled random variables, namely, ( x k , λ k , ν k ) . Hence, the LMC update (step 3) is based on a stochastic potential U and the distribution µ k of x k is now a random measure . Our analysis sidesteps this obstacle by using techniques from stochastic optimization. We also leverage techniques from primal-dual algorithms in the Wasserstein space, in the spirit of works such as [7, 8, 42] that studied the LMC (3) or alternative time-discretizations of gradient flows of the KL divergence as splitting schemes.

First, define the potential energy E and the (negative) entropy H for ( µ, λ, ν ) ∈ P 2 ( R d ) × R I + × R J as

for U as in (5). Notice from (4) that L ( µ, λ, ν ) = E ( µ, λ, ν ) + H ( µ ) -log( Z ) , where we used the KL divergence decomposition in (1). To proceed, consider the following assumptions:

Assumption 3.1. The potential energy E ( µ, λ, ν ) in (11) is m -strongly convex with respect to µ along Wasserstein-2 geodesics for m ≥ 0 and all ( λ, ν ) ∈ R I + × R J . Explicitly,

where s is an optimal coupling achieving W 2 2 ( µ, µ 0 ) (see Appendix B).

Assumption 3.2. The gradients and variances of f, g, h are bounded along iterations { µ k } k ≥ 0 , where µ k is the distribution of x k , i.e., there exists G 2 such that

Assumption 3.1 holds with m = 0 if f, g are convex and h is linear. If f is additionally strongly convex, then it holds with m > 0 [26, Prop. 9.3.2]. Assumption 3.2 is typical in (stochastic) nonsmooth optimization analyses (see, e.g., [36, 43, 44]). Notice, however, that gradients are only required to be bounded along trajectories of Algorithm 1, a crucial distinction in the case of strongly convex functions whose gradients can only be bounded locally. Assumption 3.2 can be satisfied under mild conditions on f, g, h , such as local Lipschitz continuity or linear growth.

The following theorem provides the first convergence analysis of the discrete-time PD-LMC.

Theorem 3.3. Denote by µ k the distribution of x k in Algorithm 1. Under Assumptions 2.1, 3.1, and 3.2, there exists R 2 0 such that, for η k ≤ η ,

For η k ≤ R 0 / ( G √ k ) and ¯ η k = η k / ∑ K k =1 η k , we obtain

Additionally, there exists a sequence of step sizes η k > 0 such that W 2 2 ( µ k , µ ⋆ ) ≤ R 2 0 and E [ ∥ λ k ∥ 2 ] + E [ ∥ ν k ∥ 2 ] < ∞ for all k . The same results hold (without expectations) when using exact dual gradients, i.e., if the updates in steps 4-5 are replaced by λ k +1 = [ λ k + η k E µ k [ g ] ] + and ν k +1 = ν k + η k E µ k [ h ] .

Theorem 3.3, whose proof is deferred to Appendix C, implies rates similar to those for GDA schemes in finite-dimensional Euclidean optimization (see, e.g., [43]). To recover those rates, however, we must bound the magnitudes of λ k , ν k . In [43], this is done by bounding the iterates in the algorithm itself, i.e., by projecting them onto the set D r = { ( λ, ν ) ∈ R I + × R J | max( ∥ λ ∥ 2 , ∥ ν ∥ 2 ) ≤ r } and choosing r such that Φ ⋆ ⊆ D r (Proposition 2.2(ii) ensures this is possible). We then incur a bias on the order of η in (12) that vanishes in the decreasing step size setting of (13). Though convenient, this is not necessary since there exists a sequence of step sizes such that both E [ ∥ λ k ∥ 2 ] and E [ ∥ ν k ∥ 2 ] are bounded for all k ≥ 0 . In the interest of generality, Theorem 3.3 holds without these hypotheses. It is worth noting that though faster rates and last iterates guarantees can be obtained for Euclidean saddle-point problems, they rely on more complex schemes than the GDA in Algorithm 1 involving acceleration or proximal methods [45-48].

The results in Theorem 3.3 are stated for the stochastic scheme in Algorithm 1. However, Theorem 3.3 yields the same rates (without expectations) for exact dual gradients, i.e., for the dual ascent scheme (9). In this case, the second condition in Assumption 3.2 simplifies to max k ( ∥ E µ k [ g ] ∥ 2 , ∥ E µ k [ h ] ∥ 2 ) ≤ G 2 . Not only are these milder assumptions than [24, Eq. (16)], but the guarantees hold for discrete- rather than continuous-time dynamics. Finally, (12)-(13) imply convergence with respect to the KL divergence for convex potentials ( m = 0 ) with stronger guarantees in Wasserstein metric for strongly convex ones ( m> 0 ).

The convergence rates for distributions µ k from Theorem 3.3 also imply convergence rates for empirical averages across iterates x k of Algorithm 1. This corollary is obtained by combining (12)(13) with the following proposition. By taking φ to be the constraint functions g or h from (PI) yields feasibility guarantees for PD-LMC .

Algorithm 2 (Stochastic) dual LMC

1: Inputs: N 0 k > 0 (burn-in), γ k , η k > 0 (step sizes), and λ 0 = 0 . 2: for k = 0 , . . . , K -1 3: x 0 ∼ µ 0 4: x n +1 = x n -γ k ∇ x U ( x n , λ k ) + √ 2 γ k β n , β n ∼ N (0 , I d ) , for n = 0 , . . . , N 0 k -1 5: λ k +1 = [ λ k + η k g ( x N 0 k ) ] + 6: end

Proposition 3.4. Consider samples x k distributed according to µ k and c k ≥ 0 with ∑ K k =1 c k = 1 such that ∑ K k =1 c k [ KL( µ k ∥ µ ⋆ ) + m 2 W 2 2 ( µ k , µ ⋆ ) ] ≤ ∆ K for m ≥ 0 . Then, it holds that

3.2 PD-LMC with LSI potentials

In this section, we replace Assumption 3.1 on the convexity of the potential by an LSI common in the sampling literature. We consider only inequality constraints ( J = 0 ) here and omit the function arguments ν , since accounting for equality constraints requires significant additional assumptions.

Assumption 3.5. The distribution µ λ satisfies the LSI for bounded λ , i.e., there exists σ > 0 such that 2 σ KL( ζ ∥ µ λ ) ≤ ∥∇ log ( dζ/dµ λ ) ∥ 2 L 2 ( ζ ) for all ζ ∈ P 2 ( R d ) .

The LSI in Assumption 3.5 is often used in the analysis of the standard LMC algorithm [29, 49]. It holds, e.g., when f is strongly convex and g is a (possibly non-convex) bounded function due to the Holley-Stroock perturbation theorem [50]. In fact, if f is 1 -strongly convex and | g | is bounded by 1 , then Assumption 3.5 holds for σ ≥ e -2 λ (see, e.g., [51, Prop. 5.1.6] or [52, Thm 1.1]). The LSI is akin to the Polyak-Łojasiewicz (PL) condition from Euclidean optimization [53], which supposes issues with GDA methods such as Algorithm 1. Indeed, it is not enough for the Lagrangian (4) to satisfy the PL condition in the primal variable to guarantee the convergence of GDA in Euclidean spaces. We must either modify Algorithm 1 using acceleration or proximal methods [47, 54-56] or impose the PL condition also on λ [57, 58]. Since the Lagrangian 4 is linear in λ , it is clear that Algorithm 1 will not suffice to provide theoretical guarantees in the LSI case.

We therefore consider the variant in Algorithm 2, where N 0 k LMC iterations (step 3) are executed before updating the dual variables (step 4). This is akin to using different time-scales in continuoustime, a common technique for solving saddle-point problems [54, 58]. Since it resembles a dual ascent counterpart of the LMC algorithm (3), we refer to it as (stochastic) dual LMC (DLMC). As opposed to the dual ascent algorithm from [24] in (9), however, Algorithm 2 does not require any explicit evaluation of expected values. The following theorem provides an analysis of its convergence.

Theorem 3.6. Assume that the functions f, g are M -smooth, i.e., have M -Lipschitz continuous gradients, satisfy Assumption 3.5, and that E µ [ ∥ g ∥ 2 ] ≤ G 2 for all µ ∈ P 2 ( R d ) . Let 0 < η k ≤ η , 0 < ϵ ≤ ηG 2 < 1 ,

Under Assumption 2.1, there exists B < ∞ such that the distributions { µ k } of the samples { x N 0 k } generated by Algorithm 2 satisfy

Recall from (PI) that I is the number of inequality constraints. Additionally, E [ ∥ λ k ∥ 1 ] is bounded for all k .

1

.

0

1

.

5

2

.

0

2

.

5

3

.

0

0

1

2

Density

PD-LMC

1

.

0

1

.

5

2

.

0

2

.

5

3

.

0

Proj. LMC

1

.

0

1

.

5

2

.

0

2

.

5

3

.

0

Mirror LMC

PD-LMC

Proj. LMC

Mirror LMC

Rejection

Theorem 3.6, whose proof is deferred to Appendix D, provides similar guarantees as (approximate) subgradient methods in finite-dimensional optimization (see, e.g., [22, 56]). This is not surprising seen as γ k , N 0 k in Theorem 3.6 are chosen to ensure that step 4 yields a sample x k ∼ ¯ µ k such that KL(¯ µ k ∥ µ λ k ) ≤ ϵ using [29, Theorem 1]. At this point, g ( x N 0 k ) in step 5 is an approximate, stochastic subgradient of the dual function (6). Though it may appear from (12) and (14) that Algorithms 1 and 2 have the same convergence rates, an informal computation shows that the latter evaluates on the order of dκ 2 /η as many gradient per iteration, where κ = M/σ . Note that we can once again apply Theorem 3.6 to derive ergodic average and feasibility guarantees for Algorithm 2.

4 Experiments

We now return to the applications described in Section 2.2 to showcase the behavior of PD-LMC. We defer implementation details and additional results to Appendix E. Code for these examples is publicly available at https://www.github.com/lfochamon/pdlmc .

Sampling from convex sets. Wecast the problem of sampling from a Gaussian distribution N (0 , 1) truncated to C = [1 , 3] as (PI) by taking f ( x ) = x 2 / 2 and g ( x ) = [( x -1)( x -3)] + (see Section 2.2). Fig. 1 shows histograms for the samples obtained using PD-LMC, the projected LMC (Proj. LMC) from [13], and the mirror LMC from [59], all with the same step size. Both Proj. LMC and Mirror LMC generate an excess of samples close to the boundary (between 1 . 5 and 3 times more samples than expected). This leads to an underestimation of the mean (Proj. LMC: 1 . 488 / Mirror LMC: 1 . 470 vs. true mean: 1 . 510 ). In contrast, PD-LMC provides a more accurate estimate ( 1 . 508 ). Yet, since it constrains the distribution µ rather than its samples, it is not an interior-point method and can produce samples outside of C . Theorems 3.3-3.6 show that this becomes less frequent as the algorithm progresses (in Fig. 1, only 2% of the samples are not in C ). This occurs even without using minibatches in steps 4-5 of Algorithm 1 as in [24]. In fact, our experiments show that mini-batches increase the computational complexity with no performance benefit (Appendix E). These issues are exacerbated in more challenging problems, such as sampling from a two-dimensional standard Gaussian centered at [2 , 2] restricted to an unit ℓ 2 -norm ball (Fig. 2). In this case, Proj. LMC places almost 25% of its samples on the boundary (where only 0 . 14% of samples should be), while PD-LMC only places 1 . 8% of its samples outside of the support. Mirror LMC provides a better mean estimation in this setting, although a bit more asymmetric than PD-LMC [Mirror LMC: (0 . 312 , 0 . 418) vs. PD-LMC: (0 . 446 , 0 . 444) vs. true mean: (0 . 368 , 0 . 368) ].

Rate-constrained Bayesian models. Here, we consider π to be the posterior of a Bayesian logistic regression model for the Adult dataset from [60], where the goal is to predict whether an individual

0

10

20

30

Prevalence of

>

$50k (%)

Constrained

(

δ

= 0

.

03

)

Constrained

(

δ

= 0

.

01

)

Unconstrained

Female

Male

0

1

2

3

4

5

6

Iteration

10

5

-

10

-

5

0

5

Dual variable (

ν

)

TSLA

NVDA

JNJ

AAPL

GOOG

BRK-B

LLY

×

makes more than $50 k based on socioeconomic information (details on data pre-processing can be found in [61]). We consider a standard Gaussian prior on the parameters θ ∈ R d +1 of the model, where d is the number of features. Using the LMC algorithm to sample from the posterior (i.e., no constraints), we find that while the average probability of positive predictions is 19 . 1% over the whole test set, it is 26 . 2% among males and 5% among females ('Unconstrained' in Fig. 3). To overcome this disparity, we take gender to be the protected class in (PIII), constraining both G male and G female with δ = 0 . 01 . Using PD-LMC, we obtain a Bayesian model that leads to an average probability of positive outcomes of 18 . 1% and 15 . 1% for males and females respectively. In fact, we now observe a substantial overlap of the distributions of positive predictions across genders for the constrained posterior µ ⋆ ('Constrained ( δ = 0 . 01 )' in Fig. 3). This substantial reduction of prediction disparities comes at only a minor decline in accuracy (unconstrained: 84% vs constrained: 82% ).

3. Counterfactual sampling. Though the distribution of positive predictions changes considerably for both male and female individuals, the final dual variables ( λ male = 0 and λ female ≈ 160 ) show that these changes are due uniquely to the female group [as per Prop. 2.2(iv)]. This implies that the reference model π is itself compatible with the requirement for the male group, but that reducing the disparity for females requires considerable deviations from it. By examining λ female, we conclude without recalculating µ ⋆ that even small changes in the tolerance δ for the female constraint would substantially change the distribution of outcomes [Prop. 2.2(v)]. This is confirmed by 'Constrained ( δ = 0 . 03 )' in Fig. 3. Notice that this is only possible due to the primal-dual nature of PD-LMC. This type of counterfactual analysis is even more beneficial in the presence of multiple requirements. Indeed, let π be the posterior of a Bayesian model for the daily (log-)return of a set of assets (see Appendix E for more details). Using (PIV), we consider how the market would look like if the average (log-)return of each asset were to have been (exactly) 20% higher. Inspecting the dual variables (Fig. 4), we notice that this increased market return is essentially driven by two stocks: NVDA and LLY ( ν < 0 ). In fact, the reference model π would be consistent with an even higher increase for JNJ and GOOG ( ν > 0 ). We confirm these observations by constraining only NVDA and LLY , which yields essentially the same (log-)return distribution for all assets.

5 Conclusion

We tackled the problem of sampling from a target distribution while satisfying a set of statistical constraints. Based on a GDA method in Wasserstein space, we put forward a fully stochastic, discretetime primal-dual LMC algorithm (PD-LMC) that precludes any explicit integration in its updates. We analyze the behavior of PD-LMC for (strongly) convex and log-Sobolev potentials, proving that the distribution of its samples converges to the optimal constrained distribution. We illustrated the use of PD-LMC for different constrained sampling applications. Future work include strengthening the convergence results to almost sure guarantees and improving the rates obtained using proximal and extra gradient methods, particularly in the LSI setting.

Acknowledgments and Disclosure of Funding

The work of L.F.O. Chamon is funded by the Deutsche Forschungsgemeinschaft (DFG, German Research Foundation) under Germany's Excellence Strategy (EXC 2075-390740016).

References

M. F. Faulkner and S. Livingstone, 'Sampling algorithms in statistical physics: a guide for statistics and machine learning,' arXiv preprint arXiv:2208.04751 , 2022.

R. van de Schoot, S. Depaoli, R. King, B. Kramer, K. Märtens, M. G. Tadesse, M. Vannucci, A. Gelman, D. Veen, J. Willemsen et al. , 'Bayesian statistics and modelling,' Nature Reviews Methods Primers , vol. 1, no. 1, p. 1, 2021.

Y. Song and S. Ermon, 'Generative modeling by estimating gradients of the data distribution,' Advances in neural information processing systems , vol. 32, 2019.

G. O. Roberts and J. S. Rosenthal, 'General state space Markov chains and MCMC algorithms,' Probability Surveys , vol. 1, 2004.

C. Robert and G. Casella, Monte Carlo statistical methods . Springer Verlag, 2004.

G. O. Roberts and R. L. Tweedie, 'Exponential convergence of Langevin distributions and their discrete approximations,' Bernoulli , pp. 341-363, 1996.

A. Wibisono, 'Sampling as optimization in the space of measures: The Langevin dynamics as a composite optimization problem,' in Conference on Learning Theory . PMLR, 2018, pp. 2093-3027.

A. Durmus, S. Majewski, and B. Miasojedow, 'Analysis of Langevin Monte Carlo via convex optimization,' The Journal of Machine Learning Research , vol. 20, no. 1, pp. 2666-2711, 2019.

Y. Wang and W. Li, 'Accelerated information gradient flow,' J. Sci. Comput. , vol. 90, no. 1, 2022.

L. Lang, W.-s. Chen, B. R. Bakshi, P. K. Goel, and S. Ungarala, 'Bayesian estimation via sequential Monte Carlo sampling-Constrained dynamic systems,' Automatica , vol. 43, no. 9, pp. 1615-1622, 2007.

Y. Li and S. K. Ghosh, 'Efficient sampling methods for truncated multivariate normal and student-t distributions subject to linear inequality constraints,' Journal of Statistical Theory and Practice , vol. 9, pp. 712-732, 2015.

Y.-P. Hsieh, A. Kavis, P. Rolland, and V. Cevher, 'Mirrored Langevin dynamics,' Advances in Neural Information Processing Systems , vol. 31, 2018.

S. Bubeck, R. Eldan, and J. Lehec, 'Sampling from a log-concave distribution with projected Langevin Monte Carlo,' Discrete & Computational Geometry , vol. 59, no. 4, pp. 757-783, 2018.

A. Salim and P. Richtarik, 'Primal dual interpretation of the proximal stochastic gradient Langevin algorithm,' in Advances in Neural Information Processing Systems , 2020, pp. 37863796.

K. Ahn and S. Chewi, 'Efficient constrained sampling via the mirror-Langevin algorithm,' Advances in Neural Information Processing Systems , vol. 34, pp. 28 405-28 418, 2021.

Y. Kook, Y. Lee, R. Shen, and S. Vempala, 'Sampling with Riemannian Hamiltonian Monte Carlo in a constrained space,' in Advances in Neural Information Processing Systems , A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho, Eds., 2022.

L. Sharrock, L. Mackey, and C. Nemeth, 'Learning rate free Bayesian inference in constrained domains,' in Conference on Neural Information Processing Systems , 2023.

M. Noble, V. De Bortoli, and A. Durmus, 'Unbiased constrained sampling with self-concordant barrier Hamiltonian Monte Carlo,' in Thirty-seventh Conference on Neural Information Processing Systems , 2023.

A. M ˛ adry, A. Makelov, L. Schmidt, D. Tsipras, and A. Vladu, 'Towards deep learning models resistant to adversarial attacks,' in International Conference on Learning Representations , 2018.

M. Kearns, S. Neel, A. Roth, and Z. S. Wu, 'Preventing fairness Gerrymandering: Auditing and learning for subgroup fairness,' in International Conference on Machine Learning , 2018, pp. 2564-2572.

A. Cotter, H. Jiang, M. Gupta, S. Wang, T. Narayan, S. You, and K. Sridharan, 'Optimization with non-differentiable constraints with applications to fairness, recall, churn, and other goals,' Journal of Machine Learning Research , vol. 20, no. 172, pp. 1-59, 2019.

L. F. O. Chamon, S. Paternain, M. Calvo-Fullana, and A. Ribeiro, 'Constrained learning with non-convex losses,' IEEE Trans. on Inf. Theory , vol. 69[3], pp. 1739-1760, 2023.

M. Gürbüzbalaban, Y. Hu, and L. Zhu, 'Penalized Langevin and Hamiltonian Monte Carlo Algorithms for Constrained Sampling,' 2022.

X. Liu, X. Tong, and Q. Liu, 'Sampling with trusthworthy constraints: A variational gradient framework,' Advances in Neural Information Processing Systems , vol. 34, pp. 23 557-23 568, 2021.

R. Jordan, D. Kinderlehrer, and F. Otto, 'The variational formulation of the Fokker-Planck equation,' SIAM Journal on Mathematical Analysis , vol. 29, no. 1, pp. 1-17, 1998.

L. Ambrosio, N. Gigli, and G. Savaré, Gradient flows: in metric spaces and in the space of probability measures . Springer Science & Business Media, 2005.

A. S. Dalalyan and A. Karagulyan, 'User-friendly guarantees for the Langevin Monte Carlo with inaccurate gradient,' Stochastic Processes and their Applications , vol. 129, no. 12, pp. 5278-5311, 2019.

A. S. Dalalyan, A. Karagulyan, and L. Riou-Durand, 'Bounding the error of discretized Langevin algorithms for non-strongly log-concave targets,' Journal of Machine Learning Research , vol. 23, no. 235, pp. 1-38, 2022.

S. Vempala and A. Wibisono, 'Rapid convergence of the unadjusted Langevin algorithm: Isoperimetry suffices,' in Advances in Neural Information Processing Systems (NeurIPS) , 2019, pp. 8092-8104.

D. P. Bertsekas, Convex Optimization Theory . Athena Scientific, 2009.

D. G. Luenberger, Optimization by Vector Space Methods . Wiley, 1968.

V. Jeyakumar and H. Wolkowicz, 'Generalizations of slater's constraint qualification for infinite convex programs,' Math. Program. , vol. 57, no. 1-3, pp. 85-101, 1992.

J. F. Bonnans and A. Shapiro, Perturbation Analysis of Optimization Problems . Springer, 2000.

A. J. Kurdila and M. Zabarankin, Convex Functional Analysis . Birkhäuser Basel, 2005.

C. Villani, Topics in optimal transportation . American Mathematical Soc., 2021, vol. 58.

A. P. Ruszczy´ nski, Nonlinear Optimization . Princeton University Press, 2006.

N. Nisan, T. Roughgarden, Éva Tardos, and V. V. Vazirani, Eds., Algorithmic Game Theory . Cambridge University, 2007.

D. P. Bertsekas, Constrained Optimization and Lagrange Multiplier Methods (Optimization and Neural Computation Series) , 1st ed. Athena Scientific, 1996.

B. Kloeckner, 'Approximation by finitely supported measures,' ESAIM: Control, Optimisation and Calculus of Variations , vol. 18, no. 2, pp. 343-359, 2012.

A. Nagurney and D. Zhang, Projected Dynamical Systems and Variational Inequalities with Applications . Springer, 1996.

M. Welling and Y. W. Teh, 'Bayesian learning via stochastic gradient Langevin dynamics,' in Proceedings of the 28th International Conference on International Conference on Machine Learning , ser. ICML'11. Madison, WI, USA: Omnipress, 2011, p. 681-688.

A. Salim, A. Korba, and G. Luise, 'The Wasserstein proximal gradient algorithm,' Advances in Neural Information Processing Systems , vol. 33, pp. 12 356-12 366, 2020.

A. Nedi´ c and A. Ozdaglar, 'Approximate primal solutions and rate analysis for dual subgradient methods,' SIAM Journal on Optimization , vol. 19, no. 4, pp. 1757-1780, 2009.

A. Cherukuri, E. Mallada, and J. Cortés, 'Asymptotic convergence of constrained primal-dual dynamics,' Systems & Control Letters , vol. 87, pp. 10-15, 2016.

A. Nemirovski, 'Prox-method with rate of convergence o(1/t) for variational inequalities with Lipschitz continuous monotone operators and smooth convex-concave saddle point problems,' SIAM Journal on Optimization , vol. 15, no. 1, pp. 229-251, 2004.

Y. Nesterov, 'Dual extrapolation and its applications to solving variational inequalities and related problems,' Math. Program. , vol. 109, no. 2-3, p. 319-344, 2007.

T. Lin, C. Jin, and M. I. Jordan, 'Near-optimal algorithms for minimax optimization,' in Proceedings of Thirty Third Conference on Learning Theory , 2020, pp. 2738-2779.

A. Mokhtari, A. Ozdaglar, and S. Pattathil, 'A unified analysis of extra-gradient and optimistic gradient methods for saddle point problems: Proximal point approach,' in Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics , 2020, pp. 1497-1507.

S. Chewi, M. A. Erdogdu, M. Li, R. Shen, and S. Zhang, 'Analysis of Langevin Monte Carlo: from Poincaré to Log-Sobolev,' in Proceedings of Thirty Fifth Conference on Learning Theory , ser. Proceedings of Machine Learning Research, P.-L. Loh and M. Raginsky, Eds., vol. 178. PMLR, 2022, pp. 1-2.

R. Holley and D. W. Stroock, 'Logarithmic Sobolev inequalities and stochastic Ising models,' Journal of Statistical Physics , no. 5-6, 1986.

D. Bakry, I. Gentil, M. Ledoux et al. , Analysis and geometry of Markov diffusion operators . Springer, 2014, vol. 103.

P. Cattiaux and A. Guillin, 'Functional inequalities for perturbed measures with applications to log-concave measures and to some Bayesian problems,' Bernoulli , vol. 28, no. 4, 2022.

H. Karimi, J. Nutini, and M. Schmidt, 'Linear convergence of gradient and proximal-gradient methods under the Polyak-Łojasiewicz condition,' in Machine Learning and Knowledge Discovery in Databases , P. Frasconi, N. Landwehr, G. Manco, and J. Vreeken, Eds. Cham: Springer International Publishing, 2016, pp. 795-811.

J. Yang, X. Li, and N. He, 'Nest your adaptive algorithm for parameter-agnostic nonconvex minimax optimization,' in Advances in Neural Information Processing Systems , 2022.

M. Boroun, Z. Alizadeh, and A. Jalilzadeh, 'Accelerated primal-dual scheme for a class of stochastic nonconvex-concave saddle point problems,' in American Control Conference , 2023, pp. 204-209.

M. Sanjabi, J. Ba, M. Razaviyayn, and J. D. Lee, 'On the convergence and robustness of training gans with regularized optimal transport,' in Advances in Neural Information Processing Systems , 2018.

J. Yang, N. Kiyavash, and N. He, 'Global convergence and variance reduction for a class of nonconvex-nonconcave minimax problems,' in Advances in Neural Information Processing Systems , 2020, pp. 1153-1165.

T. Fiez, L. Ratliff, E. Mazumdar, E. Faulkner, and A. Narang, 'Global convergence to local minmax equilibrium in classes of nonconvex zero-sum games,' in Advances in Neural Information Processing Systems , 2021, pp. 29 049-29 063.

K. Ahn and S. Chewi, 'Efficient constrained sampling via the mirror-Langevin algorithm,' in Advances in Neural Information Processing Systems , vol. 34, 2021, p. 26.

D. Dua and C. Graff, 'UCI machine learning repository,' 2017. [Online]. Available: http://archive.ics.uci.edu/ml

L. F. O. Chamon and A. Ribeiro, 'Probably approximately correct constrained learning,' in Advances in Neural Information Processing , 2020.

D. M. Blei, A. Y. Ng, and M. I. Jordan, 'Latent Dirichlet allocation,' J. Mach. Learn. Res. , vol. 3, pp. 993-1022, 2003.

G. Celeux, M. E. Anbari, J.-M. Marin, and C. P. Robert, 'Regularization in Regression: Comparing Bayesian and Frequentist Methods in a Poorly Informative Situation,' Bayesian Analysis , vol. 7, no. 2, pp. 477-502, 2012.

A. Lamperski, 'Projected stochastic gradient Langevin algorithms for constrained sampling and non-convex learning,' in Conference on Learning Theory . PMLR, 2021, pp. 2891-2937.

L. Li, Q. Liu, A. Korba, M. Yurochkin, and J. Solomon, 'Sampling with mollified interaction energy descent,' arXiv preprint arXiv:2210.13400 , 2022.

K. S. Zhang, G. Peyré, J. Fadili, and M. Pereyra, 'Wasserstein control of mirror Langevin Monte Carlo,' in Conference on Learning Theory . PMLR, 2020, pp. 3814-3841.

Q. Jiang, 'Mirror Langevin Monte Carlo: the case under isoperimetry,' Advances in Neural Information Processing Systems , vol. 34, pp. 715-725, 2021.

V. Srinivasan, A. Wibisono, and A. Wilson, 'Fast sampling from constrained spaces using the metropolis-adjusted mirror Langevin algorithm,' arXiv preprint arXiv:2312.08823 , 2023.

J. Shi, C. Liu, and L. Mackey, 'Sampling with mirrored Stein operators,' International Conference of Learning Representations , 2022.

M. Girolami and B. Calderhead, 'Riemann manifold Langevin and Hamiltonian Monte Carlo methods,' Journal of the Royal Statistical Society Series B: Statistical Methodology , vol. 73, no. 2, pp. 123-214, 2011.

M. Brubaker, M. Salzmann, and R. Urtasun, 'A family of MCMC methods on implicitly defined manifolds,' in Artificial intelligence and statistics . PMLR, 2012, pp. 161-172.

P. Tseng, 'On linear convergence of iterative methods for the variational inequality problem,' Journal of Computational and Applied Mathematics , vol. 60, no. 1, pp. 237-252, 1995.

N. Golowich, S. Pattathil, C. Daskalakis, and A. Ozdaglar, 'Last iterate is slower than averaged iterate in smooth convex-concave saddle point problems,' in Proceedings of Thirty Third Conference on Learning Theory , 2020, pp. 1758-1784.

V. M. Panaretos and Y. Zemel, An invitation to statistics in Wasserstein space . Springer Nature, 2020.

F. Otto, 'The geometry of dissipative evolution equations: the porous medium equation,' Communications in Partial Differential Equations , vol. 26, no. 1-2, pp. 101-174, 2001.

A. T. Schwarm and M. Nikolaou, 'Chance-constrained model predictive control,' AIChE Journal , vol. 45[8], pp. 1743-1752, 1999.

P. Li, M. Wendt, and G. Wozny, 'Robust model predictive control under chance constraints,' Computers & Chemical Engineering , vol. 24[2-7], pp. 829-834, 2000.

F. Borrelli, A. Bemporad, and M. Morari, Predictive Control for Linear and Hybrid Systems . Cambridge University Press, 2017.

A Related work

In constrained sampling, it is important to distinguish between two types of constraints: support constraints and statistical constraints. The former deals with sampling from a target distribution π that is supported on a proper subset X ⊂ R d , which arises in applications such as latent Dirichlet allocation [62] and regularized regression [63]. The latter is the problem tackled in the current work.

A first family of constrained sampling methods relies on rejection sampling: it obtains samples via any (unconstrained) method, rejecting those that violate the constraint (see, e.g., [10, 11]). Though this approach can handle constraints of any nature, it is often inefficient in terms of number of samples generated per iteration of the method (effective number of samples), especially when confronted with intricate constraints and high dimensional problems.

These drawbacks can be addressed for support constraints using techniques inspired by finitedimensional constrained optimization. Projected LMC, for instance, deals with the problem of sampling from a target distribution restricted to a convex set [13, 64]. Barrier methods have also been used to tackle the same problem [18, 65]. Similarly, mirror and proximal descent versions of popular sampling algorithms such as LMC [12, 14, 15, 66-68] and Stein Variational Gradient Descent (SVGD) [69] have been proposed. Mirror descent algorithms enforce constraints by mapping (mirroring) the primal variables to a space with a different geometry (induced by a Bregman divergence) over which the optimization is carried out. Alternatively, methods adapted to target distributions support on manifolds have also been put forward [18, 70, 71]. In practice, these methods require explicit expressions for the projections, barriers, and mirror maps describing the constraints and are therefore not adapted to handle statistical requirements such as those considered in (PI). Langevin dynamics with constraint violation penalties were considered in [23], although they cannot enforce exact constraint satisfaction.

Statistical (moment) constraints such as those considered in (PI) were investigated in [24]. As we discussed at the end of Section 2.3, this paper considers the combination of LMC and approximate dual ascent shown in (9). It also introduces a similar version of SVGD as well as algorithms based on barriers. Aside from requiring exact integration against intractable measures (namely, µ k ), convergence guarantees for these methods hold under restrictive assumptions on the constraints g i . Additionally, guarantees are derived only for continuous-time (gradient flows) dynamics.

This work is also closely related to saddle-point methods in finite-dimensional optimization. For the general problem of max x min y f ( x, y ) , the behavior of descent-ascent methods have been investigated under a myriad of scenarios, including for functions f that are (strongly) convex/(strongly) concave [43, 45-48, 72, 73] as well as non-convex/non-concave under, e.g., PL conditions [47, 54, 55, 57, 58]. In general, convergence holds for the ergodic average of iterates [43, 57, 73]. Last-iterate results often require different algorithms, involving proximal point or extra gradient methods and time scale separation [47, 48, 73]. In particular, guarantees for the GDA method used in Algorithm 1 often requires stringent conditions that are hard to enforce for dual problems such as (DI).

B Wasserstein space and discrete-time flows

In this subsection we give some background on the Wasserstein spaces and sampling as optimization. For µ, ν ∈ P 2 ( R d ) , we define the 2-Wasserstein distance as

where S ( µ, ν ) is the set of couplings between µ and ν . The metric space ( P 2 ( R d ) , W 2 ) is referred to as the Wasserstein space [74]. It can be equipped with a Riemannian structure [75]. In this geometric interpretation, the tangent space to P 2 ( R d ) at µ is included in L 2 ( µ ) and is equipped with a scalar product defined for f, g ∈ L 2 ( µ ) as

We use its differential structure as introduced in [75] and [26, Chapter 10]. For the functionals at stake in this paper (e.g., potential energies and negative entropy), the set of subgradients of a functional F : P 2 ( R d ) → R at µ ∈ P ( R d ) is non-empty if µ satisfies some Sobolev regularity [26, Theorem 10.4.13].

A Wasserstein gradient flow of F is a solution ( µ ( t )) t ∈ (0 ,T ) , T > 0 , of the continuity equation

that holds in the distributional sense, where v ( t ) is a subgradient of F at µ ( t ) . Among the possible processes ( v ( t )) t , one has a minimal L 2 ( µ ( t )) norm and is called the velocity field of ( µ ( t )) t . In a Riemannian interpretation of the Wasserstein space [75], this minimality condition can be characterized by v ( t ) belonging to the tangent space to P 2 ( R d ) at µ ( t ) denoted T µ ( t ) P 2 ( R d ) , which is a subset of L 2 ( µ ( t )) . The Wasserstein gradient is defined as this unique element, and is denoted ∇ W 2 F ( µ ( t )) . In particular, if µ ∈ P 2 ( R d ) is absolutely continuous with respect to the Lebesgue measure, with density in C 1 ( R d ) and such that F ( µ ) < ∞ , ∇ W 2 F ( µ )( x ) = ∇F ′ ( µ )( x ) for µ -a.e. x ∈ R d , where F ′ ( µ ) denotes the first variation of F evaluated at µ , i.e. (if it exists) the unique function F ′ ( µ ) : R d → R s.t.

for all ξ = ν -µ , where ν ∈ P 2 ( R d ) .

Now, we denote by T # µ the pushforward measure of µ ∈ P 2 ( R d ) by the measurable map T . We recall that the KL divergence of µ relative to π can be decomposed as (1). The distribution µ k of x k in (3) is known to follow a 'forward-flow' splitting scheme [7] of the Fokker-Planck equation in (2), namely

where I denotes the identity map in L 2 ( µ k ) and ∇ W 2 V ( µ k ) = ∇ x f ( x ) .

C Proofs from Section 3.1

Proof of Theorem 3.3. Consider the potential

where µ ⋆ is the solution of the constrained sampling problem (PI) and Φ ⋆ is the set of solutions of the dual problem (DI). Our goal is to show that V decreases (in some sense) along trajectories of Algorithm 1. We say 'in some sense' because contrary to the standard LMC algorithm, the distribution µ k of x k is now a random measure that depends on the random variables { λ k , ν k } . Explicitly, we consider the filtration F k = σ ( µ 0 , { λ ℓ , ν ℓ } ℓ ≤ k ) and show that V decreases on average when conditioned on F k . This turns out to be enough to prove Theorem 3.3.

Indeed, notice that µ k ∈ F k . Hence, the potential energy E ( µ k , λ k , ν k ) ∈ F k and the conditional law ˜ µ k = L ( x k |F k -1 ) evolves as in the regular LMC algorithm (3). That is to say, conditioned on the event F k , step 5 of Algorithm 1 follows a splitting scheme as (16), i.e.,

Notice that all distributions in (18) are now deterministic. The core of the proof is collected in the following lemma that shows that V is a non-negative supermartingale. Note that Lemma C.1 describes the gap between 'half-iterations.' This will be inconsequential for our purposes.

Lemma C.1. Under the conditions of Theorem 3.3, we have

We defer the proof of Lemma C.1 to show how it implies the bounds in Theorem 3.3. To do so, take the expectation of (19) with respect to { λ ℓ , ν ℓ } ℓ ≤ k to obtain

Using the bounds in Assumption 3.2 then yields

Then, summing the LHS of (21) over k and using the fact that V is non-negative yields

Notice that the expectation here is taken only over µ 0 given that ( λ 0 , ν 0 ) = (0 , 0) are deterministic. To proceed, we use the following proposition, whose proof we defer.

Lemma C.2. Under the hypothesis of Theorem 3.3 it holds that E [ V (˜ µ 1 / 2 , λ 1 , ν 1 )] ≤ R 2 0 for

Back in (21), we obtain that

We conclude by using the convexity of the Wasserstein distance to write

Similarly for the KL divergence. The bounds in Theorem 3.3 are then obtained for the specific choices of η k by noticing that ∑ K k =1 1 / √ k ≥ √ K and ∑ K k =1 1 /k ≤ 1 + log( K ) . Notice that all inequalities in the proof continue to hold for deterministic { λ k , ν k } . The bounds in Theorem 3.3 therefore also hold (without expectations) when using exact gradients to update the dual variables.

Finally, we show there exists a sequence of step sizes η k such that E [ V ( µ k -1 / 2 , λ k , ν k )] ≤ R 2 0 for all k ≥ 1 , where the expectation is taken over the { λ k , ν k } . This immediately implies that W 2 2 ( µ k , µ ⋆ ) ≤ R 2 0 and both E [ ∥ λ k ∥ 2 ] and E [ ∥ ν k ∥ 2 ] are bounded for all k . We proceed by induction. The base case is covered by Lemma C.2. Suppose now that there exists a sequence { η 0 , . . . , η k -1 } such that E [ V (˜ µ k -1 / 2 , λ k , ν k )] ≤ R 2 0 . From the definition of V in (17) and the fact that the ( λ ⋆ , ν ⋆ ) are bounded (Prop. 2.2), we then obtain that E [ ∥ λ k ∥ ] , E [ ∥ ν k ∥ ] are bounded. Consequently, there exists η k > 0 such that

From (21), we obtain that E [∆ k ] ≤ 0 , which together with the induction assumption yields E [ V (˜ µ k +1 / 2 , λ k +1 , ν k +1 ) ] ≤ R 2 0 . ■

Proof of Lemma C.1. The proof proceeds by combining two inequalities bounding the primal and dual terms in (17).

(i) W 2 2 (˜ µ k +1 / 2 , µ ⋆ ) . We proceed following a coupling argument. Let s k be an optimal coupling between the random variables Y ∼ ˜ µ k and Z ∼ µ ⋆ , i.e., a coupling that achieves W 2 2 (˜ µ k , µ ⋆ ) . Consider now the random variable T = Y -η k ∇ W 2 E (˜ µ k , λ k , ν k ) and observe from (18a) that it is distributed as ˜ µ k +1 / 2 . Naturally, the coupling s k is no longer optimal for ( T, Z ) , so that by the definition of the Wasserstein distance it follows that

Expanding the RHS of (23) and using the m -strong convexity of E (Assumption 3.1) yields

We can then bound the effect of the diffusion step using [8, Lemma 5] as in which yields

(ii) ∥ λ k +1 -λ ⋆ ∥ 2 + ∥ ν k +1 -ν ⋆ ∥ 2 . Notice that since λ ⋆ ∈ R I + , the projection x ↦→ [ x ] + is a contraction, i.e., ∥ [ λ ] + -λ ⋆ ∥ ≤ ∥ λ -λ ⋆ ∥ for all λ ∈ R I . Using the definition of λ k +1 , ν k +1 from Algorithm 1, we then obtain that

for all ( λ ⋆ , ν ⋆ ) ∈ Φ ⋆ . To proceed, consider the conditional expectation of (26) with respect to F k , namely,

where we used the fact that ˜ µ k , λ k , ν k ∈ F k . We conclude by using the linearity of the potential energy E from (11) in both λ and ν to get

for all ( λ ⋆ , ν ⋆ ) ∈ Φ ⋆ .

To proceed with the proof, combine (25) and (28) to get

To upper bound the term in brackets, we add and subtract log( Z ) and use the decomposition of the Lagrangian in terms of (11) to obtain

We therefore conclude that

Since (30) holds for all ( λ ⋆ , ν ⋆ ) ∈ Φ ⋆ , it holds in particular for the minimizer of the RHS, for which we can then write V (˜ µ k -1 / 2 , λ k , ν k ) . By subsequently taking the minimum of the LHS, we obtain (19). ■

Proof of Lemma C.2. From the updates in Algorithm 1, we obtain

for all ( λ ⋆ , ν ⋆ ) ∈ Φ ⋆ . Notice that since λ ⋆ ∈ R I + , the projection x ↦→ [ x ] + is a contraction, i.e., ∥ [ λ ] + -λ ⋆ ∥ ≤ ∥ λ -λ ⋆ ∥ for all λ ∈ R I . Using the triangle inequality then yields

To proceed, observe from (18a) that ˜ µ 1 / 2 = [ I d -η 0 ∇ f ] # µ 0 , which implies that W 2 2 (˜ µ 1 / 2 , µ 0 ) ≤ η 2 0 ∥∇ f ∥ 2 L 2 ( µ 0 ) . Using the bounds in Assumption 3.2, we obtain that E [ V (˜ µ 1 / 2 , λ 1 , ν 1 )] ≤ R 2 0 . ■

Proof of Proposition 3.4. Since c k ≥ 0 and ∑ c k = 1 , we can use Jensen's inequality to write

Using the relation between the ℓ 1 - and ℓ 2 -norm, we further obtain

If φ is bounded by 1 , then the summands on the right-hand side of (31) can be bounded by TV ( µ k , µ ⋆ ) . Indeed,

The total variation distance can in turn be bounded by the KL divergence using Pinsker's inequality. We therefore obtain

On the other hand, if φ is 1 -Lipschitz, the summands on the right-hand side of (31) are bounded by which implies that

as long as m> 0 .

D Proofs from Section 3.2

Proof of Theorem 3.6. The proof is based on the analysis of the stochastic dual ascent algorithm

where µ λ k is the Lagrangian minimizer from (4) and ϵ > 0 . Observe that, once again, we analyze the conditional distribution ¯ µ k = L ( ξ k | λ k ) . We collect this result in the following proposition:

Proposition D.1. Consider the iterations (32) and assume that E µ [ g 2 i ] ≤ G 2 for all µ ∈ P 2 ( R d ) . Then, for 0 < η k ≤ η and ϵ ≤ ηG 2 , there exists B < ∞ such that

The expectations are taken over the samples ξ k ∼ ¯ µ k .

We conclude by combining Prop. D.1 with [29, Theorem 1], which characterizes the convergence of the LMC algorithm (3) under Assumption 3.5. Indeed, using the γ k , N 0 k from Theorem 3.6 in Algorithm 1 guarantees that the law ¯ µ k of x k | λ k is such that KL(¯ µ k ∥ µ λ k ) ≤ ϵ , i.e., satisfies the conditions in Prop. D.1. We can then apply Jensen's inequality to get that E [KL(¯ µ k ∥ µ ⋆ )] ≥ KL( E [¯ µ k ] ∥ µ ⋆ ) = KL( µ k ∥ µ ⋆ ) , where µ k is the law of x N 0 k in Algorithm 2. ■

Proof of Prop. D.1. The proof relies on the following lemmata:

Lemma D.2. For all µ ∈ P 2 ( R d ) such that KL( µ ∥ µ λ ) ≤ ϵ , the expected value E µ [ g ] is an approximate subgradients of the dual function d in (6) at λ ∈ R + I , i.e.,

Lemma D.3. Under the conditions of Prop. D.1, it holds for all k that

Lemma D.4. The sequence ( λ k , g ( ξ k )) obtained from (32) is such that

where the expectation is taken over realizations of ξ k .

Before proving these results, let us show how they imply Prop. D.1. Start by noticing from (32) that λ i,k +1 ≥ λ i,k + ηg i ( ξ k ) . Solving the recursion and recalling that λ 0 = 0 then yields

Taking the expected value over ξ k and dividing by ηK , we obtain

where the last bound stems from the triangle inequality. Since the upper bound is non-negative for all i , we use the fact that the maximum of a set of values is less than the sum of those values to write

where we used Lemma D.3 together with ( E [ ∥ z ∥ 1 ]) 2 ≤ I · E [ ∥ z ∥ 2 ] and B 0 < B . Observe that (37) bounds the infeasibility of the ergodic average 1 K ∑ K -1 k =0 ¯ µ k for ¯ µ k as in (32).

To proceed, use the relation between the normalization of µ λ and the dual function value [see (6)] to decompose the KL divergence between ¯ µ k and µ ⋆ as

Since KL(¯ µ k ∥ µ λ k ) ≤ ϵ and d ( λ ) ≤ d ( λ ⋆ ) for all λ ∈ R I + , we get

Averaging over k and using Hölder's inequality then yields

which from Lemma D.3 and (37) becomes

where we used B 0 < B . Taking the expected value, applying Lemma D.4, and taking 1 + √ I ≤ 2 I for I ≥ 1 yields (33). ■

D.1 Proof of Lemmata D.2-D.4

Proof of Lemma D.2. From the definition of the Lagrangian (4) and the dual function (6) (with J = 0 ), we obtain L ( µ, λ ) = KL( µ ∥ µ λ ) + d ( λ ) . Using the fact that KL( µ ∥ µ λ ) ≤ ϵ , we get

Again using the definition of the dual function (6), we also obtain that d ( λ ′ ) ≤ L ( µ, λ ′ ) . Adding to (38) then gives

Notice from (4) that the first term of the Lagrangians in (39) cancel out, leading to (34).

■

Proof of Lemma D.3. Start by combining the update in (32) and the fact that, since λ ⋆ ∈ R I + , the projection x ↦→ [ x ] + is a contraction, to get

Taking the conditional expectation given λ k then yields

where we used the fact that E µ [ ∥ g ∥ 2 ] ≤ G 2 for any µ . Noticing from (32) that E [ g ( ξ k ) | λ k ] = E ¯ µ k [ g ] , we can use Lemma D.2 to get

where we used the fact that η/ 2 < η to simplify the relation.

To proceed, consider the set of approximate maximizers of the dual function

Notice that Φ ⋆ ⊆ D . Since there exists at least one λ ⋆ that achieves D ⋆ (Prop. 2.2), D is not empty. Notice that for λ k / ∈ D , we have that d ( λ k ) -D ⋆ + ηG 2 + ϵ < 0 . By the towering property, we therefore obtain from (40) that

since η > 0 .

To bound the case when λ k ∈ D we use the strictly feasible candidate µ † from Assumption 2.1. Indeed, recall that KL( µ † ∥ π ) ≤ C and E µ † [ g i ] ≤ -δ < 0 for all i . From the definition of the dual function (6), we obtain that

where we used the fact that λ ∈ R I + to write ∑ i λ i = ∥ λ ∥ 1 . Hence, it follows that

Using the fact that ∥ z ∥ 2 ≤ ∥ z ∥ 2 1 for all z and that Φ ⋆ ⊂ D , we immediately obtain that ∥ λ -λ ⋆ ∥ 2 ≤ 2 B 2 0 for all λ ∈ D . Using the towering property and the fact that D ⋆ ≥ d ( λ k ) and ϵ ≤ ηG 2 yields

To conclude, we write

Using (42) and (45) then yields

Since both (42) and (45) holds independently of λ k +1 , we can also write

Applying these relations recursively, we obtain that

Noticing from (44) that since λ ⋆ ∈ D we have ∥ λ ⋆ ∥ 2 ≤ B 2 0 < B 2 then concludes the proof. ■

Proof of Lemma D.4. To bound (35), we once again use the non-expansiveness of the projection to obtain

Taking the expectation and using the fact that E µ [ ∥ g ∥ 2 ] ≤ G 2 , we get

Applying this relation recursively from K and using the fact that λ 0 = 0 (deterministic) yields

Since E [ ∥ λ ∥ 2 ] ≥ 0 , we can divide by 2 ηK

K to obtain the desired result. ■

E Applications

In this section, we provide further details on the example applications described in Section 2.2 as well as additional results from the experiments in Section 4. In these experiments, we start all chains at zero (unless stated otherwise) and use different step-sizes for each of the updates in steps 3-5 from Algorithm 1. We refer to them as η x , η λ , and η ν . In contrast, we do not use diminishing step-sizes.

E.1 Sampling from convex sets

We are interested in sampling from target distribution

for some closed, convex set C ⊂ R d . Several methods have been developed to tackle this problem, based on projections [13, 64], mirror maps [12, 15, 16], and barriers [17, 18]. Here, we consider a constrained sampling approach based on (PI) instead.

To do so, note that sampling from (46) is equivalent to sampling from

for the unconstrained π ∝ e -f ∈ P 2 ( R d ) . Note that this is exactly (PII). Now let C be described by the intersection of the 0 -sublevel sets of convex functions { s i } i =1 ,...,I , i.e.,

Such a description is always possible by, e.g., considering the distance function d ( x, C ) = inf y ∈C ∥ x -y ∥ 2 , which is convex (since C is convex) and for which { x : d ( x, C ) ≤ 0 } = C (since C is closed). Immediately, we see that solving (PII) is equivalent to solving

where [ z ] + = max(0 , z ) . Note that (PVI) has the same form as (PI).

To see why this is the case, consider without loss of generality that i = 1 . Since [ s i ( x )] + ≥ 0 for all x by definition, it immediately holds that

By the monotonicity of Lebesgue integration, we obtain that the feasibility set of (PVI) is

In other words, the feasibility set of (PVI) is in fact F = { µ ∈ P 2 ( R d ) : µ ( C ) = 1 } = P 2 ( C ) .

To illustrate the use of (PVI), consider the one-dimensional truncated Gaussian sampling problem from Section 4. Namely, we wish to sample from a standard Gaussian distribution N (0 , 1) truncated to C = [1 , 3] . In the language of (PVI), we take f ( x ) = x 2 / 2 (i.e., π ∝ e -x 2 / 2 ) and s i ( x ) =

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

Iteration

×

10

4

0

5

10

Ergodic constraint slack

(a)

0

1

2

3

4

5

Iteration

×

10

6

0

2

4

6

8

Dual variable (

λ

)

(b)

0

1

2

3

4

5

PD-LMC iterations

×

10

5

1

.

40

1

.

45

1

.

50

1

.

55

1

.

60

Sample average

(a)

0

1

2

3

4

5

LMC evaluations

×

10

5

0

.

50

0

.

75

1

.

00

1

.

25

1

.

50

Sample average

N

b

= 1

(13 s)

N

b

= 10

(44 s)

N

b

= 100

(334 s)

N

b

= 1000

(3232 s)

True mean

(b)

( x -1)( x -3) . In order to satisfy the assumptions of our convergence guarantees (particularly 2.1), we leave some slack in the constraints by considering E µ [[ s i ( x )] + ] ≤ 0 . 005 . This also helps with the numerical stability of the algorithm. Fig. 1 shows histograms of the samples obtained using PD-LMC (Algorithm 1 with η x = η λ = 10 -3 ), the projected LMC (Proj. LMC, η = 10 -3 ) from [13], and the mirror LMC ( η = 10 -3 ) from [59]. In all cases, we take 5 × 10 6 samples and keep only the second half.

Observe that, due to the projection step, Proj. LMC generates an excess of samples close to the boundaries. In fact, it generates over three times more than required. This leads to an underestimation of the distribution mean and variance (Table 2). A similar effect is observed for mirror LMC. In contrast, PD-LMC provides a more accurate estimate. Nevertheless, PD-LMC imposes constraints on the distribution µ rather than its samples. Indeed, note from (47) that its feasibility set is such that samples belong to C almost surely , which still allows for a (potentially infinite) number of realizations outside of C . Yet, though PD-LMC is not an interior-point method , Theorems 3.3-3.6 show that excursions of iterates outside of C become less frequent as the algorithm progresses. We can confirm this is indeed the case in Fig. 5a, which shows the ergodic average of [ s ( x )] + along the samples of PD-LMC. Note that it almost vanishes by iteration 10 4 even though the dual variable λ only begins to stabilize later (Fig. 5b). This is not surprising given that it is guaranteed by Prop. 3.4. In fact, only roughly 2% of the samples displayed in Fig. 1 are not in C .

Before proceeding, we examine whether the convergence of PD-LMC could be improved by averaging more than one LMC samples when updating the dual variables, i.e., using mini-batches in steps 4-5 of Algorithm 1. Mini-batches will reduce the variance of the dual updates, although at the cost of additional LMC steps per iteration. To compensate for this fact, Fig. 6b displays the evolution of the ergodic average of PD-LMC samples as a function of the number of LMC evaluations rather than the number of iterations (as in Fig. 6a). Notice that, in this application, increasing the number of LMC samples N b does not lead to faster convergence. This illustrates that, though mini-batches could be useful in some applications (particularly when the constraints are not convex, as in Section 3.2), it

Iter. = 1e4

Iter. = 1e5

Iter. = 1e6

Iter. = 5e6

Proj. LMC

(a)

Iter. = 1e4

Iter. = 1e5

Iter. = 1e6

Iter. = 5e6

PD-LMC

(b)

is not immediate that their benefits always outweigh the increased computational cost. Oftentimes, using a single LMC sample is more than enough. It is worth noting that using PD-LMC with a large mini-batch N b was suggested in [24] to approximate the expectation needed by their continuous-time algorithm. As we see here, this is neither necessary nor always beneficial.

Wenowturn to a more challenging, two-dimensional applications. We seek to sample from a Gaussian located at [2 , 2] with covariance diag([1 , 1]) restricted to an ℓ 2 -norm unit ball (Fig. 1). Specifically, we use f ( x ) = ∥ x ∥ 2 / 2 (i.e., π ∝ e -∥ x ∥ 2 / 2 ) and s i ( x ) = ∥ x ∥ 2 -1 . Once again, we leave some slack to the algorithm by taking the constraint in (PVI) to be E µ [[ s i ( x )] + ] ≤ 0 . 001 . For reference, we also display samples from the real distribution obtained using rejection sampling.

This is indeed a challenging problem. The boundary of C is 2 standard deviations away from the mean of the target distribution, which means that samples from the target π are extremely scarce this region. Indeed, using the untruncated Gaussian as a proposal for rejection sampling yields an acceptance rate of approximately 1% . The strong push of the potential f towards the exterior of C leads Proj. LMC ( η = 10 -3 ; the last 10 6 samples are used after running 5 × 10 6 iterations) to be now even more concentrated around its boundary. In fact, almost 25% of its samples are in an annular region of radius [0 . 999 , 1) , where only 0 . 14% of the samples should be according to rejection sampling. Indeed, note from Fig. 7a, that even as iterations advance, the samples of Proj. LMC continue concentrate close to the boundary.

In contrast, PD-LMC ( η x = 10 -3 and η λ = 2 × 10 -1 ) only place 1 . 8% of its samples outside of C , mostly during the initial phase of the algorithm (Fig. 7a). Indeed, the average of the constraint function along samples from PD-LMC essentially vanishes around iteration 5 × 10 4 . Achieving this requires larger values of λ (on the order of 250 , Fig. 8b) compared to the one-dimensional case (Fig. 5b). This reflects the difficulty of constraining samples to C in this instance, a statement formalized in the perturbation results of Prop. 2.2(iv). Due to the more amenable numerical properties of the barrier function, mirror LMC ( η = 10 -3 ) performs well without concentrating samples on the boundary ( 0 . 15% of the samples on the annular region of radius [0 . 999 , 1) ).

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

Iteration

×

10

5

0

.

0

0

.

1

0

.

2

0

.

3

Ergodic constraint slack

(a)

0

1

2

3

4

5

Iteration

×

10

6

0

50

100

150

200

250

Dual variable (

λ

)

(b)

0

10

20

30

40

Prevalence of

>

$50k (%)

Constrained

(Acc.: 82%)

Unconstrained

(Acc.: 84%)

Population

(a)

0

.

0

0

.

5

1

.

0

1

.

5

2

.

0

Iteration

×

10

4

0

.

0

0

.

5

1

.

0

1

.

5

Dual variable (

λ

)

×

10

2

Female

Male

(b)

E.2 Rate-constrained Bayesian models

While rate constraints have become popular in ML due to their central role in fairness (see, e.g., [20]), they find applications in robust control [76-78] and to express other requirements on the confusion matrix, such as precision, recall, and false negatives [21]. For illustration, we consider here the problem of fairness in Bayesian classification.

Let q ( x ; θ ) = P [ y = 1 | x, θ ] denote the probability of a positive outcome ( y = 1 ) given the observed features x ∈ X and the parameters θ distributed according to the posterior π . This posterior is determined, e.g., by some arbitrary Bayesian model based on observations { ( x n , y n ) } n =1 ,...,N . Hence, E θ ∼ π [ q ( x ; θ )] denotes the likelihood of a positive outcome for x . Consider now a protected group, represented by a measurable subset G ⊂ X , for which we wish to enforce statistical parity. In other words, we would like the prevalence of positive outcomes to be roughly the same as that of the whole population. Thus, we now want to sample not from the posterior π , but from a close-by distribution of parameters θ that ensures this parity. Explicitly, for some tolerance δ > 0 , we want to sample from

Naturally, we can account for more than one protected by incorporating additional constraints.

In our experiments, we take π to be the posterior of a Bayesian logistic regression model for the Adult dataset from [60] (details on data pre-processing can be found in [61]). The N = 32561 data points in the training set are composed of d = 62 socio-economical features ( x ∈ R d , including the intercept) and the goal is to predict whether the individual makes more than US$ 50000 per year ( y ∈ { 0 , 1 } ). The posterior is obtained by combining a binomial log-likelihood with independent zero-mean Gaussian (log)priors ( σ 2 = 3 ) on each parameter of the model, i.e., we consider the

10

20

30

Prevalence of

>

$50k (%)

1

(5 s)

10

(28 s)

100

(251 s)

1000

(2464 s)

LMC iterations (

N

b

)

potential

We begin by using the LMC algorithm from (3) (i.e., we impose no constraints) to collect samples of the coefficients θ from this posterior ( η = 10 -4 ; the last 10 4 samples are used after running 2 × 10 4 iterations). We find that, while the probability of positive outputs is 19 . 1% across the whole test set, it is 26 . 2% among males and 0 . 05% among females. Looking at the distribution of this probability over the unconstrained posterior π (Fig. 9a), we see that this behavior goes beyond the mean. The model effectively amplifies the inequality already present in the test set, where the prevalence of positive outputs is 30 . 6% among males and 10 . 9% among females.

To overcome this disparity, we consider gender to be the protected class in (PVII), constraining both G female and G male . We formulate the constraint of (PVII) using an empirical distribution induced from the data. Explicitly, we consider constraints

where G female , G male ⊆ { 1 , . . . , N } partition the data set. For these experiments, we take δ = 0 . 01 . Using PD-LMC ( η x = 10 -4 , η λ = 5 × 10 -3 ), we then obtain a new set of samples from the logistic regression parameters θ that lead to a prevalence of positive outcomes (in the test set) of 17 . 1% over the whole population, 18 . 1% for males, and 15 . 1% for females. In fact, we notice a substantial overlap between the distributions of this probability across the constrained posterior µ ⋆ for male and female (Fig. 9a). Additionally, this substantial improvement over the previously observed disparity comes at only a minor reduction in accuracy. Though both distributions change considerably, notice from the value of the dual variables that these changes are completely guided by the female group. Indeed, λ male = 0 throughout the execution of PD-LMC (Fig. 9b).

Before proceeding, we once again examine the effect of using multiple LMC samples to update the dual variables, i.e., using mini-batches in steps 4-5 of Algorithm 1. Fig. 10 shows the distribution of the prevalence of positive predictions ( > $50 k) for different mini-batch sizes N b . In all cases, we collect 2 × 10 4 samples, which means that we evaluate 2 N b × 10 4 LMC updates (step 3 in Algorithm 1). Same as in the 1D truncated Gaussian case, we notice no difference between the resulting distributions. This is to be expected given our results (Theorem 3.3). The computation time, on the other hand, increases considerably with the mini-batch size. Once again, we note that PD-LMC with a large mini-batch N b was used in the experiments of [24] to overcome the challenge of computing an expectation in their dual variable updates. In turns out that this computationally intensive modification is not necessary.

Intercept

Bachelors

Adm/clerical

African-Amer.

White

Gender

Iran

Age (31.0, 37.0]

0

.

0

0

.

5

1

.

0

1

.

5

2

.

0

Odds ratio

Unconstrained

Constrained

(a)

10

20

30

Prevalence of

>

$50k (%)

0.01

0.02

0.05

Tolerance (

δ

)

(b)

E.3 Counterfactual sampling

Previous applications were primarily interested in sampling from µ ⋆ , the constrained version of the target distribution π . The goal of counterfactual sampling , on the other hand, is to probe the probabilistic model π by evaluating its compatibility with a set of moment conditions. It is therefore interested not only in µ ⋆ , but in how each condition affects the value P ⋆ = KL( µ ⋆ ∥ π ) . We next describe how constrained sampling can be used to tackle this problem.

Let π denote a reference probabilistic model, such as the posterior of the Bayesian logistic model in (48). Standard Bayesian hypothesis tests can be used to evaluate the validity of actual statements such as 'is it true that E x ∼ π [ g ( x )] ≤ 0 ? or E x ∼ π [ h ( x )] = 0 ?' Hence, we could check 'is π more likely to yield a positive output for a male than a female individual?' (from the distributions under Unconstrained in Fig. 9a, this is probably the case). In contrast, counterfactual sampling is concerned with counterfactual statements such as 'how would the world have been if E [ g ( x )] ≤ 0 ?' In the case of fairness, 'how would the model have been if it predicted positive outcomes more equitably?'

Constrained sampling evaluates these counterfactual statements in two ways. First, by providing realizations of this alternative, counterfactual world ( µ ⋆ ). For instance, we can inspect the difference between realizations of π , obtained using the traditional LMC in (3), and µ ⋆ , obtained using PD-LMC (Algorithm 1). In Fig. 11a, we show the mean of some coefficients of the Bayesian logistic models from Section E.2. Notice that it is not enough to normalize the Intercept and reduce the advantage given to males ( Female is encoded as Male = 0 ). This alternative model also compensates for other correlated features, such as education ( Bachelor ), profession ( Adm/clerical ), and age.

Second, constrained sampling evaluates the 'compatibility' of each counterfactual condition (constraint) world with the reference model. While Algorithm 1 does not evaluate P ⋆ explicitly, it provides measures of its sensitivity to perturbations of the constraints: the Lagrange multipliers ( λ ⋆ , ν ⋆ ) . Indeed, recall from Prop. 2.2 that

Hence, ( λ ⋆ , ν ⋆ ) describe the magnitude of tilts needed for the reference model π to satisfy the conditions E [ g ( x )] ≤ 0 or E [ h ( x )] = 0 . This relation is made explicit in Prop. 2.2(iv).

Concretely, observe that the dual variable relative to the constraint on the male subgroup is always zero (Fig. 9b). This implies that π is fully compatible with the condition

i.e., the statement 'the model predicts positive outcomes for males on average at least as much as for the whole population.' In contrast, accommodating statistical parity for females requires considerable deviations from the reference model π ( λ ⋆ female ≈ 160 ). Without recalculating µ ⋆ , we therefore know that even small changes in the tolerance δ for the female constraint would substantially change the

TSLA NVDA

JNJ

AAPL GOOG BRK-B

LLY

-

0

.

2

0

.

0

0

.

2

0

.

4

0

.

6

(Log)return

LMC

PD-LMC

TSLA NVDA

JNJ

AAPL GOOG BRK-B

LLY

5

10

15

20

Volatility

JNJ

AAPL

GOOG

BRK-B

LLY

2

4

0

1

2

3

4

5

6

Iteration

×

10

5

-

0

.

03

-

0

.

02

-

0

.

01

0

.

00

0

.

01

Ergodic constraint slacks

TSLA

NVDA

JNJ

AAPL

GOOG

BRK-B

LLY

(a)

0

1

2

3

4

5

6

Iteration

×

10

5

-

10

-

5

0

5

Dual variable (

ν

)

TSLA

NVDA

JNJ

AAPL

GOOG

BRK-B

LLY

(b)

distribution of outcomes. This statement is confirmed in Fig. 11b. Notice that this is only possible due to the primal-dual nature of PD-LMC.

E.3.1 Stock market model

Counterfactual analyses based on the dual variables become more powerful as the number of constraints grow. To see this is the case, consider the Bayesian stock market model introduced in Section 2.2. Here, π denotes the posterior model for the (log-)returns of 7 assets ( TSLA , NVDA , JNJ , AAPL , GOOG , BRK-B , and LLY ). The dataset is composed of 5 years of adjusted closing prices for a total of 1260 points per asset. The posterior is obtained by combining a Gaussian likelihood N ( ρ, Σ) with Gaussian prior on the mean ρ [ N (0 , 3 I ) ] and an inverse Wishart prior on the covariance Σ (with parameters Ψ = I and ν = 12 ). Using the LMC algorithm ( η = 10 -3 ; the last 3 × 10 5 samples are used after running 6 × 10 5 iterations), we collect samples from this posterior and estimate the mean and variance of the (log-)return for each stock (Table 3). In this case, Σ is initialized to 10 × I .

We might now be interested in understanding what the market would look like if all stocks were to incur a 20% increase in their average (log-)returns. To do so, we use PD-LMC ( η x = 10 -3 and η ν = 6 × 10 -3 ) to solve the following constrained sampling problem

where ¯ ρ i is the mean (log-)return of the i -th stock shown in Table 3. The distribution of ρ and the diagonal of Σ are compared to those from the unconstrained model in Fig. 12. Notice that, though we only impose constraints on the average returns ρ , we also see small changes in the stock volatilities.

Inspecting the dual variables (Fig. 13b), we notice that three dual variables are essentially zero ( TSLA , AAPL , and BRK-B ). This means that their increased returns are completely dictated by those of other

TSLA NVDA

JNJ

AAPL GOOG BRK-B

LLY

-

0

.

4

-

0

.

2

0

.

0

0

.

2

0

.

4

0

.

6

(Log)return

All stocks

NVDA, LLY

TSLA NVDA

JNJ

AAPL GOOG BRK-B

LLY

5

10

15

20

Volatility

JNJ

AAPL

GOOG

BRK-B

LLY

2

4

stocks. Said differently, their returns increasing 20% is consistent with the reference model π conditioned on the other returns increasing. Proceeding, two stocks have negative dual variables ( LLY and NVDA ). This implies that bringing their constraints down to ¯ ρ i would yield a decrease in P ⋆ (distance to the reference model π ). This is in contrast to JNJ and GOOG , whose positive λ 's imply that we should should increase their returns to reduce P ⋆ . Indeed, by inspecting the ergodic slacks (Fig. 13a) we see that all stocks approach zero (i.e., feasibility), but that JNJ and GOOG do so from above. This behavior is expected according to Prop. 3.4.

These observations show two things. First, that an increase in the average returns of LLY and NVDA is enough to drive up the returns of all other stocks. In fact, it leads to essentially the same distribution as if we had required the increase to affect all stocks (Fig. 14). Second, that the increase we would see in JNJ (and to a lesser extent GOOG ) would actually be larger than 20% . Once again, we reach these conclusion without any additional computation. Their accuracy can be corroborated by the results in Table 3.

F Example application of Prop. 2.2

In this section we illustrate the result in Prop. 2.2, i.e., we show that given solutions ( λ ⋆ , ν ⋆ ) of (DI), the constrained sampling problem (PI) reduces to sampling from µ λ ⋆ ν ⋆ ∝ e -U ( · ,λ ⋆ ,ν ⋆ ) .

Indeed, consider a standard Gaussian target, i.e., π ∝ e -∥ x ∥ 2 / 2 , and the linear moment constraint E [ x ] = b , for b ∈ R d . This can be cast as (PI) with f ( x ) = ∥ x ∥ 2 / 2 and h ( x ) = b -x (no inequality constraints, i.e., I = 0 ). Clearly, the solution of (PI) in this case is µ ⋆ = N ( b, I ) , i.e., a Gaussian distribution with mean b . What Prop 2.2 claims is that rather than directly solving (PI), we can solve (DI) to obtain a Lagrange multiplier ν ⋆ such that µ ⋆ = µ ν ⋆ for µ ν defined as in (4).

In this setting, we can see this is the case by doing the computations explicitly. Indeed, we have

Completing the squares, we then obtain

To compute the Lagrange multiplier ν ⋆ , notice from the definition of the dual function in (6) that the dual problem (DI) is in fact a ratio of normalizing factors. Explicitly,

Immediately, we obtain

Note that the dual problem is a concave program, as is always the case [30]. To conclude, we can combine (49) and (50) to get

The main advantage of using Algorithms 1 and 2 is that we do not need to determine the Lagrange multipliers ( λ ⋆ , ν ⋆ ) to then sample from µ λ ⋆ ν ⋆ = µ ⋆ . Indeed, Theorems 3.3 and 3.6 show that these stochastic primal-dual methods do both things simultaneously, without explicitly evaluating any expectations.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: the abstract and introduction list our contributions, i.e. the theoretical and experimental study of Primal-Dual Langevin Monte Carlo.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: we discuss the limitations of our theoretical results in the conclusion.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: For each theoretical result, we worked with justified assumptions, making all the dependencies of the problem clear. We provide clear and detailed proofs in the appendix.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: For each experimental result, we provide the detailed setting.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We will provide a link to a public github repository with a Python code.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/public/ guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https://nips.cc/public/ guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: the algorithms and settings of the experiments are rather simple, and detailed in the submission.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We give all precisions relative to the the evaluation of our models, either on simulated data or real data, including precising train/set proportions for instance.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: Our experiments run on a standard laptop in a few minutes, and are illustrative.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: In our opinion, this paper does not address societal impact directly, and consider the generic problem of optimization over measures and sampling.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: In our opinion the paper does not have direct positive or negative social impact.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: Our paper does not present such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [NA]

Justification: our paper does not use existing assets

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: we only have experiments in Python that will be made public.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: our experiments do not involve crowdsourcing.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: our study do not involve risk for participants.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper aims to solve the constrained sampling problem via primal-dual method. The authors proposed a new sampling method PD-LMC and provide detailed convergence analysis. Several numerical experiments were conducted to verify the sampling method.\n\nSTRENGTHS:\nThe structure of paper is easy to follow. The authors discussed detailed background and related work. The authors also established solid convergece guarantee for their algorithms under different settings.\n\nWEAKNESSES:\n1. The primal-dual method for constrained sampling problem is not novel. Similar algorithms were also studied in [1]. I notice that PD Langevin in [1] requires the estimation of expectation with respect to sampling distribution. But it can easily addressed by running many particles in parallel. Therefore I take the theoretical analysis as the most important contribution of this paper. The weaknesses are listed as following.\n\n2. In the theoretical analysis part, the authors didn't show the rate of violation of constraints, i.e. if $\\mu_k$ satisfies constraints. And this leads to the following one.\n\n3. In Theorem 3.3, when there are equality constraints, $\\mu^*$ is supported on a low dimension manifold while $\\mu_k$ is not, then $KL(\\mu_k\\\\|\\mu^*)$ is not well defined. Hence I doubt the correctness of Theorem 3.3. I suggest the authors should consider metric not depending on density ratio, like W2 distance or TV distance. \n\n4. In the numerical experiments, the paper lacks comparison with other methods, e.g. mirror Langevin related methods for sampling on convex set; PD Langevin and Control Langevin in [1] in rate-constrained Bayesian models.\n\n[1] Liu, Xingchao, Xin Tong, and Qiang Liu. \"Sampling with trusthworthy constraints: A variational gradient framework.\" Advances in Neural Information Processing Systems 34 (2021): 23557-23568.\n\nQUESTIONS:\nPlease see weakness part.\n\nLIMITATIONS:\nThe limitations are pointed out in the paper.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-08"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper studies constrained sampling schemes. The objective function is the KL divergences with a target distribution, while the constraint set is given as some expectation equations or inequalities. The authors rewrite the problem into a saddle point formulation and study the Wasserstein gradient descent and L2 gradient ascent directions. This results in a particular Langevin dynamics with a dual ascent direction. The authors study the convergence analysis of the proposed algorithm. Several numerical examples demonstrate the effectiveness of the proposed method.\n\nSTRENGTHS:\nThe authors clearly explain the constrained optimization problems. They apply the primal-dual gradient descent-ascent algorithm based on the Wasserstein space to solve the optimization problem. Convergence analysis is presented using tools in optimal transport.  \n\nSampling from a convex set is a very good application.\n\nWEAKNESSES:\nThere is a lack of literature on generalized optimization problems in Wasserstein spaces. \n\nWang et al. Accelerated Information Gradient flow. Journal of Scientific Computing. 2022. \n\nWang et al. Information Newton's flow: second-order optimization method in probability space, 2020. \n\nTan et al. Noise-Free Sampling Algorithms via Regularized Wasserstein Proximals, 2023.\n\nQUESTIONS:\n1. For the time update, the authors perform the forward Euler time step for the primal and dual variables.  Do the authors expect some advantages if one applies the proximal steps?  What is the main difference in analyzing the primal-dual algorithm in Wasserstein space and the classical Euclidean space. \n\n2. Can the authors provide some simple examples, such as Gaussian target distributions and linear moment constraints, to explain the main result? This could partially answer the first question.\n\nLIMITATIONS:\nThere are no limitations.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-08"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this work, the focus is on a constrained optimisation problem in the space of measures. Specifically, the goal is to obtain a distribution / samples from a distribution which is close in KL to a target distribution while also satisfying a set of statistical constraints. The paper discusses this somewhat atypical problem, and designs the PD-LMC method to generate approximate sample from such a distribution. Theoretical properties of PD-LMC are also presented, along with some numerical experiments to showcase the working of this method.\n\nSTRENGTHS:\n### Quality and Clarity\n\nThe paper reads well, with a clear description of the objective at hand and the necessary background to help the reader grasp the problem and the proposed solution. For someone used to the classical constrained sampling problem, the examples in Section 2.2 definitely helped contextualize the problem better, and is appreciated. The dual formulation which is central to this work is also adequately presented, which helps understanding the algorithm. It is interesting that the proposed method gets away without computing expectations and replacing them in a single particle.\n\n### Originality and Significance\n\nThe idea is novel, and looks like a clean way to adapt ideas from Euclidean optimization for the distributional constrained problem considered here. I cannot quite comment on the significance of this proposed method as I'm not familiar with this flavor of constrained sampling. Judging by the numerical experiments in section 4, the method appears to work well in low-dimensional setting.\n\nWEAKNESSES:\nThe method is interesting and novel, but there is some loss of intuition for me, which I've tried to express as a question in Q2. Essentially, the proposed method performs two levels of approximations which while practically appealing, makes it hard to follow. I also have a minor issue with laws of $x_{k}$ and the expectations appearing in Theorem 3.3. Essentially, the concern is how there's no stochasticity on the LHS, and this is where my confusion with working with laws and samples is confusing to me.\n\nQUESTIONS:\n1. Is the purpose of the $x$-updates simply to approximate the discretisation in the steps 11b and 11c? This was particularly confusing to me because problem DI doesn't involve any $\\mu$, which is what is sought out to be solved.\n2. Following-up on the previous question, to understand better, the line of reasoning is: Eq. 9 would be what is ideal, but this requires knowing $\\mu_{\\lambda_{k}, \\nu_{k}}$ and expectations under this. So, [22] proposed Eq. 10 in the equality constrainted setting, which is extended to Alg 1 in this paper. Since $\\mu_{\\lambda_{k}, \\nu_{k}}$ is a tilted version of $\\pi$, what would a proposal like approximately computing the expectations in Eq. 9 using a high-accuracy sampler like MALA do (notwithstanding the sample waste)? I suppose something similar is being done in Algorithm 2?\n3. What is preventing the dual variables from blowing up in magnitude? In other words, how large can $\\max_{k} \\mathbb{E}[\\|\\lambda_{k}\\|^{2}] + \\mathbb{E}[\\|\\nu_{k}\\|^{2}]$ be / what rate does it grow at?\n4. Why are equality constraints disregarded for PD-LMC with LSI potentials? Since $\\nu$ doesn't exist, why does it show up in Algorithm 2?\n5. Can you give a concrete example for when Assumption 3.4 is satisfied, i.e., with $g$ being a bounded perturbation, and a bound on $\\sigma$?\n\nLIMITATIONS:\nNot quite, it would be nice if the authors could comment on some limitations of their analysis / methods.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}