PaperlayerView exact revision on arXiv

Source-linked AI summary

CoCoGen: Physically-Consistent and Conditioned Score-based Generative Models for Forward and Inverse Problems

Christian Jacobsen, Yilin Zhuang, Karthik Duraisamy

arXiv:2312.10527v2cs.LG

TL;DR

The paper addresses the challenge of enforcing physical consistency in score-based models for PDE-governed problems. It incorporates PDE residual information into sampling and combines a frozen unconditional model with conditional augmentations for forward and inverse tasks. The reported approach reduces physical residuals while retaining prediction accuracy across the described applications.

  • Problem

    Score-based models have potential for physics tasks, but prior work had limited general PDE coverage and did not enforce physical behavior during sampling.

    Source: §1 · p. 2

  • Method

    The method incorporates discretized PDE residual information during sampling and uses conditional augmentations of a frozen unconditional score-based model for guided generation.

    Source: §1 · pp. 2–3

  • Results

    Physical consistency sampling systematically reduces physical residuals without compromising prediction accuracy in the described Darcy-flow applications.

    Source: §6 · pp. 19–20

  • Takeaways & Limitations

    The framework addresses both forward and inverse problems, with generated samples remaining physically plausible in the described settings.

    Source: §6 · p. 20

    The paper presents sparse-measurement inversion as an illustration of potential, not as a comparison against advanced solutions.

    Source: §5.2.1 · p. 19

Abstract

from arXiv · show

Recent advances in generative artificial intelligence have had a significant impact on diverse domains spanning computer vision, natural language processing, and drug discovery. This work extends the reach of generative models into physical problem domains, particularly addressing the efficient enforcement of physical laws and conditioning for forward and inverse problems involving partial differential equations (PDEs). Our work introduces two key contributions: firstly, we present an efficient approach to promote consistency with the underlying PDE. By incorporating discretized information into score-based generative models, our method generates samples closely aligned with the true data distribution, showcasing residuals comparable to data generated through conventional PDE solvers, significantly enhancing fidelity. Secondly, we showcase the potential and versatility of score-based generative models in various physics tasks, specifically highlighting surrogate modeling as well as probabilistic field reconstruction and inversion from sparse measurements. A robust foundation is laid by designing unconditional score-based generative models that utilize reversible probability flow ordinary differential equations. Leveraging conditional models that require minimal training, we illustrate their flexibility when combined with a frozen unconditional model. These conditional models generate PDE solutions by incorporating parameters, macroscopic quantities, or partial field measurements as guidance. The results illustrate the inherent flexibility of score-based generative models and explore the synergy between unconditional score-based generative models and the present physically-consistent sampling approach, emphasizing the power and flexibility in solving for and inverting physical fields governed by differential equations, and in other scientific machine learning tasks.

1 Introduction

The introduction identifies a gap in enforcing physical laws in score-based models for PDE problems and proposes incorporating governing equations into sampling. It applies frozen-backbone conditional models to forward and inverse Darcy-flow tasks.

1 Introduction

  • Score-based models directly learn the score function, avoiding restrictions on model or distribution form associated with likelihood-based approaches.The introduction situates score-based models among likelihood-based, implicit, diffusion, and score-based generative approaches.Source: §1 · pp. 1–2
  • Prior physics applications had limited general PDE coverage, and existing approaches described here did not enforce physical behavior during sampling.The introduction notes that residual gradients could be conditioning information without explicitly minimizing residuals during sampling.Source: §1 · p. 2
  • The sampling method uses discretized PDE information to reduce generated-sample residuals over baseline sampling without using the PDE residual during training.Unlike PINN-type methods, it uses the residual during sample generation rather than training and constrains samples rather than solving the PDE.Source: §1 · p. 2
  • The correction steps minimize residuals for general nonlinear PDEs, unlike image-reconstruction corrections based on a linear system.The paper distinguishes its approach from PINN-type methods by applying residuals during sampling and using a discrete residual based on the data.Source: §1 · pp. 2–3
  • Conditional score-based models built on a frozen unconditional backbone support Darcy-flow surrogate prediction and probabilistic reconstruction and inversion from sparse pressure measurements.The conditional models use a minimal-training augmentation approach adapted from ControlNet.Source: §1 · p. 3

2 Background

Score-based generation learns a score function and uses progressive noising followed by reverse-time sampling to produce data. The background presents SDE-based models and the reversible probability flow ODE as sampling formulations, noting different accuracy and efficiency advantages.

