Strictly Constrained Generative Modeling via Split Augmented Langevin Sampling
Authors: Matthieu Blanke
Yongquan Qu
Sara Shamekh
Pierre Gentine
Venue: ICLR 2026 Poster (2026)
Keywords: Langevin sampling
generative models
constrained sampling
duality
Lagrangian
data assimilation
optimal control
diffusion models
inverse problems
optimal control
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: We propose a novel constrained sampling algorithm for physically-constrained deep generative modeling

Abstract: Deep generative models hold great promise for representing complex physical systems, but their deployment is currently limited by the lack of guarantees on the physical plausibility of the generated outputs. Ensuring that known physical constraints are enforced is therefore critical when applying generative models to scientific and engineering problems. We address this limitation by developing a principled framework for sampling from a target distribution while rigorously satisfying physical constraints. Leveraging the variational formulation of Langevin dynamics, we propose Split Augmented Langevin (CASAL), a novel primal-dual sampling algorithm that enforces constraints progressively through variable splitting, with convergence guarantees. While the method is developed theoretically for Langevin dynamics, we demonstrate its effective applicability to diffusion models. We apply our method to diffusion-based data assimilation on a complex physical system, where enforcing physical constraints substantially improves both forecast accuracy and the preservation of critical conserved quantities. We also demonstrate the potential of CASAL for challenging feasibility problems in optimal control.

Published as a conference paper at ICLR 2026
STRICTLY CONSTRAINED GENERATIVE MODELING VIA SPLIT AUGMENTED LANGEVIN SAMPLING
Matthieu Blanke"
Yongquan Qu Columbia University LEAP NSF STC
Sara Shamekh New York University LEAP NSF STC
Pierre Gentine
New York University
Columbia University LEAP NSF STC
LEAP NSF STC
ABSTRACT
Deep generative models hold great promise for representing complex physical systems, but their deployment is currently limited by the lack of guarantees on the physical plausibility of the generated outputs. Ensuring that known physical constraints are enforced is therefore critical when applying generative models scientific and engineering problems. We address this limitation by developing a principled framework for sampling from a target distribution while rigorously satisfying mathematical constraints. Leveraging the variational formulation of Langevin dynamics and Lagrangian duality, we propose Constrained Alternated Split Augmented Langevin (CASAL), novel primal-dual sampling algorithm that enforces constraints progressively through variable splitting. We analyze our algorithm in Wasserstein space and derive explicit mixing time rates. While the method is developed theoretically for Langevin dynamics, we demonstrate its ap- plicability to diffusion models. We apply our method diffusion-based data as- similation on complex physical system, where enforcing physical constraints substantially improves both forecast accuracy and the preservation of critical con- served quantities. We also demonstrate the potential of CASAL for challenging non-convex feasibility problems optimal control.
INTRODUCTION
Generative deep learning methods have recently emerged as powerful tools to model and sample from complex data distributions, with successful applications in image synthesis (Ho et al., 2020), protein and material design (Corso et al., 2023), and probabilistic forecasting (Price et al., 2025; Morel al., 2025; Rozet et al., 2025). By learning a stochastic process from a training dataset, these models can generate arbitrarily many plausible samples conditioned partial information. They are particularly useful the physical sciences, where data often scarce and multiple states may be consistent with available observations (Epstein & Fleming, 1971; Nathaniel & Gentine, 2025). While perceptual applications mainly aim for plausibility, scientific and engineering prob- lems require samples that obey strict physical or structural constraints, such conservation laws or system dynamics (Kashinath et al., 2021). In such cases, approximate resemblance is not enough: generated samples must obey the governing physical principles. This requirement becomes even more critical when generative models are used out-of-distribution or in an autoregressive fashion, where small violations can accumulate and severely degrade long-term accuracy (Pedersen et al., 2025). Developing constrained sampling methods applicable to pre-trained generative models in zero-shot scenario (i.e., without additional training) therefore crucial.
Modern generative models, including energy- score-based, and diffusion models (Du & Mordatch, 2019; Song et al., 2020), typically rely on Langevin dynamics, where noisy gradient steps push the samples toward high-likelihood regions. Enforcing mathematical constraints during Langevin sampling remains a challenging problem. A natural idea to project each iterate onto the constraint set, leading to projected Langevin dynamics (Bubeck et al., 2015; Durmus et al., 2019; Christopher et al., 2024). While these methods offer theoretical guarantees in convex settings, they tend to perform poorly when applied to non-convex constraints, which are common in physical
"Correspondance Matthieu Blanke <matthieu war.iie..bllleeee.e bla.ke@nyu..dd.

---

Published as a conference paper at ICLR 2026
systems. In such cases, strict projections can cause the dynamics become trapped in limited gions of the constraint set, hindering exploration and introducing sigmificant sampling bias. Other approaches use soft constraint penalty functions such the barrier method (Fishman et al., 2023) and diffusion guidance (Ho & Salimans, 2022; Meunier et al., 2025), requiring a differentiable con- straint model. These methods encourage but do not enforce constraints, which is insufficient when strict satisfaction is crucial. To our knowledge, no existing approach achieves both strict constraint satisfaction and unbiased exploration.
Contributions We introduce a rigorous mathematical formalism to analyze the constrained sam- pling problem. Inspired by the variational formulation of Langevin dynamics and constrained op- timization, we propose Constrained Alternated Split Augmented Langevin (CASAL), novel sam- pling algorithm that bridges the gap between complex generative modeling and constrained sampling using variable splitting. Our method enforces hard constraints while preserving the exploration pability of Langevin dynamics. It ensures strict constraint satisfaction and benefits from convergence guarantees via duality analysis. We show that our approach generalizes deep generative modeling and diffusion models. We demonstrate the effectiveness of CASAL on complex physically con- strained sampling tasks, including data assimilation problems where maintaining physical invariants is key to reliable forecasting, and on non-convex feasibility problems in optimal control.
2 PROBLEM FORMULATION OF CONSTRAINED LANGEVIN SAMPLING
In this section, we provide mathematical formulation of constrained sampling: given a genera- tive model and a constraint set, our goal is generate samples from the conditional distribution supported on the constraint set. Such constrained distributions arise in many applications where samples must strictly satisfy known physical laws. We adopt the framework of the Langevin Monte Carlo algorithm (Rossky et al., 1978). a foundation of modern generative modeling frameworks. The application to deep generative models is discussed in Section 3.3.
Langevin Monte Carlo Consider target distribution with density p(x) = e-(()) on Rd where f=x is a differentiable potential. Markov chain Monte Carlo methods design iterative algo- rithms producing samples (ax) whose distribution a converges to p. Among them, the Langevin Monte Carlo algorithm plays a central role. It requires access to the gradient of the potential fz(),, also called the score function (Hyvärinen & Dayan, 2005), and performs noisy gradient descent updates with step size
T++1 F2 f(x))
V2Twe,
lid NO Ia).
(2.1)
Under standard assumptions, the chain converges to p, the target distribution (Durmus et al., 2019).
Constrained target distribution We now consider the case where the samples are known to sat- isfy hard constraints sampling time, in the form of a bounded measurable set C C Rd, which models prior information such as physical conservation laws. The conditional density supported on C is
pc(x)= = e f(x) (0=5(())k Ic(x), Væ E Rd
(2.2)
Ze
with 1 the indicator function of and Ze normalizing constant. Note that the conditional distri- bution (2.2) can be rewritten using modified potential: po(x) == e ff(((=) / with the constrained potential fe(x) == f(x) + Xc(x). defined with the characteristic function
if E
xc(x) == +00
(2.3)
otherwise.
We do make any assumption on the constraint set C, except that it is bounded and that pe is well- defined. Next, we provide examples of such constraints that may occur in physical applications.
Example 2.1 [Physical constraints] When describes discretized physical field, conservation of energy E can often be expressed as the non-convex set = {o E Rd ||>] = E}, while mass conservation corresponds to C {x E Rd
Σ3 2. M) for prescribed mass M.
2

---

Published as a conference paper at ICLR 2026
Objective Our objective is to design sampling algorithm that produces samples distributed ac- cording to pe for any constraint set C. It should rely on the score function Vf(x) of the uncon- strained density, and mathematical operations related to C such as constraint functions, projection operators. The method should operate in a "zero-shot" scenario, requiring no retraining or additional data.
Example 2.2 [Projected Langevin] A natural idea enforce hard constraints is to project each unconstrained update (2.1) onto C with the projection operator Pe(x) == argmin || - ==|l², as
ZEC
2++1 -
V2TW)), We 1td.
(O,
(2.4)
This projected Langevin algorithm, and its extension diffusion models, enjoy strong theoreti- cal guarantees when C is convex and is log-concave (Bubeck et al., 2015). However, with non- convex constraints, repeated projection can trap the dynamics in small feasible regions, biasing ex- ploration (Barber & Ha, 2018; Ahn & Chewi, 2021). This motivates the need for sampling methods that enforce constraints more gradually.
Example 2.3 [Soft penalty methods] Constraints can also be enforced softly by adding differen- tiable cost c(x) 2 0 to the potential, penalizing samples far from C, with tunable coefficient À E R:
ht+1 = 24 - T(Vf(x))) + AVe(x.)) + V2TWe,
iid. La). ~
(2.5)
This corresponds to guidance in diffusion models (Ho & Salimans, 2022; Huang et al., 2024c).
The cost function c(x) corresponds to a negative tog-likelihood centered on the constraint set. Such methods encourage constraint satisfaction but do not guarantee it, violations only smoothly pe- nalized.
Evaluation Assessing the performance of constrained sampling algorithms is difficult as Pe is generally intractable. In practice, rely three key performance criteria: constraint violation, sampling bias, and computational cost. Constraint violation measures the deviation samples from C. Even when samples lie within C. they must accurately follow the conditional distribution pe without bias; this is typically quantified by comparing sample statistics to known or approximated quantities under pc- Finally, the computational cost evaluated based the number of function evaluations or total runtime required to generate valid samples.
3 SPLIT-AUGMENTED LANGEVIN FOR STRICTLY CONSTRAINED SAMPLING
In this section, we introduce rigorous framework for constrained generation. We first derive variational formulation of the target distribution Pe and, using Lagrangian duality, demonstrate why standard dual methods fail to sample from it. Building on this insight, propose Constrained Alternated Split Augmented Langevin (CASAL), a primal-dual sampling algorithm that samples from Pc while ensuring strict constraint satisfaction through a variable splitting scheme.
3.1 VARIATIONAL FORMULATION OF CONSTRAINED SAMPLING
To better understand the constrained sampling problem, we formulate it an optimization problem in the space of probability measures, leveraging the connection between Langevin dynamics and optimization Wasserstein space. We let P2 (IR:) denote the set of probability measures on Rd with finite second moments. For a measure E (Te) admitting density, we identify the measure with its density function. Further details on this framework are provided in Appendix A.
Variational view of sampling Langevin Monte Carlo can be viewed as an optimization algorithm in distribution space. Let
D(allp) / log dp
dq
dg
(3.1)
the Kullback-Leibler divergence (Kullback & Leibler, 1951). It is non-negative information- theoretic quantity measuring how differs from p. Crucially, can be shown that Langevin dynamics describe the gradient flow minimizing the functional q D(gll)) in the Wasserstein
3

---

Published as a conference paper at ICLR 2026
space (2 (Jordan al., 1998; Villani, 2021). Consequently, the discrete-ting time Langevin algo- rithm (2.1) acts as a stochastic partiele approximation of gradient descent in P2(R2), driving the distribution 9t of the chain toward the minimizer p (Bernton, 2018). Formally, Langevin sampling
solves:
minimize qE
D(alp))-
(3.2)
Our approach extends this variational perspective to the constrained setting by characterizing Pe as the solution to a constrained optimization problem.
Proposition 3.1 [Information projection] Suppose that PP(()) 1. Then the conditional distribution Pe is the projection of p onto the set of distributions supported on C:
PC argmin D(qll))
(EP2(Ed)
(3.3)
subject to =C)=
This is special case of I-projection (Csiszár, 1975). To solve it, one might try to apply Lagrangian duality on (3.3). However, we show in the following result that strong duality is not attained, imply- ing that duality-basee numerical methods would fail to converge pe.
Proposition 3.2 Strong duality is not attained for (3.3).
We also show that penalty methods of Example 2.3 approximate (3.3) with a finite penalty parameter. Using Proposition 3.2, we prove that such methods cannot sample from pc-
Corollary Penalty methods (2.5) cannot enforce PC,(C)
This singularity arises because the support of the target distribution a strict subset of Rd, which prevents the constraint qualification from being satisfied. A possible relaxation to allow small violation probability (C) 2 1-8 for small > O, but this allows unphysical states and leads to poor conditioning. To overcome this, we introduce a splitting strategy that relaxes the coupling between the sample and the constraint while maintaining strict feasibility.
3.2 CONSTRAINED ALTERNATED SPLIT AUGMENTED LANGEVIN
Directly targeting Pc is challenging because the potential forces exploration in Rd while the con- straint forces the measure onto lower-dimensional manifold. To address this composite objective, we propose to split the variable x into pair (x, 2) E Rd x C. enforcing that a E C while encourag- ing and to remain close. We thus define joint probability measure g(z, 2), with marginals and 9².
Proposition 3.3 [Variable splitting] Problem (3.3) equivalent to the following problem:
minimize (RP2())))))
D(q [pp) + Egy [xc(=)]
(3.4)
subject to
Pg (x = 1.
This formulation mirrors variable splitting techniques in optimization (Gabay & Mercier, 1976), and separates the roles of x and 2 C, which are respectively maximizing likelihood and enforcing the constraint. Rather than requiring æ almost surely. relax the condition to be satisfied expectation, and penalize the variance. Specifically, we define a penalty parameter p > 0 and consider the following problem and its associated Lagrangian.
D + E [xc(=)] + LE, [|| ==1
minimize (R xRd)
subject to Ex - z] = 0.
Definition
For a primal-dual pair (q. A) P2(RR 4Rd x Rd)) x Red, the Lagrangian of (P) is
A(g.A) = D(g [p) Eg bxc [=)] + + A 2] + f f || =112).
(3.5)
Unlike the original projection Problem (3.3), this relaxed formulation admits a qualified constraint, ensuring that strong duality holds.
4

