Technical Summary: Predict-Project-Renoise (PPR)

Problem Statement

Diffusion models have established themselves as state-of-the-art frameworks for deep generative modeling, effectively approximating high-dimensional data distributions $p(x)$. However, in scientific and engineering domains such as fluid dynamics, medical imaging, and climate modeling, generating samples that merely approximate the distribution $p$ is insufficient. These applications require samples $x$ to satisfy hard constraints $x \in C$, such as exact conservation laws, boundary conditions, observational consistency, or zero divergence in incompressible flows.

Existing approaches to constrained sampling face significant limitations:

1. Training-based methods often lack flexibility, requiring retraining for different constraints, and typically enforce constraints only softly.
2. Bayesian decomposition methods (treating constraints as noisy observations) are ill-suited for hard constraints because the implicit likelihood approaches a Dirac delta. This causes gradients to become unstable and informative only in the immediate vicinity of the feasible set.
3. Projection-based methods can enforce feasibility explicitly but suffer from theoretical and practical pitfalls. They often fail to sample the true constrained distribution, leading to three specific failures: (i) constraint violations, (ii) incorrect constrained distributions, and (iii) samples lying outside the support $S$ of the original data distribution (e.g., introducing discontinuities in continuous data). Furthermore, direct projection of noisy states is often ill-posed, and repeated projection steps can systematically bias the sampler, causing mode collapse.

Methodology: Predict-Project-Renoise (PPR)

The authors propose that correct constrained sampling requires approximating the constrained marginal distributions $p^C_t$ induced by a forward diffusion Stochastic Differential Equation (SDE) starting from the constrained set $C$.

Theoretical Foundation

The paper defines a constrained forward process where the initial state $x_0$ is drawn from $p^C(x_0) \propto p(x_0)\mathbb{1}_C(x_0)$. This induces a sequence of constrained marginals $\{p^C_t\}$. The corresponding backward process involves a constrained score function:
$$\nabla_{x_t} \log p^C_t(x_t) = \nabla_{x_t} \log p_t(x_t) + \nabla_{x_t} \log P(x_0 \in C | x_t)$$
The second term represents the gradient of the log-feasibility, pointing toward regions where the clean sample $x_0$ is likely to be feasible.

The PPR Algorithm

To sample from these constrained marginals without retraining, the authors identify a corrector kernel $K_t$ whose unique stationary distribution is $p^C_t$. Since sampling from the exact constrained posterior is intractable, they approximate this kernel using a three-step iterative process at each reverse diffusion step:

1. Predict: Use a standard predictor (e.g., DDIM) to move from $x_{t+1}$ to $x_t$.
2. Project: Instead of projecting the noisy state $x_t$ or the denoised estimate $d_\theta(x_t)$ directly, PPR projects through the denoiser. It solves an optimization problem to find a noisy state $x^*_t$ such that the denoised output satisfies the constraint:
   $$x^*_t = \arg \min_{x} c(d_\theta(x))$$
   This ensures the projected clean estimate $d_\theta(x^*_t)$ remains on the data manifold $S$, avoiding artifacts like discontinuities.
3. Renoise: The denoised, feasible estimate $d_\theta(x^*_t)$ is renoised via the forward kernel to obtain a new sample $x'_t$ at the current noise level $t$.

This Predict-Project-Renoise sequence is repeated $M$ times at each diffusion step. The renoising step is critical as it restores the correct noise level and injects the stochasticity required for mixing, preventing the sampler from collapsing to fixed points.

Key Contributions

1. Constrained Forward/Backward Formulation: The paper defines the constrained forward process and derives the associated backward process and score function, identifying a corrector kernel with the constrained marginal as its unique stationary distribution.
2. The PPR Algorithm: A training-free sampling algorithm that approximates the corrector kernel by iteratively projecting through the denoiser and renoising.
3. Necessity of Components: The authors demonstrate that all three components (prediction, projection through the denoiser, and renoising) are necessary. Removing any component leads to dramatic performance degradation, higher constraint violations, or samples falling outside the data support.
4. Theoretical Error Bounds: The paper provides an analysis of the approximation error, showing it decomposes into a mixing term (decaying exponentially with iterations $M$) and an irreducible bias from the point-estimate approximation.

Experimental Results

The method was evaluated on three distinct settings:

1. DATA2D: 2D distributions (checkerboard and Gaussian mixtures) with random constraints.
2. Kuramoto-Sivashinsky (KS): A 1D partial differential equation modeling turbulence, tested on 128-dimensional trajectories.
3. WEATHER: Global weather forecasting using the Appa latent weather model (108-dimensional atmospheric state).

Performance Metrics:

• Feasibility: PPR achieved a feasibility rate (samples satisfying $c(x) \le \epsilon$) of 82.7% on the KS dataset and 79.0% on DATA2D. In contrast, the next best methods achieved feasibility rates of roughly 1% (KS) and 31% (DATA2D).
• Constraint Violation: PPR consistently demonstrated the lowest constraint violations across all datasets.
• Distributional Fidelity: PPR produced samples that remained within the data support $S$ (preserving continuity and regularity), whereas many baseline methods (e.g., CPS, PDM, HDC) produced discontinuous samples that fell outside the data manifold.
• Accuracy: On KS and Weather tasks, PPR achieved lower RMSE, better Skill scores, and lower CRPS (Continuous Ranked Probability Score) compared to baselines like DPS, MMPS, and Trust Sampling (TS).

Significance and Claims

The paper claims that PPR successfully addresses the long-standing challenge of sampling from diffusion models under hard constraints without retraining. Its primary significance lies in its ability to simultaneously achieve low constraint violations and high distributional fidelity, a combination that existing methods fail to deliver.

The authors emphasize that PPR offers a meaningful trade-off between computational cost and sample quality: increasing the number of kernel applications ($M$) consistently improves the constrained distribution approximation. Furthermore, by projecting through the denoiser, the method ensures that samples remain consistent with the learned data manifold, avoiding the artifacts (such as discontinuities) common in $x_0$-space projection methods.

The work concludes that PPR is a robust, training-free solution for scientific applications requiring exact adherence to physical laws or observational data, though it notes limitations regarding non-differentiable or discrete constraints.