2 Background

  • Score-based generation learns the log-density gradient and generates samples by reversing a progressive noising process that makes score estimation possible in broader regions of data space.Noise is added until samples approximately follow an easy-to-sample prior, then removed iteratively.Source: §2 · p. 3
  • The forward SDE adds noise from the data distribution toward a prior, and its reverse-time dynamics approximately recover data samples from that prior.The SDE drift and diffusion coefficients are selected so the terminal distribution approximately follows an easy-to-sample prior; experiments use T = 1.Source: §2 · p. 4
  • A parameterized score model approximates the time-dependent score, and denoising score matching trains it using transition-kernel scores.The background states that the transition-kernel score has a closed-form solution and that the DSM objective uses a positive weighting function.Source: §2 · p. 4
  • The probability flow ODE has the same intermediate-time marginals as the reverse SDE and can be solved backward to sample from the data distribution.The background describes the PF ODE as reversible in time.Source: §2 · p. 4
  • For physics-based applications, the PF ODE offers accuracy advantages, while the reverse SDE can be more efficient in some cases such as RePaint.The paper suggests the PF ODE's lack of noise injection may aid denoising during sampling.Source: §2 · p. 4

3 Enforcing Physical Consistency and Conditioning

The section aims to enforce governing PDEs in generated samples through sampling-time residual reduction, without changing model training. It also describes conditional score models built on pretrained unconditional models.

3 Enforcing Physical Consistency and Conditioning

  • The section frames physical consistency as generating samples that satisfy specified governing equations, with the constraints applied at inference rather than during training.This leaves the training procedure unchanged and avoids retraining to impose those constraints.Source: §3 · p. 5

3.1 Unconditional Model Training

  • The unconditional score model uses a convolutional UNet and variance-preserving SDE, with βmin = 1 × 10^-4, βmax = 10, and T = 1.These SDE settings are used across the experiments.Source: §3.1 · pp. 5–6

3.2 Learning Conditional Score Functions

  • Conditional models pair fields with conditioning information and freeze a pretrained unconditional model while training a smaller conditional augmentation.Conditions can include physical parameters, macroscopic quantities, or partial field measurements; zero-convolution initialization allows the augmentation to begin with unhindered unconditional sampling.Source: §3.2 · pp. 6–7

3.3 Sampling

  • For the last N of τ sampling steps, the method appends a negative residual-gradient update to the SDE or ODE solver to reduce the PDE residual.It can also apply M additional consistency steps after reverse-SDE sampling; ϵ, τ, N, and M balance efficiency against physical consistency, and the authors note their chosen ϵ may not be optimal for every PDE or dataset.Source: §3.3 · pp. 7–8

4 Unconditional Generation of Darcy Flow Fields

The section demonstrates unconditional generation of Darcy flow fields and sampling-time enforcement of their governing PDE. Physical-consistency steps can reduce residuals while also denoising permeability fields, without directly modifying them.

4 Unconditional Generation of Darcy Flow Fields

  • The section motivates physical-consistency sampling because unconditional samples approximate the data distribution but generally do not satisfy the PDE.The authors present this sampling approach as an alternative to relying on more advanced architectures or finely tuned training to obtain consistently small residuals.Source: §4 · p. 8

4.1 Traditional Finite Difference Solution and Datasets

  • Darcy datasets pair sampled permeability fields with pressures obtained by solving the equations; the datasets use intrinsic dimensionalities s = 16 and s = 256.Each training set contains 10,000 samples on an n = 64 grid; an additional 1,000 samples are withheld for conditional experiments.Source: §4 · p. 8, §4.1 · p. 9

4.2 Residual Computation

  • The residual measures the discretized Darcy pressure equation, while its gradient is applied only to pressure; the integral condition is satisfied by construction.Residual gradients avoid direct permeability updates.Source: §4.2 · p. 10

4.3.1 Training

  • The unconditional model is trained on Darcy flow data with intrinsic dimensionality s = 16, and τ, N, and M can be varied after training to affect physical consistency.The sampling procedure solves the reverse SDE or PF ODE with physical-consistency steps; the choice of solver and hyperparameters affects final residuals.Source: §4.3 · p. 10, §4.3.2 · p. 10