---

Published as a conference paper at ICLR 2026
Proposition 3.4 [Strong duality] Strong duality holds and is attained for (P). In particular, the Lagrangian (3.5) admits a saddle point (4c: A.), and q, is a solution of (P).
Proposition 3.4 ensures that a solution of the relaxed problem (P) can be found by searching for a saddle point of (3.5). The primal density q E P2 x Rd) is approximated with particles.
Stochastic saddle point iterations We propose stochastic primal-dual scheme solve the sad- dle point problem for (3.5). We generalize the method of Chamon et al. (2024) to our non-smooth, split setting. Given step size and noise te NO, Ta), the iterations are:
I++1 It fVx( =p(x; 24 + mc) V
(3.6a)
=++ = Pc(2:-T((z, - - 1++1 (t))
(3.6b)
Pe+1 M4 + (T/P) 2141),
(3.6c)
where / == (1/p) x À is rescaled dual variable. This scheme, summarized in Algorithm 1, is named Constrained Alternated Split Augmented Langevin (CASAL). The primal variables x and approximate the Wasserstein gradients of (3.5) with a gradient step and proximal step respectively. In particular, x follows Langevin dynamics biased toward z, while is projected onto C. The dual variable μ integrates the error to correct bias, following a stochastic approximation of the dual ascent step.
Algorithm 1 Constrained Alternated Split Augmented Langevin (CASAL)
input potential gradient Vf, projection Pe, step size > 0, coupling p O, iteration number T, initial distribution qu
output sample 2r E C
initialize xp ~ qo: 2p Po(xo), n4 E TRd
for 0 < f < T
do
draw wt ~ NO La)
T++1 = Ht - TV fV)(()) -T((x - 21 + +Me + /2Twe
2++1 Po(a Tp(2) I1+1 Me))
14411 = ft + (f/p) x (x++1 - 2+11)
end for
Connection with optimization algorithms The update formulas (3.6) resemble the Alternating
Direction Method of Multipliers (ADMM) of Glowinski & Marroco (1975) and Gabay & Mercier (1976), widely used in constrained optimization (Boyd et al., 2011). Here, the variables æ and play the role of the primal variables in ADMM and the dual. Our sampling scheme can be seen stochastic analog of the proximal ADMM (He et al., 2002) in sample space I& just like Langevin Monte Carlo parallels gradient descent. However, it differs from ADMM applied in distribution space P2 (IRd,, as our method operates directly on coupled samples.
3.3 PRACTICAL IMPLEMENTATION AND DEEP GENERATIVE MODELS
Implementation in diffusion models Our proposed algorithm is constrained variant of Langevin Monte Carlo, which plays a central role in many generative frameworks (Du & Mordatch, 2019; Song Ermon, 2019). The spli--augmented updates (3.6) can be used as a drop-in replacement for standard Langevin steps, without altering other sampler components, making constraint enforce- ment simple and modular. Leveraging the connection between Langevin dynamies and diffusion models (Ho et al., 2020), CASAL provides a training- free constrained sampling algorithm for pre- trained diffusion models. This parallel has already been exploited by Christopher et al. (2024) to introduce projected diffusion models. Details are discussed in Appendix B.
Constraint satisfaction Our algorithm applies to arbitrary constraint sets, provided that projec- tion operator (exact or approximate) is available. Unlike penalty and latent projected methods, does not require differentiable constraint model, which can be challenging to derive (Laumond, 1987). The coupling parameter p can be tuned or progressively increased along the diffusion pro- cess. This is detailed with ablation studies in Appendix C.
5

---

Published as a conference paper at ICLR 2026
Latent space constraints A significant challenge arises when sampling occurs in latent space R", while the constraint set C is defined in a distinct physical space R linked by a de- coder A R The target density then takes the form:
pc(x) = T
I Ic(A),, Vo E R.
(3.7)
Enforcing constraints in this setting is significantly harder, mapping projection penalty gra- dient from R back to I implicitly requires inverting A, which is often computationally expensive unstable due to ill-conditioning. In contrast, our framework naturally accommodates this setting by modifying the splitting constraint to Ar = 2, instead lows CASAL to perform the projection step solely in the physical space, without the need to invert the decoder or define a projection the latent space.
z. Crucially, this al-
Computational cost For learning methods to accelerate large-scale physical simulations, effi- ciency is central. Compared to unconstrained diffusion, our method adds the cost projection operation at each step. does projected diffusion. For non-convex constraints, efficient numerical methods such as augmented Lagrangian algorithms can be used to solve the projection step, and are amenable to parallelization (Boyd et al., 2011; Liang et al., 2025). In the of latent diffusion, our splitting technique allows the constraint operations to be computed directly in physical space, which can be considerably faster than propagating them through decoder in projected diffusion or diffusion guidance. More details can be found in Appendix C.
4 NON-ASYMPTOTIC CONVERGENCE ANALYSIS
We now provide theoretical analysis of CASAL. Our study is twofold: first, we quantify the re- laxation error induced by approximating the strict projection with the variable splitting problem (P). Second, we analyze the algorithmic convergence the discrete-time process (3.6) toward the solu- tion of the relaxed problem. Proofs are deferred to Appendix D.
4.1 PROPERTIES OF THE RELAXED SOLUTION
Let (g., A.) denote saddle point of the Lagrangian (3.5). We first characterize the optimum q. Proposition 4.1 The optimal distribution o.' takes the form
(x) x exp (-((z)) - P
2 p. (w))
(4.1)
where 14 (1/p) x X, and dc(2) minyec ||y =|| is the distance to the constraint set.
This result illustrates the mechanism of CASAL: the dual variable μs, shifts the effective potential to center the distribution correctly relative to the constraint, while controls the tightness of the constraint enforcement. In the limit of infinite penalty, we recover the exact conditional distribution.
Proposition 4.2 [Recovery of the projection] Let gp denote the solution (P). Then
Tp 9
Pe in distribution.
(4.2)
Thus, larger values of p bring the samples closer to C, while smaller values encourage exploration. For finite p. the following bounds quantify the approximation error relative to Pc-
Proposition 4.3 [Relaxation error] The relaxed solution q. satisfies the following conditions:
D(gfll)) < D(Pell), (C) 1, W2 g.) 1 D
(4.3)
In practice, P is finite, and Proposition 4.1 suggests bias for finite p. The role of the Lagrange multiplier X, is to correct this bias, and center the density towards the right mode.
Proposition 4.4 [Consistency] In the case of a degenerate Gaussian potential f and an affine constraint set the solution of (P) is unbiased: Eq. [x] Ey Ep [z].
This result confirms the importance of using duality. As it is classical in optimization, duality allows for convergence to the constrained optimum with finite penalty parameter (Bertsekas, 2014).

---

Published as a conference paper at ICLR 2026
4.2 CONVERGENCE AND MIXING RATE
We now analyze the convergence of CASAL to the stationary distribution q.. We rely on the follow- ing standard assumptions.
Assumptions The constraint set C is convex and bounded. The potential f is a-convex and 3- smooth. Furthermore, there exists M > such that for all 4, || - Re + 1421 S M.
We define a Lyapunov function measuring convergence in primal and dual spaces. Definition [l yapunov function] For a primal-dual pair (q, A) E P2(RR4 x Rd) x R24 we define the following Lyapunov function with respect reference pair (p, P2 Rd:
1 Vig. P. = : (f W pd) + 2
A 211
(4.4)
The following proposition establishes discrete evolution variational inequality (Pang & Stewart, 2008) for the algorithm. In particular, gives a condition on the algorithm's hyperparameters and to ensure decrease in the Lyapunov function.
Proposition 4.5 [Primal-dual Lyapunov decrement] Let 4+ be the distribution after iterations (3.6) with step size т. Then, provided 4r + 4p) S 1, for any reference pair (p, v)):
V(g, p; v) (g A :p: + W (+) 2 7 (LLg-. L(g., A)) C+? (4.5)
Summing this decrement over T iterations yields a convergence rate for the time-averaged distribu- tion, matching the standard rates for Langevin Monte Carlo (Durmus et al., 2019).
Corollary 2 [Mixing rate] Let (zx) be generated by Algorithm with step size = 1/vt+I. Let d. denote the distribution of the time-averaged iterate H, + E Ls. Then,
D(7: < O Int
(4.6)
APPLICATION TO PHYSICALLY CONSTRAINED GENERATIVE MODELING
We evaluate CASAL on three scientific generative modeling tasks where challenging non-convex physical constraints play critical role: stationary distribution sampling, data assimilation and opti- mal control. We apply our algorithm to diffusion models as described in Section 3.3. More details given in Appendix C. Although it is not the central contribution of experimental work, also test CASAL on high-dimensional partial differential equations in Appendix C.5.
Baselines Our sampling algorithm is compared with the unconstrained Langevin algorithm, the projected Langevin algorithm, constraint penalty guidance methods, and their deep diffusion model analogs (Carvalho et al., 2023; Huang et al., 2024a; Christopher et al., 2024; Zampini et al., 2025). All methods share the same score function, and differ only in how constraints are incorporated.
5.1 ENERGY-CONSTRAINED STATIONARY FIELD GENERATION
We first validate our method constrained Monte Carlo sampling of station- ary distribution, which is a critical problem in climate science and in molecular dynamics for example (Paquet & Viktor, 2015; Pedersen al., 2025). We con- sider a two-dimensional field, representing for instance a fluid (see Figure 1), discretized on 100 100 grid. The equilibrium distribution p is sampled using Langevin dynamics.
Figure 1: Sampled field snapshot.
Constraint A key macroscopic quantity is the kinetic energy, which often remains conserved or decreases in physical prediction tasks. Our task is sample from the conditional distribution under a fixed energy C = {ar E Rd | = E}. a non-convex constraint.
7

---

Published as a conference paper at ICLR 2026
Experimental setup The unconstrained distribution p bimodal in Fourier space, with asymmetric modes on the first Fourier coefficient: one positive and concentrated, the other negative and wider, allow- ing higher energy. The unconstrained distribution sampled with the Langevin Monte Carlo algorithm, and Pe is estimated via rejection sampling. The bimodal nature of makes the exploration challenging. We condition on high energy level, only achievable via the negative mode. As the positive mode cannot satisfy the energy constraint, the correct conditional distribution concentrates on the negative mode, and we can easily compare it to the generated samples. For each method, 1000 independent chains are run and the last iterate collected. We compute histograms of the first Fourier coefficient for evaluation.
target
unconstrained
penally CASAL
projected
-3
0
3
mode
Figure 2: Empirical his- tograms for the first mode.
Results Our results are shown in Figure 2. Only CASAL matches Pc closely. Projected Langevin satisfies the constraint exactly but fails to explore, yielding many samples in the wrong mode. Soft constraint penalty (Zhang et al., 2025) enforces energy conservation only on average, and there- fore the produced samples do not match PC- These results demonstrate that CASAL enforces hard constraints while retaining enough exploration to correctly sample the conditional distribution.
5.2 PHYSICS-PRESERVIE PRESERVING DATA ASSIMILATION
Data assimilation, central problem in geophysics, aims to estimate the state of dynamical sys- tem from sparse, noisy observations using prior knowledge. Recent work applies deep generative architectures to this task (Rozet & Louppe, 2023; Qu et al., 2024), but these models do not enforce physical invariants, such as energy mass conservation, which are essential for physical plausibility long-tew ang-term forecasting. We study physically constrained generative models for data assimilation on the inviscid Burgers equation, reduction of the Navier-Stokes equations with conserved mass and energy that exhibits rich dynamics and complex multiscale behaviors similar to turbulence (van Gastelen et al., 2024). A latent diffusion model is trained offline on dataset of trajectories in Fourier space, without any conditioning. Appendix C.3 gives additional background.
Constraint For the inviscid Burgers equation, states must satisfy mass and energy conservation, expressed in physical space TR= as C {= E &R= |= F, Σ, non-convex, and the projection onto C is challenging to compute. The projection operator Pe is approximated via iterative alternating- altemmating-projcctoon projection algorithm in high dimension. It is even more challenging to compute in the latent Fourier space, where pe takes the form (3.7). More details are provided in Appendix C.3.
M). This constraint set
Experimental setup We perform cyclic data assimilation on the Burgers equation discretized on a 200-point spatial grid. The ground truth trajectory evolves from random initial condition over time horizon H 8. Observations are sparse: the system is observed at 10 equally spaced times, with noisy spatial measurements at fixed, evenly spaced locations. Each method runs for cycles per trajectory, producing predicted trajectory that be compared to the ground truth. The first baseline is 3D-Var (Courtier et al., 1998), which estimates the state with a Gaussian posterior. At sampling time, the diffusion model combined with the Gaussian posterior, which conditions sampling to the available information. For each cycle, the analysis is computed as the average of diffusion posterior samples. The experiment repeated over 50 independent trajectories. We compute the average mean squared error with respect to the ground truth in the state space, in & norm, and in the constraint space, where the quadratic constraint violation error reported. All methods share the same biased linear forecast model.
Results Figure 3 shows assimilated states and averaged error curves. In this under-observed set- ting, the diffusion prior helps to regularize the structure of complex states better than the Gaussian prior, especially for longer times, where the system shows a stiffer structure. However, uncon- strained diffusion drifts away from the true trajectory, with significant deviations both mass and energy. Projected diffusion (Christopher al., 2024) strictly enforces constraints but introduces high-frrquency artifacts, leading to physically implausible states. Our algorithm CASAL achieves the best compromise: it respects conservation and guides sampling toward physically plausible
8

---

Published as a conference paper at ICLR 2026
target
Gaussian
penalty
projected
CASAL
time 0
time 3
time
82
mass
Field
energy
position
pasition
position
time
0
time
Figure 3: Left Data assimilation sampled states and reanalysis. The black crosses represent the observations. Right Averaged relative error, in terms of constraint violation and lg norm.
states, resulting in significantly lower estimation error. Wall-clock computational times reported in Appendix C.3, and show the efficiency our variable splitting approach in constraining latent diffusion model. These results highlight the potential of constrained generative modeling for robust data assimilation in physical systems.
5.3 CONSTRAINED PRIORS FOR FEASIBILITY PROBLEMS IN OPTIMAL CONTROL
As final application, we evaluate CASAL on a feasibility problem in optimal control: find trajecto- ries that satisfy both system dynamics and non-convex obstacle avoidance constraints. These prob- lems are hard due to the non-convexity of obstacle regions. We consider a dynamical system with state ys)) and control u(s), with s the physical time, and define a trajectory as x == (y(s), u(s)).
Constraint Dynamics are encoded via the constraint set C. := fo 5 Fly, u), |u| S Umax Obstacle constraints define the potentially non-convex set C. (x y(s) & O Vs). for obstacle regions O,, The goal is to find trajectories in the intersection Ca n Co. For this task, ADMM is classical solver alternating projections onto Ca and Co, but its convergence can be compromised when C. is non-convex. Instead, we propose to guide ADMM with samples from generative prior: a diffusion model trained on trajectories, with constraints enforced at sampling. This approach has seen promising results in control and robotics with diffusion penalty guidance and projected diffusion (Carvalho al., 2023; Shaoul et al., 2025; Zampini al., 2025), which we implement and compare with CASAL.
Experimental setup We consider planar quadrotor system, controlled in acceleration an- gle (Tedrake, 2009). A latent diffusion model is trained on a dataset of obstacle-free trajectories, obtained with variety of random periodic excitations. At test time, non-convex obstacles are intro- duced. The corresponding constraint is imposed during sampling. In order to avoid the obstacles, the algorithm needs to find a swinging trajectory. Each sampled trajectory then used to initialize an ADMM solver for feasibility problems (Bílkoáá & Šorel, 2021), and we record the fraction of samples for which feasible solution is found.
Results Figure 4 shows some sampled tra- jectories and success rates as the obstacle size increases, computed over 10000 sam- ples. Constraint penalty guidance favors ob- stacle avoidance, but some sampled trajecto- ries penetrate the obstacles. Projected dif- fusion avoids obstacles but suffers projec- tion bias, producing distorted and unphysical paths. Our algorithm balances both aspects: it produces obstacle-avoiding trajectories that remain dynamically feasible, leading to signif- icantly higher success.
obstacle
penalty
projected
CASAL
success
yi
2
Figure 4: Left Dashed lines are sampled trajecto- ries, solid lines are the projections onto the feasi- bility set. Right Success rates for various obstacle sizes of the rightmost obstacle.

