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

TL;DR

This paper introduces Predict-Project-Renoise (PPR), a novel algorithm that enables diffusion models to sample under hard constraints, maintaining data fidelity and satisfying physical laws in complex applications.

Contribution

The work develops a corrector kernel and an iterative projection method to enforce hard constraints in diffusion model sampling, a capability previously unavailable.

Findings

PPR achieves low constraint violations in complex models.

PPR maintains high distributional fidelity in constrained sampling.

Effective in applications like weather forecasting and physical sciences.

Abstract

Diffusion models cannot enforce hard constraints, yet applications in the physical sciences demand exact satisfaction of conservation laws, boundary conditions, and observational consistency. In this work, we identify a corrector kernel whose unique stationary distribution is the constrained marginal at each noise level, and approximate it by iteratively projecting through the denoiser and renoising via the forward kernel. The resulting Predict-Project-Renoise (PPR) algorithm enables sampling from pretrained diffusion models under hard constraints. Its three components are each necessary: projecting through the denoiser keeps samples close to the data manifold, while renoising and iterating drive samples toward the constrained marginal. On 2D distributions, the Kuramoto-Sivashinsky equation, and global weather forecasting with a $10^8$-dimensional atmospheric model, PPR simultaneously achieves low constraint violations and high distributional fidelity, a combination that existing methods fail to deliver.