4.3.2 Physically-Consistent Sampling

  • With sufficiently large τ, physical-consistency steps bring generated residuals to training-set levels or lower, while the PF ODE generally yields smaller residuals than reverse-SDE sampling.Reverse-SDE sampling requires more consistency steps to approach training-set residuals; its added noise may favor diversity when fidelity is not paramount.Source: §4.3.2 · p. 11, Figure 3 · p. 11
  • 4 PDE operator evaluations suffice to reach a generated-dataset residual no greater than the training-set residual, using the PF ODE with τ = 2000, M = 1, and N = 3.The study prioritizes residual reduction over sampling efficiency.Source: §4.3.2 · pp. 10–11
  • Applying consistency steps to pressure indirectly denoises permeability through the jointly predicted score functions, while reducing physical residuals.The steps are applied to pressure rather than permeability, which the authors leave unmodified.Source: §4.3.2 · pp. 11–12

5 Conditional Generation

Conditional score-based models reuse an unconditional model for forward prediction and inverse reconstruction, with physical-consistency sampling guiding outputs toward the governing PDE. The section illustrates these applications on Darcy flow and the 1D Burgers equation.

5 Conditional Generation

  • A single pretrained unconditional model is reused through separate conditional augmentations for parameter-based prediction and sparse-pressure reconstruction.This avoids training separate models for each condition.Source: §5 · p. 13

5.1 Conditional Field Generation from Generative Parameters

  • On Darcy-flow test parameters, average L2 errors are 7.0 × 10^-4 for pressure and 6 × 10^-4 for permeability predictions.Sampling uses physical-consistency enforcement, and each test parameterization is evaluated with one sample.Source: §5.1 · p. 13, Figure 5 · p. 14

5.2 Conditional generation for temporal PDE

  • For 1D Burgers, more physical-consistency steps reduce residual and slightly improve RMSE, while model RMSE is comparable to autoregressive baselines.The model generates all time slices simultaneously; residual error is mainly due to shock-wave position, with no error propagation across the evaluated time steps.Source: §5.2 · pp. 15–16, Table 1 · p. 16, Table 2 · p. 16

5.2.1 Conditional Field Inversion and Reconstruction with Sparse Measurements

  • With sparse pressure measurements, CoCoGen outperforms POD on reconstruction error and physical residual, though POD pressure reconstruction approaches perfect with full-field measurements.The comparison is intended to illustrate potential, not performance against advanced inverse-problem methods.Source: Figure 11 · p. 19, §5.2.1 · p. 19

6 Conclusions

CoCoGen combines physical-consistency sampling with conditional score-based models to address forward and inverse physics problems. The conclusions report reduced physical residuals without compromising prediction accuracy, while noting configuration caveats and possible benefits from more advanced solvers.

6 Conclusions

  • CoCoGen produces small residuals across applications; probability-flow ODE sampling yields smaller residuals than reverse-SDE sampling, while more consistency steps reduce residuals and noise.An excessively small τ can impede achieving small physical residuals.Source: §6 · p. 19
  • More advanced reverse-SDE or probability-flow ODE solvers may further reduce residuals, especially when combined with physical-consistency sampling.Source: §6 · p. 20
  • For 1D Burgers, time-slab predictions achieve low error against benchmark methods, and the last time step can condition unrolled long-time predictions.Source: §6 · p. 20
  • Physical-consistency sampling reduces physical residuals without compromising prediction accuracy, and the Darcy example shows pressure-field effects de-noising permeability while lowering residuals.The authors report these outcomes relative to baseline sampling methods.Source: §6 · p. 20
  • The framework is presented as a flexible basis for applying score-based generative models to physics-based and other scientific machine-learning tasks.Source: §6 · p. 20

Appendix A Solving Darcy Flow with a Linear System

The appendix constructs an overdetermined finite-difference linear system for Darcy pressure on a discretized domain. Its equations cover boundary and interior nodes plus an integral condition, and the system is solved by least squares.

Appendix A Solving Darcy Flow with a Linear System

  • The system has n^2 + 1 equations for pressure at n^2 spatial locations, partitioned into corner, edge, interior, and integral-enforcement equations.Source: Appendix A · p. 23
  • Corner and non-corner edge equations encode finite-difference relations among permeability K, pressure p, and source f_s at boundary nodes.The appendix gives four corner equations and 4(n − 2) edge equations.Source: Appendix A · pp. 23–24
  • Interior nodes have (n − 2)^2 equations, while the final equation approximates the integral condition with a two-dimensional trapezoidal rule; the overdetermined system is solved by least squares.Source: Appendix A · p. 24

Appendix B Unconditional Model and Training Details

The unconditional Darcy score models share a U-Net-type architecture and training setup across dataset sizes. Training uses a variance-preserving SDE and runs for 200,000 iterations on four GPUs.

