Diffusion models as probabilistic neural operators for recovering unobserved states of dynamical systems

TL;DR

This work shows that diffusion-based generative models can serve as probabilistic neural operators for PDEs, capable of forward prediction, inverse mapping, and reconstruction from partial observations. By introducing mixed conditional training, a single diffusion model learns multiple tasks, outperforming traditional neural operators on several dynamical systems and providing multiple plausible solutions when identifiability is incomplete. The approach leverages conditioning on observed data, partial masking, and PDE-consistency cues (PDE residuals) to improve accuracy and interpretability. While inference is slower due to sampling, the probabilistic framework offers a principled way to quantify uncertainty and select plausible reconstructions using additional information or PDE-based priors.

Abstract

This paper explores the efficacy of diffusion-based generative models as neural operators for partial differential equations (PDEs). Neural operators are neural networks that learn a mapping from the parameter space to the solution space of PDEs from data, and they can also solve the inverse problem of estimating the parameter from the solution. Diffusion models excel in many domains, but their potential as neural operators has not been thoroughly explored. In this work, we show that diffusion-based generative models exhibit many properties favourable for neural operators, and they can effectively generate the solution of a PDE conditionally on the parameter or recover the unobserved parts of the system. We propose to train a single model adaptable to multiple tasks, by alternating between the tasks during training. In our experiments with multiple realistic dynamical systems, diffusion models outperform other neural operators. Furthermore, we demonstrate how the probabilistic diffusion model can elegantly deal with systems which are only partially identifiable, by producing samples corresponding to the different possible solutions.

Figures (5)

• Figure 1: Conditional and mixed conditional training of diffusion models demonstrated with the Shallow-Water Equation (SWE-orig). The system has two variables (channels) represented by the colored rectangles with time ($t$) and spatial coordinate ($x$) on the x- and y-axes. The 'desired output' is the full state which we train the model to reconstruct from partial information. The clean parts of the input represent the conditioning information (at training and inference time) and the noisy parts are reconstructed by denoising. Each conditional model is trained with conditioning information for a single task. The mixed conditional training yields a single model for all defined tasks by sampling one task for each mini-batch during training.
• Figure 2: Model input for mixed conditional training.
• Figure 3: Example simulations from two of the studied systems.
• Figure 4: State reconstruction results for non-identifiable system SWE-init for a test sample with multiple possible outcomes.
• Figure 5: Detailed results for the non-identifiable system SWE-init. Left: PDE error vs. MAE for 100 samples for a single test case. Middle: Histogram of correlations between MAE and PDE error across the test set (red line shows the case on the left). Right: PDE errors of different models in the test set.