FunDiff: Diffusion Models over Function Spaces for Physics-Informed Generative Modeling

TL;DR

FunDiff addresses the challenge of generative modeling over continuous function spaces governed by physics. It combines a physics-aware function autoencoder (FAE) with a latent diffusion transformer (DiT) to generate continuous functions evaluable at arbitrary points while enforcing physical priors through PDE residuals and architectural constraints. The framework comes with a minimax-theoretic foundation for density estimation in infinite-dimensional spaces and demonstrates strong empirical performance across fluid dynamics and solid mechanics with robustness to noisy and low-resolution data. By separating constraint enforcement from latent-space diffusion and enabling exact differentiation in continuous space, FunDiff offers a principled, scalable approach to physics-informed function generation with practical impact for simulation, design, and inverse problems.

Abstract

Recent advances in generative modeling -- particularly diffusion models and flow matching -- have achieved remarkable success in synthesizing discrete data such as images and videos. However, adapting these models to physical applications remains challenging, as the quantities of interest are continuous functions governed by complex physical laws. Here, we introduce $\textbf{FunDiff}$, a novel framework for generative modeling in function spaces. FunDiff combines a latent diffusion process with a function autoencoder architecture to handle input functions with varying discretizations, generate continuous functions evaluable at arbitrary locations, and seamlessly incorporate physical priors. These priors are enforced through architectural constraints or physics-informed loss functions, ensuring that generated samples satisfy fundamental physical laws. We theoretically establish minimax optimality guarantees for density estimation in function spaces, showing that diffusion-based estimators achieve optimal convergence rates under suitable regularity conditions. We demonstrate the practical effectiveness of FunDiff across diverse applications in fluid dynamics and solid mechanics. Empirical results show that our method generates physically consistent samples with high fidelity to the target distribution and exhibits robustness to noisy and low-resolution data. Code and datasets are publicly available at https://github.com/sifanexisted/fundiff.

Figures (12)

• Figure 1: Overview of the function generative framework. The framework comprises three main steps. (1) We first train a function autoencoder (FAE) that maps discretized function data into a continuous latent space while preserving physical constraints. The encoder employs a Perceiver module to handle variable discretizations and project them into a unified latent representation. A Vision Transformer (ViT) processes the discretized inputs, and the decoder enables continuous function evaluation via cross-attention between query coordinates $x$ and encoded features. When known physical priors are available, the decoder is modified accordingly. (2) Next, we train a Diffusion Transformer (DiT) in the learned latent space to generate diverse function samples. For conditional generation, partial observations are encoded using the pretrained transformer and incorporated into the diffusion model inputs. The diffusion model is trained using rectified flow. (3) During inference, latent codes are sampled by solving a reverse-time ODE from Gaussian noise, and then decoded into continuous functions via the decoder. The framework naturally supports various physical constraints through its continuous latent representation, enabling applications in super-resolution, flow reconstruction, data purification, and non-intrusive prediction.
• Figure 2: Damped sinusoidal functions. Analysis of generative quality across models trained with different numbers of samples ($n=16,32,64,128,256$). (Top) Comparison between the worst-case generated samples (blue dots) and their reconstructed signals (red dashed lines) obtained through parameter estimation. (Bottom) The first four panels show the estimated parameter distributions for amplitude ($A$), damping coefficient ($\gamma$), angular frequency ($\omega$), and vertical shift ($b$). Red dashed lines indicate the true parameter bounds from the underlying uniform distributions. The rightmost panel displays the MSE of the reconstructed signal. The results demonstrate that models trained with more samples generate functions that better align with the true distributions and achieve lower reconstruction error.
• Figure 3: Kolmogorov flow. (a) Unconditional generation with enforced divergence-free constraint. (b) Comparison of spectral energy distributions between models constrained with and without divergence-free conditions. (c) Flow field reconstruction from noisy, low-resolution observations. Visualization of reconstructed velocity fields using different methods with input data of 64 × 64 resolution corrupted by 20% Gaussian noise. (d) Comparative analysis of reconstruction errors across methods at various downsampling factors. (e) Relative $L^2$ error of reconstructed velocity fields across varying noise levels ($\sigma$) for different spatial resolution reductions (1x, 2x, and 4x downsampling). Dashed lines show mean performance with (blue) and without (red) PDE constraints, while shaded regions indicate one standard deviation.
• Figure 4: Burgers' equation.(a) Solution reconstruction from noisy, low-resolution observations. Visualization of reconstructed solutions using different methods with input data of 40 × 40 resolution corrupted by 20% Gaussian noise. (b) Comparative analysis of reconstruction errors across methods at various downsampling factors. (c) Relative $L^2$ error of reconstructed solutions across varying noise levels ($\sigma$) for different spatial resolution reductions (1x, 2x, and 5x downsampling). Dashed lines show mean performance with (blue) and without (red) PDE constraints, while shaded regions indicate one standard deviation. (d) Histogram of PDE residual norms, comparing the physical consistency of solutions with and without PDE constraints, demonstrating improved physical consistency when PDE constraints are incorporated.
• Figure 5: Linear elasticity. (a) Comparison of diagonal strain concentration tensor distributions (Voigt notation) between provided coarse measurements ($32 \times 32$) and reconstructions from the trained FunDiff model. The generated samples maintain enforced symmetry constraints and achieve a relative $L^2$ error of $5.12\%$ when compared to ground truth data at $256 \times 256$ resolution. (b) Comparison of reconstruction errors under different noise types and downsampling rates.

Theorems & Definitions (13)

• Theorem 1: Density estimation over RKHS space
• proof
• Definition 2: Path norm neyshabur2015norm
• Definition 3: Spectral norm
• Lemma 4
• Theorem 5: Density estimation over Barron space
• proof
• Theorem 6: Minimax lower bound for Wasserstein estimation over an RKHS ball
• proof
• Theorem S1: Lower Bound, Theorem 3 of Ref. niles2022minimax