---

Published as a conference paper at ICLR 2026
6 RELATED WORK
Constrained diffusion Early approaches adapted optimization methods to Langevin dynamics, including Projected Langevin Monte Carlo (Bubeck et al., 2015), proximal Monte Carlo (Salim et al., 2019), Mirrored Langevin (Hsieh et al., 2018), and penalized Langevin (Gurbuzbalaban et al., 2024). Extensions to diffusion models have also been explored (Fishman et al., 2023; Liu et al., 2023: Christopher et al.. 2024). These methods offer convergence guarantees in convex settings, where constraints do not hinder exploration, but are less effective in non-convex physical problems. Carvalho et al. (2023), Huang et al. (2024a) and Meunier et al. (2025) impose soft constraint penal- ties in the diffusion process. Crucially, these methods rely on a differentiable constraint penalty function. Therefore, the sampling objective is inherently different from (3.3), as the almost-sure constraint defines a non-smooth potential. Our approach tackles this non -on-smoothness using a prox- imal operator and projections, thereby ensuring strict constraint satisfaction, rather gradient steps on a smooth approached loss. The variational formulation of Langevin sampling has been used by Chamon al. (2024) to enforce constraints average.
Variable splitting Variable splitting, inspired by ADMM, has been applied to Bayesian posterior sampling (Vono al., 2019), plug-and-play samplers for inverse problems (Bouman Buzzard, 2023; Wu et al., 2024; Martin et al., 2024), and guided diffusion (Zhang et al., 2025). These works apply variable splitting to a smooth maximum a posteriori optimization problem, where the aux- iliary variable is updated by gradient descent. Crucially, our framework enforces exact constraint satisfaction through non-smooth constraint potential, without requiring differentiable constraint model. Moreover, formalize sampling optimization problem in density space, which is key to obtain our probabilistic sampling guarantees. Our algorithm also extends to latent diffusion, enabling computational savings.
Physically constrained neural networks Physical constraints have also been imposed on de- terministic neural networks (Négiar et al., 2023; Hansen al., 2023). In related sampling ap- proaches, Cheng et al. (2024) and Utkarsh et al. (2025) integrate projection into flow-mator dow-matching. Our approach differs in targeting strict satisfaction a diffusion framework.
CONCLUSION
In this work, we addressed the complex problem of constrained sampling from a novel theoretical perspective. While prior methods often rely on heuristics, we introduced the first mathematical framework for this task based on Lagrangian duality in the space of probability measures. This formalism not only provides rigorous foundation but also offers diagnostic tool to understand the limitations of projected and penalty-based approaches.
Our proposed algorithm, CASAL, offers principled bridge between constrained sampling the- ory and modern generative modeling. By leveraging a variable splitting scheme, we ensure strict constraint satisfaction while maintaining the exploration properties of Langevin dynamics. The laxation error and the mixing times are analytically quantified in Wasserstein distance. Crucially, framework naturally extends to latent diffusion, where it remarkably decouples the physical constraint space from the generative latent space. Our applications to complex physical systems illustrate the practical importance of informing data-driven methods with physical constraints.
While CASAL naturally extends to deep generative models, several promising directions remain. natural extension would be the derivation of natively time-dependent consrtained variational formulation of the sampling problem for diffusion models and other modern deep generative models such as stochastic interpolants (Albergo & Vanden-Eijnden, 2023). Finally, while the reliance on projection operator is the necessary price for strict feasibility, further research into approximate or parallelized projections could broaden the applicability of CASAL to more high-dimensional, constrained sampling problems in science and engineering
10

---

