Constrained Sampling with Primal-Dual Langevin Monte Carlo

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.

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.

The 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.

Furthermore, 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.

That 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.

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.

Consider the results in Theorem 3.3, namely (13): $\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.$ Since $x\_k \sim \mu\_k$ and $\mathbb{E}\_{\mu^\star}[g] \leq 0$, we obtain that $\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},$ where we used the $\ell\_1/\ell\_2$-norm relation.

If $g$ is bounded by 1, the summands are bounded by $TV(\mu\_k,\mu^\star)^2$. Using Pinsker's inequality, we therefore obtain $\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}}$ If $g$ is $1$-Lipschitz (and $f$ is $m$-strongly convex), we have $\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),$ which implies $\mathbb{E} \bigg[ \frac{1}{K} \sum\_{k=1}^K g(x\_k) \bigg]\leq \sqrt{\frac{2 \Delta\_K}{m}}.$ The same arguments hold for $h$.

Hence, 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.

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.

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.

We thank the reviewer for their enthusiastic opinion of our paper. We next address their questions one-by-one.

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.

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.

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.

Indeed, 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)$.

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).

We can show this indeed the case here by doing computations explicitly. Indeed, we have

Completing the squares we obtain

From the definition of the dual function in (7) (line 139), we can write (DI) explicitly as a ratio of normalizing factors, namely

Immediately, we obtain

Note that, as we mention in the text, the dual problem is indeed a concave program. To conclude, observe that as we indeed have

Our 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.

Please 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.

We thank the reviewer for their positive comments on our paper.

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.

Indeed, note that (13) comes from the inequality (line 617, Appendix C) $\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),$ where $\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.

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].

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).

The 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).

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.

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.

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$.

I feel that the constrained setting in the submission deviates from the traditional constrained sampling setup in that the constraints are expressed as functionals, and the dual formulation extends the formulation in the Euclidean space to space of probability measures. The above paper appears more relevant to when the potential to sample from is a composite function (smooth and non-smooth parts). But I'm also interested to hear about the authors' thoughts (hopefully they can respond).

This manuscript and [Salim and Richtarik] use primal-dual methods in distinct ways. As reviewer d9K2 points out, they tackle the problem of sampling from composite potentials that can contain a non-smooth part restricting the support of the distribution. In contrast, we study (PI) which imposes statistical (moment) constraints on the distribution (see Table 1 and Section 2.2 for a detailed discussion on the differences between these constraint types). While (PI) can be used to tackle support constraints (as we show in Section 2.2 and illustrate in our experiments), it can also handle statistical restrictions for which proximal operators cannot be constructed (see examples in Section 2.2). Though our focus on computational aspects means we did not consider proximal updates in this paper, we do believe that studying proximal versions of PD-LMC is an interesting future direction (see more details in our response to reviewer 4nfu).

Additionally, [Salim and Richtarik] consider stochastic "primal" updates [see (3) in their paper] as in stochastic gradient Langevin [33]. In constrast, PD-LMC (Algorithm 1) uses stochastic gradients only for its dual updates (steps 4-5). Indeed, note from (6) that $U$ is a deterministic function that depends only on parameters of the sampling problem (namely, $f,g,h$). Hence, its evaluation does not involve computing any expectation. Combining stochastic gradient Langevin with PD-LMC is also an interesting research direction but outside the scope of this work (see comment line 198).

That being said, as with other papers mentioned in these reviews, [Salim and Richtarik] is certainly relevant related literature and we will definitely include these clarifications in the revised version of the manuscript.