Technical Summary: Softly Constrained Denoisers for Diffusion Models Applied to Partial Differential Equations

1. Problem Statement

Diffusion models have emerged as powerful generative priors for solving Partial Differential Equations (PDEs) in fields such as fluid dynamics and material science. However, a central challenge remains: effectively incorporating governing PDE constraints into the generative model without compromising the model's ability to learn the true data distribution.

Existing approaches typically enforce constraints in two ways:

1. Training-time Regularization: Adding PDE residuals as loss terms (e.g., Physics-Informed Diffusion Models).
2. Inference-time Adjustments: Modifying the sampling trajectory via guidance (e.g., using constraint gradients during generation).

The paper argues that these methods suffer from critical limitations, particularly when the governing PDE is misspecified (i.e., the mathematical model is an approximation that does not perfectly match observed data).

• Regularizers bias the generative distribution away from the true data distribution. They force the model to satisfy the constraint even when the constraint is incorrect, leading to a loss of data fidelity.
• Inference-time guidance introduces approximation errors during sampling and can similarly bias the output distribution.

The core problem is that PDE models are inherently approximate (e.g., Darcy's law assumes laminar flow; wave numbers may be measured with noise). A robust generative model must exploit these constraints where they are informative but retain the flexibility to deviate from them when the constraints are misspecified, thereby converging to the true underlying data distribution $p_{data}(x)$.

2. Methodology: Softly Constrained Denoisers (SCD)

The authors propose Softly Constrained Denoisers (SCD), an architectural modification that embeds soft inductive biases directly into the denoiser network, rather than enforcing constraints via external loss terms or sampling adjustments.

Core Architecture

The SCD decomposes the denoiser into two components:

1. Base Denoiser ($D^{orig}_{\theta}$): A standard network (e.g., U-Net) that learns the conditional mean $E[x_0 | x_t]$ from the data.
2. Constraint-Informed Correction: A term derived from the PDE residual, modulated by a learnable scaling factor.

The final denoiser output is defined as:
$$D^{SCD}_{\theta}(x_t, t) = D^{orig}_{\theta}(x_t, t) + \gamma_{\theta}(x_t, t)\sigma(t)^2 \nabla_{D^{orig}_{\theta}} \log l_c(D^{orig}_{\theta}(x_t, t))$$

Where:

• $l_c(x) = \exp(-\|R(x)\|)$ is a relaxed constraint function based on the PDE residual $R(x)$.
• $\gamma_{\theta}(x_t, t)$ is a learnable scaling factor (a small network) that determines how much the model relies on the constraint at each noise level $t$.
• The correction term uses the gradient of the constraint evaluated at the base denoiser's output, avoiding the computationally expensive vector-Jacobian products required by standard guidance methods.

Theoretical Guarantees

The paper proves that this parameterization preserves the fundamental guarantees of diffusion models:

• Optimality: At convergence, the SCD still approximates the optimal conditional mean $E[x_0 | x_t]$. The correction term effectively learns to project the deviation of the base denoiser from the optimum onto the constraint gradient.
• Unbiasedness: Unlike regularized approaches, SCD does not shift the optimal denoiser away from the true conditional expectation. Consequently, the model retains the ability to converge to the true data distribution $p_{data}(x)$ even if the constraint is misspecified.
• ELBO Preservation: The method does not degrade the Evidence Lower Bound (ELBO) of the diffusion model, whereas regularized objectives strictly degrade it.

Training

The model is trained using the standard Denoising Score Matching (DSM) loss:
$$L = E_{t, x_0, x_t} [w(t) \| D^{SCD}_{\theta}(x_t, t) - x_0 \|^2]$$
No additional regularization terms are added to the loss function. The network learns when and how much to utilize the constraint information through the parameter $\gamma_{\theta}$.

3. Key Contributions

1. Architectural Innovation: A method to transform any standard denoiser into a "softly constrained" version with minimal computational overhead, embedding PDE knowledge directly into the network structure.
2. Theoretical Analysis: Formal proofs demonstrating that previous regularizer-based methods break the distributional guarantees of diffusion models and are sensitive to constraint misspecification. Conversely, the paper proves that SCDs preserve these guarantees.
3. Empirical Validation: Demonstration of superior performance in misspecified settings across toy examples and PDE benchmarks (Darcy Flow and Helmholtz Equation), showing better constraint satisfaction than existing methods while maintaining sample quality and robustness.

4. Experimental Results

The authors evaluated SCD against Vanilla Diffusion, Physics-Informed Diffusion Models (PIDM), and Diffusion Posterior Sampling (DPS) guidance.

Illustrative Examples (Toy Data)

• Setup: Generating points on a circle with introduced "misspecifications" (a "dent" or a "chopped" segment).
• Findings:
  • PIDM failed to learn the true distribution when constraints were misspecified, placing no samples on the "dent" and producing high Wasserstein-1 ($W_1$) distances.
  • DPS showed increased $W_1$ distance as misspecification increased.
  • SCD maintained low $W_1$ distances comparable to Vanilla diffusion, successfully learning the true data distribution even when the constraint was partially incorrect.

Darcy Flow Benchmark

• Task: Generating permeability ($K$) and pressure ($p$) fields satisfying $\nabla \cdot (K\nabla p) + f = 0$.
• Misspecification: The flow measurement $f$ was scaled incorrectly to simulate calibration errors.
• Findings:
  • PIDM residuals increased significantly with misspecification, performing worse than the Vanilla baseline.
  • SCD preserved the Negative Log-Likelihood (NLL) (indicating high distributional fidelity) while maintaining low residuals. It successfully adapted to the correct regions of the constraint while ignoring the misspecified parts.

Helmholtz Equation Benchmark

• Task: Solving $\nabla^2 u + k^2 u - a = 0$ with noisy wave number measurements.
• Comparison: SCD vs. FunDPS (a guidance-based method using reconstruction loss and residuals).
• Findings:
  • FunDPS heavily favored reconstruction error, sacrificing physical consistency (high residuals).
  • SCD achieved the most competitive residual values while preserving NLL, demonstrating robustness to noisy gradients in the constraint function.

5. Significance and Claims

The paper claims that Softly Constrained Denoisers offer a "drop-in upgrade" for diffusion-based generative modeling in scientific applications. The primary significance lies in the balance between data fidelity and constraint compliance:

• SCDs exploit constraint knowledge to improve physical plausibility without forcing the model into a biased distribution when the physics model is imperfect.
• They resolve the trade-off where regularizers force the model to obey incorrect physics, and inference-time guidance introduces sampling errors.
• The approach is particularly valuable for inverse problems and scientific simulations where governing equations are approximations of complex real-world phenomena.

The authors note limitations, including the requirement for differentiable constraints and the assumption of clean training data. They suggest future work could explore deriving differentiable constraints from non-differentiable ones and tuning the trade-off between distributional fidelity and constraint compliance along a Pareto front.