Published as a conference paper at ICLR 2026
Reproducibility statement The proofs of the theoretical results included in this paper available in Appendix D. The code the proposed algorithm available online at
and
Implementation details and comparison with other algorithms and diffusion models are discussed in Appendix B. Ablation studies, discussion about algorithm hyperparameters and experimental details are available in Appendix C.
Acknowledgments We thank Carla Roesch, Luiz Chamon and Anna Korba for their insightful feedback on this work. The authors acknowledge funding, computing, and storage resources from the NSF Science and Technology Center (STC) Learning the Earth with Artificial Intelligence and Physics (LEAP) (Award #2019625). The authors acknowledge the MILInES organization for its support.

---

Published as a conference paper at ICLR 2026
REFERENCES
Kwangjun Ahn and Sinho Chewi. Efficient constrained sampling via the mirror-Langevin algorithm.
Advances in Neural Information Processing Systems, 14-28405--2818, 2021. 3
Michael Samuel Albergo and Eric Vanden-Eijnden Building Normalizing Flows with Stochastic Interpolants. In The Eleventh International Conference on Learning Representations, 2023. URL https: fttts://openrevee
net forum :id= 17ggiicee...tt.... 10
Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media, 2008. 17
Rina Foygel Barber and Wooseok Ha. Gradient descent with non-convex constraints: local concavity determines convergence. Information and Inference: A Journal of the IMA, 7(4):755-806, 2018.
Espen Bernton. Langevin Monte Carlo and JKO splitting. In Sébastien Bubeck, Vianney Perchet, and Philippe Rigollet (eds.). Proceedings of the 31st Conference On Learning Theory, volume 75 of Proceedings of Machine Learning Research, pp. 1777-1798, PMLR, 06-09 Jul 2018. URL https: https://procee /ttp///oeeeede//// wwweediingg... ngs ml.. wail..press press// /tee/v/5///// v75/berntonl 8a html.
Dimitri P Bertsekas. Constrained optimization and Lagrange multiplier methods. Academic press, 2014. 6
Zuzana Bílková and Michal Sorel. Projection methods for finding intersection of two convex sets and their use in signal processing problems. Electronic Imaging, 33:1-6, 2021. 9, 23
Matthieu Blanke, Ronan Fablet, and Mare Lelarge. Neural Incremental Data Assimilation. In ICML 2024 Al for Science Workshop, 2024. 22
Charles A Bouman and Gregery T Buzzard. Generative plug and play: Posterior sampling for inverse problems. In 2023 59th Anmual Allerton Conference on Communication, Control, and Computing (Allerton), pp. 1-7. IEEE, 2023. 10
Stephen Boyd, Neal Parikh, Eric Chu, Borja Peleato, Jonathan Eckstein, et al. Distributed optimiza- tion and statistical learning via the alternating direction method of multipliers. Foundations and Trends@ in Machine learning, 3(1):1-122, 2011. 6,
Sebastien Bubeck, Ronen Eldan, and Joseph Lehec. Finite-time analysis of projected Langevin Monte Carlo. Advances in Neural Information Processing Systems, 28, 2015. 1, 3, 10
Joao Carvalho, An Le, Mark Baierl, Dorothea Koert, and Jan Peters. Motion planning diffusion: Learning and planning of robot motions with diffusion models. In 2023 IEFF/RS International Conference on Intelligent Robots and Systems (IROS), pp. 1916-1923. IEEE, 2023. 7, 9, 10, 23
Luiz Chamon, Mohammad Reza Karimi Jaghargh, and Anna Korba. Constrained Sampling with Primal-Dual Langevin Monte Carlo. Advances in Neural Information Processing Systems, 37: 99285-29323, 2024. 5, 10, 17, 28
Chaoran Cheng, Boran Han, Danielle C Maddix, Abdul Fatir Ansari, Andrew Stuart, Michael W Mahoney, and Bernie Wang. Gradient- Generation for Thirteenth International Conference on Learning Representations, 2024. 10
Systems. In The
Jacob K Christopher, Stephen Back, and Nando Fioretto. Constrained synthesis with projected diffusion models. Advances in Neural Information Processing Systems, 37:89307-88333, 2024. 1, 5, 7, 8, 10, 19
Gabriele Corso, Hannes StÅ, Bowen Jing. Regina Barzilay, Tommi Jaakkola, et al. DifiDock: Dif- fusion Steps. Twists, and Turns for Molecular Docking. In International Conference Learning Representations (ICLR 2023), 2023.
Philippe Courtier, Andersson, W Heckley, D Vasiljevic, M Hamrud, A Hollingsworth, Rabier, M Fisher, and Pailleux. The ECMWF implementation of three- dime-simmalio variational assimila- tion (3D-Var). I: Formulation. Quarterly Journal of the Royal Meteorological Society, 124(550): 1783-1807, 1998. 8
12

---

Published as a conference paper at ICLR 2026
Thomas M Cover. Elements of information theory. John Wiley & Sons, 1999. 38
Imre Csiszár. I-divergence geometry of probability distributions and minimization problems. The annals of probability, pp. 146-158, 1975. 4
Yilun Du and goor Mordatch. Implicit generation and modeling with energy based models. Advances in neural information processing systems, 32, 2019. 5, 18
Alain Durmus, Szymon Majewski, and Biažej Miasojedow. Analysis of Langevin Monte Carlo via convex optimization. Journal of Machine Leurning Reseurch, 20(73):1 2019. 1, 2, 36
Edward S Epstein and Rex Fleming. Depicting stochastic dynamic forecasts. Journal of Atmo-
spheric Sciences, 28(4):500-511 1971.
Geir Evensen. The ensemble Kalman filter: Theoretical formulation and practical implementation. Ocean dynamics, 53:343-361 2003.
Nic Fishman, Leo Klarner, Valentin De Bortoli, Emile Mathieu, and Michael John Hutchinson. Diffusion Models for Constrained Domains. Transactions Machine Leurning Research, 2023. ISSN 2835-8856. URL https:ttp/// /ttp::ss///peeeeiiee.... wtwpt///oooooodduuutooeoere
Expert
Certification. 2, 10
Daniel Gabay and Bertrand Mercier. A dual algorithm for the solution of nonlinear variational problems via finite element approximation. Computers & mathematics with applications, 2(1): 17-40, 1976. 5,5
Roland Glowinski and Americo Marroco. Sur I' l'approximation, par éléments finis d' d'ordre un, et la résolution, par pénalisation-dualit d' T'uee classe de problèmes de dirichlet non linéaires. Revue française d' a'automaaique,, informatique, recherche opérationnelle. Analyse numérique, 9(R2):41- 76, 1975.
Ian J Goodfellow, Jean Pouget-Abadic Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Advances in neural information processing systems, 27, 2014. 18
J J Gordon, David J Salmond, and Adrian FM Smith. Novel approach to Sonlincarnon-Gauusiann Gaussian Bayesian state estimation. In IEE proceedings (radar and signal processing), volume 140, pp- 107-113. IET, 1993. 21
Mert Gurbuzbalaban, Yuanhan Hu, and Lingjiong Zhu. Penalized Overdamped and Underdamped Langevin Monte Carlo Algorithms for Constrained Sampling. Journal of Machine Learning Re- search, 25(263): 2024. 10
Derek Hansen, Danielle C Maddix, Shima Alizadeh, Gaurav Gupta, and Michael W Mahoney. Learning physical models that can respect conservation laws. In International Conference on Machine Learning, pp. 12469-12510. PMLR, 2023. 10
Bingsheng He, Li-Zhi Liao, Deren Han, and Hai Yang. A new inexact alternating directions method for monotone variational inequalities. Mathematical Programming, 92(1):103-118, 2002. 5
Geoffrey E Hinton. Training products of experts by minimizing contrastive divergence. Neural
computation, 14(8) 2002. 18
Jonathan Ho and Tim Salimans.
Classifier-free diffusion guidance.
arXiv preprint
arXiv:2207. 12598, 2022. 2,3
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-688 6851, 2020. 1,5, 19
Ya-Ping Hsich, Ali Kavis, Paul Rolland, and Volkan Cevher. Mirrored Langevin dynamics. Ad- vances in Neural Information Processing Systems, 31, 2018. 10
Jiahe Huang. Guandao Yang, Zichen Wang, and Jeong Joon Park. DiffusionPDE: Generative PDE- solving under partial observation. Advances Neural Information Processing Systems, 37: 130291-130323 2024a. 7, 10,
13

---

Published as a conference paper at ICLR 2026
Langwen Huang, Lukas Gianinazzi, Yuejiang Yu, Peter D. Dueben, and Torsten Hoefter. DiffDA: diffusion model for weather-scal scale data assimilation. Proceedings of the 41st International Con- ference Machine Learning, 2024b. 22
William Huang, Yifeng Jiang, Tom Van Wouwe, and Karen Liu. Constrained diffusion with trust sampling. Advances in Neural Information Processing Systems, 37:93849-93873, 2024c. 3
Aapo Hyvärinen and Peter Dayan. Estimation of non-normalizze statistical models by score match- ing. Journal of Machine Learning Research, 6(4), 2005. 2
Richard Jordan, David Kinderlehrer, and Felix Otto. The variational formulation of the Fokker-
Planck equation. SIAM journal on mathematical analysis, 29(1):1-17, 1998. 4, 17
Karthik Kashinath, M Mustafa, Adrian Albert, JL Wu, C Jiang Soheil Esmacilzadeh, Kamyar Aziz- zadenesheli, R Wang, Ashesh Chattopadhyay, A Singh, et al. Physics-informed machine learning: studies for weather and climate modelling. Philosophical Transactions of the Royal Society A,
2021.
Diederik P Kingma. Max Welling, et al. Auto-encoding Variational Bayes, 2013. 18
Solomon Kullback and Richard A Leibler. On information and sufficiency. The annals of mathe- matical statisties, 22(1):79-886 1951.
Jean-Paul Laumond. Finding Collision-Free Smooth Trajectories for a Non-Holonomic Mobile Robot. In IJCAI, volume 87, pp. 1120-1123, 1987. 5
Théotime Le Hellard, Franki Nguimatsia Tiofack, Quentin Le Lidec, and Justin Carpentier. Sobolev Diffusion Policy. 2025. 22
Jinhao Liang, Jacob K Christopher, Sven Koenig, and Ferdinando Fioretto, Simultaneous Multi- Robot Motion Planning with Projected Diffusion Models. arXiv preprint arXiv:2502.03607, 2025. 6
Guan-Horng Liu, Tianrong Chen, Evangelos Theodorou, and Molei Tao. Mirror diffusion models for constrained and watermarked generation. Advances Neural Information Processing Systems, 36:42898 42917, 2023. 10
Andrew C Lorene. Analysis methods for numerical weather prediction. Quarterly Journal of the Royal Meteorological Sociery, 1174414 1986. 21
Ségolène Martin, Anne Gagneux, Paul Hagemann, and Gabriele Steidl. PnP-Flow: Plug-and-play image restoration with flow matching. arXiv preprint arXiv:2410.02423, 2024. 10
Etienne Meunier, David Kamm, Guillaume Gachon, Redouane Lguensat, and Julie Deshayes. Learning to generate physical ocean states: Towards hybrid climate modeling, arXiv preprint 254.02900..02499 2025. 2, 10
Rudy Morel, Francesco Pio Ramunno, Jeff Shen, Alberto Bietti, Kyunghyun Cho, Miles Cranmer, Siavash Golkar, OLEXANDR GUGNIN, Geraud Krawezik, Tanya Marwah, Michael McCabe, Lucas Thibaut Meyer, Payel Mukhopadhyay, Ruben Ohana, Liam Holden Parker, Helen Qu, François Rozet, K.D. Leka, Francois Lanusse, David Fouhey, and Shirley Ho. Predicting par- tially observable dynamical systems via diffusion models with a multiscale inference scheme. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025. URL forrmm?rrmmm d= 3
https: httpss////ppeeeeeee...oo
Juan Nathaniel and Pierre Gentine. Generative emulation of chaotic dynamics with coherent prior.
arXiv preprint 250iv:2504 /4264, 2025.
Geoffrey Négiar, Michael W. Mahoney, and Aditi Krishnapriyan. Learning differentiable solvers for systems with hard constraints. In The Eleventh International Conference Learning Represen- tations, 2023, URL https:///// /ttp///ppp/eeeeeeee///ee wtt.eviwwwwwwweeeeeeeeeeeeeeeeee forum? vid 6wnGksrO.... 10
Jong-Shi Pang and David Stewart. Differential variational inequalities. Mathematical program- ming, 113(2):345 4224, 2008. 7
14

---

Published as a conference paper at ICLR 2026
Eric Paquet and Herna L Viktor. Molecular Dynamics, Monte Carlo Simulations, and Langevin Dynamics: a Computational Review. BioMed research international, 2015(1):183918, 2015. 7
Chris Pedersen, Laure Zanna, and Joan Bruna. Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos. arXiv preprint arXiv:2503.18755 2025. 7,
Ilan Price, Alvaro Sanchez-Gonzalez, Ferran Alet, Tom R Andersson, Andrew El-Kadi, Dominic Masters, Timo Ewalds, Jacklynn Stott, Shakir Mohamed, Peter Battaglia, et al. Probabilistic weather forecasting with machine learning. Nature, 637(80444:884 90, 2025.
Yongquan Qu, Juan Nathaniel, Shuolin Li, and Pierre Gentine. Deep generative data assimilation in multimodal setting. In Proceedings the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 449-459, 2024. 8, 22
P. J. Rossky, J. D. Doll, and H. L. Friedman. Brownian dynamics as smart Monte Carlo simulation. The Journal of Chemical Physies, 69(10) 4633, 11 1978. ISSN 0021-9606. doi: 10.1063/ 1,436415, URL https: /doi ozg/ 1063// 43 64155. 2
Litu Rout, Negin Raoof, Giannis Daras, Constantine Caramanis, Alex Dimakis, and Sanjay Shakkot- tai. Solving linear inverse problems provably via posterior sampling with latent diffusion models. Advances in Neural Information Processing Systems, 36:49960 199990
49990, 2023. 22
François Rozet and Gilles Louppe. Score-hased data assimilation. Advances in Neural Information
Processing Systems, 36:40521 2023. 22
François Rozet, Ruben Ohana, Michael McCabe, Gilles Louppe, François Lanusse, and Shirley Ho. Lost in Latent Space: An Empirical Study of Latent Diffusion Models for Physics Emulation. arXiv preprini arXiv:2507.0266 2025.
Adil Salim, Dmitry Kovalev, and Peter Richtárík. Stochastic proximal Langevin algorithm: Potential splitting and nonasymptotic rates. Advances in Neural Information Processing Systems, 32, 2019. 10
Yoshikazu Sasaki. Some basic formalisms in numerical variational analysis. Review, 98(12):875-883, 1970. 21
Monthly Weather
Yorai Shaoul, Itamar Mishani, Shivam Vats, Jiaoyang Li, and Maxim Likhachev. Multi-Robot Mo- tion Planning with Diffusion Models. In The Thirleenth International Conference on Learning Representations, 2025. URL httpp::////....... http///ppeeeeiieeeee /
net/ forum AUCAAI
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019. 5, 18, 19
Yang Song, Jascha Sohl- Bohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. International Conference on Learning Representations, 2020. 1, 18, 19
Russ Tedrake. Underactuated robotics: Learning, planning, and control for efficient and agile ma- chines course notes for MIT 6.832. Working draft edition, 3(4):2. 2009. 9, 22
Utkarsh Utkarsh, Pengfei Cai, Alan Edelman, Rafael Comee-Bombarelli, and Christopher Vincent Rackauckas. Physics-constrained flow matching: Sampling generative models with hard con- straints. arXiv preprint 0rX112250600471,, 2025. 10
T. Gastelen, W. Edeling, and B. Sanderse. Energy-conserving neural network for turbulence closure modeling. Journal of Computational Physics, 508:113005 2024. ISSN 0021-9991. doi: http:://6o..or//1..016//.......p http://ccineeeearaccc/aa/////211
URL https: /wwww sciencedrrectt..
8, 22
Cédric Villani. Topics in optimal transportation, volume 58. American Mathematical Soc., 2021. 4, 17
Maxime Vono, Nicolas Dobigeon, and Pierre Chainais. Aplie-and-Aummented Gibbs Sam- pler- Aopp-icataiioon to Large-Scale Inference Problems. IEEE Transactions on Signal Processing, 67(6):1648-166) 2019. doi: 10. 2tt.TTSS..009.884492222 10
15

---

Published as a conference paper at ICLR 2026
Max Welling and Yee Whye Teh. Bayesian learning via stochastic gradient Langevin dynamics. In Proceedings of the 28th International Conference on International Conference on Machine Learning. ICML' 11, pp. 681-688, Madison, WI, USA, 2011. Omnipress. ISBN 9781450306195. 18
Zihui Wu, Yu Sun, Yifan Chen, Bingliang Zhang, Yisong Yue, and Katherine Bouman. Princi- pled probabilistic imaging using diffusion models as plug-and-play priors. Advances in Neural Information Processing Systems, 37:118389-118427, 2024. 10
Stefano Zampini, Jacob K Christopher, Luca Oneto, Davide Anguita, and Ferdinando Fioretto. Training-fre constrained generation with stable diffusion models. arXiv preprint arXiv:2502.005625, 2025. 7, 9, 23
Youyuan Zhang, Zehua Liu, Zenan Li, Zhaoyu Li, James Clark, and Xujie Si. Decoupling Training- Free Guided Diffusion by ADMM. In Proceedings of the Computer Vision and Pattern Recogni- fion Conference, pp. 23292-23302, 2025. 10, 21
16

---

Published as a conference paper at ICLR 2026
A
VARIATIONAL FRAMEWORK FOR LANGEVIN MONTE CARLO
Consider the functional
F(g) P(lll)) = glog(1/D).
(A.1)
The Wasserstein gradient flow is defined as the following differential system
dq
OF 0q
vv
(A.2)
a1
For functional (A.1), this differential system becomes
dq
V (IVS(x)) Aq(x,
(A.3)
Ōt
which is found be the Lokker-Planckk equation for the Langevin dynamics
dr V dB.
(A.4)
More details can be found in (Jordan et al., 1998; Ambrosio et al., 2008; Villani, 2021; Chamon ct al., 2024).
17

---

Published as a conference paper at ICLR 2026
B DEEP GENERATIVE MODELS AND LANGEVIN SAMPLING
Many modern generative frameworks, from energy-based models to state-of-the-art diffusion mod- els, rely Langevin dynamics for sampling (Hinton, 2002; Du & Mordatch, 2019; Song & Ermon, 2019; Song et al., 2020). In Appendix B, we review how key classes of generative models relate to Langevin updates.
For these generative models, sampling takes the form
I++1
Tt ( 22n we, We (O T)-
(B.1)
We interpret these steps as the discretization of a Wasserstein flow for a time-dependent func- tional Fqg t), which summarized in Appendix A. We can then identically apply our constrained sampling algorithm, as a time-dependent variation of Algorithm 1, detailed in Algorithm 2. From variational point of view, this results in framing the constrained sampling time-varying con- strained optimization problem.
Algorithm 2 Time-dependent CASAL
input time dependent potential gradient V fa t), iteration number T. time-dependent step sizes Tr, projection Pc. cou- pling p > O, intial distribution qu
output sample 27 C
initialize Ho ~ 40 = Pe(xo) = 0 Rdd
for 0 < S T - do
draw W; [O, 14)
2x+1 Σς Ti f (It, Txp(1x 2t + Me) + /2
2++1 = Pe ( TeP 2t 2t+1 -14)
Mt41 14 T/ (1141 2(41)
end for
B. ENERGY-BASED MODELS
An energy-based model (EBM) defines density
1 p(z) N exx((_f(xx))
(B.2)
where fo is learned energy function. Sampling from typically relies on Langevin dynam- ies 12.1) or stochastic gradient Langevin dynamics (Welling Teh, 2011). Energy-based models with Langevin sampling have demonstrated strong performance across range of tasks(Du & Mor- datch, 2019), and offer distinct advantages over methods such as Variational Autoencoders (Kingma al., 2013) and Generative Adversarial Networks (Goodfellow et al., 2014), A particularly valuable property energy- based models is their flexibility in incorporating constraints via summing up the corresponding energies. From this perspective, our algorithm, when applied to energy energy-based based models, can be interpreted as providing stronger constraint enforcement through an augmented Lagrangian potential and corresponding proximal Langevin updates- going beyond the simple addition of straint energies.
B.2 SCORE-BASED GENERATIVE MODELS
Score-based generative models aim learn the score function V (x() of a family of progres- sively noised data distributions {pt hejo, rather than modeling the data density directly. Once the score is learned, typically via denoising score matching, samples can generated by Langevin-type updates.
Proposed by Song & Ermon (2019), this method generates samples by applying Langevin dynamics > σ A score model sa(z, o)) is trained to
at sequence of decreasing noise levels oT approximate the noise- deieedeepnndnnt score V, log σ) the perturbed data distribution p(z; or), which is obtained by convolving p(x) with a Gaussian of various noise level oe Then update step takes the form
Ft+1 St (x) 2 we, wt N00 I),
(B.3)
18

---

Published as a conference paper at ICLR 2026
where To of are tary-ng step sizes. The update takes the form of (B.1) with Vf(x, f) ss, ot). This can be seen as an unadjusted Langevin algorithm with tem- perature ot, gradually refining the sample noise decreases. In this directly applied at each noise level to impose constraints. It is worth noting that the projected diffu- sion model (Christopher et al., 2024) also falls into this category hard projection following each Langevin update within the annealed Langevin dynamics framework. Note that this covers the case where several Langevin steps are taken fixed noise level, as in the work of Song Ermon (2019), by choosing Tę to be constant for number of steps t.
algorithm can be
B.3 DIFFUSION MODELS
Denoising diffusion probabilistic models (DDPMs), introduced by Ho et al. (2020), define a forward process that gradually corrupts data point Yo by adding Gaussian noise through a fixed Markov chain:
(ive 1t-1) = (wei / 3, Vt- Brl),
(B.4)
where Br E (0, 1) a small noise schedule. This leads to a closed-form expression for q( ro). with the following definitions:
at
Bc,
0g
II a2..
(B.5)
The reverse process is parameterized by a neural network c0(x,,,) t), which predicts the noise com- ponent. The sampling procedure follows:
D++1
Ft
V1 -1-0x co( + oww, wy ~(0,
1 0z
(B.6)
Vas
where og is typically set to match the forward variance Be. As noted by Ho et al. (2020), this step corresponds Fuler-Maruyama discretization of variant of Langevin dynamies, and the learned noise predictor eg implicitly estimates the score V pt up to a scaling factor. Hence, the sampling formula (B.6) takes the form (B.1) with T₁ of/2 and
v log (ara)) so
VT- a co(
(B.7)
The DDPM can be regarded as discrete score-based model under the variance preserving stochastic differential equation (VP-SDE) interpretation (Song et al., 2020), and thus our CASAL sampling is valid for DDPM sampling.
B.4 SCORE-BASED DIFFUSION MODELS
Score-based diffusion models (Song et al., 2020) directly learn the score function of perturbed data distributions and generate samples by simulating the reverse-time stochastic dynamics.
Forward SDE Define forward Itô SDE that gradually adds noise to data zo ~ Pdata:
dz = a(x, t) dt + b(t) aw.
(B.8)
where for the variance-preservinng (VP) choice,
a(x, t) = 1 (t) x, b(t) = B(t).
(B.9)
This yields marginal distributions (z) that interpolate between the data and near- Gear-Gaussian noise increases.
Reverse SDE The time-reversed process follows dx = [a(x, f) b(1)2 V, (o + b(t) aww,
(B.10)
where W/ is reverse-time Wiener process. A neural network so f) is trained by score matching to approximate V p(z, t).
Predictor-Corrector sampling. Once the score network is trained, CASAL sampling is ap- plicable. CASAL can also be integrated seamlessly into the predictor-correttt sampling scheme proposed by Song al. (2020). The pedictor--coreette sampler interleaves the following steps.
19

---

Published as a conference paper at ICLR 2026
Predictor: an Euler-Mel Marlyymaa step of the reverse SDE,
I++1 Ht Ta [a( t) b(t)² so(1,, t)] + b(t) 2 N( La).
(B.11)
Corrector: a few steps of Langevin MCMC refine samples.
I++1 It (0 t) v2Twu, wy (O,
(B.12)
Similar to the previous sections, these formulas take the form of (B. 1), with different time-varying potential gradients V fx t) and step sizes.
Summary Across EBMs, diffusion models, and hybrid schemes, the core sampling formula is an overdamped Langevin update, possibly annealed through noise scales. This makes our constrained sampling algorithm CASAL compatible with all these approaches zero-shot plug-in.
20

---

Published as a conference paper at ICLR 2026
C EXPERIMENTAL DETAILS
C.I PARTICLE ON A CIRCLE
3-
0-
0-
-1-
-3-
-0-
-3
-3
-3
potential
Figure 5: True posterior density, projected Langevin samples and CASAL
Figure 5 shows the exploration issue arising with Projected Langevin Monte Carlo in the case of non-convex constraints and a bi-modal distribution. Here, projecting on the constraint set
C -=]
(C.1)
leads to poor exploration, as the samples are stuck on the positive side of the likelihood landscape, while the only high-likelihood zone compatible with the constraint is on the other side. Our algo- rithm solves this exploration issue leveraging variable splitting
C.2 FIELD GENERATION
compared to projected Langevin Monte Carlo, primal-dual Monte Carlo and constraint-penalized Langevin Monte Carlo. For the latter, implement the variable- splitting algorithm of Zhang et al. (2025), and the penalty parameter is dual variable that adapted and updated following the same scheme as CASAL.
Baselines Our algorithm
Sampling Langevin Monte Carlo is iterated over 1000 steps, and we set p to follow linear inter- polation schedule between 2 and 20.
Constraints The field subject to energy conservation (Example 2.1). The projection in closed
forms. For the penalty method, the penalty cost is c(x) = -E MY². The primal-dual Langevin
Monte Carlo algorithm enforces the constraint function h(x) Σ x4 -M on average.
C.3 DATA ASSIMILATION
Context In many geophysical and engineering applications, one relies on numerical simulation to predict the time-dependent evolution of complex system, whose state at physical time / is denoted by x E Rd. But these models are inherently imperfect-eithee because of computational constraints or incomplete knowledge of the true dynamics. When real-world observations E Rm become available (for example in digital- -win settings), we assume a statistical model of the form
y = h(x)) +
(C.2)
where R -> R is an observation operator and the measurement error. The imperfect simulation yields prior forecast E dd the background estimate, which must be adjusted using to produce more accurate estimate of the true state, usually referred to as the analysis, as the initial condition for the next simulation. Equivalently, seeks samples from the posterior
p(z|b, y) p(y(z) p(z(6).
(C.3)
This estimation problem is formulated sequentially for each new observation, by propagating the obtained posterior analysis with forecast model, and repeating the process. Classically, this is achieved by one of three approaches: sequential Monte Carlo methods (e.g. particle filters (Gor- don et al., 1993)), ensemble- based filters (c.g. the Ensemble Kalman Filter (Evensen, 2003)), or variational methods that solve for the MAP estimate (e.g. 3D-Var/4D-Var (Sasaki, 1970; Lorene,
21

---

Published as a conference paper at ICLR 2026
1986)). The 3D- Var algorithm assumes that the background error distribution and observation error
distribution are Gaussian,
z/6 ~ (b,, B), E~ NN, P), then taking negative logarithm of (C.3) yields the following optimization target:
(C.4)
J(x) = 2 ( + 211 bll R-1
(C.5)
Deep learning represents promising tool to learn more complex priors for data assimilation (Huang al., 2024b; Rozet & Louppe, 2023; Qu et al., 2024; Blanke et al., 2024)
Data For simulating the Burgers equation, we implemented the same method as van Gastelen et al. (2024), but we added an extra constant linear advection term. We work in Fourier space with the first 20 Fourier modes. The field evolves according to the Burgers equation for 4 time units. We generate 1,000 trajectories, with the field recorded at 10 timesteps for each trajectory, with the initial state drawn at random in Fourier space with power-law decay of the coefficient magnitude.
Learning architecture We implemented a DDPM diffusion model, using the formalism detailed in Appendix B. Diffusion is learned in latent space, defined as the first 10 Fourier modes. The neural network involved is a fully connected network with depth 3 and width 128, using cosine time embedding. is trained for 200 epochs. At sampling time, 1000 diffusion steps are used with p 10.
Baselines Our algorithm is compared unconstrained latent diffusion, latent penalty diffusion and latent projected diffusion, incorporating observations using diffusion posterior sampling Rout al. (2023) or as strict constraint, and propagating the projection and the penalties through the decoder.
Constraints The field is sampled subject to energy and mass conservation constraints (Exam- ple 2.1). The projection is computed by alternating projections on the two constraint set, which have closed forms.
The initial conditions are drawn random following the same distribution of the training data.
Additional results Figure shows the evolution of key estimation metrics for a data assimilation
trajectory, for the various methods compared.
target
Ganssian
penalty
projected
CASAL
enror mass energy
x
x
1.0
0.5
time
Figure 6: Mass conservation, energy conservation and &a error.
C.4 FEASIRILITY PROBLEM
Trajectory planning for quadrotor obstacle avoidance is an imoprtant problem (Le Hellard et al., 2025). We implement linearized version of the planar quadrotor dynamics described in (Tedrake, 2009).
Data The trajectories are discretized in time as (y1,..., US, U1,... ..,...us) E R22S with S = 200 and a time interval As = 0.01. The dynamics constraint
Ca {x (0) (s)(s) Ay(s) Bu(s), [u(s) S Umnax)
(C.6)
22

---

Published as a conference paper at ICLR 2026
is described by linear equality constraint, discretized into linear system, and inequality straint on the control inputs. The projection on this convex constraint set is obtained by Dykstra's double projection algorithm (Bilkoáá Sorel, 2021), and used within the ADMM solver.
Learning architecture We implemented DDPM diffusion model, using the formalism detailed in Appendix B. The trajectories signals are learned in a latent Fourier space encoding the first 10 modes of the input signal. The neural network involved is fully connected network with depth and width 128, using a cosine time embedding. It is trained for 200 epochs. At sampling time, 1000 diffusion steps are used with p 100.
Baselines We implement the latent projected diffusion algorithm (Zampini et al., 2025), and dif- fusion guidance with constraint penalties (Carvalho al., 2023). and propagate the penalty function through the decoder.
The obstacles are segments, and projecting onto the feasible region is performed by moving the penetrating trajectory portions trajectory either directly above or directly underneath the obstacle. For the penalty method, the constraint penalty is the quadratic distance between the trajectory and the obstacle, which is simple enough differentiate through in this case.
Constraint
C.5 PIIYSIALLY--CNSTTRANEE EMULATION OF PARTIAL DIFFERENTIAL EQUATIONS
To confirm the applicability of CASAL to a wider range of high-dimensiooal PDEs with complex constraints, we evaluate our method on three additional systems: the Darcy flow, the Navier-Stokes equations, and the 2D Poisson equation, following the experiment setting of Huang et al. (2024a). All methods these experiments share the same network architecture and observational data.
Darcy Flow We experiment on the Darcy flow setting of DiffusionPDE, discretized on 32 x 32 grid. We use latent diffusion in Fourier space and sample both the coefficient field and the solution field from 25 sparse observations. The samplers aim at enforcing the Darcy equation connecting the coefficients and the solution. The reconstructed fields are posterior averages over 10 samples.
As shown in Table 1, CASAL significantly reduces the reconstruction error compared to the base- lines while satisfying the physical constraints to a high degree of precision. All approaches demon- strated similar wall-clock computational times.
Table 1: Reconstruction error and constraint violation for the Darey flow experiment.
Algorithm
Reconstruction error
Constraint violation
Projected diffusion
22 %
1 10 10-6 %
Penalty
23 %
25 %
CASAL
10 %
2 x 10 5
Navier-Stokes Following the approach of DiffusionPDE,, we evaluate generative data assimilation on the Navier-Stokes equations. We constrain the generated fields match sparse observations while simultaneously strictly enforcing the physical constraints of enstrophy and total circulation.
Table 2 demonstrates that CASAL achieves the lowest reconstruction error for the initial and termi- nal states, while strictly satisfying the physical constraints at a computational cost comparable to the baselines.
Table Reconstruction error, constraint violation, and runtime for the Navier-Stokes experiment.
Algorithm
Reconstruction error
Constraint violation
Wall-clock time (s)
Projected diffusion
10 %
10 10-b %
1.07
Penalty
11 %
20% %
1.00
CASAL
8
10 %
1.05
Poisson Equation
To provide clear demonstration of transfer beyond the Burgers equation to
purely elliptic setting, we benchmark CASAL on the 2D Poisson equation. We consider a 100 x 100 discretization and train single latent diffusion model on solutions with random boundary conditions
23

---

Published as a conference paper at ICLR 2026
in Fourier space. At inference time, 10% of the grid points observed, and the global mass is known.
All samplers enforce the mass constraint (in this setting, penalty-guided diffusion reduces to latent analogue DiffusionPDE). Each sampler produces 50 posterior samples. The results in Table show that CASAL substantially improves the sampling accuracy over baseline methods while strictly enforcing the mass constraint.
Table 3: Reconstruction error, constraint violation, and runtime for the Poisson equation.
Algorithm
Reconstruction error
Constraint violation
Wall-clock time (s)
Projected diffusion
15 %
x 10 10-
0.34
Penalty
14 %
25 %
0.33
CASAL
5 %
2 x 10 10-7%
0.32
C.6 COMPUTATIONAL COST
To assess the computational cost of CASAL, we summarize below the costs for the constraint sets used in our experiments. We provide below a breakdown the projection cost for different con- straint types, and compare them to the overall runtime.
In our experiments, we use three types of base projections:
projection onto a sphere for energy conservation,
projection onto intervals for obstacle avoidance,
projection onto a linear subspace for mass conservation,
We also combine these projections in the where the constraint is the intersection of such sets. All of these admit efficient implementations.
For the sphere and the interval, cost is (dd) with closed-form formulas.
For a linear subspace,
C= (x Ax b 0)
(C.7) (T( Ar). which requires
with mdd The projection is given by Pe(x) a precomputed pseudoinverse at cost (ddm) and a matrix-vector product at cost O(m), which remains small compared to neural network evaluations.
In all tested settings, the runtime overhead from projections was small compared to the cost of score evaluations in diffusion models. Furthermore, CASAL is compatible with approximate projections, allowing further savings. For more complex constraint sets, iterative solvers such as ADMM can be employed with limited number steps, trading accuracy for speed in early iterations, where perfect constraint enforcement is not yet required.
To validate this point, we conducted the following additional runtime experiment. We measure the average wall-clock time for the different sampling algorithms in the data assimilation problem, where each sampled state is projected on the intersection of 2 constraint sets: one for mass and one for energy. All times are in seconds per 106 sampling steps, measured
Apple M1 setup.
Experiment
unconstrained Langevin
penalty
projected Langevin CASAL
Fluid generation Data assimilation
0.37
0.41
0.43
0.45
0.17
0.25
2.46
0.31
Table 4: Comparison of computational times.
We also note that diffusion guidance can be computationally costly as it can require multiple penalty gradient steps per sampling step to enforce constraints. In the experiment 5.3, we found this method to be of the order of 3 times slower than CASAL and projected diffusion.
In summary, CASAL has comparable runtime to projected Langevin. For a number of use cases, its additional cost is modest, especially in the context of deep generative models where the com- putational budget is dominated by score evaluations. This cost can also be adjusted in practice, by computing approximate projections in the early steps of sampling.
24

---

Published as a conference paper at ICLR 2026
C.7 ABLATION STUDIES FOR P AND À
For the field generation task of Section 5.1, we generate 1000 samples of energy-constrained fields, using different schedules for p. For each schedule, evaluate the samples with the following measure error: we report the proportion of samples that fall near the unlikely positive mode of the bimodal potential, implying that the sampled distribution deviates from the target
We experiment with different schedules: two schedules use constant p throughout the iterations. The other schedules are linear and logarithmic interpolation between these two values. We run the experiment for different numbers of Langevin iterations. The results are reported in the following table.
number of steps / schedule constant p Pmin 1000
constant
Pmax
linear
logarithmic
42.4%
22%
0.04%
3.8%
5000
44.7%
2.8%
0.001%
0.0002%
Table 5: Proportion of samples in the wrong mode for different schedules of p
We observe that allowing to vary across iterations substantially improves sample quality. With too small the deviation between and is too large. With too largep, the chain fails to explore the energy landscape. When the number of steps is limited, only annealed schedules manage to recover the correct mode. This experiment highlights the importance of adaptive schedules in practice.
In the other experiments, we found that the time-varying step size induce by diffusion models, which also scales p, was sufficient to balance exploration and constraint satisfaction
We conduct an ablation study on both the field generation experiment (Section 5.1) and the Burg- data assimilation task (Section 5.2) to investigate the influence of the initial value of the dual variables.
In the first experiment, the fields sampled with the Langevin Monte Carlo algorithm, with fixed potential. In the second experiment, a diffusion model is used, so the score function is time-varying function throughout iterations.
We run our sampling algorithm with Gaussian initialization of X, with different sizes a. For each value of a, run 100 independent chains. For the flow sampling experiment, report the maxi- mum number of sampling steps required to converge. For the data assimilation experiment, because the dual problem changes over time, do evaluate the convergence of the dual variables. In- stead, we report the average reconstruction accuracy.
Steps to convergence
Reconstruction eπror
0
10
0.56
1
200
0.56
5
800
0.56
10
1000
0.57
20
1600
0.58
50
2500
0.79
100
N/A
1.34
Table 6: Influence of
25

---

Published as a conference paper at ICLR 2026
D PROOFS
In the following proofs, if E (2((R has a density, we identify and its probability density function.
D.1 PROOF OF PROPOSITION 3.1
Proposition 3.1 [Information projection] Suppose that 0 < Py(C) < 1. Then the conditional
distribution Pc is the projection of p onto the set distributions supported on
PC
argmin
D(gll))
(EP2(R4)
(3.3)
subject to
Proof. For simplicity, assume the existence of probability distributions. For q E P2(R)).
/ ((x) log
g(z)
D(gllp)
dr
p(x)
J( ((x)) log
g(x) Ze
dr
(D.1)
pc(ar) Z
Ze
D(gllpc)= +
7
where Ze satisfies
1 [ pe
Rd
Z
Ic p(z)ddx
(D.2)
Ze
Z PC(
Ze
Therefore,
(gllp)) D(allee) + (C).
(D.3)
This quantity minimized for g = Pc, and the minimal value (C).
L
D.2 PROOF OF PROPOSITION 3.2
Proposition 3.2 Strong duality is not attained for (3.3).
Proof. The Lagrangian of (3.3) is
L(g, A) = ((dlp)) + + P
(D.4)
D(llp)) + XEE]((())]
For all À IR, the dual function is
g(A) = min D(gllp) AEg[c(x)]
(D.5)
(EP2(R
and this minimum is attained by
(x(x)
Z e -(())
(D.6)
7
p(x)e Ac(a) -Ac()))
Z
The Lagrangian evaluated at Px equals
7
((2) log
(D.7)
Z
26

---

Published as a conference paper at ICLR 2026
To compute Zy, note that
PA
Rd
(D.8)
Z
plax)da ( e>>c(t))
Z
Za cc
Let
s(A) 1(e e (-Ar(m)p(I)dk
(D.9)
Then,
2 1 7x [P)(C) + +(())
(D.10)
By assumption, for all A IRd, 0 s(A) 1. Furthermore, obtain by combining (D.7) and (D.10), that
g( = log P2(C) + =(A)
(D.11)
This value is always strictly lower than its limit:
AY, ((A) log
lim g(A), x++00
(D.12)
P. Pp(C)
which is preciscly the optimal value of Problem (3.3), attained by = Pc. Indeed,
D(pellp) = /
Z
Ze p(x) log dr
Zc
(D.13)
Z
Z
(C)
log
Zc
Z
where Ze satisfies
Pc
Rd
7
[ p(z)dr
(D.14)
Ze
Z
Ze (C)
follows that
D(pellp) = log P,(C)
(D.15)
This value is found to be the minimizer of Problem (3.3) using Gibbs' inequality.
Corollary Penalty methods (2.5) cannot enforce (C(CC 1.
Proof. Penalty methods sample from Px. with finite X. By the previous proof,
(x(x) x exp (=f(x) Ac(x))
(D.16)
minimizes the Al, A). By definition of the dual function,
g(A) L(gx, A).
(D.17)
For all densities E (2d) satisfying the constraint P,(C) 1, the positive duality gap implies
L(gx, A) < D(pellp) & L(g. A)-
(D.18)
Therefore, Px does not satisfy (C) 1.
27

---

Published as a conference paper at ICLR 2026
D.3 PROOF OF PROPOSITION 3.3
Proposition 3.3 [Variable splitting] Problem (3.3) is equivalent to the following problem:
minimize, (EPa(R4XR4)
(4 [Ip) B8 [xc(=)]
(3.4)
subject to Pg (x 2) 1.
Proof. Given q(2,2) the solution of Problem (3.4), the marginal qt gives the solution of Prob-
lem (3.3). Given (ax the solution Problem (3.3), the solution of Problem (3.4) can be obtained by defining copy of x.
D.4 PROOF OF PROPOSITION 3.4
Proposition 3,4 [Strong duality] Strong duality holds and is attained for (P). In particular, the Lagrangian (3.5) admits a saddle point (q,,A.) and q, is solution of (P).
Proof. The expectation constraint is satisfied by the couple density of
~~ No z|x ~ (x,
(D.19)
Hence, there exists feasible point. Furthermore, the constraint function is surjective, JE, = v can be satisfied with æ (O, Ta) and z|ær N(x v, la). Therefore, Slater's constraint qualification condition is satisfied, and implies that strong duality holds and is attained. Relevant references can be found in Chamon al. (2024).
D.5 PROOF OF PROPOSITION 4.1
Lemma [Integral formulation of the Lagrangian] The Lagrangian (3.5) of (P) is equal to
[ [9(s)) f11x + - =|2] do (x)dg(21<) ||ll
L(g, D(g"llp) with μ = (1/p) x A.
(D.20)
Proposition 4.1 The optimal distribution of takes the form
(x) exp (f((x)) -20---20 _db(x+ uw)).
(4.1)
where p. (1/p) X, and dc(=) minyee |yy =|| the distance the constraint set.
Proof. By Lemma 1 and by definition of the saddle point, μs+ is such that minimizes
L(g, = D(g" (lp) + Eg | [1 = 11 + xc(2),
P 1ll
2
(D.21)
D + 1 11 - + + xc(=) dg (z)dg(=)=) P Ill..1?
Therefore, among the different conditionals g(=(zr) (z). the optimal (=(x)) minimizes
/ 15 n..1 xc(+)) dg (x))g(=1(x)
(D.22)
For each x, the function brackets is minimized by
(xx = Pe(x + 14))
(D.23)
and hence, by integration against (<2(x).
I Ar - z 12 + xc(=) dg(=lx) > 6 -1Ax- (x) + +.,l xc(x,
(D.24)
This value is achieved by the deterministic conditional concentrated at (x)) Hence
Vr TRd
(=(x²) (2)
(D.25)
28

---

Published as a conference paper at ICLR 2026
By definition of 2, (z) and 9o= this implies that, for all 2.
/ 1 Ar 2 p.ll + xc(2) (z|2) = L - + p.)
(D.26)
Now replace this expression in the Lagrangian, seen as function of g"
L(q. = D(p= + / (z + (x)dg((x) - || P 2 + p.)] / logg" q" dq" Ex
(D.27)
P 2
||
This functional is equal, up to constant to the Kullback-Leibler divergence between q² and
(a) exp (-((x) - 2 Le +X.)).
(D.28)
and is hence minimized by the latter.
D.6 PROOF OF PROPOSITION 4.2 Proposition 4.2 [Recovery of the projection] Let 9p denote the solution to (P). Then
fp 1 p-000
Pe distribution.
(4.2)
Proof of Proposition 4.2. This result follows from Proposition 4.1.
D.7 PROOF OF PROPOSITION 4.3
Lemma [Optimal value bound] Let q. be solution Problem (P). Then
D(glllp) + L Fy [| =11²] < D(pellp)
(D.299
2
Proof. Let q P2 x R defined by x ~ p. and 2/2 8g. Then æ 2 P., P(R Rd),, and Ex 2] = 0, is in the feasible set of Problem (P). This distribution provides the stated upper bound on the optimal value.
Proposition 4.3 [Relaxation error] The relaxed solution q. satisfies the following conditions:
1 D(pelle).
D(g)llp) < D(pellp), (C) = 1, W2 d.) ) S P
(4.3)
Proof. The result follows from the previous lemma, noticing that W2'gg" q S Ey. ||| ==
(D.30)
D.8 PROOF OF PROPOSITION 4.4
Proposition 4.4 [Consistency] In the case of a non-degenerate Gaussian potential f and an affine constraint set C, the solution of (P) is unbiased: E,,. [x]] E,, [2] Epc [x].
Proof. Let C fi Rd Ar b) with A xdx of full rank, b E R and k 1. We may parameterize C = + WR with = A E R and W E Bd a full column rank matrix gathering basis vectors of ker A. Let W IR - ker denote the associated bijective operator.
We may assume that is standard normal Gaussian, since operating linear transform to pre- serves the problem with a different transformed affine constraint space. We thus assume that
f(x) 2
x.
(D.31)
29

---

Published as a conference paper at ICLR 2026
Let (= v). and let h C R be bounded function. Then,
E[h(y)] /
h(W 'o00-'(x---
h(W-~( -)))pe((x)
Z
/
h(W v))e -
Ze
(D.32)
Z
W/ J/L. h(y)e
s(utWodd
Zc
7
| h(y) exp (-2 (o++((++))) dy.
Wi
Zc
This shows that y is Gaussian vector, with quadratic potential
(w) T (u + Wy)-
(D.33)
In partieular, the expectation of maximizes this quadratic form:
W (WEly) + 2) = 0.
(D.34)
Because x 1 Wy.
E Epe [x] = v + WEyy
(D.35)
Therefore,
pe [x] = O.
(D.36)
which means exactly that ippe [x]] E C is the orthogonal projection of [xp[ 0 onto C. Recall that the projection onto equal to
Pa(a:) = W
(D.37)
Second, note from Proposition 4.1 that o, is Gaussian, the projector is linear and
( = [Pc(z) =||²- = [[w - (x - e)1l.
(D.38)
The mean m= Eg [x]] under vf maximizes the likelihood, and thus solves
+ p(WW Ia) (WW+ Im) + p. v) O,
(D.39)
which simplifies to
m + p(Ta - (WW+)mm + p4 - v) = 0.
(D.40)
Left-multiplying by WT, we obtain that
W vmm=
(D.41)
obtain another equation by combining the equality constraint Eg, [x]] E,, [=]. and the char- acterization of Proposition 4.1 implying ql-mmsstt surely that Pe(æ + u.). By taking the expectation, we obtain
m = =m [x] [x] o wW+ (m + - v).
(D.42)
Substituting in (D.40), we obtain
(u 1 - + 0.
(D.43)
implying that m E C. The two conditions W m = 0 and m C characterize m Ey, [x]. and thus prove that
Eg. [x] = E.g [z] = Epe [x]].
(D.44)
30

---

Published as a conference paper at ICLR 2026
D.9 PROOF OF CONVERGENCE
We use the two following standard descent lemmas: one for the gradient step, and one for the projection step.
Lemma 3 [Smooth gradient descent inequality] Let f be a a-convex and ß-smooth function. Let E I2 and let F4 := -VVf((). Then, for all y RR..
2T 2/2 - v/1²
2T ||+ - y11 2 f(x+) f(y) +
2T (1 3T) ||, - x=1² + 1|1 y/1 (D.45)
Lemma [Projection inequality] Let C be a convex set of R Let 2 E R Then, for all w E R
|| w| 2 Pc(z) w|| [Pc(2)) - =1|?
(D.46)
We now study the effect of the alternated descent steps. Lemma 5 [Alternated gradient Lyapunov decrement] Let x & E Rd and 2, E R Define the iterations
1+ = r - AAr + μ (0--)) 2)) = 2 T(z Ai pt)
(D.47)
Let 7 = (1A|]2 + 1). Then for all (y.w) Rdd R²t,
1
1
27 (1 y. - w)ll
[1 - y, w)]l2
2T
((t, - cly, w)
1 = -1 (2 I+, 2+) || - 2,, 2, - 211
(D.48)
||A|| x [| -1)ll || E||.
2T [xx+ - -,, 2)| 2 |1(2 - x, 2)11
Proof. The quadratic decrement factors as follows
[1x - y. w)11 2
1
||( - y, 24 - w)ll
1
(x - - - -2+) T
b (x ++++++++++++)))++))+ - y, (= $ + 2+) - w
(x x =- --) T
(x+ - y. 2+ w))
1
1 [r x+),
+ (a 2- -4) T
2 2+)
2
(x - 2 24) T (a y, 8 w) + (a x4, - 2+) (x4 2,,+ - -))+ |1(x+ - x, 2, - =)
(D.49)
Since (x+, 24) = (I, V (b 2) + 7 (0, A(r - -)),
Vola, 2) (y) ã. =) - (1r x+, 2+) T (yr - -. - + (a 4) A AT (w 2). (D.50)
Substituting in the quadratic decrement expression, we obtain
2T ( y, - w)ll 1 11 - y, w)]2
1
2T
-V(( =){(v= i. w -2)
1
(D.51)
+ (x - 2+, 2 -=2) T (a+ - -. 24 - 2)
1
+(x - 8) A (w 2) +
[1( - x, 24 - 2)11
2T
31

---

Published as a conference paper at ICLR 2026
By convexity of o,
Φly, w) 2 0(0,2) + Vo(z,z) 1, (
(D.52)
The function (br,2) + A1l z + 1=| is smooth. Letting M == (A Ia) (dx()))))
its Hessian is M T M, which has the same eigenvalues as MM = AAT + Ia- Therefore, its smoothness constant is y || + 1, with ||A|| the largest singular value of A.
By smoothness,
(i, S ((z, 2) + VØ(z, 2) (e xx - 2) + 21(( 2, 8 - =2)ll?
(D.53)
2
Combining with the convexity inequality,
oly, w) 2 ((; Vo(zz (w ã, w - ;) 2 11( x, 2 =)
-
(D.54)
This implies the following quadratic decrement bound
1
1
27 (x y w)ll
2r 11 v. w)ll
6 2) oly, w)
1
(x 2+, - 2+) T (x+ - -,,24 - 2)
(D.55)
(r - T)_ A (10 2)
2r [1+ Ir 2, =)1 ? l1(: -x, - =)11
1
Finally, by the
inequality
24-2 2+) (x++ -2, - =2) 2 - (x 24, 2 1| 2, -- -11. (D.56) (D.56)
and
(x - σ) AT (w- 2) 2 - |All x [| - +i]l x || - 2||-
(D.57)
Lemma 6 [Alternated proximal Lyapumov decrement] Let x E IR and 2, E TR... Define the
iterations
2+ = prox,, - Tp(Ax Ax + p 2))
(D.58)
24 = Pel= - Tp(= - Ax+ - H).
Let = 2 + 1). Then for all (y, E Rdd x C,
2T [2 y, - w)ll 2r
[1x+ -, - -w]l2
f(x+) - [ly] + p 6 Ar + + μt - 2+ 2 - Ay + 4 - w||2
1
27 |100 = - =](|
(1 B (x) x',4- 2
2r
(xx4 U1 2 w
(D.59)
1
11 (x - x',2 x²) || 24,2 z'' -2+1
PT
|||A|| x ||0 -x+111 || De
[2+ - I. zし - =3113.
Proof. Let x' := x --p(pr=------ + -p-22
and 2 =p(z Ax. - -p). Applying Lemma 5 to x' and v'
with step size TP, and t x. we obtain for all y Rdd z RR
32

---

Published as a conference paper at ICLR 2026
1
1
[x y. w)ll
11 y. z w)ll z
2r
2r
1
pl| Ar + 11 2+11
plllAy + F w|l2
2
2
(2 2. - =2)l
2r
(D.60)
|| (o x',2 2- 2') || x || - x+, 2 2+1l
pr
M|| x ||x x+4l x ||:0 2+1
? (x 2, 2+ 2)
with 7 - A| + 1).
We bound
|| -w=2+1 2+1l De.
(D.61)
and
(xr - ,',== z²) = |1(:c u4, xp4) S M.
(D.62)
pT
Furthermore, applying and summing the descent inequalities for f and Pe at x' and
2r | y/12
| y11
1
+ 2T
1| =ul2
1| - w112
2T
(D.63)
f(x+) - f(v)= + 2T
(1 3T]11 - 2211 + a
2 1 - y11
|| 2/122
2T
The result obtained by adding these inequalities.
D.10 PROOF PROPOSITION 4.5
Lemma [Primal-dual I Lyapunov decrement] Let x E Rd and 2 E R Define the iterations
x+ = prox,,(x Tp(Ar + 4 1x-2)) 2))
2+ = = T(+ - -24).
Pelz Tp(2 Ar - p).
(D.64)
Let 2(||A]2 1). Then for all (y, w. E Rdd C x A provided 4T + 2|[ 1)) $
1
1
1
2T [1x y, w)1l
(25
y, 2か w)ll
a 111
||
u112 >
2r
2r
f(x+) f(y) 6 $1A 2+11
p 6 Al w||
+v T (Ar+ - -24) T(Ay- - -w)
(D.65)
+ +2
|1( - - 2+ 1?
- 4 - 2M 2
pT ||Ar - 12.
2
33

---

Published as a conference paper at ICLR 2026
Proof. By Lemma 6,
1
2T [ V, 2 w)ll
(1x 2/ w)ll
27
1| + F 2+11 2/2 I1 + μr w|l
f(x+) f(y) + 2
11 - I, x'' - 2) + - (1 ( x'.2 - 2
2T
2r
(D.66)
+ 2
(x - V. 24 w)
(xr - x²,== x0'2-22' =z')l || x || I4, - =+
P|| x || - x+|l ||W 2+11 II
2 11( - r, -))1
We are going to absorb the negative terms of the right-hand side into the positive quadratic terms, up to terms of order T.
First decompose
11(+ - I. 2, -0,24-2)11- -)]1 s 2/ (I+ n', 24 =0)112 211((a x, 2' -))13
(D.67)
so that
- (2+ 2,24 =2]11 2 -1 21)12 - 711(2' x, x2,''= 2))1 2)1
(D.68)
By absorbing these contributions, the remaining quadratic terms are
2r (1 (y-) - x,, - = 2)11++ + 2r (1 - (3 + 2 =1)11
(D.69)
By combining the Cauchy-Schwartz inequality and Young's 's inequality, for all >> 0,
1
p||A||
| x+
|| ||w =+ 2 7 -1 x -/2 11
(D.70)
In order to absorb these terms in the quadratic terms, we choose o = 1/(4T) and obtain
|||A|| x ||I - x+ || x w 2+1l 2
x - x+ 2+12 4 4TA2\|AA|2D??
87
|1xx x4, 2+)1
4T @2||AA|/DD
8T
1 11 x4,2' - -+)11 - ATp²"|A|1 DŹ.
[1x 2 21)112
47
(D.71)
After absorbing the first term, the remaining quadratic terms are
1
1
2r -1/2-277)l('' - x, - 2]1 + (1/2 (B 2,))))((++++ - -2':2----211 =2)111 (D.72)
2r
By a similar argument,
P (x - x', 2- -2)=21)1 x 11xx 24,2 2+)l 2 [1 x4,2 - - 222
T
8T
1
x', 2- 2²)
11
-2)12
20²
87
(D.73)
After absorbing the first two terms, the remaining quadratic terms are
1
2r (1/4- - 2y - xx, 2,,''222 2) + 2r
(1/4 - (B + 27)+ - -2, =2)112
(D.74)
34

---

Published as a conference paper at ICLR 2026
Finally, provided 4 3)T 1,
11 - y, - w)ll [ v - ww]l? 2
2r
2r
f(x+) - -(())= + 6 _1 + - 2+ 11
-/2 =11 + + 1 - ww1
(D.75)
a 11 - y, - =w)1
+ 2
4TT211A|222 - 2p M
Now note that || + 11 2+11 11 11 w112 = | - 2+12 -lAlylyy -lAy-yll'' [0]]=+1++ (Ax
(Ax+ - 24 - (Ay - w))
(D.76)
where PP = X.
By the dual iteration x+ A + 7 -Axt+ z+),
1
11+ - 111
-
1
/| - 111 + I1lAE -
-++12=++12- +-+00 (A -1) T T(A++ -
27
2T
(D.77)
_
1
||
1112
+ T 1/2 5|Ar+ - 2+112 + (A v)T(Ar+ - 2+),
2T
implying
1
w/12
1
11 - w11 - 24 12 + (v x) - -24).
(D.78)
2r
2T
Gathering these equations,
1
1
27 (<x y, = - w)]l 2 |1( - y; 24 w))l |
1
> u11 b a1 - 1112+ N
27
27
[xx+) - ((y) + G11 f1A 2+12 L Al - 10
tu T [Ax+ - 2+) - pd AT (Ay - w)
a |1( 2, 10--w) w)ll
2
47 ||| D2- - 2222222222
PT || - 12.
2
(D.79)
Proposition 4.5 [Primal-dual Lyapunov decrement] Let 4+ be the distribution after iterations (3.6) with step size т. Then, provided 4T + 4p) S 1, for any reference pair (p, v)):
V(g.x:p,v) EV | + a W 2 ((L(g-, I[g., A)) + C+2 (4.5)
2
Proof. Let (I4, 2+) 9+. Note that x4 Σ+ + 22ww, where (#+, is defined in the previous lemma, and w NO0 Ia) is an independent Gaussian noise. We let @+ be the distribution of (++,2+).
The previous lemma is valid for all (x, 2), (y, w). Choose y(z) to be the optimal transport plan between o² and p², and w(z) to be the optimal transport plan between of and p². Then
E2 v|l Wz .p))
(D.80a)
Fy w||2| W2(g .p)).
(D.80b)
Furthermore, by definition of the Wasserstein distance,
E y1 > W (gī Ey112 - w0l > W g p²),
(D.81a)
(D.81b)
35

---

Published as a conference paper at ICLR 2026
and
(" + W2 97) 1 By [[[ - y. - -w)l
(D.82)
By integration of Lemma against
1
1
1
1
W2
(q" " p²) +
W (q p²)
W 2 (q" (T))
W2 (4))
2T
2T
27
2T
1
1
Bgg [l| w/|3]
Ry [l] - m
2r
2r
a
2
(D.83)
[f((((+) / 11A 2+ 11 - (1p(s(())))))) +
2 211 1013
ppv By [AØ+ - 2+] - pv Ep |Ay - w)
47 1 2+ 112 - 22p 22 PT
2 | - 2+11
Note that
| 2+112 = ||Ar+ 2+ - //TAAul]
(D.84)
= || 2+ 11? + 2T Aurl1 22w (Ax+ 2-);
whose expectation is
E2 [] 2+11% [1 = B,y [11 2+11 11 + 2
(D.85)
The other quadratic terms in 2+ are handled similarly. Expectations of linear terms in ū are not affected by the Gaussian noise. We thus obtain
,p") + (2 p²)
W2'0",
W3 oi)
2T
2T
2T
2r
1
+ Eg [1|x
w11
Eg [11+ m
2r
2r
+
2
W2 (1++4))
(D.86)
[ + 6 11 - -+112]-1pppp +
-11 - =013]
+pv TTTEE, Eg [Ax+ - -+1-p=) Epp [Ay w]
- 4r ||A|||||u 2+12
2p² 2 -
PT 2
||Ax+ 2+11 + (2p||A]] + da)r.
In order to control the value of By [s]. we use the 3-smoothness assumption and apply Lemma 3 of Durmus et al. (2019):
E [f(x)] Eg, [f(x)] S Bdr
(D.87)
We finally handle the entropy terms. For E P2(RRd), let
H(p) / y logg.
(D.88)
Since x+ = 3. + /2 Lemma 5 of Durmus et al. (2019) applies and yields
1
W p²)
W2 p²) 2 H(g:) H())).
(D.89)
2r
2r
Gathering the inequalities above, obtain
1 + W2
1
1
W3 (g
,²'
(a
W2 .4.
2r
2T
27
2T
1
1
Egg [I|A -v11
E.y [[]
=
2r
2T
(D.90)
(W2(g=:0) 2
2
L(qr,v) 1(p,1)
-CT.
36

---

Published as a conference paper at ICLR 2026
with
C= -==AA||A|/11ue - 2+ 2+112 1Pp-2p2 G1AA, - + da) + 3d.
(D.91)
D.11 PROOF OF PROPOSITION D.1
Lemma [Saddle gap bound] Let (q., A.) be the saddle point of Problem (P). Then, for all (g, A) P2(R)) x R
L(g, L(g,, A) 2 Dog")
(D.92)
Proof. By definition of (g+: A,), T((+, A) = T( X.). Furthermore, by Lemma 1,
L(g, X.) - I(q., A.) = (p) + /
P 11 2 1 xc(2) (g"(r)dg(21x))
_
D(gillp) + /
P [|Axx || 1112 + xc(x) (z)d
2
= D(d""lp)) - -D(
(D.93)
+ / [2 ||A - = mel2 xe(=) do (x) (dg(=)>) - dg.(2(x))
/ ? -1 + xe(2) (d) (1) - dg (z)) (=(xr)
From the characterization of 9. in Proposition 4.1, (21xx) (m) and
/ 161AAs=++ - + m-12 112 + xc(=)) dr (x) (dg(=)x) d=/()/=)) 2 0.
(D.94)
Therefore,
+ + l + +. (dd dxt
L(g. L(g,,
(D.95)
that is
L(q, L(q., / dd + Egl((x) + P
2 +
- / log do. - By, [f(x) + f l(
(D.96)
2
L(g, A.) 2 D(9")la") P( D l9.))
Proposition D.1 [Mixing rate] Then, there exists such that
T-1
T-1
DDG + E(V (gT, < C Σ
-i + Vigo,
(D.97)
t=
Proof. Let Tt. be the probability measure of (xa, 2)) conditionally on Ao. A and note that (q++1, Ae+1) is obtained from Ar A, by applying updates (3.6). We apply Proposition 4.5 qr and (p, v) (g., ).), and apply Lemma 8 to bound
L(g. A.) - L(g., 2 D( 144).
(D.98)
We obtain
W 2 (1t 4..) + W2 (.)
W Tet1)
W2 vit1)
2T
2T
2T
2r
+ E, [[|A, xx
2T E, [1>c+1 >x
27
(D.99)
+
a
W (1)+1, 2
D( - Cr.
37

---

Published as a conference paper at ICLR 2026
Summing these inequalities for TTs 1, 1,
Σ Ty D (6+1 11
W2( ,.) + W3 + 11
1
2012
(D.100)
T-1
2.12 12 (6 )f) + (16. + -
+ C
Σ
2
We finish the proof by taking the expectation over Ar = 2o....,Ar7 Ar. To compute the average of the mutual information, note that
E [[] = =(f(()))
(D.101)
and, by the entropy
=
= [[H(g(Z]|A)] 2 H())).
(D.102)
We finally obtain
Σ (ef+++l144) 1 ((W x1) = (0z [5.. +001++1111110 -
(D.103)
T-1 2 W (00 Φ.) + (96 (96. 9.) + =1100 2 + C
1
Σ
:t
1=0
Corollary 2 [Mixing rate] Let (zt) be generated by Algorithm with step size = 1/vt+1. Let q. denote the distribution of the time-averaged iterate ã, := 7 E.= Zz. Then,
Inl V
D (1.)) O
(4.6)
Proof of Corollary 2. By convexity of the Kullback-Leibler divergence (Cover, 1999),
D 11
D
(D.104)
Applying Proposition D.1 to T = 1/V1,
D + + Ar)
==0
(D.105)
C
V(go,o)
8=0
This implies
[-1
H-1
Σ
D
D
V(go, Au),
(D.106)
t
Vi ==
Vt
and
Σ
logt.
(D.107)
8 t++
==0
Therefore,
o (logt/vi)
D (gg.)
(D.108)
38

---

Published as a conference paper at ICLR 2026
Usage of Large Language Models We used large language models at the sentence level to correct English writing and avoid word repetition.
39


Reviews:
[{"review_id": "LaRGunK0zg", "reviewer": "Reviewer_zhci", "summary": "The paper proposes Split Augmented Langevin (SAL), a method for strictly enforcing equality constraints during sampling in generative methods such as score-based diffusion models. It introduces a primal–dual Langevin framework that couples unconstrained and projected variables to guarantee feasibility at every step.", "questions": "**Q1.** How does the proposed SAL method compare quantitatively to prior constrained generation frameworks such as PCFM [2], D-Flow [3], DiffusionPDE [4], or ECI [1] that also enforce hard or functional constraints?\n\n**Q2.** What is the average magnitude of the constraint violation observed in Figure 3, and does the method guarantee strict satisfaction (i.e., ($||h(x)|| = 0$)) or only approximate feasibility?\n\n**Q3.** How would the projection operator behave for complex or nonlinear constraints where analytical projectors are unavailable? Would this require solving an inner constrained least-squares problem at every step, and how does that impact computational cost?\n\n**Q4.** Given that the initialization step projects Gaussian noise to obtain ($z_0$), what guarantees exist that this projection produces feasible or representative samples, especially for high-dimensional or nonconvex manifolds?\n\n**Q5.** How does your method perform in high-dimensional or multi-modal constrained settings? Does it exhibit the slow mixing behavior often observed in Langevin-based samplers, and what evidence supports stable convergence in such regimes?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "**S1.** The paper effectively highlights the central difficulty in constrained generation, balancing feasibility with sample diversity, and offers a thoughtful critique of simple projection-based solutions.\n\n**S2.** It presents a theoretically grounded primal–dual Langevin formulation that unifies ideas from constrained optimization and generative sampling in a principled way.\n\n**S3.** The exposition is clear and precise, with mathematical reasoning and well-structured derivations that make the approach easy to follow.", "weaknesses": "**W1.** The paper does not include enough baselines or adequately discuss prior work on constrained generation. Comparable methods such as PCFM [2], D-Flow [3], DiffusionPDE [4], and ECI [1] are only briefly mentioned or omitted. ECI, for instance, also enforces hard constraints, while PCFM demonstrates constrained sampling under nonconvex conditions: both directly relevant for benchmarking.\n\n**W2.** The empirical analysis is limited, lacking standard quantitative metrics such as MMSE, SMSE, or FID reported routinely in this literature. Reported results rely largely on qualitative visualizations and histograms, and even show non-zero constraint violations (e.g., Fig. 3) without quantifying their scale or assessing the degree of enforcement softens your claim of hard constraint.\n\n**W3.** The evaluated constraints are relatively simple, and the computational claims appear optimistic. The method assumes analytic projectors, but more complex nonlinear or PDE-based constraints would require solving constrained least-squares problems at every step, which could become computationally expensive and limit scalability.\n\n**W4.** The algorithm assumes that projecting Gaussian noise to obtain ($z_0$) yields feasible initial states. This assumption may hold for simple quadratic or linear constraints but lacks guarantees for challenging manifolds, such as PDE solutions with shocks and discontinuities. Testing on more complex physical systems, as in PCFM [2] or guided functional diffusion methods [5], would better support claims of generality.\n\n**References**\n\n[1] Cheng et al., *Gradient-Free Generation for Hard-Constrained Systems (ECI)*, arXiv:2412.01786 (2024).\n\n[2] Utkarsh et al., *Physics-Constrained Flow Matching: Sampling Generative Models with Hard Constraints*, arXiv:2506.04171 (2025).\n\n[3] Ben-Hamu et al., *D-Flow: Differentiating through Flows for Controlled Generation*, arXiv:2402.14017 (2024).\n\n[4] Huang et al., *DiffusionPDE: Generative PDE-Solving under Partial Observation*, NeurIPS 37 (2024).\n\n[5] Yao et al., *Guided Diffusion Sampling on Function Spaces with Applications to PDEs*, arXiv:2505.17004 (2025).\n\n[6] Chamon et al., *Constrained Sampling with Primal-Dual Langevin Monte Carlo*, NeurIPS 37 (2024).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "hDF8co6OXu", "reviewer": "Reviewer_7Rvb", "summary": "The paper proposes Split Augmented Langevin (SAL), a training-free sampler that enforces hard constraints by splitting variables and projecting a companion variable onto the constraint set each step, with the final output $z_T \\in C$. It provides a variational/duality view, consistency to the strictly constrained law as coupling increases, and a drop-in use within diffusion models.\n\nContributions:\n1) A split augmented Langevin sampler with projection that enforces samplewise constraints while preserving exploration.\n2) A relaxed duality analysis with attained strong duality and recovery of the strictly constrained law.\n3) A practical integration recipe for diffusion/latent diffusion requiring only a projection operator $P_C$, demonstrated on three tasks.", "questions": "1) Problem setting and constraints. For physics, equality-in-expectation $\\mathbb{E}_q[h(x)]=0$ seems insufficient. In many cases, the idealized requirement $\\mathbb{E}_q[|h(x)|]=0$ would force $h(x)=0$ almost surely (i.e., strict feasibility), but this is typically impractical to enforce via smooth penalties. Please clarify your recommended formalization: average constraints (Section 3) versus support constraints $x\\in C$ almost surely (Section 4). Do you envisage $h(x)\\ge 0$ and $\\lambda\\ge 0$ anywhere, or do you recommend set membership with projection as the primary approach?\n\n\n2) Finite-time behavior and stopping. Can you relate finite $\\rho$ and step sizes to discrepancy from the target conditional $p_C$? A practical stopping rule based on $\\|x_t - z_t\\|$ or constraint residual thresholds would help practitioners decide when sampling is accurate enough.\n\n3) Adaptive schedules for $\\rho$ and $\\eta$. Do you recommend an adaptive policy driven by observed coupling gaps or residuals (e.g., increase $\\rho$ when violations plateau, decrease when exploration slows)? A short ablation comparing fixed, annealed, and adaptive schedules would be valuable.\n\n4) Inexact projections and robustness. When $P_C$ is computed approximately (few inner iterations), what error conditions preserve convergence and strict feasibility in practice? Please include an experiment sweeping projection tolerance/iterations and reporting feasibility rates and sampling bias.\n\n5) Intersections of multiple constraints. For $C=\\cap_i C_i$, do you prefer a single projection onto the intersection, alternating projections onto $C_i$, or multiple split variables $z^{(i)}$ with separate dual updates? A small comparison on stability, wall-clock, and feasibility would inform users.\n\n6) Classical PDE benchmarks. Consider adding concise tests where projections are standard and informative: Poisson with linear constraints $Au=b$, Darcy flow with mass/flux balance, and incompressible Navier–Stokes via Hodge projection. Reporting $\\|Au-b\\|_2$, divergence norms, spectra, and feasibility histograms would strengthen generality claims.", "limitations": "", "rating": 2, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1) A split augmented Langevin formulation that enforces per-sample constraints while retaining exploration; a practical combination of variable splitting and projection for sampling.\\\\\n2) Coherent variational/duality setup with clear algorithms; sensible baselines demonstrating feasibility and distributional fidelity.\n3) Motivation is explicit; notation and procedures are readable; appendices provide proofs, variants, and implementation notes.\n4) Training-free and model-agnostic; integrates with diffusion pipelines, making strict constraints more accessible in scientific applications.", "weaknesses": "1) Projection robustness: many constraints need iterative or approximate $P_C$. Quantify how projection error affects feasibility and sampling bias; include experiments sweeping inner-iteration counts or projection tolerances.\n2) Hyperparameter tuning: performance depends on $\\rho$ and dual step $\\eta$. Propose adaptive schedules driven by observed coupling gaps $\\|x_t-z_t\\|$ or constraint residuals, and compare against fixed/annealed baselines.\n3) Experimental breadth: add compact, classical PDE tests where projections are standard and informative (Poisson $Au=b$, Darcy with mass/flux balance, incompressible NS via Hodge projection), reporting PDE residuals, divergence norms, spectra.\n4) Latent diffusion details: when enforcing constraints in decoded space, discuss scaling $\\rho$ through the decoder Jacobian and show sensitivity analyses to ensure stability and consistent enforcement.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "KJcAPiIdnN", "reviewer": "Reviewer_gUbm", "summary": "The paper introduces SAL, a principled framework for strictly constrained sampling. By combining a variable-splitting formulation with stochastic primal–dual updates, SAL attains strict feasibility under nonconvex constraints while preserving the exploratory behavior. Theoretical analysis provides strong guarantees, and experiments on physical systems and control tasks demonstrate the method’s potential. However, the experiment results are somewhat limited, leaving some unconvincibility about the practical robustness of the proposed method. I would consider raising my rating if additional experimental results were provided.", "questions": "1. I understand that the SAL combines the strengths of several existing baselines. Since prior work such as Primal–Dual Sampling already performs well, to what extent is SAL an incremental improvement that integrates ideas from Projected LMC, variable splitting, and augmented Lagrangian methods? It would be helpful to include a summary table comparing different baselines across key aspects (feasibility, stricticity, etc.) to make the distinctions clearer.\n2. The paper mentions that a projection operator is required for SAL as well as for other methods. However, do such projection operators always exist for arbitrary physical constraints or control objectives? If not, could the authors discuss possible strategies or approximations?\n3. Given that the authors are trying to model constrained sampling over physical field functions rather than simple random variables, and physical conservation laws also apply to fields, is it formally appropriate to represent the problem directly using random variables $x$ and $z$? Some clarification on this abstraction would make the formulation more convincing.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The SAL algorithm achieves strict feasibility without sacrificing exploration in constrained generation, addressing the weakness of previous projected or penalty-based approaches through a well-motivated formulation.\n2. The paper provides a clear variational interpretation of constrained sampling and offers theoretical guarantees supporting the soundness and effectiveness of the proposed method.\n3. The experiments cover a broad range of tasks and settings, demonstrating that SAL maintains physical constraints while preserving sampling diversity and quality.", "weaknesses": "1. None of the three experiments report quantitative metrics or summary tables comparing different methods.\n2. Lines 329–333 mention several baselines intended for comparison, yet the most relevant one, the Primal–Dual Langevin method, does not appear in the reported experimental results.\n3. Lines 420–421 note that ADMM is a classical solver for obstacle-avoidance problems and highlight its limitations, but the experiments result does not provide a direct comparison between ADMM and the SAL, which is unreasonable.\n4. Lines 320–322 discuss computational cost, claiming that SAL adds only the cost of a projected diffusion step compared to the unconstrained baseline. However, the additional variable-splitting updates also may introduce extra overhead. Additionaly, a quantitative runtime comparison with baselines would be necessary to substantiate the efficiency claim.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "LCEop7uT2L", "reviewer": "Reviewer_1KBT", "summary": "This paper tackles the problem of sampling from generative models while strictly satisfying physical constraints like conservation laws. Existing approaches either fail to enforce constraints exactly (penalty methods) or become trapped in limited regions under non-convex constraints (projection methods). The authors propose Split Augmented Langevin (SAL), which uses variable splitting and augmented Lagrangian dynamics to progressively enforce constraints while preserving exploration capability. The method operates in a training-free manner on pre-trained diffusion models and provides convergence guarantees through duality analysis. Experiments on energy-preserving field generation, data assimilation for the Burgers equation, and optimal control problems demonstrate that SAL achieves both strict constraint satisfaction and accurate conditional sampling, outperforming existing baselines in maintaining physical plausibility and forecast accuracy.", "questions": "While the theoretical contributions (Propositions 3-6) are rigorous, the connection between the mathematical advantages and the observed performance improvements in Section 5 is not clearly articulated. Specifically, it remains unclear how the duality gap in Proposition 3 directly explains the constraint violations in Figure 3, or how Proposition 5's attained duality resolves both the exploration bias (Figure 2) and projection artifacts (Figure 4) simultaneously. I recommend adding explicit statements or a summary that maps each theoretical result to the corresponding experimental observations, making the paper more accessible to readers outside convex optimization. A comparative table contrasting penalty methods, projected Langevin, and SAL across key mathematical properties (duality, constraint satisfaction, exploration) would significantly improve clarity.", "limitations": "", "rating": 4, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "**Principled framework with theoretical guarantees:** The variational formulation and duality analysis provide strong theoretical foundations, proving that the relaxed problem recovers the target distribution as coupling strength increases.\n\n**Training-free and modular:** The method works as a drop-in replacement for standard Langevin steps in pre-trained diffusion models, requiring no retraining or additional data.\n\n**Strong empirical validation:** Experiments on diverse physical systems demonstrate practical effectiveness, showing improved constraint satisfaction and forecast accuracy compared to existing baselines.", "weaknesses": "**Limited Complexity of Constraints.** While the experiments effectively demonstrate the method's core capabilities, the constraint types considered are relatively simple. The energy conservation constraint in Section 5.1 is a quadratic sphere, and the mass conservation in Section 5.2 is linear. Real-world physical systems often involve more complex constraints, such as coupled higher-order nonlinear conservation laws (e.g., multi-component energy functionals with nonlinear operators), or systems governed by multiple interacting PDEs with intricate coupling structures. It remains unclear how the projection step and the coupling parameter would perform when the constraint manifold has significantly more complex geometry or when the projection itself becomes computationally expensive. Evaluating SAL on such challenging constraints would strengthen confidence in its applicability to complex physical modeling tasks.\n\n**Gap between Theoretical Analysis and Actual Implementation.** While the paper presents a practical and effective algorithm, there is a notable gap between the theoretical analysis and the actual implementation. The convergence analysis in Section 4.3 focuses on the optimization problem (Eq 4.4) in the space of probability measures, establishing strong duality (Proposition 5) and asymptotic recovery (Proposition 6). However, the actual algorithm (Eq 4.5-4.6) is a stochastic sampling procedure operating in sample space with Langevin noise. The missing piece is a rigorous convergence analysis of the proposed iterative scheme (4.6) itself, which would require establishing: (1) Non-asymptotic convergence rates in Wasserstein distance, (2) Finite-sample guarantees for fixed $\\rho$ values, (3) Guidance on hyperparameter selection ($\\rho, \\tau, \\eta$) based on convergence theory. The current theoretical results only guarantee that as $\\rho \\to \\infty$, the solution to (4.4) approaches $p_C$, but do not characterize how the finite-step stochastic algorithm (4.6) approximates this solution. This is acknowledged by the authors as a limitation, but addressing it would significantly strengthen the contribution. Despite this gap, the empirical results convincingly demonstrate the method's effectiveness.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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:

Decision: Accept (Poster)