Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Predict-Project-Renoise: Sampling Diffusion Models under Hard Constraints

Omer Rochman-Sharabi, Gilles Louppe

TL;DR

This work tackles the challenge of sampling diffusion models under hard constraints by defining a constrained forward process and corresponding constrained marginals, then introducing Predict-Project-Renoise (PPR) to sample from these marginals. PPR alternates between enforcing feasibility through a denoiser-based projection and restoring the diffusion noise via renoising, thereby approximating the constrained distribution $p_t^{\mathcal{C}}$ at each step. Across Data2D, Kuramoto-Sivashinsky, and global weather tasks, PPR substantially reduces constraint violations and yields samples that closely match the constrained ground truth, outperforming existing projection-based and posterior-sampling baselines. The approach offers a principled, training-free mechanism to incorporate hard constraints into score-based diffusion models, with a controllable compute-quality tradeoff and broad applicability to scientific domains such as weather forecasting and PDE-driven systems.

Abstract

Neural emulators based on diffusion models show promise for scientific applications, but vanilla models cannot guarantee physical accuracy or constraint satisfaction. We address this by introducing a constrained sampling framework that enforces hard constraints, such as physical laws or observational consistency, at generation time. Our approach defines a constrained forward process that diffuses only over the feasible set of constraint-satisfying samples, inducing constrained marginal distributions. To reverse this, we propose Predict-Project-Renoise (PPR), an iterative algorithm that samples from the constrained marginals by alternating between denoising predictions, projecting onto the feasible set, and renoising. Experiments on 2D distributions, PDEs, and global weather forecasting demonstrate that PPR reduces constraint violations by over an order of magnitude while improving sample consistency and better matching the true constrained distribution compared to baselines.

Predict-Project-Renoise: Sampling Diffusion Models under Hard Constraints

TL;DR

This work tackles the challenge of sampling diffusion models under hard constraints by defining a constrained forward process and corresponding constrained marginals, then introducing Predict-Project-Renoise (PPR) to sample from these marginals. PPR alternates between enforcing feasibility through a denoiser-based projection and restoring the diffusion noise via renoising, thereby approximating the constrained distribution at each step. Across Data2D, Kuramoto-Sivashinsky, and global weather tasks, PPR substantially reduces constraint violations and yields samples that closely match the constrained ground truth, outperforming existing projection-based and posterior-sampling baselines. The approach offers a principled, training-free mechanism to incorporate hard constraints into score-based diffusion models, with a controllable compute-quality tradeoff and broad applicability to scientific domains such as weather forecasting and PDE-driven systems.

Abstract

Neural emulators based on diffusion models show promise for scientific applications, but vanilla models cannot guarantee physical accuracy or constraint satisfaction. We address this by introducing a constrained sampling framework that enforces hard constraints, such as physical laws or observational consistency, at generation time. Our approach defines a constrained forward process that diffuses only over the feasible set of constraint-satisfying samples, inducing constrained marginal distributions. To reverse this, we propose Predict-Project-Renoise (PPR), an iterative algorithm that samples from the constrained marginals by alternating between denoising predictions, projecting onto the feasible set, and renoising. Experiments on 2D distributions, PDEs, and global weather forecasting demonstrate that PPR reduces constraint violations by over an order of magnitude while improving sample consistency and better matching the true constrained distribution compared to baselines.
Paper Structure (45 sections, 34 equations, 24 figures, 4 tables, 2 algorithms)

This paper contains 45 sections, 34 equations, 24 figures, 4 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Contributions.
  3. Score-based generative modeling
  4. Forward process.
  5. Backward process.
  6. Training and sampling.
  7. Method
  8. Constrained forward process.
  9. Constrained backward process.
  10. Predict-Project-Renoise (PPR) sampling.
  11. Related Work
  12. Projection-based constrained sampling.
  13. Limitations of existing projection methods.
  14. Posterior sampling under observation processes.
  15. Experiments
  16. ...and 30 more sections

Figures (24)

  • Figure 1: Data2D results. Wasserstein and $k$-NN cross-edge rate between the ground truth $p_t^\mathcal{C}$ and the approximations at various $t$. The thick lines show the means while the shaded areas show the $25$th and $75$th quantiles. PPR has a $k$-NN cross-edge rate closer to $0.5$ than any of the baselines, and is close to MMPS in Wasserstein. MMPS is performant because its Gaussian assumption for $p(x_0 | x_t)$ is a good approximation in this setting as discussed in Appendix \ref{['app:metrics']}.
  • Figure 2: Data2D results. Distribution of the constraint violation (log scale). To avoid $\log(0)$, $10^{-6}$ was added. PPR is an order of magnitude better than the best baselines. The lower constraint violation for the unconverged PDM than for DPS and TS is explained by $x_t$ getting stuck in local minima of the constraint.
  • Figure 3: Data2D results. Samples from the true marginal $p_t^\mathcal{C}$ (blue) and approximations (orange) at various diffusion times $t$. The samples are shown in the square $[-2, 2]^2$ for better visibility. PPR's marginals closely resemble the ground truth. PDM gets stuck in local minima far from the square, and is omitted for visibility, while DPS and TS sample the checkerboard prior.
  • Figure 4: KS results. The skill, RMSE, and CRPS measure the accuracy of the prediction. PPR achieves consistently lower RMSE, skill, spread, and CRPS, while reducing constraint violations by orders of magnitude. In DPS and $x_0$-space methods, there is no clear way to allocate more compute other than by running more diffusion steps.
  • Figure 5: KS results. Mean $\|u_0^i - u_0^{i-1} \|$ measures the mean norm of the difference between two timesteps, while the max is computed over all timesteps. In the base system, both the mean and max are close to $1$. Higher values represent strong discontinuities. PPR has the smoothest continuity, followed by TS. The discontinuities appear at the conditioning times, and are sharper for methods in $x_0$-space and for MMPS. In DPS, the samples are continuous because it is heavily biased towards the prior.
  • ...and 19 more figures