Spatially-Aware Diffusion Models with Cross-Attention for Global Field Reconstruction with Sparse Observations

TL;DR

This study develops and enhances score-based diffusion models in field reconstruction tasks, and introduces a condition encoding approach to construct a tractable mapping mapping between observed and unobserved regions using a learnable integration of sparse observations and interpolated fields as an inductive bias.

Abstract

Diffusion models have gained attention for their ability to represent complex distributions and incorporate uncertainty, making them ideal for robust predictions in the presence of noisy or incomplete data. In this study, we develop and enhance score-based diffusion models in field reconstruction tasks, where the goal is to estimate complete spatial fields from partial observations. We introduce a condition encoding approach to construct a tractable mapping mapping between observed and unobserved regions using a learnable integration of sparse observations and interpolated fields as an inductive bias. With refined sensing representations and an unraveled temporal dimension, our method can handle arbitrary moving sensors and effectively reconstruct fields. Furthermore, we conduct a comprehensive benchmark of our approach against a deterministic interpolation-based method across various static and time-dependent PDEs. Our study attempts to addresses the gap in strong baselines for evaluating performance across varying sampling hyperparameters, noise levels, and conditioning methods. Our results show that diffusion models with cross-attention and the proposed conditional encoding generally outperform other methods under noisy conditions, although the deterministic method excels with noiseless data. Additionally, both the diffusion models and the deterministic method surpass the numerical approach in accuracy and computational cost for the steady problem. We also demonstrate the ability of the model to capture possible reconstructions and improve the accuracy of fused results in covariance-based correction tasks using ensemble sampling.

Figures (20)

• Figure 1: Schematic of the proposed condition encoding block with the UNet-based diffusion model, $F_\theta$, and two ways of encoding sensor information: (A) cross-attention and (B) classifier-free guidance.
• Figure 2: Schematic of the proposed condition encoding block. For CFG, mean-pooling is performed to reduce the dimensionality and to combine it with the noise scale embedding.
• Figure 3: Comparison of the generated permeability fields for the Darcy flow problem with 1.37% observed data points. Reverse sampling process of the diffusion models is configured with 20 steps, using a predictor-corrector scheme and a single trajectory. The black crosses denote the observed data points.
• Figure 4: Bar plot of nRMSE for the PDEs with 1% observed data points (1.37% for the Darcy flow) and various observation noise levels. The red dashed line denotes the error of reconstructing the field using the mean of the training data. The diffusion models are configured with 20 steps, with a predictor-corrector scheme and an ensemble of 25 trajectories.
• Figure 5: nRMSE of the PDEs with 1% observed data points (1.37% for the Darcy flow) for different numbers of reverse steps. The diffusion models are configured with a predictor-corrector scheme and an ensemble of 25 trajectories.