Strictly Constrained Generative Modeling via Split Augmented Langevin Sampling
Matthieu Blanke, Yongquan Qu, Sara Shamekh, Pierre Gentine
TL;DR
The paper tackles the challenge of enforcing physical constraints in generative modeling by formulating constrained Langevin sampling and introducing Split Augmented Langevin (SAL), a primal-dual, split-variable method that achieves strict constraint satisfaction while preserving exploration. SAL builds a relaxed, variational framework with augmented Lagrangian potentials and stochastic primal-dual updates, yielding convergence guarantees and a drop-in replacement for diffusion-based samplers. Through theoretical results and three physics-focused applications—energy-prescribed bimodal fields, Burgers-data assimilation, and non-convex trajectory feasibility in control—SAL demonstrates improved constraint adherence and predictive fidelity over traditional constrained sampling and diffusion baselines. The approach is modular, compatible with pre-trained diffusion models, and offers practical pathways to integrate physical laws into scientific generative modeling with zero-shot deployment. This has significant implications for climate forecasting, fluid dynamics, and robotics, where preserving conserved quantities and feasibility constraints is essential for trustworthy simulations and planning.
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 (SAL), 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. In particular, we use constrained diffusion models to generate physical fields satisfying energy and mass conservation laws. 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 SAL for challenging feasibility problems in optimal control.