Constrained Sampling with Primal-Dual Langevin Monte Carlo

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.

Paper

Similar papers

Peer review

Reviewer ZMkL5/10 · confidence 3/52024-07-08

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.

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. 2. 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. 3. 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. 4. 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. [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.

Questions

Please see weakness part.

Rating

5

Confidence

3

Soundness

2

Presentation

3

Contribution

2

Limitations

The limitations are pointed out in the paper.

Reviewer 4nfu7/10 · confidence 5/52024-07-08

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.

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. Sampling from a convex set is a very good application.

Weaknesses

There is a lack of literature on generalized optimization problems in Wasserstein spaces. Wang et al. Accelerated Information Gradient flow. Journal of Scientific Computing. 2022. Wang et al. Information Newton's flow: second-order optimization method in probability space, 2020. Tan et al. Noise-Free Sampling Algorithms via Regularized Wasserstein Proximals, 2023.

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

Rating

7

Confidence

5

Soundness

3

Presentation

4

Contribution

3

Limitations

There are no limitations.

Reviewer d9K25/10 · confidence 4/52024-07-11

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.

Strengths

### Quality and Clarity The 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. ### Originality and Significance The 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.

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. 2. 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? 3. 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? 4. Why are equality constraints disregarded for PD-LMC with LSI potentials? Since $\nu$ doesn't exist, why does it show up in Algorithm 2? 5. 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$?

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

Not quite, it would be nice if the authors could comment on some limitations of their analysis / methods.

Reviewer ZMkL2024-08-08

Thanks for the feedback. Most of my concerns are addressed and I'm happy to raise the score. The rate of violation of constraints is an important proposition that should be added to the main text. I would like to mention that the proposed PD LMC algorithm is the same with PD Langevin in [22] when particle number $N=1$, while the theoretical analysis in this paper is solid and novel.

Reviewer 4nfu2024-08-08

Reply to authors

The authors carefully address my questions. I also like the example of Gaussian distributions. I suggest the publication of this paper.

Reviewer d9K22024-08-10

Thank you for your rebuttal

Q3: I feel like this is an important point to cover. Specifically if the quantity increases faster than $K$, then the analysis would yield a vacuous bound. In your method, 1) how would you obtain a bound, 2) how would you enforce a bound on these variables, and 3) how would this affect the guarantees for the method? I have no further questions aside from the ones above.

Authorsrebuttal2024-08-12

We agree with the reviewer that these are important points and address them in the discussion after Theorem 3.3 in the current version of the manuscript (lines 235-244). There, we argue for the form of Theorem 3.3 rather than modifying Algorithm 1 because Theorem 3.3 also shows there exists a sequence of step sizes such that $(\lambda\_k,\nu\_k)$ are bounded. Additionally, none of our experiments enforce bounds (they implement PD-LMC exactly as described in Algorithm 1). Yet, we always observe the dual variables converging (see, e.g., Fig. 2 in the manuscript), suggesting that pathological cases are not common. That being said, the theory would not be complete without addressing these cases, which we discuss in the remarks of lines 235-244. We expand our treatment of these points, including the discussion below, in the revised version of the manuscript. If any point remains unclear, we are happy to provide further details. To the reviewer's points: 1) **Bounds on $(\lambda^\star,\nu^\star)$**: Under Assumption 2.1, we provide such a bound on $\\|\lambda^\star\\|\_1$ in Lemma D.3. Explicitly, suppose there exists a $\mu^\dagger$ such that $\mathrm{KL}(\mu^\dagger \\| \pi) \leq C$, $\mathbb{E}\_{\mu^\dagger} [g\_i] \leq -\delta < 0$, and $\mathbb{E}\_{\mu^\dagger} [h\_j] = 0$ (Assumption 2.1). Then, by definition $$ D^\star = d(\lambda^\star,\nu^\star) \leq \mathrm{KL}(\mu^\dagger \\| \pi) + \sum\_i \lambda^\star\_i \mathbb{E}\_{\mu^\dagger} [g\_i] + \sum\_j \nu^\star\_j \mathbb{E}\_{\mu^\dagger} [h\_j] \leq C - \\|\lambda^\star\\|\_1 \delta $$ for the dual function in (7). By strong duality, $D^\star = P^\star \geq 0$, which yields $$ \\|\lambda^\star\\|\_1 \leq \frac{C}{\delta}. $$ A bound on $\nu^\star$, though more complicated, can be obtained under a similar constraint qualification (see, e.g., the general [Gauvin, "A necessary and sufficient regularity condition to have bounded multipliers in nonconvex programming," 1977]). 2) **Enforcing the bound**: Suppose $\\|\lambda^\star\\|\_1,\\|\nu^\star\\|\_1 \leq B$. Suffices it to replace steps 4 and 5 in Algorithm 1 by $$ \lambda\_{k+1} = \Big[ \lambda\_k + \eta\_k g(x\_k) \Big]\_0^{B} \quad \text{and} \quad \nu\_{k+1} = \Big[ \nu\_k + \eta\_k h(x\_k) \Big]\_{-B}^{B}, $$ where $[z]\_{L}^U = \min(\max(z,L),U)$. This is a common solution used in the Euclidean case (see, e.g., [35]). 3) **How does this affect the guarantees?** It does not. Indeed, note that proof of Theorem 3.3 is based on the fact that (17) is a Lyapunov function. This is proved in Lemma C.1. When bounding the Euclidean norm of the dual variables (line 642), the first step is realizing that the projection $\lambda \mapsto [\lambda]\_0^B$ and $\nu \mapsto [\nu]\_{-B}^B$ are contractions, since $\lambda^\star\_i \in [0,B]$ and $\nu^\star\_j \in [-B,B]$. Hence, we can use $$ \\| [\lambda]\_0^B - \lambda^\star \\|^2 \leq \\| \lambda - \lambda^\star \\|^2 \quad \text{and} \quad \\| [\nu]\_{-B}^B - \nu^\star \\|^2 \leq \\| \nu - \nu^\star \\|^2 $$ to obtain (26). The proof then proceed as is and yields the exact same results as in Theorem 3.3, except that (13) [and similarly (14)] can now be rewritten as $$ \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 ( 3 + B^2 ). $$

Reviewer d9K22024-08-12

Thank you for the further commentary

I have no further questions, and would like to maintain my score. I think that parts 2 and 3 from above would benefit the exposition.

Area Chair oRkn2024-08-13

Please comment on the relation of your work to ...

Please comment on the relation of your work to ... Salim and Richtarik, **Primal Dual Interpretation of the Proximal Stochastic Gradient Langevin Algorithm**, NeurIPS 2020, https://papers.nips.cc/paper/2020/hash/2779fda014fbadb761f67dd708c1325e-Abstract.html This seems relevant due to the primal-dual interpretation of Langevin described here. Moreover, the work directly addresses the issue of constrained sampling. This work is not cited, but seems very relevant. Is it not? If either case, I believe, this should be clearly explained here to the reviewers and in the paper. Apologies for posting this so late; I hope you can still respond. Reviewers: What do you think? Best regards, AC

Reviewer d9K22024-08-13

My understanding of the relation

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

Authorsrebuttal2024-08-13

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.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC