Constrained Sampling with Primal-Dual Langevin Monte Carlo
Authors: Luiz F. O. Chamon
Mohammad Reza Karimi Jaghargh
Anna Korba
Venue: NeurIPS 2024 poster (2024)
Keywords: Sampling
optimization
constrained optimization
duality
Primary area: optimization
TL;DR: We consider the problem of sampling from a target probability distribution whose density is know up to a normalization constant while satisfying a set of statistical constraints specified by general nonlinear functions.

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
Luiz F. O. Chamon
Mohammad Reza Karimi
University of Stuttgart
ETH Zürich mkarimi minff ethz. ch
luiz. chamon@simtech. uni mti-stuttgart.de
Anna Korba CREST, ENSAE, IP Paris anna. konakokbbaeenaee..
Abstract
This work considers the problem of sampling from a probability distribution known up to normalization constant while satisfying 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 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 ascent dynamics in Wasserstein space to put forward a discrete- time primal-dus Langevin Monte Carlo algorithm (PD-LMC) that simultaneously constrains the target distribution and samples from We analyze the convergence of PD-LMC under standard assumptions 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 several applications.
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 [+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 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), moments (e.g., portfolio return) constraints. This limitation often addressed by post- post-processing, transforming variables, or by introducing penalties in the target distribution. Though successful in specific settings, these approaches have considerable downsides. Post-processine techniques such as rejection sampling (sec, c.g., [10] [T]) 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., (8) only accommodate (deterministic) support constraints and are not suited
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Table 1: Type and target of sampling constraints.
Soft constraint
Hard constraint
Mirror/proximal LMC [22
15 projected LMC [13], barriers 16 18
Sample (xr ju)
Distribution (p)
Penalized LMC 23
PD-LMC
Though modifying the target distribution directly offers more flexibility (see. e.g., (23),, it does not guarantee that constraints are satisfied (Table T We refer the reader to Appendix Afror a more detailed literature review.
for statistical requirements such as robustness or fairness 19 22
This paper overcomes these issues by directly tackling the constrained sampling problem. Explicitly, it seeks to sample not from target distribution и on Rd, but from the distribution p² that solves
p A min (EP2(B4)
subject to
Egen [e (r)) $ 0,
I
(PI)
Ex (x)))) = 0, j
where P2(R") denotes the set of probability measures Rd with bounded second moments and the functions gi, hú represent the requirements. Note that (PI) only considers measures μ against which gi, hy are integrable. Otherwise, the expectations are taken be +00, making the correspond- ing measure infeasible. Observe that (PI) is more general than the support-constrained sampling problem considered in, e.g., 12 Indeed, it constrains the distribution μ rather than its sam- ples (Tblle Tp. Algorithms based on projections, barrier, 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 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 pre- viously proposed to tackle (PI). but it requires the exact computation of expectations with respect intractable distributions 24 This paper not only overcomes this limitation, but also provides convergence guarantees for broader class of constraint functions.
The main contributions of this work include:
a discrete-time constrained sampling algorithm (PD-LMC, Algorithm T for solving precludes any explicit integration (Section 3);
PI that
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 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 Sobolo inequality (LSI) for variant of PD-LMC (Algorithm Seection 3.2];
numerical examples illustrating applications of (PI) and the effectiveness of PD-LMC Section 4).
Problem formulation
2.1 Background on Langevin Monte Carlo
Consider target distribution и E P2(Rd) absolutely continuous with respect to Lebesgue measure whose density (also denoted w) can be expressed as T<x) <-((x)/ /2 for some normalization constant Z. Define the Kullback-Leibler (KL) divergence of with respect to
(L(H||=) = / |og
dp dm
log(()dd--oogl du = fdw+ / [log(H)du - log(Z) V(H) + H(pp) log(Z), (1)
where is the Ridon-Ni derivative, V is the potential energy, and the negative entropy is absolutely continuous with respect n; and too otherwise. For wide class of functions f (c.g..
di

---

smooth and strongly convex), samples from x can be obtained from the path of the Langevin diffusion process. whose instantaneous values x(t) have distributions μ(t) evolving according the Explicitly.
Fokker-Planck equation 25
dz(1)= = (V)( + (1)(() and @u((t) V.
{T(x(())=))
Ōt
for d-dimensional Brownian motion (t). where V 9 denotes the divergence of q and
Vw, (L(ulI(T) denotes the Wasserstein-2 gradient of KL(-Im) at [26] Theorem 10.4.17] (see Appendix B for more details). Indeed, the Langevin diffusion (2) brings the distribution pu() of srl)
progressively closer to the target In fact, the Fokker-Planck equation can be interpreted
gradient flow of the KL divergence with respect to the Wasserstein-2 distance [5).
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
Fk+1 = Ik - 7% Vzxx)) + 22k 3x iid NO, (2).
for a step size 7 > O, where 1a 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 biased dime-discrttizations of the Langevin diffusion in 2 77, rates of convergence of LMC have been obtained for smooth
and strongly convex 8 27 Or convex 28 potentials 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 set of statistical requirements. Explicitly, we wish to sample from a distribution µ² that solves (PI). Since PI constrains the distribution rather than its samples 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 Seection 4 and more details on their formulations are provided in Appendix E
1. 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 (I(H]]TT
n² E argmin #EPa(8d)
(PII)
subject to P Pr~pll E 1]
for a closed convex set c Rdd Indeed, let be the intersection of the 0-sublevel sets of convex functions {s, / Such description always exists (see (Appendix Ep. Then, (PII can be cast as (PI) using g. (x) = (x)]+ for [z]+ = max(0, 2). Notice that the gi are convex and that although they not everywhere differentiable, T(s;(x) 0) ();); is a subgradient of gi, where I(E)
on the event and 0 otherwise. Observe that support constraints can also be imposed using projections, mirror/proximal maps, and barriers
[12188. These methods, however, constrain the samples æ rather than their distribution μ as in (PII.
Rate-constrained Bayesian models: rate constraints have garnered attention in ML due to their central role in fairnsss [20] 21). Consider data pairs (ar, y). where E X are features and E (0, labels, and a protected (measurable) subgroup G C X. Let ग be a Bayesian posterior the parameters A of model q(:: 8) denoting the probability of positive outcome (based, e.g., binomial model). We wish to enforce statistical parity. i.c., we wish the prevalence of positive outcomes within the protected group G be close to higher than in the whole population. We cast this problem as (PI by constraining the average probability of positive outcome as in
p
min (RP)(R")
(L(u||=)
(PIII)
subject to
Ex,0~ [q(x; 8) [ 2 En, [g(x; 0)]
where 5 >> O denotes our tolerance (22 Naturally, multiple protected groups can be accommo- dated by incorporating additional constraints. Hence, constrained sampling provides a natural way to encode faimess in Bayesian inference.
3

---

3. Counterfactual sampling: rather than imposing requirements on probabilistic models, constrained sampling can also be used probe them by evaluating counterfactual statements. Indeed, let denote reference probabilistic model such that sampling from и yields realizations of the "real world." Consider the counterfactual statement "how would the world have been [[](())) s 0?" Constrained sampling not only gives realizations of this alternative world, but it also indicates its "compatibility" with the reference model, namely the value Ja (PI).
More concretely, consider a Bayesian stock market model. Here, т is posterior model for the (log-)ret Jreturns of / assets, e.g., distributed Gaussians (p, E). Here, the vector p 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 to change by solving
P
min MEPa(R-)
[L(H|TT)
subject to
E (1 [Ai] Pi. i= I.....I
(PIV)
(( [xd S oî
Due to correlations in market, certain choices pi or o; may be more "unrealistic" than others. Additionally, it could be that some of these conditions vacuous conditioned on the others. As we show next, our approach to tackling (PI) effectively isolates the contribution cach requirement in the solution μ², thus enabling 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 convenient character- ization based on convex duality that amenable to be sampled using the LMC algorithm (3). Indeed, let d - Re and h IR + R be vector-valued functions collecting the constraint functions yi and hy respectively. The Lagrangian [PI then defined as
Z Zzv
L() & KL(A]TT + T Ea[gl + T En[h] = KL(H]|)X)) + log
(4)
for A E R, and E RR,, where
U((,dd)) for U(z,A,v) 9((()) + vT)((r)
(x)
(5)
Zxu
and normalization constant ZA.. Notice that p* = min, maxx20, Li,, V), which is why Pi referred to as the primal problem.
To obtain the dual problem of PT), define the dual function
v(,,v)&
min MEPa(R()
I(,,A,D)
Notice from (4p that the minimum in (6 is achieved for PA from (5). the Lagrangian minimizer, that d(x, v) log( log(Z (Z/ The solution of (6) is therefore tilted version of म, whose tilt is controlled by the dual variables (A, p). Since (6) is relaxation of (PI), it yields lower bound on the
primal value, i.e., d(A,v) v) < p* for all (A, R! IR The dual problem seeks the tilts (X* v*) that yield the best lower bound, i.c.,
D max AERS.VER
d(A,v) v).
(DI)
The set x+ argmax A20, d(A, v) of solutions of (DI) called the set of Lagrange multipliers. depends on the distributions μ and through its objective d.
Note from O that DI
The dual problem DI has several advantageous properties. Indeed, while the primal problem (PD is an infinite dimensional, smooth optimization problem in probability space, the dual problem (DI a finite dimensional, non-smooth optimization problem in Euclidean space. What is more, it
concave problem regardless of the functions f,9. h, since the dual function (6) is the minimum of
set of affine functions in (X, v) 30 Prop. 1. These properties are the more attractive given
that, under mild conditions stated below, DD can be used to solve PP]

---

Assumption 2.1. There exists n² E (2(R)) with KL( ||T)
S C <
oo such that Byy [9: - - < 0
and E [hy] 0 all j.j.
Proposition 2.2. Under Assumption 2.7 the following holds:
(i) D*;
(ii) there exists a finite pair (X+ v*) x*; (iii) for any solution p² of (I and (X+ v*) of DI it holds that
L(p(,,,) S L(p X - Lp, X", v/")
for all (H, X, 1) P2(B*) x R4 x R,; (7)
(iv) the solution of (I) is nT = px.v for (X*, 1²) =x;
(v) consider the perturbation of (PI)
v), -
min (RP2(Rd)
(L(uITT) subject to Exz [g((x)] S Ug Banp [h,(x)) = Ugj (PV)
Then, (X,,"*) are subgradients of PM (0,0) Pt, i.e, P* (*( 2 P - :Tu v Vt and if P (u, v) is differentiable at (0, 0), then V, -x and V Pu, -V at ().
Proof. In finite dimensional settings, (i)((v) are well-known duality results (see, e.g., 300). 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 convex program. Under Slater's condition (Assumption 2.1)
it is (i) strongly dual (P* = D*) and (ii) there exists at least one solution (A*, v²* of (DI) (see B1
Sec. 8.6, Thm. 1] 32 Cor. 4.1)). This implies (iii) the existence of the saddle- pointt 33
Prop. 2. 156], (iv) that μ E argmin, L(u, X", ,*) = {px-v-) }. since the KL divergence is strongly
convex and its minimizer is unique [34] Thm. 7.3.7], and (v) that (A*, 2/*) are subgradients of the perturbation function p v) [33 Prop. 4.27].
shows that given solutions (A*, v²) of DI), the constrained sampling problem (P) reduces to sampling from px*** α e U-((AA".') (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)
Proposition 2.2
convex in the standard T,2 geometry, i.e., along mixtures of the form lpo + (1 f)ja for / E [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 X 2 0 and v E IR,, i.e., for f, to be convex and h to be linear.
Hence, Proposition 2.2 reduces the constrained sampling problem (P)) to that of finding the Lagrange
multipliers (A*,**) "²). Despite their finite dimensionality. however, computing these parameters is intricate. Indeed, since (DI) is a concave program, we could obtain (A", v²) E th using
Ax+1
[>v + 17s Ex,,, [9]] and Mh+1 V4 + mp Eux, [h].
(8)
for m. > 0. where we used the fact that Eps, [o] and Eys. [h] are (sub)gradients of the dual function (6) (X, v) [36] Thm. 2.87]. This procedure known optimization and game theory as dual ascent or best response 377 38 Notice, however, that (8) is not a practical algorithm as it requires explicit
integration with respect to the intractable distribution pxv from (5)
This issue was partially addressed in (24] (in continuous time and without equality constraints, i.c.. 0) by replacing the Lagrangian minimizer MALr in (8) by the distribution of LMC samples, as in
Tk+1 = Tk 7h VUKk, Ac) - 2T 8k
iid
~ (0,
AR+1 [Ax + The Eyuls]]
Law(xx).
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 μ4 (namely, using the x&)., empirical approximations of [99) may not only require an exponential (in the dimension d) number of samples (e.g.. [39 Thm. 1.2)), but it introduces enrors 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.
5

---

Algorithm 1 Primal-dual LMC
1: Inputs: 7) > O (step size), 2n ~ Pos and (Ao, vb) = (0.
2: for = 0.. :.....P0 K
3: F6+1 = Fkk 7/2 V U. Ae. vx) + 2 Sk. for BB (0,
4: Ax+1 [>x + +=gg
5: Mk+1 = Us + +nh(Tx))
6: end
3 Primal-dual Langevin Monte Carlo
Consider the GDA dynamics for the saddle-point problem Di in Wasserstein space. Explicitly,
0u(t)
(10a)
Øt [v,,((((t), (())
(10b)
Ot
0v, (t) = Ve, [(((t), v())
(10c)
Ot
for the Lagrangian L defined in 4.. where [z] = for X > 0 and [ max(a, 0) other- wise [40] Sec. 2.2]. Observe that V x [1,, 1. and V., Li,
1. Hence,
the algorithm from 24 described in (9) involves a deterministic implementation of (TOb) that fully
integrates over u(t). contrast, we consider a stochastic, single-parte particla implementation of (TO) that
leads to the practical procedure in Algorithm
Euler-Maruyama dineediscretization of the Langevin diffusion associated to T0a) (step but replace the expectations in (TOb)- (1cc) by single-sample approximations (steps 5). Algorithm 1 can therefore be seen as a particle implementation of the deterministic Wasserstein GDA algorithm (TO). As such, it resembles a primal-dual counterpart of the LMC algorithm in (3). which is why we dub primal-dual LMC (PD-LMC). Alternatively, Algorithh I can be interpreted as stochastic approximation of the dual ascent method in (9). This suggests that the gradient approx- imations steps 43 could be improved using mini-batches, which is fact how 224 approximates the expectation in Our theoretical analysis and experiments show that these mini batches are neither necessary nor always worth the additional computational cost (see Seettonn 3.1 and /Section 4p. Note that the "stochastic approximations" in Algorithm| [ refer to the dual updates (steps 5) rather than the LMC update (3tep3) as in stochastic gradient Langevin 4T. Though these methods could be combined, it is beyond the scope of this work.
Explicitly, also
The remainder of this section is dedicated to analyzing the convergence properties of PD-LMC for both stochastic dual gradients (as in Algorithm [ and exact dual gradient (as in (9)). For the latter, we obtain guarantees for the discrete implementation (9) under weaker assumptions than the continuous-timg analysis of (24] We consider strongly log-concave target distributions in Section 3.1
and those satisfying an LSI Section 3.2)
3.1 PD-LMC with (strongly) convex potentials
As opposed to the traditional LMC algorithm 3 or the deterministic updates in ( Algorithm involves three coupled random variables, namely, (Ix, As, vx). Hence, the LMC update (step 3) is based on a stochastic potential and the distribution μ4 of Øy is now 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 works such as [7] that studied the LMC ( or alternative time- discereittiiininnns of gradient flows of the KL divergence as splitting schemes.
First, define the potential energy E and the (negative) entropy H for (1,, A, v) P2(RR) xR x R
as
E(H,A,))= U(z, v)ddu(() and H((1) log(H)du.
(11)
for U as in (5). Notice from 4 that L(u, A, v) E(u, A, v) + H(H)) log Z), where we used the KL divergence decomposition in (I. To proceed, consider the following assumptions:

---

Assumption 3.1. The potential energy Ep, X, v) in (II is m-strongly convex with respect to
along W geodesics for m 2 0 and all (A. ( R x F. Explicitly,
A, 2 E() + / (Vw; E X.
x. v) 2 (1,
where s is an optimal coupling achieving W (u, no) (see (see Appendix B Assumption 3.2. The gradients and variances of f. g, are bounded along iterations fuxta20 where μ4 is the distribution of zy. i.e., there exists G such that
max (1 V fll V (e)))))) [1 S G² and max ( (B [1 E. [1]]]]] S G
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 (26] Prop. 9.3.2]. Assumption 1.22is typical in (stochastic) non- smooth optimization analyses (see, e.g., 36 43
44 ). Notice, however, that gradients are only
required to be bounded along trajectories of Algorithm T a crucial distinction in the case of strongly
convex functions whose gradients can only be bounded locally. Assumption (3. 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 μt the distribution of Ix in Algorithm
7
Under Assumptions 2.7 3.7
and 3.2 there exists Ri such that, for Ti S 1).
1
[T 11+**+++++++ T W2 (Mx. 1*) Gr
R n
G
Σ
[E[I 13)- (12)
K
K
For m S Ro/(GVk) and πkK mk/ Ek=1 Mk- obtain
K
Th KL(Uk II=*)=++++++)* + m 2 (m): RoG(1 (ogg(K))
(3+
121))
VK
k
(13)
Additionally, there exists sequence of step sizes 1]s > such that W2(ps: n²) S Rg and TB[I] >te El|ppk 12) < oc for all k. The same results hold (without expectations) when using exact dual gradients, ie., if the updates in steps 5 are replaced by Ax+1 A>x- +7p E,g, [9]] and Vi41 1/ t Epp [t].
Theorem 3.3
whose proof deferred to Apppeddix C implies rates similar those for GDA schemes in dimen-dimensional Euclidean optimization (see, e.g., (3). To recover those rates, however, we must bound the magnitudes of λx. Vg. itself, i.e., by projecting them onto the set D, {(x, E IR_ RJ maxx[])\]² || S rr and choosing such that h- C DD (Proposition 2.2 ii) ensures this is possible). We then incur bias on the order of y (2) that vanishes in the decreasing step size setting of (T3). Though convenient,
4 this is done by bounding the iterates in the algorithm
this is not necessary since there exists a sequence of step sizes such that both Ei||> and |[ bounded for all & 2 O. In the interest of generality, Theorem 3.3 holds without these hypotheses. is worth noting that though faster rates and last iterates guarantees can be obtained for Euclidean gaddle-poing problems, they rely more complex schemes than the GDA in Algorithm involving
acceleration or proximal methods [4 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.c., for the dual ascent scheme (9) In this case, the second condition in Assumption 3.2 simplifies
to maxx (Il Epp, [g]]|3, EE,,, [ l2] < G². Not only are these milder assumptions than 124 Eq. (16)), but the guarantees hold for discrete- rather than continuous-time dynamics. Finally, 12
im-
ply 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 μl from Theorem 3.3 also imply convergence rates for empirical averages across iterates Ią Afgorrim This corollary is obtained by combining (12
(13) with the following proposition. By taking &o to be the constraint functions g or from (PI) yields
fcasibility guarantees for PD-LMC.
7

---

Algorithm 2 (Stochastic) dual LMC
1: Inputs: N2 0 (burn-in), Tks T7 0 (step sizes), and Ào 0.
2: for = 0...., K -1
3: 30 140
4: In+1 In
Yk (Zn Aa) 2 Bur Bn No), for ,,...,Ng
5: Ak+1 >x me)(#Ng))
6:
end
Proposition 3.4. Consider samples Ta distributed according to and ck 2 0 with Ek=1 Cje 1
such that Σ k=1 ck
KT,(ux 1o (1 + (1s, I A for m 2 0. Then, it holds that
K
V2AK. K.
if is bounded by 1,
E
(x))(()))
Byr [sp] 3
2Ak
if is 1-Lipschitz and >> O.
m
Proof. See|Appendix C
3.2 PD-LMC with LSI potentials
In this section, we replace Assumption 3.1 on the convexity of the potential by an L.II common in the sampling literature. We consider only inequality constraints (J 0) here and omit the function arguments v, since accounting for equality constraints requires significant additional assumptions. Assumption 3.5. The distribution M4 satisfies the LSI for bounded X, i.e., there exists >> 0 such that 20 (L(Sl)μx) V log (dC/dux)lliacs for all < P2(
The LSI in Assumption 3.3 is often used in the analysis of the standard LMC algorithm 29 49 It holds, e.g., when f is strongly convex and g (possibly non-convex) bounded function due the Holley-Stroock perturbation theorem [50] In fact, if is 1-strongly convex and [g| is bounded by 1,
then Assumption 3.5 holds for a 2 e 22 (see, e.g., [51 Prop. 5.1.6] or [52] Thm 1.1])- The LSI is akin the Polyak-Lojasiewiez (PL) condition from Euclidean optimization (3). which supposes issues with GDA methods such as Algorithm Indeed, it is not enough for the Lagrangian (4) to satisfy the PL condition in the primal variable guarantee the convergence of GDA in Euclidean spaces. We must either modify Algoriihm using acceleration proximal methods 47 54 56 or impose the PL condition also on [37 Since the Lagrangian
linear in X, it is clear that Algorithm T will not suffice to provide theoretical guarantees in the LSI case.
We therefore consider the variant in Algorithm 2 where No LMC iterations (step
executed before updating the dual variables (). This is akin to using different time-scales in continuous- time, common technique for solving saddle-point problems
(54 58). Since resembles dual ascent counterpart of the LMC algorithm (3) we refer to it as (stochastic) dual LMC (DLMC). As opposed the dual ascent algorithm from [24 ()), 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 -mooth, i.e., have M -ipishhhit continuous gradients, satisfy Assumption 3.5 and that Ellll9l11 < G2 for all μ E P2(RR)) Let 0 < m S n.
S G2 1,
1 log
2 (2KT(Mollua.
16dM2 and Nº
7
70
Under Assumption 2.1 there exists B < do such that the distributions {pu} of the samples {x }
generated by Algorithm satisfy
G
21B²
(L(ullla") set
(14)
K R=0
2
Recall from PT that is the number of inequality constraints. Additionally, E[[lAR]) is bounded for all k.
8

---

PD-LMC
Proj. LMC
Mirror LMC
Density
0
1.5
2.0
2.5
3.0
1.0
2.5
3.0
1.0
2.0
2.5
3.0
Figure I: Sampling from a ID truncated Gaussian (ground truth displayed as dashed lines).
PD-LMC
Proj. LMC
Mirror LMC
Rejection
Figure 2: Sampling from a 2D truncated Gaussian (true mean in red and sample mean in orange).
Theorem 3.6.
whose proof is deferred to Appendix D provides similar guarantees as (approxi- mate) subgradient methods in kiite-dimensionall optimization (see, c.g., [22 [56).. This is not surprising seen as 7k. NV in Thhorem 3.6 are chosen to ensure that 4ields sample 23 Pa such that (T (1x))) S c using (29] Theorem 1J. At this point, g(zNg) in step 5 is an approximate, stochastic subgradient of the dual function (6). Though it may appear from (12) and (14) that Al- gorithms
[ and 2 have the same convergence rates, an informal computation shows that the latter
evaluates on the order of dr² /n) as many gradient per iteration, where k = M/o. 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 Code for these examples is publicly available at https: //ww. gi thub. com/ wfcchaoon/mmdmmmmcmc
1. Sampling from convex sets. We cast the problem of sampling from a Gaussian distribution (O. 1) truncated to C = [1, 3] as (PI) by taking f(x) x² 2 and g(x) = [( ( (z 3) (see Section 2.2). Fig. I 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 constrains the distribution j rather than its samples, it is not an interior-poin/ 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. T only 2% of the samples are not in C). This occurs even without using mini- batches in steps 5 of [Algoiithm I 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 fy-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 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 the posterior of Bayesian logistic regression model for the Adult dataset from (60), where the goal is predict whether an individual

---

TSLA
JNJ
GOOG
LLY
NVDA
AAPL
BRK- B
Unconstrainel
Dual variable (w)
Constrained
(6 0.01)
Constrained 0.03)
Female
Male
10
3
Prevalence of > $50k (%)
Iteration
x10
Distribution of the probability of pre- dicting >> 50k under the Bayesian logistic model (black lines indicate the mean across genders).
Figure 3:
Figure Counterfactual sampling of the stock market: dual variables.
makes more than $50k based on socioeconomic information (details on data pre-processing can be found in [6]). We consider a standard Gaussian prior the parameters 0 E R2+11 of the model, where d 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. over the whole test set, it 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 Gmate and Gremale with 8 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 the distributions of positive predictions across genders for the constrained posterior μe ("Constrained (8 0.01)" in Fig. Bb. 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 (Amale and Aeemale 160) show that these changes are due uniquely to the female group [as per Prop. iv This implies that the reference model i itself compatible with the requirement for male group. but that reducing disparity for females requires considerable deviations from it. By examining Afemale, we conclude without recalculating ja" that even small changes in the tolerance for the female constraint would substantially change the distribution of outcomes [Prop. [22 This is confirmed by "Constrained (8 0.03)" in Fig. 3 Notice that this is only possible due to the primal-dual nature PD-LMC. This type of counterlactual analysis is Indeed, let æ be the posterior of a Bayesian model for the daily (log- return of a set of assets (see Appendix for more details). Using PIV)
more beneficial in the presence of multiple requirements.
consider how the market would look like if the 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 (u < 0). In fact, the reference model π would be consistent with an even higher increase for JNJ and GODG (v 0). We confirm these observations by constraining only NVDA and LLY, which yields essentially the same (log-)return distribution for all assets.
average (log-)retune
5 Conclusion
We tackled the problem of sampling from target distribution while satisfying a of statistical constraints. Based on a GDA method in Wasserstein space, we put forward a fully stochastic, discrete- time 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 distribution its samples converges to the optimal constrained distribution. We illustrated 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.
10

---

Acknowledgments and Disclosure of Funding
The work of L.F.O. Chamon is funded by the Deutsche Fosschngsgeemeinschatt (DFG, German Research Foundation) under Germany's ' Excellence Strategy (EXC 2075-390740016).
References
[1] M. F. Faulkner and S. Livingstone, "Sampling algorithms in statistical physics: guide for statistics and machine learning," arXiv preprint arXiv:2208.0477 04751, 2022.
[2] 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 statisties and modelling," Nature Reviews Methods Primers, vol. 1, no. p. 2021.
[3] Y. Song and S. Ermon, "Generative modeling by estimating gradients of the data distribution." " Advances in neural information processing systems, vol. 32, 2019.
[4] G. O. Roberts and J. S. Rosenthal, "General state space Markov chains and MCMC algorithms," Probability Surveys, vol. 1, 2004.
[5] C. Robert and G. Casella, Monte Carlo statistical methods. Springer Verlag, 2004.
[6] G. O. Roberts and L. Tweedie, "Exponential convergence of Langevin distributions and their discrete approximations," Bernoulli, pp. 341-363, 1996.
[7] A. Wibisono, "Sampling optimization in the space of measures: The Langevin dynamics as
a composite optimization problem," in Conference on Learning Theory. PMLR, 2018, pp. 2093-3021
[8] A. Durmus, Majewski, and B. Miasojedow, "Analysis of Langevin Monte Carlo via convex optimization," The Journal of Machine Learning Research, vol. 20, no. 1, pp. 2666-217 266622711, 2019.
[9] Y. Wang and W. Li, "Accelerated information gradient flow." J. Sci. Comput., vol. 90, no. 1, 2022.
[10] 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. pp. 1615-1622, 2007.
[11] 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. pp. 712-732, 2015.
[12] Y.-P. Hsieh, A. Kavis, Rolland, and V. Cevher, "Mirrored Langevin dynamics," Advances in Neural Information Processing Systems, vol. 31, 2018.
[13] S. Bubeck, R. Eldan, and J. Lehec, "Sampling from log-concave distribution with projected Langevin Monte Carlo," Discrete & Computational Geometry, vol. 59, no. 4, pp. 757-783, 2018.
[14] A. Salim and Richtarik, "Primal dual interpretation of the proximal stochastic gradient Langevin algorithm," in Advances in Neural Information Processing Systems, 2020, pp. 3786- 3796.
Advances in Neural Information Processing Systems, vol. 34, pp. 28 405-28 418, 2021. [15] K. Ahn and S. Chewi, "Efficient constrained sampling via the mirror- Langevin algorithm,"
[16] Y. Kook, Y. Lee, 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.
[17] I.. Sharrock, Mackey, and C. Nemeth, "Learning rate free Bayesian inference constrained domains," in Conference Neural Information Processing Systems, 2023.
11

---

[18] M. Noble, V. De Bortoli, and A. Durmus, "Unbiased constrained sampling with self-concordant barrier Hamiltonian Monte Carlo," in Thirty-sevent Conference on Neural Information Pro- cessing Systems, 2023.
[19] A. Mądry, A. Makelov, L. Schmidt, D. Tsipras, and A. Vladu, "Towards deep learning models resistant to adversarial attacks," in International Conference on Learning Representations, 2018.
[20] M. Kearns, S. Neel, A. Roth, and Z. S. Wu, "Preventing faimess Gerrymandering: Auditing and learning for subgroup fairness," " in International Conference on Muchine Learning, 2018, pp. 2564-2572.
[21] 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.
[22] 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.
[23] M. Gürbüizalaaaa, Y. Hu, and Zhu, "Penalized Langevin and Hamiltonian Monte Carlo Algorithms for Constrained Sampling," 2022.
[24] 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.
[25] R. Jordan, D. Kinderlchrer, and Otto, "The variational formulation of the Fokker-Plance equation." SIAM Journal on Mathematical Analysis, vol. 29, no. 1, pp. 1-17, 1998.
[26] L. Ambrosio, N. Gigli, and G. Savaré, Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media, 2005.
[27] 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-55 1118-,3112220 2019.
[28] A. S. Dalalyan, A. Karagulyan, and L. Riou-Durand, "Bounding the error of discretized Langevin algorithms for non-strongly log-con con-cce targets," Journal of Machine Learning Research, vol. 23, no. 235, pp. 1-38, 2022.
[29] S. Vempala and Wibisono, "Rapid convergence of the unadjusted Langevin algorithm: Isoperimetry suffices." in Advances in Neural Information Processing Systems (NeurIPS), 2019, pp. 8092-8104.
[30] D. P. Bertsekas, Convex Optimization Theory. Athena Scientific, 2009.
[31] D. G. Luenberger, Optimization by Vector Space Methods.
Wiley, 1968.
[32] 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.
[33] J. F. Bonnans and A. Shapiro, Perturbation Analysis of Optimization Problems. 2000.
Springer,
[34] A. J. Kurdila and M. Zabarankin, Convex Functional Analysis.
Birkhäuser Basel, 2005.
[35] C. Villani, Topics in optimal transportation. American Mathematical Soc., 2021, vol. 58.
[36] A. P. Ruszczyíski, Nonlinear Optimization. Princeton University Press, 2006.
[37] N. Nisan, T. Roughgarden, Êva Tardos, and V. V. Vazirani, Eds., Algorithmic Game Theory. Cambridge University, 2007.
[38] D. Bertsekas, Constrained Optimization and Lagrange Multiplier Methods (Optimization and Neural Computation Series), 1st ed. Athena Scientific, 1996.
12

---

[39] B. Kloeckner, "Approximation by finitely supported measures," ESAIM: Control, Optimisation and Calculus of Variations, vol. 18, no. 2, pp. 343-359, 2012.
[40] A. Nagurney and D. Zhang, Projected Dynamical Systems and Variational Inequalities with Applications. Springer, 1996.
[41] M. Welling and Y. W. Teh, "Bayesian learning via stochastic gradient Langevin dynamics," in Proceedings of the 28th International Conference International Conference Machine Learning, ser. ICML'11. Madison, WI, USA: Omnipress, 2011. p. 681-688.
[42] A. Salim, A. Korba, and G. Luise, "The Wasserstein proximal gradient algorithm."" Advances in Neural Information Processing Systems, vol. 33. pp. 12 223-1112 366, 2020.
[43] A. Nedić and A. Ozdaglar, "Approximate primal solutions and rate analysis for dual subgradient methods," SIAM Journal Optimization, vol. 19, no. 4, pp. 1757-1780, 2009.
[44] A. Cherukuri, E. Mallada, and J. Cortés, "Asymptotic convergence of constrained primal- dynamics," Systems & Control Letters, vol. 87, pp. 10-15, 2016.
[45] A. Nemirovski, "Prox-method with rate of convergence o( 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.
[46] 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.
[47] 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. 1738-2779)
[48] A. Mokhtari, A. Ozdaglar, and S. Pattathil, "A unified analysis extra 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.
[49] S. Chewi, M. A. Erdogdu, M. Li, Shen, and Zhang, "Analysis 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.
[50] R. Holley and D. W. Stroock, "Logarithmic Sobolev inequalities and stochastic Ising models." Journal of Statistical Physics, no. 5-6, 1986.
[51] D. Bakry, I. Gentil, M. Ledoux al., Analysis and geometry of Markov diffusion operators. Springer, 2014, vol. 103.
[52] 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.
[53] H. Karimi. J. Nutini, and M. Schmidt, "Linear convergence of gradient and proximal proximal-gradient methods under the Dalyak-Lojasiewizz condition." " in Machine Learning and Knowledge Discov- ery in Databases, P. Frasconi, N. Landwehr, G. Manco, and J. Vreeken, Eds. Cham: Springer International Publishing. 2016, pp. 795-811.
[54] J. Yang, X. Li, and N. He, "Nest your adaptive algorithm for parameter-agnostic nonconvex minimax optimization." in Advances Neural Information Processing Systems, 2022.
[55] M. Boroun, Z. Alizadeh, and A. Jalilzadeh, "Accelerated primal-dual scheme for a class of stochastic noncomvex-concave saddle point problems," in American Control Conference, 2023, pp. 204-209.
[56) M. Sanjabi, J. Ba, M. Razaviyayn, and J. D. Lee, "On the convergence and robustness of training gans with regularized optimal transport," Advances Neural Information Processing Systems, 2018.
13

---

[57] J. Yang. N. Kiyavash, and N. He, "Global convergence and variance reduction for a class of nonconvex-nowe nonconccaneeeceeeeeeeeeeeee minimax problems." in Advances in Neural Information Processing Systems, 2020, pp. 1153-1165.
[58] T. Fiez, L. Ratliff, E. Mazumdar, E. Faulkner, and A. Narang, "Global convergence to local min- max equilibrium classes of nonconvex zero-sum games, in Advances Neural Information Processing Systems, 2021, pp. 29 049-29 063.
[59] K. Ahn and Chewi, "Efficient constrained sampling via the mirror- algorithm,"" in Advances in Neural Information Processing Systems, vol. 34, 2021, p. 26.
[60] D. Dua and C. Graff, "UCI machine learning repository," 2017. [Online]. Available: http://archive.ios. huw.rciive..c...uu.eum//lm
[61] L. F. O. Chamon and A. Ribeiro, "Probably approximately correct constrained learning," in Advances Neural Information Processing, 2020.
[62] D. M. Blei, A. Y. Ng, and M. Jordan, "Latent Dirichlet allocation," Mach. Learn. Res., vol. 3, pp. 993-1022, 2003.
[63] G. Celeux, M. E. Anbari, J.-M. Marin, and C. P. Robert, "Regularization in Regression: Com- paring Bayesian and Frequentist Methods in Poorly Informative Situation," Bayesian Analysis, vol. 7, no. 2. pp. 477-502, 2012.
[64] A. Lamperski, "Projected stochastic gradient Langevin algorithms for constrained sampling and non-convex learning." in Conference on Learning Theory. PMLR, 2021, pp. 2891-2937.
[65] L. Li. Q. Liu, A. Korba, M. Yurochkin, and Solomon, "Sampling with mollified interaction energy descent," arXiv preprint arXiv:2210.1344 13400, 2022.
[66] 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.
[67] Q. Jiang, "Mirror Langevin Monte Carlo: the case under isoperimetry." Advances in Neural Information Processing Systems, vol. 34, pp. 715-725, 2021.
[68] V. Srinivasan, A. Wibisono, and A. Wilson, "Fast sampling from constrained spaces using the metropolis-adjustee mirror Langevin algorithm," arXiv preprint arXiv:2312.08822 2023.
[69] Shi, C. Liu, and L. Mackey, "Sampling with mirrored Stein operators," International Confer- ence of Learning Representations, 2022.
[70] 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.
[71] 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.
[72] 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.
[73] 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.
[74] V. M. Panaretos and Y. Zemel, An invitation to statistics in Wasserstein space. Springer Nature, 2020.
[75] F. Otto, "TTh geometry of dissipative evolution equations: the porous medium equation," Communications in Partial Differential Equations, vol. 26, no. 1-2, pp. 101-174, 2001.
[76] A. T. Schwarm and M. Nikolaou, "Chance-constrained model predictive control," AIChE Journal, vol. 45[8], pp. 1743-1752, 1999.
14

---

[77] Li. M. Wendt, and G. Wozny, "Robust model predictive control under chance constraints," Computers Chemical Engineering, vol. 24[2-7]. pp. 829-834, 2000.
[78] Borrelli, A. Bemporad, and M. Morari, Predictive Control for Linear and Hybrid Systems. Cambridge University Press, 2017.
15

---

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' C RRd which arises in applications such as latent Dirichlet allocation 62] and regularized regression 63 The latter is the problem tackled in the current work.
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] I1). Though this approach can handle constraints of any nature, is often inefficient 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 finite- dimensional constrained optimization. Projected LMC, for instance, deals with the problem of sampling from target distribution restricted to a convex set [[3] 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 122 14
15 66-688 and Stein Variational Gradient Descent (SVGD) [69] have been proposed. Mirror descent algorithms enforce constraints by map- ping (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 target distributions support on manifolds have also been put forward [18] 770] 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 those considered in Pl were investigated in [24]. As 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 similar version of SVGD well algorithms based on barriers. Aside from requiring exact integration against intractable measures (namely, px),
convergence guarantees for these methods hold under restrictive assumptions on the constraints gi- Additionally, guarantees are derived only for continuous-tine (gradient flows) dynamics.
This work is also closely related to saddle-point methods in mnite-dimensional optimization. For the general problem of max, ming f(x, y). the behavior of descent-ascent methods have been investigated under myriad of scenarios, including for functions f that are (strongly) con- as well as non-convex/non-concaa under, e.g., PL condi- 58 In general, convergence holds for the ergodic average of iterates [43 57 73)
vex/(strongly) concave 43
48 72 73
tions 47 54 55
57
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 often requires stringent conditions that are hard to enforce for dual problems such
DI..
16

---

B Wasserstein space and discrete-time flows
In this subsection we give some background on the Wasserstein spaces and sampling as optimization. For n, V E )2,(R"), we define the 2- Wasserstein distance as
/ ||= ds(,))
W (),
inf SES(U.)
where S(u,v) is the set of couplings between μ and The metric space (P2( W2) is referred to as the Wasserstein space [74]. can be equipped with Riemannian structure [75]. In this geometric interpretation, the tangent space to P2(Rd) is included in (²( and is equipped with scalar product defined for f.g E (n) as
(f 12() = f(z)g(z)dp(())
We use its differential structure as introduced in 75 and 226 Chapter 10]. For the functionals
stake in this paper (e.g., potential energies and negative entropy), the set of subgradients of functional F P2(Rd) + IR at μ E P(R4) is non-empty if satisfies some Sobolev regularity [26 Theorem 10.4.13].
A Wasserstein gradient flow of F is a solution (u(t)))((,TT) T > O, of the continuity equation
0p(t) Ōt
+ + (u(t)u(t)) =
that holds in the distributional sense, where v(t) is subgradient of at μ(t). Among the possible processes (u(t)),, one has minimal L (u())) norm and is called the velocity field of (a(t)), In a Riemannian interpretation of the Wasserstein space 775,, this minimality condition can be characterized by vt) belonging to the tangent space to P2 ( at put) denoted Tu(e) P2( which is subset of (,(())) The Wasserstein gradient is defined as this unique element, and is denoted Vw, (()()), In particular, if E P2(1))) is absolutely continuous with respect to the Lebesgue measure, with density in C' (IR) and such that F(u) < oo, F()(()) = (p) for p-a.e. I Rdd, where F'(u) denotes the first variation of F evaluated at μ, ie. (if it exists) the unique function F'()): Rdd s.t.
lim (F(u+ F(N)) -FWW) / F(()(=)dE(())
(15)
for all & = -- μt, where P2(RR).
Now, we denote by Twp the pushforward measure of P2 by the measurable map T. We recall that the KI. divergence of μu relative to ıt can be decomposed as (T. The distribution /4 of mx in is known to follow a "forward-flow splitting scheme 7 of the Fokker-Plance equation in 2 namely
Pk+ 1/2 [I - Y V w. ((Hx)) Px
(forward discretization for V)
(16)
#k+1 = Pk+1/2 * N (0, Txx
(exact flow for H),
where denotes the identity map in L (1x) and Vw, V(u)) Vz)
17

---

C Proofs from Section 3.1
Proof of Theorem 3.3 Consider the potential
V(n, & (1,))))
|| 112 + ||I v²1 (A (AT,")
min
(17)
where n² is the solution of the constrained sampling problem (P) and o* is the set of solutions of the dual problem (DI). Our goal is show that decreases (in some sense) along trajectories of Algorithm T We say "in some sense" because contrary to the standard LMC algorithm, the distribution μ* of 24 now random measure that depends on the random variables {Ak+ M Explicitly, we consider the filtration F& = oHo. {de, besk))) and show that V decreases on average when conditioned on Fk. This turns out be enough to prove Theorem 3.3
Indeed, notice that P4 E Fk. Hence, the potential energy E(ju, Ae, vkk) E Fi and the conditional law f L(xx (k--) evolves in the regular LMC algorithm (3). That is to say, conditioned on the event Fs. step 5 of Algorithm Aggorthmm[oolll fooloos a splitting scheme as (16), i.e.,
Px+1/2 [T - Th V W2 Ejuss An. [k)] M
(1 (88)
PK+1 PR+1/2 * 2TIa)
(18b)
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 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
((V((k41/2 A 4 ((ix de, vk) 27k KL 11) + m
W (iks
2
+ n
Vigg Au, uk) uk) (p) + Ennpy [lx()))1 Ez-pa [1xx))
(19)
We defer the proof of Lemma C.I to show how it implies the bounds in Theorem 3.3 the expectation of (19) with respect to (Ae, veless obtain
To do take
[[A+] A E VV (px+1/2. ARt1, uk+1)
V(ix-1/2: de, ux)] S
2n EE
KL ll (x*) + m
W.2 (3
2
+ nh
EV,,E((is, Av. uk) 11))ll + Evaur [|(())) + Eampe [1(())]]]. (20)
Using the bounds in Assumption 3.2 then yields
Ax] E V (Ā A8+1, kk+1) (iix-1/22 Ae, ux)l $
KL( || u²)) + m W (1t (1') + (3 E[IA& E[ll (21)
27k E
2
Then, summing the LHS of (1) over k and using the fact that V is non-negative yields
K
(ii1/2, 21, - -B[[(1/22 (22)
Notice that the expectation here taken only over μin given that (Na, (0,0) 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 [ (111,/,,111 A1, S Ror
R₁ W3 n²*) [ [V/1]2 + E E [19911] + E [1]]11]]] + ||**1| 1100112
18

---

Back in 21 we obtain that
E |l p² 21) E(W2 (M p²) S Σ
R
Th 3d +
Σ n ( 11 1x 111)
k=1
We conclude by using the convexity of the Wasserstein distance to write
E
/² 2 (B w²)] W2 (14+1/2,
Similarly for the KL divergence. The bounds in Theorem 3.3 are then obtained for the specific choices of Tis by noticing that Σm= 1/v/ 2 VK and Σ 1/4 < + log(K), Notice that all inequalities in the proof continue to hold for deterministic (Ax, m. 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 sequence of step sizes The such that ([ 1/2, AR: vx)l S R for all k 2 1, where the expectation is taken over the {AR, This immediately implies that W3 (pRR, n²*) I RH and both [[[/>=1]] and R[[|2
[|'] are bounded for all k. We proceed
by induction. The base case covered by Lemma C.2, Suppose now that there exists a se- quence (
such that (1xx1221122 Ak. 1kk
Ri. From the definition of V in (17)
and the fact that the (X", 1/*) are bounded (Prop. 2.2 we then obtain that Elll> III E Eilps |I are
bounded. Consequently, there exists TiK > 0 such that
13 + EII 11 + Elll G < 2m E
m W (Tx, (ix, m*)] KL 11 +
2
From 21.. we obtain that E A < 0, which together with the induction assumption
yields E
[ (Ax+1/22 7++1) S R
Proof of Lemma C.7
The proof proceeds by combining two inequalities bounding the primal and
dual terms in 17).
(i) W2 (P f*). We proceed following a coupling argument. Let s be an optimal coupling between the random variables Y ~ jux and Z ~ n², i.e., a coupling that achieves ( p*). Consider now the random variable T = Y n V Vw,& E Mk) and observe from (8aa) that it is distributed as M Naturally, the coupling s is no longer optimal for (T, 2), so that by the definition of the Wasserstein distance it follows that
W2 (#x+1/2, < fIle nm E Au, vk) yyl2ds (z, v).
(23)
Expanding the RHS of 23 and using the m-strong convexity of E (Assumption 3.1
yields
w? (Px+11/22 1²) S w (ik 2²) - mmm (Fx.
+ 2Th [E(x", " Av. vk) ETirs Ake vb)]
+ n V Ww3&C & ( AEe V/. )ll
We can then bound the effect of the diffusion step using (8 Lemma 5] as in W.2 (Px-1/2, )²)= S 2n [H(x*) - H((Ax)))
W (Pkk
(24)
which yields
W (/4 < W2 (1 (/A-1/2--11) m, (jtk,
+ 27 [E(,", " vx) + H(p²) - E E(i,, (K vk) - H(((w))
(25)
+ n V E AR, vx 1122(pp)))
(ii) || - 1+ + || - n* 1| Notice that since X E RR, the projection x +> [x] is a
contraction, i.c. [[[}| - X < a| X+ for all X E R². Using the definition of AR+1, M441 from
we then obtain that Algorithm 0
||e+1 11 + || v² 2 I 12122 + 2meeAe x²) (T()))) ná [lo(
(26)
+ | - 20 112 + 27p(vr (M))T(((r)))) + ná 11(x=))1
19

---

for all (A", ,²) E o- To proceed, consider the conditional expectation of 26 with respect to Fk. namely,
B B[1 - [= + B[llua+1 - - [=(] - 11 - 2-11 + 11 - v²12
+ (Aee - -2)T Bxwax [o(x)] + (ux- =²) Bewax [h((r)]
2 Th
Bi |19(1))11 + Bazaa [1 (27)
+
where we used the fact that jik, Ak. u E Fk. We conclude by using the linearity of the potential
energy & from TT in both À and v to get
(AA )²) Epppp [g(x)] + (vx= - v²) Expii [h(x)] Efics )k: vb) - E 1,
Back in
27 we obtain
E AR+1
| + EE [| [=[[[[F]]]
S A 1-12² + 1| - -v11
+ 2nd [E (Jis, )E, uk) -E(j,, X w*)]
(28)
+ n Bewiie [19(z)11 + Bonrge [1(()))]].
for all (*" E d².
To proceed with the proof, combine 25 and (8) to get
W2 (#*+1/2, R²) + EE [[ - X [F.] E [|l - v0 [1(5p[a]]] S
W2 (1*=1//;1=*) + | - 2²122 1|ux - v²11
+ 27 E(u" Ak. + H(x²)) - E( X,, 1²) H((i(e)
m. ((x,(())
+ mh Vi Au. Mk) ux) (12= + EEwRy [1)( + Re-aa [+] (29)
To upper bound the term in brackets, we add and subtract log(2) and use the decomposition of the Lagrangiam in terms of [1] to obtain
E(4",, Au. 2k) + H((²) - E(Kx, ,²*) - H((x)) = 1(n",Ax.U1 Ax. L(p 2t, v*).
Using the saddle-pois poidte property (7). we then get
L(p", AR. M/x) - L( A,, v²) < L(m", ,x,, - L( ,,) S - -KL(is 11),
We therefore conclude that
W WPx+1/22 + [[1 - X 12 + Bllle4+1 - -0 [1]
W3 1²) + |1 - - + l|ex- 1z 1| 27T KL(Øx 11=)) + m (Mkks n²)
n VV (ūis, Au; uk (12: + Bee-Ri [1)(
2 + Bznaa l1 (30)
Since (30) holds for all (x*,v*)= E x², it holds in particular for the minimizer of the RHS, for which we can then write (f 1/2, v1) By subsequently taking the minimum of the LHS, we obtain 19
Proof of Lemma
From the updates in Algorithm D we obtain
( \1, mn) W2(0112, n²) + B[[[no](z))] x+ + ((llmah((z)) =*1
for all (X*, x² t² Notice that since A* E RR,, the projection + [x]+ is contraction, i.e..
[[A]+ - X+|| < a| 1+ for all X R Using the triangle inequality then yields
E[V A1, ) S W (21 10) + (2(Ho, 11*) + Blllmg((a)) - X 1= ([llnoh(an) -2(12
< W2 (1o, n²) + || + || 112
+ W33 (M1 Po) n E E (zo) 11 + v [E|| h<a))())))))))) l1.)
To proceed, observe from 18a that /1/2 [L
no V
# Mo. which implies that W:3 (M1/2: Mo) S
n V Villaa
Using the bounds in Assumption
3.2
we obtain that ([V(22/2 )1, S R2.
20

---

Proof of Proposition
3.4
Since 2 0 and Σ
1, we can use Jensen's inequality to write
I
Σ
cx)((T))))
Ex [sp]
Ck
J podjue
f
pody'
Using the relation between the 1 and f2-norm, we further obtain
K
pduux - podut S
K
/ sodlu
Cie
Cs
(31)
If po is bounded by 1, then the summands on the right-hand side of (31 can be bounded by TV μ²). Indeed,
If podpu - /ody"
svldp du'1 / Idpl dp'l 2TV(2;:))))
The total variation distance can in turn be bounded by the KI. divergence using Pinsker's inequality. We therefore obtain
E
(up((x))
E [p]
S
2
c
KL(MalAr) S 22A..
k=1
On the other hand, if op is 1- 1-pischiitz, the summands on the right-hand side of (31 are bounded by
|f poduu somue-Ja
-f ypdyr
I WW (²) S W (2 H*),
which implies that
24
E
Σ ck
(xx((x))
Ep- $
m
as long as m > 0.
21

---

D Proofs from Section 3.2
Proof of Theorem 3.6 The proof is based on the analysis of the stochastic dual ascent algorithm
Ak+1
[>x + m=o(Ex)]++ for E7- ~ Fi such that (L(Hxllex)) S.
(32)
where μλg is the Lagrangian minimizer from 4 and O. Observe that, once again, analyze the conditional distribution fig = [[(x We collect this result in the following proposition:
Proposition D.1. Consider the iterations (32) and assume that Eglgå S G for all μ E P2(RR).
Then, for 0 Uk S n and & n/G², there exists B 0o such that
B ) + 2
7G2
Σ
(33)
K
n
R=0
The expectations are taken over the samples En. Pa-
We conclude by combining Prop. D.I with 29 Theorem 1]. which characterizes the convergence
of the LMC algorithm 33) under Assumption 3.5 Indeed, using the 7k, Ng. from Theorem 3.6
in Algorithm guarantees that the law i of T is such that [T (1 I c, i.e., satishes
the conditions in Prop. D.1 We can then apply Jensen's inequality to get that B[KL(AAlle")
(L(EEūū)]l*** KL( [ ), where pa the law of ng Algorithm
Proof of Prop. D. The proof relies the following lemmata:
Lemma D.2. For all E P2(RR) such that {I(μ))m)) I c, the expected value Ey191 is an approximate subgradients of the dual function in (6) at E RT. ie.,
d(x) 2 d(X) + (A ( Ealgl for all X E RT.
(34)
Lemma D.3. Under the conditions of Prop. D.7 holds for all that
C - D' n2 + and E 1 S B 2BB + +
|A S By
Lemma D.4. The sequence (Ak. (Ex)) obtained from (32 such that
K=1
T/2 2
Σ E [T
(35)
K
where the expectation taken over realizations of Ex-
Before proving these results, let us show how they imply Prop. D.I Start by noticing from that Ai,++1 2 λiA + (Ex)(Ex). Solving the recursion and recalling that Åo 0 then yields
32
K-1
A
Σ 9.(E))
k==
Taking the expected value over & and dividing by njK, we obtain
K-1
[[Ack] Ba
EE x1I +
(36)
K
nK
k=
where the last bound stems from the triangle inequality. Since the upper bound is non-negative for all we use the fact that the maximum of a set of values is less than the sum of those values to write
1 K-
B[lAA
1 + | S
(1+ VTJB
max
(37)
K
k
nK
k=0
where we used Lemma bounds the infcasibility of the ergodic average J Ek=o th for Bi in 32
D.3 together with (1[[|21|[])² < 1 ||||=|| and Bo B. Observe that
37
K-1
22

---

To proceed, use the relation between the normalization of foy and the dual function value [see 6 to decompose the KL divergence between Pa and μ² as
(L((xlll>"> KT | (X* 2x)
En, [g] d(Ax)) d(2²).
Since (L S e and d(A) - d(x²) for all X R4,, we get
KL ||=²) s (X* AR) Emg Lo]
Averaging over & and using Hölder's inequality then yields
K
Σ
KL(Amllu") Se ||^I|]
max i=1.....
1 K
Σ
Ep
loo]
Σ
Epa Ln].
K
K k=
k=0
k=
which from Lemma
D.3
and
37 becomes
(1+ VT)BB
KL Set
Set
x B [g].
K
K
k=0
k=
where we used By < B. Taking the expected value, applying Lemma D.4 and taking VI S 21
for I 2 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 O (with J
0), we obtain L(,, X) (L(a]lux + d(A). Using the fact that KL(ullaA) S e, we get
0 S d(A) L(,, A) +
(38)
Again using the definition of the dual function (6), we also obtain that d(X) S L(, Adding to (38 then gives
a(x)= S +d(A) + Muu X) - [(,, A) +
(39)
Notice from (4) that the first term of the Lagrangians in 339) cancel out, leading to (34
Proof of Lemma D.3
Start by combining the update in 32 and the fact that, since A E R., the
projection
x is contraction, to get
| 21 < A x1 + 2n X²) g(Ex) + n 11((5))1?
Taking the conditional expectation given A then yields
E| 2 Ax1 S A 111 2m(Ax X²) E(g(Ex) Axl + G²²
where used the fact that E, Ez, [g], we can use Lemma D.2 to get
c ||||| 1
< G2 for any μ. Noticing from 32 that Elg(Ex) As]
B|I 11 Xx) S | x112 + 2n [d(As) D* + G2 +c].
(40)
where we used the fact that n/2 77 to simplify the relation.
To proceed, consider the of approximate maximizers of the dual function
D (
d(A) 2 D n -c].
(41)
Notice that th DD. Since there exists at least one X that achieves D* (Prop. 2.2 D not empty. Notice that for D & D, we have that d(Ax)) D + G2 + < O. By the towering property, we therefore obtain from (40p that
- A $ DJ < EIJ - X1I2 Au $ D).
(42)
since 7) 0.
To bound the case when A D we use the strictly feasible candidate n² from Assumption 2.1 Indeed, recall that (L(aa]| & C and E [9s] < -8 < 0 for all i. From the definition of the dual function (6) we obtain that
d()) < L(t x) || 6.
(43)
23

---

where we used the fact that À E RR4 to write E, A 1 Hence, it follows that
|| Bo _ C D' + n + for all X D.
(44)
8
||=]| for all and that o D, immediately obtain that || 2011 v 2B3 for all X D. Using the towering property and the fact that D d(Ax)) and S n2 yields
Using the fact that ||= &
B - 1²1 A DJ s 2B 2n + G²² B².
(45)
To conclude, we write
B1| - - AK11 ~011 A E P|A, E E
1 X $ D] $ D).
Using and 45) then yields
B Ak x-1
B PlAx D] B X-1 A P| DD
max( B², |I|| x A & D)).
Since both 42 and
45 holds independently of AI+1, also write
E A
2 Au D) max(B², |[[|A=11 XM11 A $ DD)
Applying these relations recursively, we obtain that
+| - |+||]] < max B ,11A0 x²133) = max(BB 1|A*
Noticing from 44 that since X D have | < B B then concludes the proof.
Proof of Lemma D.4 To bound 35 we once again use the non-expansiveness of the projection to obtain
||x++ 1 < || + r² [1 + 2n (EEx).
Taking the expectation and using the fact that 1 |lll|l 1 G², we get
A [| s TB[| [ || + G + 2n E g(Ex)]
Applying this relation recursively from K and using the fact that A = (deterministic) yields
K-1
E Ax
1 < K2 + 2n
Σ EAA g(x)))
Since E(|l> (12]
0, we can divide by 2n K to obtain the desired result.
24

---

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 zero (unless stated otherwise) and use different step-sizes for each of the updates steps 3 from Algorithm
We refer to them as Ya, 7A, and M/.. In contrast, we do not use diminishing step-sizes.
E.1 Sampling from convex sets
We are interested in sampling from target distribution
π (2) x E e
- I(x E).
(46)
for some closed, convex set C C R Several methods have been developed to tackle this problem,
based projections [3 64], mirror maps [
15 [6], and barriers [17 [8] Here, consider
constrained sampling approach based on Pl instead.
To do so, note that sampling from
46 is equivalent to sampling from
μd E argmin KL(UI]T)
MEPx(C)
for the unconstrained x e E P2 ". Note that this is exactly PII.. Now let be described by the intersection of the O-sublevel sets of convex functions (s.).e..... iec.
1 t s;(x) 53.
=
Such a description is always possible by, e.g., considering the distance function d(ær, C) = inf ||x- yll2. which is convex (since convex) and for which {a d(x,, C) s 0} (since C closed). Immediately, we see that solving (II)) is equivalent to solving
min (EP2(Rd)
[L(p]TT
(PVI)
subject to
Exzp [[s(()]) < 0, i == 1,...,,,
where [=]+ = max(0, =). Note that PVII has the same form Pl
To why this is the case, consider without loss of generality that i 1. Since [s:(*)]+ 2 0 for all by definition, it immediately holds that
C (a: sx) s 0} = (x [s(x)) s 0) = (r: [s(x)] = 0}.
By the monotonicity of Lebesgue integration, we obtain that the feasibility set of (PVI)
F= {ur P2(IR)) E [[s())]+] s
{N E P2(RR)) EN [[s(x)]+] 0] {pr E P2(RR)) [s(x)]+ = 0, p-a...
(47)
In other words, the feasibility set of (PVI) is in fact F= {u E (A(Rd)
= 11)
To illustrate the of (PVI), consider the one-dimensional truncated Gaussian sampling problem from Section 4 Namely, we wish to sample from a standard Gaussian distribution A/ 1) truncated
to C 1.33. In the language of (PVI). we take fxx = x²/2 (i.e., -= and s;(
Table 2: Mean and variance estimates
True mean
Proj. LMC
Mirror LMC
PD-LMC
ID truncated Gaussian
1.510
1.508
1.470 [0.112, 0.418)
1.488 [0.446, 0.444]
2D truncated Gaussian
00.668, 0.368]
[0.611, 0.610)
25

---

Ergodic constraint slack
Dual variable (A)
5
0.0
0.2
0.4
0.8
1.0
2
Iteration
x10
Iteration
x10
(a)
(b)
Figure 5: the-dimensional truncated Gaussian sampling: (a) Ergodic average of the constraint function (slack) and (b) Evolution dual variable A.
1.60
1.50
Sample average
1.55
Sample average
1.25
1.50
1.00
No 1(3s))
No 10 (44s)
0.75
No 1334
0.50
mean
1.40
1
1
2
PD-LMC iterations
10
LMC evaluations
10
(a)
(b)
Figure 6: The effect of the mini-batch size N₁ on PD-LMC for sampling from a ID truncated Gaussian: Estimated mean vs. (a) iteration and (b) LMC evaluations.
1xx 3). In order to satisfy the assumptions of our convergence guarantees (particularly 2.1 we leave some slack in the constraints by considering (x))+++ 0.005. This also helps with the numerical stability of the algorithm. Fig. [ shows histograms of the samples obtained
using PD-LMC (Algorithm with 7a nl 10 the projected LMC (Proj. LMC, 10-3)
from [13], and mirror LMC (n = 10 from In all cases, we take 5 10 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 iu 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 (potentially infinite) number of realizations outside of C. Yet, though PD-LMC excursions of iterates outside of C become less frequent as the algorithm progresses. We can confirm this is indeed the case in Fig. 5p, which shows the ergodic average of [s(=)]+ along the samples of PD-LMC. Note that it almost vanishes by iteration 10 even though the dual variable À only begins
an interior-pon oatert method, Theorems [3.3] 3.6 show that
to stabilize later (Fig. 5p). This is not surprising given that is guaranteed by Prop. 3.4 In fact, only roughly 2% of the samples displayed in Fig. are not in C.
Before proceeding, we examine whether the convergence of PD-LMC could improved by averaging more than one LMC samples when updating the dual variables, i.e., using mini mini-batches in steps 4 of Algorithh
will reduce the variance of the dual updates, although the cost of
additional LMC steps per iteration. To compensate for this fact, Fig. 6p displays the evolution of the ergodic average PD-LMC samples as function of the number LMC evaluations rather than number of iterations (as in Fig. 64). Notice that, in this application, increasing the number of LMC samples N₂, does not lead faster convergence. This illustrates that, though mini-batches could be useful in some applications (particularly when the constraints
convex, in Section 3.2
26

---

Proj. LMC
PD-LMC
Iter. = le4
Iter. = le5
Iter. 1c4
Iter. lc5
Iter. le6
Iter. 5e6
Iter. leb
Iter. Se6
(b)
Figure 7: Density estimate of two-dimensional truncated Gaussian using samples from (a) Proj. LMC and (b) PD-LMC.
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 large mini-batch N₂ 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.
We now turn to more challenging, two-dimensional applications. We seek sample from Gaussian located at 2, 2] with covariance diag((1, restricted to an f2-norm unit ball (Fig. D Specifically, f(x) ||= /2 (i.c.,
e 1I=11//2) and s,(x) || 1. Once again, we leave some slack to the algorithm by taking the constraint in (PVI) to be [,,[[s; S 0.001. For reference, we also display samples from the real distribution obtained using rejection sampling.
This is indeed challenging problem. The boundary of 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 an acceptance rate of approximately 1%. The strong push of the potential f towards the exterior of C leads Proj. LMC (n) = 10-3; the last 10 samples are used after running 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. 0.14% the samples should be according to rejection sampling. Indeed, note from Fig. Zh. that even iterations advance, the samples of Proj. LMC continue concentrate close to the boundary.
proposal for rejection sampling yields
10 iterations) to
In contrast, PD-LMC (n, = 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. 1). Indeed, the average of the constraint function along samples from PD-LMC essentially vanishes around iteration 5 x 104, Achieving this requires larger values À (on the order of 250, Fig. Kp) compared to the one-dimensional case (Fig. 5p). This reflects the difficulty of constraining samples to C in this instance, statement formalized in the perturbation results of Prop. 2.2 iv). Due to the more amenable numerical properties of the barrier function, minor LMC (1 10 1) performs well without concentrating samples on the boundary (0.15% of the samples on the annular region radius [0.999, 1)).
27

---

250
Ergodic constraint slack
0.3
Dual variable (A)
150
0.2
100
0.1
50
0.0
0.0
0.2
0.4
0.6
0.8
1.0
Iteration
x10
Iteration
x10
(a)
(b)
Figure 8: Two-dimensional truncated Gaussian sampling: (a) Ergodic average of the constraint function (slack) and (b) Evolution dual variable X.
X10
Popularion
1.5
Dual variable (A)
Unconstrained (Ace:: 84%)
1.0
Femall
Male
Constrained (Acc.: 82%)
20
30
0.0
0.5
1.0
1.5
2.0
Prevalence of S50k (%)
Iteration
x10*
(a)
(b)
Figure 9: Fair Bayesian logistic regression on the Adult dataset: (a) prevalence of positive outputs and (b) dual variables.
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 upplications in robust control (76 and to express other requirements 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; 0) = Ply 1|z, denote the probability of a positive outcome (y = 1) given the observed features E X' and the parameters 0 distributed according to the posterior This posterior is determined, e.g., by some arbitrary Bayesian model based on observations {(r, }/) Hence, [q:=(q(x;; 0)] denotes the likelihood of positive outcome for Consider group, represented by measurable subset G X, for which we wish to enforce statistical parity. In other words, would like the prevalence of positive outcomes be roughly the same that of the whole population. Thus, we now want to sample not from the posterior и,, but from a close-by distribution of parameters 0 that ensures this parity. Explicitly, for some tolerance 8 >> O, want to sample from
protected
μ¹* E min (EPa(Rd)
KL(H|(=)
(PVII)
subject to
Ex, (x;; 0) G > Ex, [g(x; 0) -
Naturally, we account for more than protected by incorporating additional constraints.
In our experiments, we take be the posterior of Bayesian logistic regression model for the Adult dataset from [60] (details on data rrepprocessing can be found in 61]). The N 32561 data points the training set are composed
= 62 socio-economical features (a E Rdd including the intercept) and the goal is to predict whether the individual makes more than US$ 50000 per year (y (0, 1)). The posterior obtained by combining binomial likeeeelddood with independent zero-mean Gaussian (log)priors (o² 3) each parameter of the model, i.c., consider the
28

---

LMC iterations (Na)
(3
(28 s)
100
(251
1000
(2464 s)
10
20
30
Prevalence of S50k (%)
Figure 10: The effect of the mini-batch size N₂ on PD-LMC in fair Bayesian classification.
potential
((3) = Σ log(1 + (2yn 1)2 T
d
03
(48)
2x
We begin by using the LMC algorithm from (i.c., we impose no constraints) to collect samples of the coefficients 0 from this posterior
10-1; the last 10 samples are used after running 2 x 10 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 the distribution of this probability over the unconstrained posterior π (Fig. 9h), 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, consider gender to be the protected class in (PVII., constraining both Gremate and Gmate- We formulate the constraint of (PVII) using an empirical distribution induced from the data. Explicitly, consider constraints
Ean [g(2n; 8)]
2
E
Ea~ [g(
Gremale
N
ne Grenalee
n=1
N
Σ n nEGual
Eann (((2((())))) 6)]
2
Ea-p (g(2x))) -
1
N
where Gremalo- Gmale C ,l,...,,N partition the data set. For these experiments, we take ô = 0.01. 10 10-3), we then obtain a new of samples from the logistic
Using PD-LMC (n) 10-4. 17 regression parameters 0 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 jı" for male and female (Fig. h). Additionally, this substantial improvement over the previously observed disparity comes at only 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, Amale 0 throughout the execution of PD-LMC (Fig. (2))
Before proceeding. we once again examine the effect of using multiple LMC samples to update the dual variables, i.c., using mini-batches in steps 455 of Algorithh I Fig. [h]shows the distribution of the prevalence of positive predictions (> $50k) for different mini-batch sizes N₁. In all cases, we collect 2 x 10tt samples, which means that we evaluate 2N, x 10 LMC updates (step 3 in Algorithm l}. Same as the ID 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₂ was used in the experiments of [4 to overcome the challenge of computing an expectation their dual variable updates. turns out that this computationally intensive modification is not necessary.
29

---

2.0
Unconstrained Constrained
0.01
Odds ratio
1.5
Tolerance (6)
1.0
0.02
0.5
0.0
Intercept
Bachelors
Adnislorical
African Aneer
White
Gonder
Age (3.0. 37.01
Iran
0.05
10
20
Prevalence > (%)
(a)
(b)
Figure 11: Counterfactual sampling in fair Bayesian logistic regression: (a) Selected (mean) coeffi-
cients and (b) different tolerances.
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 set of moment conditions. It therefore interested not only in μ²,, but in how each condition affects the value P = KL(p" [T)). We next describe how constrained sampling
be used tackle this problem.
Let denote a reference probabilistic model, such as the posterior of the Bayesian logistic model in (88.. Standard Bayesian hypothesis tests can be used to evaluate the validity of actual statements such as "is true that (g<=( S 0? or Pz [x))] = 0?" Hence, we could check "is a more likely to yield positive output for a male than a female individual?" (from the distributions under Unconstrained in Fig. Fig. this is probably the ease). In contrast, counterfactual sampling concerned with counterfactual statements such as "how would the world have been if [=]] < 0?" In the case of fairness, "how would the model have been 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 differ- ence between realizations of iT, obtained using the traditional LMC in (3). and μt" obtained using PD-LMC (Algorithm T). In Fig. TTp 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" each counterfactual condition (con- straint) world with the reference model. While AlgoriAAm I dooes not evaluate /* explicitly, it provides measures of its sensitivity to perturbations of the constraints; the Lagrange multipliers (X*, .v"). In- deed, recall from Prop. 2.2 that
e Mig
II
1
x
Hence, (X* )*) describe the magnitude of filts needed for the reference model π to satisfy the conditions [x)( & 0 [x)( = 0. This relation is made explicit in Prop. 2
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
Σ
Eonyi [a 0)] 2 N
Σ n=1
Egrn [9(=n; -
1mmate|
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 (Aeeaaee 160). Without recalculating n², we therefore know that even small changes in the tolerance ò for the female constraint would substantially change the
30

---

20
LMC
PD LMC
0.6
15
(Log)retum
0.4
Volatility
0.2
10
1N3
AAPL.
G00G
BRK-8
L.v
0.0
-0.2
INJ
Figure 12: Counterfactual sampling of the stock market under a 20% average return increase on each
stock: mean (p) and variance [diag(E) [E)] distributions.
TSLA
JNJ
GOOG
LLY
NVDA
AAPI.
BRK-B
Ergodic constraint slacks
0.00
Dual variable
-0.01
TSLA
AAPL GOOG
BRK-B LLY
0.02
NVDA
-10
-0.03
0
1
2
5
lteration
x10
Teeration
100
(a)
(b)
Counterfactual sampling of the stock market under 20% average return increase on each stock: (a) ergodic average of constraint functions (slacks) and (b) dual variables.
Figure 13:
distribution of outcomes. This statement is confirmed Fig. IIp. Notice that this is only possible due to the primal-dual nature of PD-LMC.
E.3.1 Stock market model
Counterfactual analyses based the dual variables become powerful the number of con- straints 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 assets (TSLA, NVDA, JNJ, AAPL, GDOG, BRK-B, and LLY). The dataset is composed of years of adjusted closing prices for total of 1260 points per asset. The posterior obtained by combining Gaussian likelihood (p,,, 2) with Gaussian prior on the mean p (O 3/)] and an inverse Wishart prior on the covariance Σ (with parameters / and V 12). Using the LMC algorithm (1) 10 3, the last 3 x 10S samples are used after running 6 x 10t iterations), we collect samples from this posterior and estimate the mean and variance of the (log- -return for each stock (Table In this case, Σ is initialized to 10
We might now be interested in understanding what the market would look like if all stocks were to incur 20% increase in their average (log-)returns. To do so,
PD-LMC (12 10-3
and 7) = 6 x 10-3) to solve the following constrained sampling problem
minimize (R()
(L(u]]T)
(PVIII
subject to E(p 1.2p: i I,
where is the mean (log- return of the i-th stock shown in Table 3 The distribution of p and the diagonal Σ are compared to those from the unconstrained model in Fig. [2] Notice that, though we only impose constraints the average returns p. also see small changes in the stock volatilities.
Inspecting the dual variables (Fig. 13p), we notice that three dual variables are essentially zero (TSLA, AAPL, and BRK-B). This means that their increased returns completely dictated by those of other
31

---

20
stocks
NVDA.
0.6
4
15
(Log)retum
0.4
Volatility
2
0.2
10
INd
AAPL
GO0G
BRK-B
L.Y
00.2
AAPL
Figure 14: Distribution of mean (p) and variance [diag( (E) of the stock market constrained to a 20% average return increase all stocks only LLY and NVDA.
Table 3: Mean + standard deviation of mean (log- )returnsttss (p).
Reference model (T)
+20% (all stocks)
20% (LLY and NVDA)
TSLA
0.20 0.12
0.24 = 0.12
0,23 = 0.12
NVDA
0.25 + 0.09
0.31 + 0.09
0.31 + 0.09
JNJ
0.02 + 0.07
0.02 + 0.07
0.03 + 0.07
AAPL
0.11 + 0.06
0.14 + 0.06
0.14 + 0.06
GOOG
0.09 0.06
0.11 = 0.06
0.11 0.06
BRK-B BB
0.06 0.07
0.07 + 0.07
0.07 = 0.07
LLY
0.16 + 0.06
0.19 + 0.06
0.19 0.06
stocks. Said differently, their returns increasing 20% is consistent with the reference model ditioned on the other returns increasing. Proceeding, two stocks have negative dual variables (LLY and NVDA). This implies that bringing their constraints down Pi would yield decrease in P* (dis- tance to the reference model т). This is in contrast to JNJ and GOOG, whose positive X's imply that we should should increase their returns reduce p*, Indeed, by inspecting the ergodic slacks (Fig. 13 we see that all stocks approach zero (i.e., feasibility), but that JNJ and GOOG do so from above. This behavior expected according 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. [4). Second, that the increase we would in JNJ (and to lesser extent GODG) would actually be larger than 20%. Once again, reach these conclusion without any additional computation. Their accuracy can be corroborated by the results in Table
32

---

Example application of Prop. 2.2
In this section we illustrate the result in Prop. 2.2 i.e., we show that given solutions (X* v*) of (DI. the constrained sampling problem reduces sampling from pA & e-Uti-A".p
Indeed, consider a standard Gaussian target, i.c., x e e-||==1//2 and the linear moment con- straint [B|æc] = 6, for E R This can be cast as (PI) with f(x) || and h(x) b - r (no
inequality constraints, i.e., I 0). Clearly, the solution of (PI) this case is μ [b,, I), i.e.,
Gaussian distribution with mean b. What Prop 2.2 claims is that rather than directly solving (P) we
can solve (DI) to obtain Lagrange multiplier v such that p pa. for Pu defined as in
In this setting, we can see this is the case by doing the computations explicitly. Indeed, we have
||01 2
(x) xc कt (x)e
[((a) exp
v (b x)
Completing the squares, we then obtain
| u11
|011
(x) x exp
T
(49)
To compute the Lagrange multiplier v², notice from the definition the dual function in (6) that the dual problem (DI) is in fact ratio of normalizing factors. Explicitly,
||a|||
( (
exp
dz
1/ argmaxx log DERd
argmax log PERd
eex([|u|13/2 - TTb) J exp
||2- 1/1
do
2
Immediately, we obtain
argmax - 1||u|12/2
(50)
MERn
Note that the dual problem a concave program, as is always the case combine (49 and (50) get
30 To conclude, we can
|
b12
|0|1
Py (x)
exp
- My N(b,)
=m=b
The main advantage of using Algorithms| 1 and is that we do not need to determine the Lagrange multipliers (X*, v²" then sample from #x*=** n²" Indeed, Theorems 3.3
3.3 and 3.6 show that these
stochastic primal-dual methods do both things simultaneously, without explicitly evaluating any expectations.
33

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification: the abstract and introduction list contributions, i.e. the theoretical and experi- mental 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 is clear that these goals are not attained by the paper.
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 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 these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
violations of
The authors should reflect on the scope of the claims made, c.g., if the approach was only tested few datasets with few runs. In general, empirical results often depend implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, 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 grounds for rejection, worse outcome might be that reviewers discover limitations that aren't it acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play important role in developing norms that preserve the integrity of the community. Reviewers will 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 complete (and correct) proof?
34

---

Answer:
[Yes]
Justification: For each theoretical result, worked with justified assumptions, making all the dependencies of the problem clear. We provide clear and detailed proofs 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 encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided 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 experi- mental 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, No answer this question will not be perceived well by the reviewers: Making the paper reproducible important, regardless of whether the code and data are provided not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary either make it possible for others replicate the model with the same dataset, access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility results, access to a hosted model (e.g., in the case of a large language model), releasing of model checkpoint, other means that appropriate the research performed.
provide
also be provided via detailed instructions for how replicate the
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 new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily architecture clearly and fully.
model architecture, the paper should describe the
(c) If the contribution is 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 open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may tricky in some cases, in which case authors are welcome 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., registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
35

---

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?
Question:
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 guides/Codes CooodeSummssiionPPP..
for more details.
While encourage the release of code and data, 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 the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to reproduce the results. See the NeurIPS code and data submission guidelines (https: //tiss.///ippeeeeeeeeeeeeeee cttppsscccppbciicc Staides//CodeSummsssse ssiiissiccccccy 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 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 data and code is permitted.
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 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 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.
Experiment Statistical Significance
Question: Does the paper report eтror 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 emorr 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).
36

---

The method for calculating the error bars should be explained (closed form formula, call to 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 I-sigma error bars, but 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 reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
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 reproduce the experi- ments?
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 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 't make into the paper).
Code Of Ethies
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: http:://pp..p........./.cc.c..aa
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 NeurlPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there a special consideration due to laws regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal 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.
the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address socictal impact.
37

---

Examples of negative societal impacts include potential malicious unintended uses (e.g.. disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploy- ment 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 par- ticular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point out. For example, an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
legitimate to point out that
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 intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
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
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g.. pretrained language models, image generators, or scraped datasets)?
Question:
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 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 challenging, and many papers do not require this, but we encourage authors to take this into account and make 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 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 used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should 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, wwwpppppppphaeee.c thcode. com/ datasets has curated licenses for some datascts. Their licensing guide can help determine the license of a dataset.
38

---

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 Python that will be made public. Guidelines:
The answer NA means that the paper does not release assets.
Researchers should communicate the details of the dataset/code/moodl as part of their sub- missions 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
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Question:
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 the paper involves human subjects, then much detail possible should be included in the main paper.
supplemental material is fine, but if the main contribution of
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should 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 Sub- jects
Does the paper deseribe 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 requirements of your country or institution) were obtained?
Question:
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 NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applica- ble), such as the institution conducting the review.
39


Reviews:
[{"review_id": "eqZHxUWvmi", "reviewer": "Reviewer_d9K2", "summary": "In 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.", "questions": "1. 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$?", "limitations": "Not quite, it would be nice if the authors could comment on some limitations of their analysis / methods.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "### 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.", "weaknesses": "The 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive comments on our paper.\n\n**Weakness** The reviewer has a point that to obtain a practical algorithm we perform stochastic updates in both the primal (step 3) and dual (steps 4-5), which complicates things. This in fact marks an important distinction with the analyses of the traditional LMC or when the dual updates are performed using exact expectations as in [22], in that we now have to handle the full joint distribution $(x,\\lambda,\\nu)$. To overcome this challenge, we work with conditional laws to obtain \"in expectation\" guarantees. That is why the bounds in Theorem 3.3 are not stochastic.\n\nIndeed, note that (13) comes from the inequality (line 617, Appendix C)\n$$\\frac{1}{K} \\sum\\_{k = 1}^K \\mathbb{E}\\_{\\lambda,\\nu} \\Big[\\text{KL}(\\tilde{\\mu}\\_{k} \\| \\mu^\\star) + \\frac{m}{2} W\\_2^2(\\tilde{\\mu}\\_k,\\mu^\\star)\\Big] \\leq\\frac{R\\_0^2}{\\eta K}+ \\eta G^2 \\bigg( 3 +\\sum\\_{k = 1}^K \\frac{\\mathbb{E}\\_{\\lambda}[\\|\\lambda\\_k\\|^2]+ \\mathbb{E}\\_{\\nu}[\\|\\nu\\_k\\|^2]}{K}\\bigg),$$\nwhere $\\tilde{\\mu}\\_k$ is the conditional law of $x\\_k \\vert \\\\{\\lambda\\_{\\ell}, \\nu\\_{\\ell}\\\\}\\_{\\ell<k}$. Our proof analyzes the evolution of this conditional law [as in (18)] and then apply Jensen's inequality using the fact that $\\mathbb{E}\\_{\\lambda,\\nu} [\\tilde{\\mu}\\_k] = \\mu\\_k$, the law of $x\\_k$ (line 618). Due to space constraints, we did not include this technical discussion in the main text but we plan to do so in the final version.\n\n**Q1** Note that (DI) involves $\\mu$ through its objective, the dual function $d$ from (7). But the reviewer is correct that, in contrast to (PI) that directly seeks $\\mu^\\star$, the dual problem (DI) seeks $(\\lambda^\\star,\\nu^\\star)$ that can then be used to obtain $\\mu^\\star$ by $\\mu\\_{\\lambda^\\star\\nu^\\star}$ defined in (5)--(6) (line 134). In fact, finding $\\mu^\\star$ and $(\\lambda^\\star,\\nu^\\star)$ are equivalent [Prop. 2.2(iv)]. Hence, the goal of the $x$-update is in fact twofold: sampling from $\\mu\\_{\\lambda^\\star\\nu^\\star}$, i.e., $\\mu^\\star$, and compute the $(\\lambda^\\star,\\nu^\\star)$ [using the discrete versions of (11b)--(11c) in steps 4-5 of Algorithm 1].\n\n**Q2** The reviewer is correct. Computationally, (9) is intractable; (10) proposed by [22] is better, but still intractable (due to the expectations); but Algorithm 1 is practical and can be implemented without approximations. From the point-of-view of convergence, (9) performs \"dual ascent\" and would converge to $(\\lambda^\\star,\\nu^\\star)$; the continuous-time counterpart of (10) was shown to converge in [22] (we extend these results in Theorem 3.3 by showing the discrete version itself converges under milder conditions on $g$, as we discuss in line 247); and we prove that Algorithm 1 also converges (again, under milder assumptions).\n\nThe reviewer's proposal would indeed yield Algorithm 2 (with an extra Metropolis acceptance step). But while this approach has advantages in the non-convex case (Section 3.2), its benefits do not necessarily outweigh the increased computational cost. For instance, consider the sample mean estimate of the one-dimensional truncated Gaussian from Fig. 1 (main paper) as a function of the number of LMC iterations (step 3) taken in Algorithm 1 (see Fig. 1 in the pdf attached to the global response). I.e., for $N=1$ we have Algorithm 1 and for $N > 1$ we follow the reviewer's proposal. Note that the additional computational cost does not improve convergence in this case (especially when considering the number of \"LMC steps\" $N$ which is proportional to the computational complexity).\n\n**Q3** This is hard to guarantee in general, even in the Euclidean convex case (see, e.g., [35]). As we mention in the paper (line 235), a common solution is to clip the iterates $(\\lambda\\_k,\\nu\\_k)$ in Algorithm 1 by some upper bound on $(\\lambda^\\star,\\nu^\\star)$. Such bounds are well-known and can be obtained under a variety of constraint qualifications such as Assumption 2.1, although explicit bounds for equality constraints require additional assumptions (see, e.g., [26,35] or the more general [Gauvin, \"A necessary and sufficient regularity condition to have bounded multipliers in nonconvex programming,\" 1977]). Alternatively, a decreasing \"weight decay\" regularization can be used (see, e.g., [62]). In some cases, it is even possible to directly bound the iterates $\\lambda\\_k$ (see Lemma D.3). For these reasons, we chose to present our results in the form of Theorem 3.3, which we find to be more applicable, and then comment on potential solutions after the theorem. We will expand our remark on this point in the camera-ready version.\n\n**Q4** The reviewer is correct: there is a typo in step 4 of Algorithm 2. We omit the equality constraints in the non-convex case to keep the analysis manageable. Deriving results similar to Lemma D.3 and D.4 for equality constraints requires additional assumptions (see response to Q3) and substantial derivations that we felt obscured the results. We note that it is sometimes possible to cast equality constraints as two tight inequalities, although this might lead to numerical issues. We will expand on this remark in the camera-ready version.\n\n**Q5** Let $f(x) = x^2/2$ ($1$-strongly convex) and $g(x) = \\sin(x)$ (bounded). We therefore have $\\mu\\_\\lambda \\propto e^{-x^2/2-\\lambda\\sin(x)}$. For $\\lambda = 0$, Assumption 3.4 holds with $\\sigma = 1$. As $\\lambda$ increases, $\\sigma$ decreases but $\\sigma > 0$ for all finite $\\lambda$. In fact, $\\sigma \\geq e^{-2\\lambda}$ (see, e.g., [42, Prop. 5.1.6] or Theorem 1.1 in [Cattiaux and Guillin, \"Functional inequalities for perturbed measures with applications to log-concave measures and to some Bayesian problems,\" 2022]). Since $\\lambda^\\star$ is finite (see bound in Lemma D.3), we could explicitly clip the iterates $\\lambda\\_k$ and get an *a priori* bound on $\\sigma$. But even without any explicit limit, we show in Theorem 3.5 that $\\mathbb{E} \\\\| \\lambda\\_k \\\\|\\_1$ is bounded for all $k$.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "sj63XBM4Td", "reviewer": "Reviewer_4nfu", "summary": "The 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.", "questions": "1. 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.", "limitations": "There are no limitations.", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The 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.", "weaknesses": "There 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their enthusiastic opinion of our paper. We next address their questions one-by-one.\n\n**Weakness** The reviewer has a point. We focused on regular Langevin dynamics and its literature, but this paper is indeed part of the large line of work of sampling as an optimization in the space of probability measures. It is therefore connected (or could be extended) to many other schemes. We appreciate the suggested references and will include them in the manuscript.\n\n**Q1.1** Since our main focus was on deriving and studying a computationally-friendly algorithm for the constrained sampling problem (PI), we did not consider proximal updates that are more expensive than the fully explicit Euler updates used in this paper (for both primal and dual variables). However, we acknowledge that investigating proximal steps is an interesting future direction and we believe that it could in fact lead to stronger theoretical guarantees (under additional assumptions) as is the case for gradient descent-ascent in Euclidean space (see, e.g., [40]). We will make note of this point in the revised version of the paper.\n\n**Q1.2** To a large extent, our approach is akin to that used to study any optimization algorithm/dynamical system: we construct a Lyapunov function to bound the duality gap $L(\\mu,\\lambda^\\star,\\nu^\\star) - L(\\mu^\\star, \\lambda, \\nu)$ [note from the saddle-point relation (8) that this is the right optimality measure]. The techniques used to bound this gap in Wasserstein space, however, are different.\n    \nIndeed, while the dual (DI) remains a (finite dimensional, non-smooth) optimization problem in Euclidean space, the primal (PI) is an (infinite dimensional, smooth) optimization problem in Wasserstein space. The proof must therefore account for the mixed nature of the Lagrangian. Additionally, Algorithm 1 performs stochastic updates in both the primal (step 3) and dual (steps 4-5). Hence, the potential $U$ used to update $x$ in step 3 is a random variable and we have to deal with the joint distribution $(x\\_k,\\lambda\\_k,\\nu\\_k)$. This is an important distinction with the traditional LMC or when the dual updates (steps 4-5) are performed using exact expectations as in [22]. To address these issues, we instead use conditional laws to obtain our guarantees in expectation over realizations of $(\\lambda\\_k,\\nu\\_k)$.\n\n**Q2** Consider a standard Gaussian target, i.e., $\\pi \\propto e^{-\\\\|x\\\\|^2/2}$, and the linear moment constraint $\\mathbb{E}[x] = b$, for $b \\in \\mathbb{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 $\\mu^\\star = \\mathcal{N}(b,I)$. What Prop 2.2 claims is that rather than directly solving (PI), we can solve (DI) to obtain a Lagrange multiplier $\\nu^\\star$ such that $\\mu^\\star = \\mu\\_{\\nu^\\star}$ for $\\mu\\_{\\nu}$ defined as in (5) (line 134).\n\nWe can show this indeed the case here by doing computations explicitly. Indeed, we have\n$$\n\\mu^\\star(x) = \\mu\\_{\\nu^\\star}(x) \\propto \\pi(x) e^{-(\\nu^\\star)^\\top h(x)}\n    = \\exp\\big[ -\\\\|x\\\\|^2/2 -(\\nu^\\star)^\\top (b-x) \\big].\n$$\nCompleting the squares we obtain\n$$\n\\mu^\\star(x) \\propto \\exp\\big[ -\\\\|x - \\nu^\\star\\\\|^2/2 + \\\\|\\nu^\\star\\\\|^2/2 -(\\nu^\\star)^\\top b \\big].\n$$\nFrom the definition of the dual function in (7) (line 139), we can write (DI) explicitly as a ratio of normalizing factors, namely\n$$\n\\nu^\\star\n    = \\text{argmax}\\_{\\nu}\\ \\log\\Big( \\frac{\\int \\pi(x)dx}{\\int \\mu\\_\\nu(x) dx} \\Big)\n    = \\text{argmax}\\_{\\nu}\\ \\log\\Big( \\frac{\\int \\pi(x)dx}{\\exp(\\\\|\\nu\\\\|^2/2 -\\nu^\\top b) \\int \\pi(x) dx} \\Big).\n$$\nImmediately, we obtain\n$$\n\\nu^\\star = \\text{argmax}\\_{\\nu}\\ -\\\\|\\nu\\\\|^2/2 +\\nu^\\top b = b.\n$$\nNote that, as we mention in the text, the dual problem is indeed a concave program. To conclude, observe that as we indeed have\n$$\n\\mu^\\star(x) = \\mu\\_{\\nu^\\star}(x) \\Bigm\\vert\\_{\\nu^\\star = b} \\propto \\exp\\big[ -\\\\|x - b\\\\|^2/2 - \\\\|b\\\\|^2/2 \\big]\n\\text{, i.e., $\\mu^\\star(x) = \\mathcal{N}(b,I)$.}\n$$\n\nOur main results (Theorems 3.3 and 3.5) show that it is not necessary to first determine the Lagrange multipliers $(\\lambda^\\star,\\nu^\\star)$ to then sample from $\\mu\\_{\\lambda^\\star\\nu^\\star} = \\mu^\\star$. Indeed, the stochastic primal-dual Langevin methods in Algorithms 1--2 simultaneously determine $(\\lambda^\\star,\\nu^\\star)$ and sample from $\\mu\\_{\\lambda^\\star\\nu^\\star} = \\mu^\\star$. Additionally, they do so without explicitly evaluating any expectations.\n\nPlease do not hesitate to let us know if something is not clear. If this explanation indeed clarifies some of our duality derivations, we consider including it in the appendix as a warm-up example.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Tk64TrBE2u", "reviewer": "Reviewer_ZMkL", "summary": "This 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.", "questions": "Please see weakness part.", "limitations": "The limitations are pointed out in the paper.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The 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.", "weaknesses": "1. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive comments on our results, particularly our theoretical guarantees. As the reviewer correctly notes, in contrast to [22] ([1] in the reviewer's comments), our analysis handles two levels of approximation: time- and space-discretization. We next address their questions.\n\n**Q1** Indeed, as we clearly acknowledge in the paper, Algorithms 1 and 2 can be seen as approximations of the PD Langevin from [22]. But there are substantial computational and theoretical differences between these methods.\n\nThe dual variable updates in the algorithms from [22] require computing an expectation over the current particle distribution [see (10)], which is intractable. While we agree that it can be *approximated* by running many particles in parallel, this would result in a substantially higher computational cost than the stochastic updates in our paper. Indeed, we show PD-LMC (Algorithm 1) converges using a single-particle approximation.\n\nFurthermore, it is not clear from [22] how these approximation errors impact the convergence of the algorithms since it only provides guarantees for exact dual updates (computing expectations). In contrast, our theoretical analyses exactly characterize the bias introduced by approximation errors (Proposition D.1) and show that they do not affect the convergence in the convex setting (Theorem 3.3). This is confirmed by our experiments. These results hold for discrete-time, expectation-free algorithms.\n\nThat being said, we acknowledge that we could have better illustrated the relative performance of PD Langevin and our stochastic version PD-LMC. We have addressed this oversight and provided an illustration of the performance of both schemes as the number of LMC particles $N$ grows for the one-dimensional truncated Gaussian and the rate-constrained Bayesian model (see Fig. 1 and 2 in the pdf attached to the global response). This hopefully illustrates that one can achieve essentially the same sampling performance as PD Langevin at a considerably lower computational cost.\n\n**Q2** The reviewer has a point: our results do not *directly* analyze the decrease rate of the constraint violation. Yet, our theoretical guarantees show that PD-LMC converges to $\\mu^\\star$ that is feasible by definition. Hence, Theorems 3.3 and 3.5 do address the issue of feasibility. Additionally, it is possible to use our results to explicitly control the constraints violation along iterations.\n\nConsider the results in Theorem 3.3, namely (13):\n$$\\frac{1}{K} \\sum\\_{k = 1}^K \\mathrm{KL}(\\mu\\_{k} \\| \\mu^\\star)+ \\frac{m}{2} W\\_2^2(\\mu\\_k,\\mu^\\star)\\leq\\frac{R\\_0^2}{\\eta K}+ \\eta G^2 \\bigg( 3 +\\sum\\_{k = 1}^K \\frac{\\mathbb{E}\\_{\\lambda}[\\|\\lambda\\_k\\|^2]+ \\mathbb{E}\\_{\\nu}[\\|\\nu\\_k\\|^2]}{K}\\bigg) \\triangleq \\Delta\\_K.$$\nSince $x\\_k \\sim \\mu\\_k$ and $\\mathbb{E}\\_{\\mu^\\star}[g] \\leq 0$, we obtain that\n$$\\mathbb{E} \\bigg[ \\frac{1}{K} \\sum\\_{k=1}^K g(x\\_k) \\bigg]\\leq \\frac{1}{K} \\sum\\_{k=1}^K \\int g d\\mu\\_k - \\int g d\\mu^\\star\\leq \\frac{1}{K} \\sum\\_{k=1}^K \\Big| \\int g d\\mu\\_k - \\int g d\\mu^\\star \\Big|\\leq \\sqrt{\\frac{1}{K} \\sum\\_{k=1}^K \\Big| \\int g d\\mu\\_k - \\int g d\\mu^\\star \\Big|^2},$$\nwhere we used the $\\ell\\_1/\\ell\\_2$-norm relation.\n\nIf $g$ is bounded by 1, the summands are bounded by $TV(\\mu\\_k,\\mu^\\star)^2$. Using Pinsker's inequality, we therefore obtain\n$$\\mathbb{E} \\bigg[ \\frac{1}{K} \\sum\\_{k=1}^K g(x\\_k) \\bigg]\\leq \\sqrt{\\frac{1}{2K} \\sum\\_{k=1}^K \\mathrm{KL}(\\mu\\_{k} \\| \\mu^\\star)^2}\\leq \\sqrt{\\frac{\\Delta\\_K}{2}}$$\nIf $g$ is $1$-Lipschitz (and $f$ is $m$-strongly convex), we have\n$$\\Big| \\int g d\\mu\\_k - \\int g d\\mu^\\star \\Big|^2\\le W\\_1^2(\\mu\\_k,\\mu^\\star) \\le W\\_2^2(\\mu\\_k, \\mu^\\star),$$\nwhich implies\n$$\\mathbb{E} \\bigg[ \\frac{1}{K} \\sum\\_{k=1}^K g(x\\_k) \\bigg]\\leq \\sqrt{\\frac{2 \\Delta\\_K}{m}}.$$\nThe same arguments hold for $h$.\n\nHence, our results can be used to control (in expectation) the constraint violation of ergodic averages along the PD-LMC trajectories. These behaviors are also observed in our experiments. For instance, Fig. 3 of the pdf attached to the global response shows the ergodic constraint slacks $\\sum h(x\\_k)/K$ for the equality constraints of the Bayesian stock market problem (Fig. 2 in the paper). Notice that by the end of the simulation their values are within $3 \\times 10^{-3}$ of zero. We thank the reviewer for raising this point and will include these corollaries and discussions in the camera-ready version.\n\n**Q3** We respectfully, but strongly disagree with this statement. We believe there may have been a misunderstanding arising from the fact that the constraints in (PI) are *distribution constraints* rather than *constraints on the support of $\\mu$*. We refer the reviewer to Section 2.2 for a detailed discussion of the differences between these constraint types. In our case, imposing a moment constraint (even as an equality constraint) does not result in $\\mu^\\star$ being supported on a lower-dimensional manifold. Another way to see this is by noticing that $\\mu^\\star \\propto \\pi \\exp^{-\\nu^\\star h}$ [see Prop. 2.2(iv)]. Hence, for $\\pi$ full-dimensional and $(\\nu^\\star, h)$ finite [which is the case, see Prop. 2.2(ii)], $\\mu^\\star$ will also be full-dimensional. For a concrete example, we refer the reviewer to the response to *Q2 of reviewer 4nfu*, where we derive an explicit solution for (DI) under an equality constraint.\n\n**Q4** We have included results for the Mirrored Langevin in the two-dimensional truncated Gaussian as well as the PD Langevin from [22] for the rate-constrained Bayesian problem (Figs. 4 and 2 respectively in the pdf attached to the global response). We note again that the latter requires an explicit integration that is intractable and that using LMC chains to approximate it reduces to the settings of Algorithms 1-2, whose convergence is guaranteed by our results (Theorems 3.3 and 3.5). As expected, we therefore obtain the same results albeit at a considerable increase in computational complexity.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the editors and reviewers for their time and for the positive comments on our paper.

In the sequel, we respond to each of their questions individually. Throughout our responses, we refer to references and equations as numbered in the submitted version of the manuscript. We also refer to figures contained in the pdf attached to this response as "pdf attached to the global response."

Meta-review: The focus of this paper is on a constrained optimization problem in the space of measures. The goal is to obtain 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 the authors design the PD-LMC method to generate an approximate sample. Theoretical properties of PD-LMC are studied. Numerical experiments showcase the workings of this method.

All reviewers were positive about this work (scores 5, 5, and 7), indicating that the strengths outweigh the weaknesses. Some of the strengths of the paper include:

- The paper is well written and reads well 
- The objective of the research and the background help the reader grasp the problem and the proposed solution
- The dual formulation central to this work is also adequately presented
- The idea seems to be novel, and looks like a clean way to adapt ideas from Euclidean optimization for the distributional constrained problem considered here
- The authors establish solid convergence guarantee for their algorithms in various settings

I fully expect the authors to address all the observed weaknesses in the camera-ready version of the paper. On this condition, based on the above reasons, I propose the paper to be accepted.

AC
Decision: Accept (poster)