Appendix B Unconditional Model and Training Details

  • Both Darcy datasets use identical unconditional U-Net-type models with 17.7M trainable parameters, trained for 5,000 epochs at learning rate 10^-4 and batch size 128.Source: Appendix B · p. 24
  • The models use a variance-preserving SDE with linear β(t), β_min = 10^-4, β_max = 10, and T = 1.Source: Appendix B · p. 24
  • Training took 24.5 hours for 200,000 iterations on four NVIDIA A6000 GPUs, including periodic sampling of eight fields per 1,000 optimization steps.Each batch of sampled fields took an average of 31 seconds.Source: Appendix B · p. 24

C.1 Darcy flow and ControlNet

The conditional augmentations use ControlNet-style architectures with condition encoders adapted to each task, while the unconditional backbone remains frozen. The appendix reports separate parameter counts and training times for surrogate modeling and field reconstruction/inversion.

C.1 Darcy flow and ControlNet

  • Conditional augmentations follow a ControlNet-based design, with condition encoders adapted to the input form and training otherwise matching the unconditional model.Source: §C.1 · p. 24
  • For surrogate modeling, a three-linear-layer encoder and reshaping/convolutional layers add 5.8M trainable parameters to a frozen 17.7M backbone; training converged in 11.6 hours.The combined model has 23.5M parameters and used 100,000 optimization iterations.Source: §C.1 · p. 25
  • For field reconstruction and inversion, a three-convolutional-layer encoder adds 5.7M trainable parameters to the frozen backbone; training converged in 18.1 hours.The combined model has 23.4M parameters and used 160,000 optimization iterations.Source: §C.1 · p. 25

C.2 1D Burgers equation and EDM

The section adapts classifier-free guidance to condition score-based generation on an initial velocity profile. It describes the sampling schedule and training setup for 1D Burgers experiments.

C.2 1D Burgers equation and EDM

  • For 1D Burgers, the initial velocity profile conditions the score, with γ = 1; the condition vector can be applied to the 2D channel through linear modulation with learnable coefficients.The conditional score replaces the original score function in the loss.Source: §C.2 · p. 25
  • Sampling solves the PF ODE for 1,000 steps, adding physical-consistency steps from 505 to 1,000 at intervals of five and varying their number from zero to five.The diffusion model matches the Darcy-flow model's depth and width.Source: §C.2 · p. 25
  • Training runs for 1,000 epochs with batch size 32 and learning rate 1 × 10^-4, using cosine annealing and warm-up.Source: §C.2 · p. 25

Appendix D Analytic Approximation to Conditional Sampling - Field Reconstruction and Inversion

The appendix approximates conditional sampling from an unconditional score model to reconstruct pressure fields and infer permeability from sparse pressure measurements. RePaint improves reconstruction and physical residuals, with added sampling cost.

Appendix D Analytic Approximation to Conditional Sampling - Field Reconstruction and Inversion

  • The method approximates conditional sampling without additional training by using a pretrained unconditional score to reconstruct pressure from measured locations and infer permeability.It approximates the conditional score using a full vector containing known and unknown dimensions; physical-consistency steps remain applicable.Source: Appendix D · pp. 25–26
  • Data imputation can be orders of magnitude more expensive with a reverse PF ODE than a reverse SDE; the forward SDE's closed-form sampling makes the procedure more feasible.Source: Appendix D · p. 26
  • At m = 500, RePaint r = 3 versus r = 0 lowers average physical residual from 1970.2 to 8.7, average reconstruction error from 9.1 × 10^-3 to 3.1 × 10^-5, and field inversion error from 6.4 × 10^-3 to 3.8 × 10^-3.Figure 12 qualitatively shows more accurate, less noisy reconstructions with RePaint.Source: Table 4 · p. 26, Figure 12 · p. 27, Appendix D · pp. 27–28
  • With fewer measurements, physical residuals remain high despite physical-consistency enforcement, and measured points may mismatch reconstructed points.Source: Appendix D · p. 27
  • RePaint alternates forward and reverse SDE steps to mix measured and unmeasured dimensions, but score evaluations grow linearly with resampling count r.Resampling only at selected time steps can partly reduce cost but requires choosing those steps.Source: Appendix D · p. 27
  • RePaint reconstruction outperforms ControlNet-based conditional generation, but the best results require around 5× more compute time.Source: Appendix D · p. 28
Loading 2312.10527v2…