FunDiff: diffusion models over function spaces for physics-informed generative modeling

Abstract

Recent advances in generative modeling-particularly diffusion models and flow matching-have been widely used for 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. To address this, we introduce FunDiff, an efficient and robust 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, generates continuous functions that can be evaluated 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, demonstrating that diffusion-based estimators achieve optimal convergence rates under suitable regularity conditions. We further demonstrate the practical effectiveness of FunDiff across diverse applications in fluid dynamics and solid mechanics. Empirical results indicate that our method can generate physically consistent samples with high fidelity to the target distribution, and exhibit robustness to noisy and low-resolution data.

Generative AI excels at images and video, but many scientific problems involve continuous physical fields. Here, authors introduce FunDiff, a function-space generative model that produces high-fidelity, physics-consistent function samples.

Introduction

Generative modeling, a cornerstone of modern machine learning, aims to learn and sample complex data distributions P(x) by constructing models that can generate new instances consistent with the learned distribution^1,2. Among recent advances in this domain, diffusion models and flow matching have emerged as powerful frameworks that achieve state-of-the-art performance. Diffusion models^3,4 learn to reverse a forward noising process by gradually denoising data through a Markov chain, which iteratively transforms random noise into high-quality samples. In parallel, flow matching^5,6 provides an alternative approach by learning continuous normalizing flows via optimal transport theory, directly modeling probability paths that transform a tractable prior distribution into the target data distribution.

Diffusion models have shown strong performance across a wide range of applications, including high-resolution image synthesis^7–9, audio generation¹⁰, protein structure prediction¹¹, and molecular design¹². Despite these advances, the focus has largely remained on discrete data types, such as images, text, and molecular graphs, whereas many scientific domains, especially in physics, demand the generation of continuous function data, where each sample represents a continuous function (e.g., a velocity or temperature field) defined over spatial and temporal domains. For instance, generating velocity fields in turbulent fluid dynamics, simulating deformation patterns in complex materials under external loadings, and forecasting temperature variations in climate models all require handling functions defined over continuous spatial and temporal domains. These problems inherently involve infinite-dimensional function spaces, a challenge that recent advances in operator learning frameworks have begun to address effectively^13–19. This presents unique challenges:

How can we extend generative modeling to generate continuous functions that can be evaluated at arbitrary points, while preserving important structural or physical properties inherent in scientific problems?

Prior work in computer vision and graphics has extensively explored the concept of function-space generation through implicit neural representations and generative models, such as DeepSDF²⁰, NeRF²¹, From Data to Functa²², and SIREN²³, as well as their diffusion-based extensions including 3DShape2VecSet²⁴, Shap-E²⁵, and HyperDiffusion²⁶. These methods treat signals such as shapes, radiance fields, or signed distance functions as continuous functions parameterized by neural networks, establishing an important foundation for generative modeling in continuous domains. Furthermore, recent works have attempted to extend the framework of generative models to function space, either by adapting existing algorithms to operate on discretized function data in ${\mathbb{R}}^n$ or by transforming input functions into a finite-dimensional latent space^22,27–31. Dupont et al.³² pioneered this direction by treating data points as continuous functions and learning distributions over them using adversarial training. Multiple theoretical frameworks have emerged for function diffusion models, including Gaussian measures on Hilbert spaces³³, dimension-independent bounds for infinite-dimensional models³⁴, generalized diffusion processes in function spaces³⁵, and denoising diffusion operators with rigorous Gaussian processes and Langevin dynamics³⁶.Although these studies have demonstrated promising results for fixed discretizations, they generally struggle to scale to higher resolutions or to handle variable discretizations consistently, as those encountered in scientific computing, partial differential equation (PDE) modeling, or operator learning^14,37–40.

In parallel, diffusion models have emerged as powerful tools for tackling inverse problems and modeling dynamical systems governed by PDEs. Recent work has developed methods for reconstructing spatiotemporal fields from sparse observations and partial supervision^41–43. To address resolution limitations, several approaches have integrated neural operators with diffusion models^44–46, demonstrating improved performance in turbulence modeling and capture of multiscale phenomena. Furthermore, various strategies have attempted to incorporate physical priors directly into diffusion processes^42,47–52. For instance, DiffusionPDE⁴² introduces physics-based guidance during the posterior sampling stage, whereas PIDM⁴⁹ embeds PDE residuals directly into the training objective. However, enforcing PDE residuals is highly sensitive to solution accuracy; small numerical errors can destabilize both training and inference. Moreover, physical constraints are inherently defined in continuous space, but these diffusion models operate on discretized representations. This mismatch introduces discretization artifacts, often leading to inaccurate or ill-posed PDE residuals, especially in inverse problems where key state variables are only partially observed. These challenges underscore a fundamental gap between current data-driven generative models and the objective of producing physically consistent, continuous functions.

To address these challenges, we propose FunDiff, a machine learning framework for physics-informed generative modeling over function spaces. Our contributions are as follows. We introduce a unified architecture that combines latent diffusion processes with function autoencoders, enabling high-quality data generation in a function-space representation. This design supports evaluation at arbitrary spatial locations while seamlessly incorporating physical priors via the function-space representation. We establish theoretical guarantees for our approach by proving that diffusion-based models attain minimax-optimal rates for density estimation in infinite-dimensional spaces under mild smoothness assumptions. We demonstrate the practical effectiveness of FunDiff across diverse applications in fluid dynamics and solid mechanics. Our experiments show that the method generates physically consistent samples with high fidelity to target distributions while maintaining robustness to noisy and low-resolution input data. Taken together, our framework bridges the gap between theoretical guarantees in function spaces and practical neural network algorithms, establishing a principled approach to physics-informed generative modeling.

Results

In this section, we present theoretical results on generative models in function spaces, which establish the mathematical foundation of our work. Building on this theory, we introduce FunDiff, a practical framework for generating functions while incorporating physical priors. We then demonstrate the effectiveness of our proposed FunDiff framework through a series of experiments across diverse physical systems. Specifically, we present four case studies with increasing complexity: (1) the Burgers’ equation, showcasing the incorporation of PDE constraints; (2) Kolmogorov flow, highlighting divergence-free field generation; (3) linear elasticity, illustrating symmetry preservation; and (4) turbulent mass transfer, demonstrating conditional generation capabilities. For notations, interested readers are encouraged to refer to the Supplementary Table S1. It is important to emphasize that in all examples, FunDiff does not act as a neural PDE solver; rather, it learns to generate physically consistent continuous functions that approximate the distributions of solutions to selected PDE-governed systems when trained on such data.

Generative models over function spaces

In this section, we introduce a framework for generative modeling in function spaces. Given the inherently challenges of handling directly with infinitely-dimensional spaces, we begin by reducing the problem to a finite-dimensional setting, with preliminaries in Supplementary Section B. By embedding functions into a finite-dimensional vector space, the task of generating new functions can be reformulated as a well-studied vector generation task.

To formalize this idea, let us begin with a vector space generative oracle $\mathcal{O}\left(d,n\right)$, which, given n i.i.d. samples ${\mathbf{x}}^1,{\mathbf{x}}^2,\dots ,{\mathbf{x}}^n\in {\mathbb{R}}^d$ from an underlying distribution $\mu \in \mathcal{P}\left({\mathbb{R}}^d\right)$, produces an estimate $\widehat{\mu }\in \mathcal{P}\left({\mathbb{R}}^d\right)$ of μ. This oracle is assumed to achieve the minimax statistical estimation rate

$$\mathbb{E}\left[W_1\left(\widehat{\mu },p\right)\right]\lesssim \epsilon \left(d,n\right),$$

where the expectation is taken over the randomness of the samples, and ε(d, n) aligns with the known minimax estimation rates^53,54. For the continuous density class $\mathcal{C}\left({\mathbb{R}}^d\right):=\left\{\mu :{\mathbb{R}}^d\to \left[0,\infty \right):\mu \mathrm{is\; continuous}\right\}$, we have a minimax density estimation rate

$${\sup }_{\widehat{\mu }}{\inf }_{\mu \in \mathcal{C}\left({\mathbb{R}}^d\right)}\mathbb{E}\left[W_1\left(\widehat{\mu },\mu \right)\right]\asymp n^{-1/d}$$

for d > 2, where $\widehat{\mu }$ is generated by diffusion models in this paper. It is a special case for the $n^{-\frac{\alpha +1}{2\alpha +d}}$ minimax rate when the smoothness α = 0. Therefore, for the entire continuous density class $\mathcal{C}\left({\mathbb{R}}^d\right)$, we have ε(d, n) = n^−1/d.

Building upon this oracle, we seek to develop a generative modeling method that operates over a space of functions. Suppose that we have n functions $f^{\left(1\right)},f^{\left(2\right)},\dots ,f^{\left(n\right)}:{\mathbb{R}}^d\to \mathbb{R}$, drawn i.i.d. from an unknown distribution $P\in \mathcal{P}\left(\mathcal{F}\right)$, where $\mathcal{F}$ is a subset of the continuous function space $\mathcal{C}\left({\mathbb{R}}^d\right)$. Our goal is to devise a theoretically sound methodology for generating new function samples from P, given only these observed samples.

To achieve this, we introduce a function autoencoder (FAE) framework. We define an encoder $\mathcal{E}:\mathcal{F}\to {\mathbb{R}}^D$, which maps a function to a finite-dimensional vector representation, and a decoder $\mathcal{D}:{\mathbb{R}}^D\to \mathcal{F}$, which maps these vectors back to functions. Applying $\mathcal{E}$ to the given functions f⁽¹⁾, f⁽²⁾, …, f⁽ⁿ⁾ yields vector samples ${\theta }^{\left(1\right)},{\theta }^{\left(2\right)},\dots ,{\theta }^{\left(n\right)}\in {\mathbb{R}}^D$. We then use the vector-space generative oracle $\mathcal{O}\left(D,n\right)$ on these vectors to obtain an estimated distribution $\widehat{p}\in \mathcal{P}\left({\mathbb{R}}^D\right)$. Finally, we can sample new vectors $\theta \in {\mathbb{R}}^D$ from $\widehat{p}$ and apply $\mathcal{D}$ to generate new functions $f=\mathcal{D}\left(\theta \right)\in \mathcal{F}$, the distribution of these newly generated functions is indicated by $\widehat{P}$ and serves as an estimation of P. This construction effectively leverages the theoretical foundations of vector-space generative modeling and extends them to the realm of infinite-dimensional function spaces.

In order to measure the distribution loss, we apply an analogy of the Wasserstein distance $W_1\left(P,\widehat{P}\right)$. Notice that for distributions over the function space, the Wasserstein distance is defined similarly as

$$W_1\left(P,\widehat{P}\right):={\sup }_{\mathrm{Lip}\left(\mathcal{L}\right)\le 1}\left[{\mathbb{E}}_{f~P}\mathcal{L}\left(f\right)-{\mathbb{E}}_{f~\widehat{P}}\mathcal{L}\left(f\right)\right],$$

where $\mathcal{L}$ is a functional and its Lipschitz constant means:

$$\mathrm{Lip}\left(\mathcal{L}\right)={\sup }_{f,g}\frac{\mathcal{L}\left(f\right)-\mathcal{L}\left(g\right)}{{∥f-g∥}_2}.$$

In the next part, we propose an upper bound of $W_1\left(P,\widehat{P}\right)$ and explain all its factors.

Density estimation loss decomposition

Suppose the vector representations $\left\{{\theta }^{\left(k\right)}=\mathcal{E}\left(f^{\left(k\right)}\right)\mid k\in \left[n\right]\right\}$ are drawn from a distribution p = E_♯P, which stands for the pushforward of P under the encoder $\mathcal{E}$. Applying the vector-space generative oracle $\mathcal{O}\left(D,n\right)$ to these samples yields an estimated distribution $\widehat{p}\in \mathcal{P}\left({\mathbb{R}}^D\right)$. We then feed $\widehat{p}$ into the decoder $\mathcal{D}$, producing $G_{♯}\widehat{p}$, an estimate of the original distribution P in the function space. For the estimation loss, we have the bound

$$W_1\left(P,\widehat{P}\right)\le {\mathbb{E}}_{f~P}\parallel f-\mathcal{D}\circ \mathcal{E}\left(f\right){\parallel }_2+\mathrm{Lip}\left(\mathcal{D}\right)\cdot W_1\left(p,\widehat{p}\right).$$ 1

Here, the first term ${\mathbb{E}}_{f~P}\parallel f-\mathcal{D}\circ \mathcal{E}\left(f\right){\parallel }_2$ represents the generalization error for the autoencoder $\mathcal{D}\circ \mathcal{E}:\mathcal{F}\to \mathcal{F}$. The second term $\mathrm{Lip}\left(\mathcal{D}\right)$ stands for the Lipschitz constant of the decoder $\mathcal{D}$, defined by the inequality

$${∥\mathcal{D}\left(\theta \right)-\mathcal{D}\left({\theta }^{\prime }\right)∥}_2\le \mathrm{Lip}\left(\mathcal{D}\right)\cdot {∥\theta -{\theta }^{\prime }∥}_2.$$

This quantity measures the robustness of $\mathcal{D}$ with respect to perturbations in the input vector space. The third term, $W_1\left(p,\widehat{p}\right)$, reflects the accuracy of the density estimation for vector space, highlighting the performance of the oracle $\mathcal{O}\left(D,n\right)$ in estimating the true distribution p. The dimension of the feature space D must be chosen judiciously to strike a balance between ensuring accurate autoencoder reconstruction and mitigating the effects of the curse of dimensionality for the vector-space generative oracle. The proof of (1) can be seen in Supplementary Section C.1. Next, we propose some simple examples to help us understand the theoretical framework.

Illustrative example 1: reproducing kernel Hilbert space (RKHS)

For a reproducing kernel Hilbert space (RKHS) $\mathcal{H}$, we denote its eigenvalues as μ₁ ⩾ μ₂ ⩾ …, and their corresponding eigenfunctions as ψ₁, ψ₂, …. For any function $f\in \mathcal{H}$, its eigen-decomposition has the form $f={\sum }_{i=1}^{\infty }{\theta }_i\sqrt{{\mu }_i}{\psi }_i$ and its RKHS norm is defined as $\parallel f{\parallel }_{\mathcal{H}}^2={\sum }_{i=1}^{\infty }{\theta }_i^2$. Define the RKHS unit ball as $\mathcal{B}:=\left\{f\in \mathcal{H}:\parallel f{\parallel }_{\mathcal{H}}\le 1\right\}$.

Given n sample functions $f^{\left(1\right)},f^{\left(2\right)},\dots ,f^{\left(n\right)}~P\in \mathcal{P}\left(\mathcal{B}\right)$, assume their eigen-decomposition is

$$f^{\left(k\right)}={\sum }_{i=1}^{\infty }{\theta }_i^{\left(k\right)}\cdot \sqrt{{\mu }_i}{\psi }_i,\left(k=1,2,\dots ,n\right).$$

We define the encoder $\mathcal{E}:\mathcal{B}\to {\mathbb{R}}^D$ as the truncation of D coefficients and the decoder $\mathcal{D}:{\mathbb{R}}^D\to \mathcal{B}$ as the linear combination:

$$\mathcal{E}:f\mapsto {\left({\theta }_1,{\theta }_2,\dots ,{\theta }_D\right)}^{\top },\mathcal{D}:{\left({\theta }_1,{\theta }_2,\dots ,{\theta }_D\right)}^{\top }\mapsto {\sum }_{i=1}^D{\theta }_i\cdot \sqrt{{\mu }_i}{\psi }_i.$$

Hence, the autoencoder $\mathcal{D}\circ \mathcal{E}$ effectively serves as a projection (truncation) onto the subspace spanned by the first D eigenfunctions. Next, in order to upper bound the density estimation loss $W_1\left(P,\widehat{P}\right)$, we introduce two different assumptions corresponding to the polynomial and exponential eigenvalue decay of the RKHS $\mathcal{H}$.

Assumption 1

(2β-polynomial decay) Assume that the RKHS $\mathcal{H}$ is embedded to a β-th order Sobolev space where $\beta >\frac{1}{2}$. It leads to the property that the eigenvalues with respect to the kernel function K( ⋅ , ⋅ ) have the property of 2β-polynomial decay as μ_i ≍ i^−2β for all i ⩾ 1. To be more specific, we assume that

$$c_{\mathcal{H}}\cdot i^{-2\beta }\le {\mu }_i\le C_{\mathcal{H}}\cdot i^{-2\beta }$$

holds for  ∀ i ⩾ 1, where $c_{\mathcal{H}}$ and $C_{\mathcal{H}}$ are constants.

Assumption 2

(γ-exponential decay) Assume that the eigenvalues with respect to the kernel function K( ⋅ , ⋅ ) have the property of γ-exponential decay as ${\mu }_i\asymp \cdot \exp \left(-C_1\cdot i^{\gamma }\right)$ for all i ⩾ 1 where C₁ > 0 is a positive constant. To be more specific, we assume that

$$c_{\mathcal{H}}\cdot \exp \left(-C_1\cdot i^{\gamma }\right)\le {\mu }_i\le C_{\mathcal{H}}\cdot \exp \left(-C_1\cdot i^{\gamma }\right)$$

holds for  ∀ i ⩾ 1, where $c_{\mathcal{H}}$ and $C_{\mathcal{H}}$ are constants.

Then, the reconstruction loss ${\mathbb{E}}_{f~P}\parallel f-\mathcal{D}\circ \mathcal{E}\left(f\right){\parallel }_2$ represents the truncation loss under the RKHS eigen-decomposition, and the robustness $\mathrm{Lip}\left(\mathcal{D}\right)$ is directly related to the RKHS norm bound of $\mathcal{B}$.

Theorem 1

(Density estimation over RKHS space) Assume P is a distribution supported on the RKHS unit ball $\mathcal{B}:=\left\{f\in \mathcal{H}:\parallel f{\parallel }_{\mathcal{H}}\le 1\right\}$. By using the generative oracle $\mathcal{O}\left(D,n\right)$ and the function autoencoder (FAE) framework introduced above, the density estimation loss over RKHS space can be bounded as follows: Under the polynomial decay setting:

$$W_1\left(P,\widehat{P}\right)\lesssim {\left(\frac{\log \log n}{\log n}\right)}^{\beta }.$$

Under the exponential decay setting:

$$W_1\left(P,\widehat{P}\right)\lesssim \exp \left(-{\left(C_1/2\right)}^{\frac{1}{1+\gamma }}{\left(\log n\right)}^{\frac{\gamma }{1+\gamma }}\right),$$

after we choose an optimal latent dimension D.

Proof

The proof is provided in Supplementary Section C.2.

Illustrative example 2: Barron space

First, we introduce the concept of Barron space⁵⁵. For function $f:\mathrm{\Omega }:={\left[0,1\right]}^d\to \mathbb{R}$, we aim to recover f using a fully connected two-layer neural network with ReLU activation:

$$f\left(\mathbf{x};\theta \right)={\sum }_{k=1}^m{a}_k\sigma \left({\mathbf{b}}_k\cdot \mathbf{x}+c_k\right).$$

Here, σ( ⋅ ) is the component-wise ReLU activation $\sigma \left(t\right)=\max \left(0,t\right)$, ${\mathbf{b}}_k\in {\mathbb{R}}^d,a_k,c_k\in \mathbb{R}$ for k ∈ [m], and the whole parameter set $\theta ={\left\{\left(a_k,{\mathbf{b}}_k,c_k\right)\right\}}_{k=1}^m$ is to be learned, and m is the width of the network. As we can see, $\theta \in {\mathbb{R}}^{m\left(2+d\right)}$. To control the magnitude of the learned network, we use the following scale-invariant norm.

Definition 1

(Path norm⁵⁶) For a two-layer ReLU network, the path norm is defined as

$$\parallel \theta {\parallel }_{\mathrm{P}}={\sum }_{k=1}^m\mid a_k\mid \cdot \left(\parallel {\mathbf{b}}_k{\parallel }_1+\mid c_k\mid \right).$$

Definition 2

(Spectral norm) Given f ∈ L²(Ω), denote by $F\in L^2\left({\mathbb{R}}^d\right)$ an extension of f to ${\mathbb{R}}^d$. Let $\widehat{F}$ be the Fourier transform of F, then

$$f\left(\mathbf{x}\right)={\int }_{{\mathbb{R}}^d}{e}^{i⟨\mathbf{x},\omega ⟩}\widehat{F}\left(\omega \right)\mathrm{d}\omega ,\forall \mathbf{x}\in \mathrm{\Omega }.$$

We define the spectral norm of f by

$$\gamma \left(f\right)={\inf }_{F\in L^2\left({\mathbb{R}}^d\right),F{\mid }_{\mathrm{\Omega }}=f{\mid }_{\mathrm{\Omega }}}{\int }_{{\mathbb{R}}^d}\parallel \omega {\parallel }_1^2\cdot \mid \widehat{F}\left(\omega \right)\mid \mathrm{d}\omega .$$

Assumption 3

We consider all functions with a bounded value and a bounded spectral norm.

$${\mathcal{F}}_s=L^2\left(\mathrm{\Omega }\right)\cap \left\{f\left(\mathbf{x}\right):\mathrm{\Omega }\to \mathbb{R}\mid \gamma \left(f\right)\le 1,\parallel f{\parallel }_{\infty }\le 1\right\}.$$

We let the function class $\mathcal{F}={\mathcal{F}}_s$, which makes the distribution P supported on Barron space ${\mathcal{F}}_s$.

In the following analysis, we mainly rely on the approximation property for two-layer ReLU networks over the Barron space ${\mathcal{F}}_s$, which are proposed in Refs. ^55,57,58.

Lemma 1

For any $f\in {\mathcal{F}}_s$, there exists a two-layer network $f\left(\mathbf{x};\tilde{\theta }\right)$ of width m such that:

$${∥f\left(\mathbf{x};\tilde{\theta }\right)-f\left(\mathbf{x}\right)∥}_{\infty }^2\le \frac{16{\gamma }^2\left(f\right)}{m}\le \frac{16}{m}.$$

Furthermore, the parameter $\tilde{\theta }$ can be bounded as:

$$\parallel {\mathbf{b}}_k{\parallel }_1=1,\mid c_k\mid \le 1,\mathrm{and}\mid a_k\mid \le \frac{2\gamma \left(f\right)}{m}\le \frac{2}{m}\mathrm{holds\; for}\forall k\in \left[m\right].$$

In the context of the FAE framework, we set our encoder as $\mathcal{E}:f\mapsto \tilde{\theta }$. If there exist multiple valid $\tilde{\theta }$ for f ∈ F_s, we can choose any one of them as the encoder’s output. As we see, the encoder $\mathcal{E}$ maps the function of the Barron space to a parameter set with bounded path norm. The decoder $\mathcal{D}:\theta \mapsto f\left(\mathbf{x};\theta \right)$ is simply a parameter injection of a two-layer ReLU network. Hence, the autoencoder $\mathcal{D}\circ \mathcal{E}$ effectively serves as a two-layer ReLU network approximation. Next, we upper-bound the density estimation loss $W_1\left(P,\widehat{P}\right)$.

Theorem 2

(Density estimation over Barron space) Assume P is a distribution supported on the Barron space ${\mathcal{F}}_s$. By using the generative oracle $\mathcal{O}\left(D,n\right)$ and the function autoencoder (FAE) framework introduced above, the density estimation loss over Barron space can be bounded as

$$W_1\left(P,\widehat{P}\right)\lesssim \sqrt{\frac{d\log d}{\log n}}.$$

Proof

The proof is provided in Supplementary Section C.3.

Now we have proved the statistical upper bound of density estimation over the RKHS ball and Barron space. To further show the optimality of such an upper bound, we provide a statistical lower bound of density estimation over the RKHS ball, which matches its upper bound with only a negligible $\mathcal{O}\left(\mathrm{poly}\left(\log \log n\right)\right)$ gap.

Theorem 3

(Minimax lower bound for Wasserstein estimation over an RKHS ball) Assume $\mathcal{H}$ is an RKHS following 2β-polynomial eigenvalue decay (Assumption 1)

$$c_{\mathcal{H}}\cdot i^{-2\beta }\le {\mu }_i\le C_{\mathcal{H}}\cdot i^{-2\beta }$$

and define the RKHS unit ball $\mathcal{B}:=\left\{f\in \mathcal{H}:\parallel f{\parallel }_{\mathcal{H}}\le 1\right\}$. We observe n i.i.d. samples X₁, X₂, …, X_n ~ P, where the unknown distribution P is supported on the RKHS ball $\mathcal{B}$. Let the statistical minimax rate

$${\mathcal{R}}_n:={\inf }_{{\widehat{P}}_n}{\sup }_{P:\mathrm{supp}\left(P\right)\subseteq \mathcal{B}}\mathbb{E}\left[W_1\left(P,{\widehat{P}}_n\right)\right].$$

Then there exists a constant K(β, c, C) > 0, depending only on pure constants β, c, C in the eigenvalue decay assumption, such that for all sufficiently large n,

$${\mathcal{R}}_n\ge K\left(\beta ,c,C\right)\cdot {\left(\frac{1}{\log n}\right)}^{\beta }.$$

Proof

The proof is provided in Supplementary Section C.4.

Notice that its corresponding statistical upper rate is ${\left(\frac{\log \log n}{\log n}\right)}^{\beta }$ according to our results in Theorem 1. It means that our upper and lower bounds only have a $\mathcal{O}\left(\mathrm{poly}\left(\log \log n\right)\right)$ gap, which is negligible to our statistical rate.

FunDiff in practice

To realize the framework of function generative models outlined above, we propose FunDiff (Fig. 1), which consists of a FAE and a latent diffusion model to generate continuous functions. The FAE learns the mapping from discrete to continuous function representations while also enabling the enforcement of prior physical constraints. The latent diffusion model is then trained directly on the latent manifold learned by the pretrained FAE encoder. Our generation process follows a simple two-step procedure: we first sample latent codes by integrating the reverse-time ODE from Gaussian noise to the target data distribution, and then transform these codes into continuous functions via the FAE decoder.

Fig. 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.

Physical prior integration

The FAE framework enables flexible integration of physical priors through multiple mechanisms of either soft or hard constraints. For general PDE constraints, we augment the loss with a residual term:

$${\mathcal{L}}_{\mathrm{PDE}}={∣\mathcal{R}\left(\widehat{f}\right)∣}_2^2,$$

where $\mathcal{R}$ represents the PDE residual and $\widehat{f}$ is the reconstructed function. The continuous representation of FAE enables automatic differentiation to compute all required derivatives, allowing constraint enforcement across the entire domain. It is worth noting that these physics-informed terms serve as regularizers rather than numerical solvers of the governing equations; they guide the learned FAE to produce output functions that are physically consistent.

Furthermore, some physical properties can be incorporated through architectural modifications. For periodic boundary conditions, one approach is to modify the coordinate embedding with Fourier bases of specific frequencies⁵⁹. When handling divergence-free fields, particularly important in fluid dynamics applications, the FAE can be configured to learn a stream function ψ, from which velocity field components are derived as $u=\frac{\partial \psi }{\partial y},v=-\frac{\partial \psi }{\partial x}$, thereby ensuring divergence-free flow. Symmetry constraints can be incorporated when needed through modifications to the decoder’s final layer.

Related work

We emphasize that our approach fundamentally differs from existing physics-informed diffusion models^42,47–52, which impose physical constraints either by modifying the diffusion model’s training objective or by guiding the posterior sampling process. These methods operate on discretized outputs, where physical laws-typically enforced through PDE residuals-are subject to discretization errors. This not only leads to inaccurate constraint enforcement but also introduces additional optimization challenges. Moreover, in inverse problems where key physical variables are unobserved, such residuals may be undefined at inference time, further limiting the applicability of these approaches. In contrast, our method enforces physical constraints within the FAE, which is fully decoupled from the training and inference processes of the latent diffusion model. Since the FAE provides a continuous neural representation, we can compute PDE residuals by automatic differentiation when needed. This design enables both exact and soft enforcement of constraints in continuous space, mitigating discretization-induced errors, and enabling more stable and physically consistent generation. Details on how to impose these constraints are provided in the following examples through comprehensive numerical experiments, demonstrating that our framework can generate diverse function samples while naturally preserving the underlying physical principles.

Learning 1D damped sinusoidal functions

To validate our theoretical analysis and demonstrate the training pipeline of the proposed framework, we first present a pedagogical example using a simple distribution of one-dimensional damped sinusoidal signals. The target functions are defined as

$$f\left(x\right)=A\exp \left(-\gamma x\right)\sin \left(\omega x\right)+b,x\in \left[0,1\right],$$ 2

where the parameters are sampled from uniform distributions: amplitude $A~\mathcal{U}\left[0.5,1\right]$, decay rate $\gamma ~\mathcal{U}\left[2,4\right]$, angular frequency $\omega ~\mathcal{U}\left[6\pi ,8\pi \right]$, and vertical offset $b~\mathcal{U}\left[-0.5,0.5\right]$.

We employ an FAE with a 6-layer encoder and a 2-layer decoder, each with a hidden dimension of 256 and 8 attention heads. The FAEs are trained on datasets of varying sizes, with function samples discretized at 128 uniform grid points. We train the FAE by minimizing the loss using the Adam optimizer for 5 × 10⁴ iterations. The learning rate follows a warmup and exponential decay schedule: it linearly increases from 0 to 10⁻³ over the first 1,000 steps, then exponentially decays every 1,000 steps with a decay rate of 0.9. Following the FAE training, we train the latent diffusion model by minimizing (9) using an identical learning rate schedule.

For evaluation, we generate 8,192 sample functions using our inference procedure with 100 denoising steps. We assess the quality of these generated samples through parameter estimation, leveraging the analytical form of damped sinusoidal functions. Specifically, for each generated sample, we perform a least-squares fitting using (2) to obtain the inferred parameters and reconstruct the corresponding signal. This approach enables a rigorous quantitative assessment of our model’s ability to capture the underlying function space using two metrics: comparisons between reconstructed and generated signals, and analyses of inferred parameter distributions against target uniform distributions.

We provide a comprehensive evaluation of our model’s generative performance across varying training dataset sizes (Fig. 2). The top row illustrates the worst-case generated samples alongside their reconstructions via parameter estimation. As the number of training samples increases, the generation quality improves accordingly. The bottom row analyzes the parameter distributions, revealing the evolution of our model’s capacity to capture the underlying function space. With limited training data (16 or 32 samples), the estimated parameter distributions exhibit narrower ranges or noticeable deviations from the target uniform distributions, indicating insufficient coverage of the parameter space. As the size of training dataset increases, the estimated parameter distributions progressively converge to the target uniform distributions, with the model trained on 256 samples achieving notably better alignment with the true distributions. This improvement is quantitatively validated by the decreasing mean squared error (MSE) between the generated samples and their reconstructions. These empirical results strongly support our theoretical analysis.

Fig. 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 (γ), angular frequency (ω), 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.

Kolmogorov flow

In this example, we use the proposed FunDiff to generate divergence-free velocity fields, a fundamental requirement for ensuring the physical consistency of incompressible flow simulations in computational fluid dynamics. We consider two scenarios: unconditional flow generation and flow reconstruction from partial observations. To evaluate our approach, we use the two-dimensional Kolmogorov flow, a well-established benchmark for analyzing transition to turbulence and the stability of incompressible fluid systems. The governing equations are formulated in the vorticity-velocity representation as follows:

$$\frac{\partial \omega \left(x,t\right)}{\partial t}+\mathbf{u}\left(x,t\right)\cdot \nabla \omega \left(x,t\right)=\frac{1}{\mathrm{Re}}{\nabla }^2\omega \left(x,t\right)+f\left(x\right),\nabla \cdot \mathbf{u}\left(x,t\right)=0,\omega \left(x,0\right)={\omega }_0\left(x\right),$$

where ω is the vorticity, u = (u, v) is the velocity field, and $\mathrm{Re}$ is the Reynolds number, set to 1000. The forcing term f(x) drives the flow and x = [x₁, x₂]. The equations are defined with periodic boundary conditions in the spatial domain x ∈ (0, 2π)² and time interval t ∈ (0, T]. The forcing term is defined as

$$f\left(x\right)=-4\cos \left(4x_2\right)-0.1\omega \left(x,t\right),$$

where the additional drag force term 0.1ω(x, t) prevents energy accumulation on large scales.

To strictly impose periodic boundary conditions and ensure the divergence-free condition, we modify the forward pass of the FAE decoder as follows:

$$\left[x,y,{\mathcal{E}}_{\theta }\left(f\right)\right]$$ 3

where ${\mathcal{D}}_{\theta }$ denotes the forward pass of the decoder, and the coordinate embedding γ is defined as

$$\gamma \left(x,y\right)=\left[\cos \left(x\right),\sin \left(x\right),\cos \left(y\right),\sin \left(y\right)\right].$$ 4

This embedding guarantees that the output functions have a periodicity of 2π within [0, 2π]². An incompressible velocity field is then computed as (u, v) = (∂ψ/∂y, − ∂ψ/∂x).

Unconditional flow generation

Following the training procedure in the Method section, we present our results in Fig. 3a, which displays representative generated samples along with their corresponding divergence fields, computed via automatic differentiation. The predicted divergence values are negligible, confirming the effectiveness of our approach. We also compare the spectral energy distributions between the training dataset and 1,000 generated samples (Fig. 3b). The spectra show good agreement up to wavenumbers of $\mathcal{O}\left(10^2\right)$. However, the model leads to some artificial high-frequency components at larger wavenumbers.

Fig. 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. Error bars indicate the standard deviation across test samples. e Relative L² error of reconstructed velocity fields across varying noise levels (σ) 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.

To demonstrate the importance of the divergence-free constraint, we conduct a comparative study using an unconstrained FAE trained directly on the velocity field. In this case, the forward pass is formulated as

$$\left[x,y,{\mathcal{E}}_{\theta }\left(f\right)\right]{\to }^{{\mathcal{D}}_{\theta }}\left[u\left(x,y\right),v\left(x,y\right),\omega \left(x,y\right)\right].$$

This unconstrained approach results in higher divergence values (Supplementary Fig. S1). Additionally, the spectral energy shows some discrepancies compared to the model constrained by the exact divergence-free condition.

Flow reconstruction from partial observations

For our second case study, we evaluate the model’s capability to reconstruct flow fields from corrupted partial observations. We use 90% of the data for training and the rest for testing. To improve resolution-invariant reconstruction capabilities, we implement a dynamic downsampling strategy during the training of the FunDiff components. Specifically, at each training iteration, the input velocity field is randomly downsampled by factors of 1, 2, or 4. We first examine the impact of the divergence-free constraint by comparing FunDiff models with and without this condition (Fig. 3e). Both variants demonstrate robust performance across different downsampling factors and noise levels. Notably, the divergence-constrained FunDiff model consistently achieves lower test errors, validating the importance of physical constraints in the reconstruction process.

Furthermore, we benchmark our approach against several state-of-the-art neural operators and PDE surrogate models, including FNO¹³, ViT⁸, and UNet⁶⁰, as well as more recent architectures such as CViT⁶¹, DPOT⁶², CNextU-Net⁶³, and AViT⁶⁴. To further evaluate performance under inverse problem settings, we also include three diffusion-based baselines: Diffusion Posterior Sampling (DPS)⁶⁵, DiffusionPDE⁴², and Spatially-Aware Diffusion⁵¹. Detailed implementation specifications for these baseline models are provided in Supplementary Section D.1. Fig. 3c presents reconstructed flow fields from various models, where the input consists of 64 × 64 resolution velocity fields corrupted by 20% Gaussian noise. It can be observed that neural operator baselines tend to produce highly discontinuous or spatially inconsistent patterns, whereas diffusion-based models recover coherent large-scale structures but often lose fine-scale features. The results demonstrate that our model uniquely achieves both smooth and accurate reconstructions. Additional visualizations of our model’s predictions, with and without the divergence-free condition imposed in FAE decoder, are shown in Supplementary Fig. S2 and Fig. S3, respectively. The quantitative comparison using relative L² errors across different downsampling factors is presented in Fig. 3e. These results confirm that our FunDiff model maintains consistent performance across all tested resolutions while exhibiting strong resilience to noise.

Burgers’ equation

As our next example, we demonstrate the effectiveness and flexibility of our framework for generating physically consistent samples by incorporating physical laws as soft constraints during FAE training. To this end, we consider the Burgers’ equation, which serves as a fundamental prototype for conservation laws that can develop discontinuities. The governing equation is given by

$$\frac{du}{dt}+u\frac{du}{dx}-\nu \frac{d^{2}u}{d{x}^2}=0,\left(x,t\right)\in \left(0,1\right)\times \left(0,1\right],$$ 5

$$u\left(x,0\right)=a\left(x\right),x\in \left(0,1\right)$$ 6

with periodic boundary conditions, where ν = 0.001 denotes the viscosity coefficient, and the initial condition a(x) is sampled from a Gaussian random field (GRF) satisfying periodic boundary conditions.

Let u_θ denote the FAE output. We define the PDE residual as

$$\mathcal{R}\left[u_{\theta }\right]\left(x,t\right)=\frac{d{u}_{\theta }\left(x,t\right)}{dt}+u_{\theta }\left(x,t\right)\frac{d{u}_{\theta }\left(x,t\right)}{dx}-\nu \frac{d^2{u}_{\theta }\left(x,t\right)}{d{x}^2}.$$

The FAE is trained by minimizing a composite loss

$$\mathcal{L}\left(\mathit{\theta }\right):{\underbrace{w_{data}\frac{1}{NM}\sum _{i=1}^N\sum _{j=1}^M{∣u_{\mathit{\theta }}^i\left(x_j,t_j\right)-u^i\left(x_j,t_j\right)∣}^2}}_{{\mathcal{L}}_{data}\left(\mathit{\theta }\right)}+{\underbrace{w_{physics}\frac{1}{NM}\sum _{i=1}^N\sum _{j=1}^M{∣\mathcal{R}\left[u_{\theta }^i\right]\left(x_j,t_j\right)∣}^2}}_{{\mathcal{L}}_{physics}\left(\mathit{\theta }\right)}$$

Here, for each sample uⁱ, we randomly sample query coordinates ${\left\{\left(x_j,t_j\right)\right\}}_{j=1}^M$ from the uniform grid during each gradient descent iteration. All network derivatives are computed using automatic differentiation. Specifically, we set w_data = 100, w_physics = 1 and randomly downsample uⁱ by a factor of 1, 2, 5 at each iteration of gradient descent.

Similar to the Kolmogorov flow example, our objective here is to train a FunDiff model to reconstruct the solution field from sparse observations. We evaluate our model against established baseline methods. We show three examples in Fig. 4a. Neural operator baselines generally produce blurry and discontinuous reconstructions, whereas diffusion-based approaches yield smoother results but often exhibit small oscillations near sharp gradients. By contrast, FunDiff produces sharp, smooth reconstructions that agree well with the reference numerical solutions, with relative L² errors below 10%. Fig. 4b presents a quantitative comparison of model performance across inputs with varying downsampling factors and 20% Gaussian noise corruption. Notably, FunDiff outperforms all baselines. Additional examples of our model’s predictions are shown in Supplementary Fig. S4. Their physical consistency is confirmed by low PDE residuals, remaining on the order of $\mathcal{O}\left(10^{-2}\right)$ across the entire domain. We also systematically examine the effect of noise and data sparsity (Fig. 4c). As expected, reconstruction error increases with higher noise levels, but FunDiff demonstrates strong robustness to data sparsity, maintaining stable performance even under aggressive downsampling (factors up to 5).

Fig. 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. Error bars indicate the standard deviation across test samples. c Relative L² error of reconstructed solutions across varying noise levels (σ) 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.

To rigorously evaluate the impact of physical constraints, we conduct an ablation study by training the function encoder without the physics loss term ${\mathcal{L}}_{physics}$. The absence of PDE-based regularization leads to higher predictive errors and PDE residuals in the generated solutions (Fig. 4d). This degradation in physical consistency is further visualized through the generated samples in Supplementary Fig. S5, underscoring the crucial role of physics-informed training in our framework.

Linear elasticity

Symmetries play an important role in physics, providing a framework for understanding conservation laws and invariances that govern natural phenomena. This example demonstrates how to embed prior symmetries into FunDiff models. To this end, we consider linear elasticity in computational micromechanics, which involves solving the Lippmann-Schwinger equation⁶⁶ in a 2D periodic representative volume element (RVE) microstructural domain $\mathrm{\Omega }\subset {\mathbb{R}}^2$ with spatial coordinates x = (x, y) under plane strain conditions. The problem is formulated as

$$\left\{\begin{array}{c}{\epsilon }_{ij}\left(\mathbf{x}\right)+{\int }_{\mathrm{\Omega }}{G}_{ijkl}^{\left(0\right)}\left(\mathbf{x},{\mathbf{x}}^{\prime }\right):\left[{\sigma }_{ij}\left({\mathbf{x}}^{\prime }\right)-C_{ijkl}^0:{\epsilon }_{kl}\left({\mathbf{x}}^{\prime }\right)\right]\mathrm{d}{\mathbf{x}}^{\prime }-{\overline{\epsilon }}_{ij}=0,\mathbf{x}\in \mathrm{\Omega },\hfill \\ {\sigma }_{ij}\left(\mathbf{x}\right)=C_{ijkl}\left(\mathbf{x}\right):{\epsilon }_{kl}\left(\mathbf{x}\right).\hfill \end{array}\right)$$ 7

These relations represent balance and constitutive equations, respectively, which are written in a component-wise manner (1 ≤ i, j, k, l ≤ 2) using the Voigt notation. x and ${\mathbf{x}}^{\prime }$ denote a pair of target and source points inside the RVE domain Ω. σ_ij, ε_ij and C_ijkl denote stress, strain, and fourth-order elastic tensor, respectively. C_ijkl(x) is expressed in terms of local Young’s modulus E(x) and Poisson ratio ν(x) as

$$C_{ijkl}\left(\mathbf{x}\right)=\frac{E\left(\mathbf{x}\right)\nu \left(\mathbf{x}\right)}{\left(1+\nu \left(\mathbf{x}\right)\right)\left(1-2\nu \left(\mathbf{x}\right)\right)}{\delta }_{ij}{\delta }_{kl}+\frac{E\left(\mathbf{x}\right)}{2\left(1+\nu \left(\mathbf{x}\right)\right)}\left({\delta }_{ik}{\delta }_{jl}+{\delta }_{il}{\delta }_{jk}\right),$$

where δ_ij denotes the Kronecker delta function. $G_{ijkl}^{\left(0\right)}$ denotes Green’s operator corresponding to the homogeneous reference material with elastic tensor $C_{ijkl}^0$⁶⁷. ${\overline{\epsilon }}_{ij}$ denotes the far-field strain used to apply the strain-controlled boundary condition, equal to the average strain over the RVE domain Ω. In this example, we model the point-wise strain concentration tensor $A_{ijkl}\left(\mathbf{x}\right)\in {\mathbb{R}}^{3\times 3}$, which linearly maps ${\overline{\epsilon }}_{ij}$ to ε_ij(x) independent of loading conditions:

$${\epsilon }_{ij}\left(\mathbf{x}\right)=A_{ijkl}\left(\mathbf{x}\right):{\overline{\epsilon }}_{kl}.$$

A_ijkl(x) can be computed by triply solving (7) subject to three orthogonal unit far field strain ${\overline{\epsilon }}_{ij}$ as

$${\overline{\epsilon }}_{ij}=\left\{{\left[1,0,0\right]}^T,{\left[0,1,0\right]}^T,{\left[0,0,1\right]}^T\right).$$ 8

Due to the intrinsic symmetry of the RVE configuration, the resulting strain concentration tensor field A_ijkl(x) exhibits a symmetric pattern with respect to both $x=\frac{1}{2}$ and $y=\frac{1}{2}$. We aim to train a FunDiff model to reconstruct the strain concentration tensor field from sparse observations while respecting the inherent symmetry priors. To this end, we modify the forward pass of the decoder to enforce these constraints. Let z denote the FAE encoder output latents. The symmetry-preserved decoder ${\widehat{\mathcal{D}}}_{\theta }$ is then constructed as

$${\widehat{\mathcal{D}}}_{\theta }\left(\mathbf{z},x,y\right)={\mathcal{D}}_{\theta }\left(\mathbf{z},x,y\right)+{\mathcal{D}}_{\theta }\left(\mathbf{z},1-x,y\right)+{\mathcal{D}}_{\theta }\left(\mathbf{z},x,1-y\right)+{\mathcal{D}}_{\theta }\left(\mathbf{z},1-x,1-y\right).$$

It can be easily shown that this formulation ensures that ${\widehat{\mathcal{D}}}_{\theta }$ generates functions with perfect symmetry about the lines x = 1/2 and y = 1/2. Although here we only focus on such a relatively simple symmetry for illustration purposes, it is worth noting that more general symmetries can be embedded in a similar manner⁶⁸.

Figure 5 a presents reconstructed samples generated by our FunDiff model from a coarse strain concentration field corrupted by 10% Gaussian noise. The results show that our model preserves underlying symmetries and accurately resolves sharp material interfaces. To assess robustness against different types of noise, we report test errors in Fig. 5b. The consistently low errors across all noise configurations highlight the framework’s reliability in practical applications.

Fig. 5. Linear elasticity.

a Comparison of diagonal strain concentration tensor distributions (Voigt notation) between provided coarse measurements (32 × 32) and reconstructions from the trained FunDiff model. The generated samples maintain enforced symmetry constraints and achieve a relative L² error of 5.12% when compared to ground truth data at 256 × 256 resolution. b Comparison of reconstruction errors under different noise types and downsampling rates. Dashed lines represent the mean reconstruction error, and shaded bands denote one standard deviation across test samples.

To further demonstrate the critical role of the continuous decoder in our framework, we implement a baseline using a discrete autoencoder with a Vision Transformer encoder and an MLP decoder projecting to the target pixel space. We train a diffusion transformer under identical conditions, with representative generated samples shown in Supplementary Fig. S6. Unlike our approach, the baseline fails to preserve symmetry or enforce physical continuity in strain fields, instead producing piecewise discontinuous predictions.

Turbulence mass transfer

In the last example, we examine the model's capacity to infer non-intrusive quantities, such as pressure, in turbulent problems. Specifically, we consider turbulent mass transfer in a rectangular channel. This phenomenon plays a critical role in the chemical engineering process, directly influencing reaction kinetics, conversion efficiency, and product purity. Our study focuses on a computational domain Ω = {(x, y)∣0.006 < x < 0.018, 0 < y < 0.01}. Within this domain, we are interested in training a FunDiff model to infer the obstacle topology and pressure field given the sparse and noisy velocity measurements.

The simulation employs a two-equation model for turbulent mass transfer based on the Reynolds-averaged Navier-Stokes (RANS) equations⁶⁹, which can be expressed as

$$\begin{array}{c}\frac{\partial \left(\rho u_i\right)}{\partial x_i}=0,\hfill \\ u_i\frac{\partial \left(\rho u_j\right)}{\partial x_i}=-\frac{\partial P}{\partial x_j}+\frac{\partial }{\partial x_i}\left[\left({\mu }_m+{\mu }_t\right)\frac{\partial u_j}{\partial x_i}\right],\hfill \\ u_i\frac{\partial \left(\rho C\right)}{\partial x_i}=\frac{\partial }{\partial x_i}\left[\rho \left(D_m+D_t\right)\frac{\partial C}{\partial x_i}\right],\hfill \end{array}$$

where p, u_i, x_i are the pressure, velocity and location, respectively; the subscript i represents the coordinate direction which can take x and y for the two-dimensional system. Here, C is the time-averaged concentration; ρ, μ_m and D_m are density, laminar viscosity, and laminar mass diffusivity, which are physical property constants and fixed to 1000 kg/m³, 1.003 × 10⁻⁶ m²/s, 1.003 × 10⁻⁸ m²/s, respectively. μ_t and D_t are the turbulent viscosity and turbulent mass diffusivity, respectively.

We begin by training an FAE to reconstruct the key physical variables: velocity field, pressure field, and the signed distance function (SDF) representing obstacle geometry. To enhance the capacity of the model to handle inputs of varying resolutions, we augment training by randomly downsampling inputs by factors of 1, 2, and 4 while preserving the target resolution. The latent diffusion model is then trained with latent representations of the pressure and SDF, conditioned on the encoded representations of the velocity field, obtained from the trained autoencoder’s latent space.

Representative generated samples (Fig. 6a) validate the model’s ability to infer pressure fields and reconstruct obstacle geometries despite heavily corrupted input velocity measurements. The reconstructions align closely with numerical estimations. Fig. 6b presents the test error across various downsampling factors, noise levels, and noise types. The results indicate that our model remains robust even under severe noise corruption and low-resolution inputs. Our results demonstrate the model’s effectiveness in preserving critical flow features and geometric details, and highlight the framework’s potential for real-world applications where high-fidelity reconstructions are required from sparse or noisy measurements.

Fig. 6. Turbulent mass transfer.

a Left: Input velocity fields at low resolution with 10% noise corruption. Middle: Comparison between reference and generated pressure fields. Right: Comparison between reference and generated obstacle geometry. b Relative L² error comparison between generated samples and reference solutions across varying input resolutions and noise levels. Dashed lines represent the mean reconstruction error, and shaded bands denote one standard deviation across test samples.

Discussion

In this work, we propose FunDiff, a flexible and efficient framework for generative modeling in function spaces that addresses fundamental challenges in generating continuous function data while preserving physical constraints. Our primary contributions are twofold: (1) a rigorous theoretical foundation establishing conditions for function distribution approximation via autoencoders, and (2) a constrained function autoencoder (FAE) that inherently enforces physical principles. Through extensive experimentation across diverse physical systems, we have demonstrated that our approach successfully generates physically valid samples while maintaining high fidelity to the target PDE solution distributions.

While our current work lays a strong foundation for physics-informed generative modeling, several promising directions for future research emerge. One particularly interesting direction is to develop not purely data-driven, but physics-informed diffusion-based PDE solvers. This could be achieved by incorporating PDE residuals directly into the score function of diffusion models and by using additional initial and boundary conditions to guide sampling, thereby steering generated samples toward the corresponding PDE solutions. Another natural extension would be to adapt our framework to handle more complex geometries, thereby enabling broader applicability in real-world scientific and engineering problems. Handling complex domains necessitates incorporating geometric inductive biases, such as mesh-based representations, implicit neural surfaces, or graph-based function approximations. Moreover, extending our approach to multiphysics applications represents a critical and exciting frontier. Many engineering and scientific challenges involve the interplay of multiple physical processes, such as fluid-structure interactions, thermo-mechanical coupling, and electromagnetic-thermal phenomena, which require models that can simultaneously capture and reconcile diverse physical constraints. These advancements will expand the scope of generative modeling for physical systems and ultimately contribute to more accurate and efficient simulation-based design across disciplines.

Methods

To realize our theoretical framework for generative modeling in function spaces, we develop FunDiff, a practical implementation that addresses the key challenges of handling infinite-dimensional function distributions while preserving physical constraints. Our approach must accommodate discretized functions at varying resolutions, enable continuous evaluation at arbitrary coordinates, and flexibly incorporate domain-specific physical priors without compromising the generative modeling process. The architecture is designed to decouple the enforcement of physical constraints from diffusion training, ensuring stable optimization while maintaining mathematical rigor in constraint satisfaction. We detail the FunDiff network architecture, training procedures, and mechanisms for integrating constraints that enable robust function generation across diverse scientific domains.

Network architecture

The FAE processes discretized functions $\mathbf{f}\in {\mathbb{R}}^{n\times d}$, where n denotes the spatial discretization resolution and d the codomain dimension. The FunDiff architecture comprises three core components: (1) a resolution-invariant encoder $\mathcal{E}$ combining a Vision Transformer (ViT)⁸ and a Perceiver module⁹, (2) a latent diffusion model operating on the compact latent space, and (3) a physics-aware decoder $\mathcal{D}$ that reconstructs continuous functions.

Encoder

The encoder $\mathcal{E}$ employs a ViT backbone to extract features and a Perceiver module to handle functions of different discretizations. Input functions $\mathbf{f}\in {\mathbb{R}}^{H\times W\times C}$ are first processed by ViT-style patchification, resulting in flattened patches ${\mathbf{f}}_p\in {\mathbb{R}}^{L\times D}$, where $L=\frac{H}{P}\times \frac{W}{P}$ represents the length of the sequence, P is the size of the patch, and D denotes the embedding dimension. To accommodate varying resolutions, trainable positional embeddings (PE) are dynamically interpolated to match the input sequence length L, yielding position-aware patch embeddings f_pe = f_p + Interp(PE, L).

A perceiver-style cross-attention module maps these variable-length inputs to fixed-dimensional latent representations. Specifically, we learn a predefined number of latent input queries ${\mathbf{z}}_q\in {\mathbb{R}}^{N\times D}$. These queries serve as learnable parameters in a cross-attention module, which processes the visual features of our input as follows

$${\mathbf{z}}^{\prime }={\mathbf{z}}_q+\mathrm{MHA}\left(\mathrm{LN}\left({\mathbf{z}}_q\right),\mathrm{LN}\left({\mathbf{f}}_{pe}\right),\mathrm{LN}\left({\mathbf{f}}_{pe}\right)\right),{\mathbf{z}}_{agg}={\mathbf{z}}^{\prime }+\mathrm{MLP}\left(\mathrm{LN}\left({\mathbf{z}}^{\prime }\right)\right),$$

where MHA, LN, and MLP represent multi-head self-attention, layer normalization, and a multi-layer perceptron, respectively. We further process the aggregated tokens z_agg using a sequence of L pre-norm Transformer blocks^70,71,

$${\mathbf{z}}_0=\mathrm{LN}\left({\mathbf{z}}_{agg}\right),{\mathbf{z}}_{\mathcal{l}}^{\prime }=\mathrm{MSA}\left(\mathrm{LN}\left({\mathbf{z}}_{\mathcal{l}-1}\right)\right)+{\mathbf{z}}_{\mathcal{l}-1},\mathcal{l}=1,\dots ,L,{\mathbf{z}}_{\mathcal{l}},=\mathrm{MLP}\left(\mathrm{LN}\left({\mathbf{z}}_{\mathcal{l}}^{\prime }\right)\right)+{\mathbf{z}}_{\mathcal{l}}^{\prime },\mathcal{l}=1,\dots ,L,$$

where MSA is a standard multi-head self-attention module used to model long-range dependencies across the sequence.

Decoder

The decoder module draws inspiration from the Continuous Vision Transformer (CViT)⁷², which enables continuous evaluation at any query coordinate through a cross-attention mechanism between query coordinates and encoded features. We start by embedding coordinates into the latent space, which can be efficiently achieved through either random Fourier features⁷³ or grid-based embeddings⁷².

The core of our decoder uses embedded coordinates ${\mathbf{x}}_0\in {\mathbb{R}}^D$ as queries, with the encoder output z_L serving as both keys and values in a series of K cross-attention Transformer blocks:

$${\mathbf{x}}_k^{\prime }={\mathbf{x}}_{k-1}+\mathrm{MHA}\left(\mathrm{LN}\left({\mathbf{x}}_{k-1}\right),\mathrm{LN}\left({\mathbf{z}}_L\right),\mathrm{LN}\left({\mathbf{z}}_L\right)\right),k=1\dots K,{\mathbf{x}}_k={\mathbf{x}}_k^{\prime }+\mathrm{MLP}\left(\mathrm{LN}\left({\mathbf{x}}_k^{\prime }\right)\right),k=1\dots K.$$

Finally, a small feedforward network projects the features x_K to the desired output dimension. This architecture naturally extends to other function representations (e.g., point clouds via PointNet⁷⁴ or meshes via graph neural networks⁷⁵) through appropriate input processing modules.

Latent diffusion with conditional diffusion transformer

We implement the diffusion process in the learned latent space using a Diffusion Transformer (DiT)⁷⁶. We condition the model on coarse solution measurements c by encoding them through the pretrained FAE encoder ${\mathcal{E}}_{\theta }$. The resulting latent-space guidance vectors are additively combined with the diffusion inputs:

$${\tilde{\mathbf{z}}}_1={\mathbf{z}}_1+{\mathcal{E}}_{\theta }\left(\mathbf{c}\right).$$

The DiT backbone consists of a sequence of Transformer blocks that utilize AdaLN-Zero modulation⁷⁶, an enhanced variant of Adaptive Layer Normalization designed for stable time-dependent conditioning. For each block, the modulation parameters are computed as

$$\mathit{\gamma },\mathit{\beta }=\mathrm{MLP}\left(t\right),{\mathbf{z}}^{\prime }=\mathrm{LN}\left(\mathbf{z}\right)\odot \mathit{\gamma }+\mathit{\beta },$$

where t represents the timestep. Following the approach in ref. ⁷⁶, we incorporate element-wise scaling factors α, initialized to zero, in the residual pathways:

$${\mathbf{z}}_{\mathrm{out}}=\mathit{\alpha }\odot \mathrm{MSA}\left({\mathbf{z}}^{\prime }\right)+\mathbf{z}.$$

By initializing these scaling factors to zero, each Transformer block initially acts as an identity transform. This strategy enables gradual learning of temporal and conditional modulations while maintaining stable gradient flow during the early stages of training.

Training

Our training procedure consists of two stages. In the first stage, we train an FAE to learn a compact and expressive latent representation of functions. We consider a set of ground truth functions ${\left\{f_i\right\}}_{i=1}^N$, where each function $f_i:\mathcal{X}\to \mathbb{R}$ maps input coordinates $\mathbf{x}\in \mathcal{X}\subset {\mathbb{R}}^d$ to scalar outputs. For training, each function f_i is discretized at a fixed set of input locations ${\left\{{\mathbf{x}}_j\right\}}_{j=1}^M$, resulting in a vector of evaluations ${\mathbf{f}}_i=\left[f_i\left({\mathbf{x}}_1\right),\dots ,f_i\left({\mathbf{x}}_M\right)\right]\in {\mathbb{R}}^M$. The reconstruction loss is given by

$${\mathcal{L}}_{\mathrm{FAE}}=\frac{1}{NM}{\sum }_{i=1}^N{\sum }_{j=1}^M{∣f_i\left({\mathbf{x}}_j\right)-\mathcal{D}\left(\mathcal{E}\left({\mathbf{f}}_i\right)\right)\left({\mathbf{x}}_j\right)∣}^2,$$

where $\mathcal{E}$ and $\mathcal{D}$ denote the encoder and decoder, respectively and f_i(x_j) represents the value of a ground truth function f_i at query point x_j. During training, we randomly sample functions {f_i} and query coordinates ${\left\{{\mathbf{x}}_j\right\}}_{j=1}^M$ at each iteration.

In the second stage, we freeze the pre-trained FAE and train the DiT model in the induced latent space. Following the rectified flow framework, the training objective is formulated as

$${\mathcal{L}}_{\mathrm{DiT}}={\int }_0^1{\mathbb{E}}_{\mathbf{f}~p_{\mathrm{data}}}\left[{∣\left({\mathbf{z}}_1-{\mathbf{z}}_0\right)-\mathbf{g}\left({\mathbf{z}}_{\tau },\tau \right)∣}^2\right]d\tau ,\mathrm{with}{\mathbf{z}}_{\tau }=\left(1-\tau \right){\mathbf{z}}_1+\tau {\mathbf{z}}_0,{\mathbf{z}}_1=\mathcal{E}\left(\mathbf{f}\right),$$ 9

where z₁ is obtained by encoding the input function through the pre-trained FAE encoder, ${\mathbf{z}}_0~\mathcal{N}\left(0,\mathbf{I}\right)$ is sampled from the standard normal distribution, and g(z_t, t) is parameterized by our DiT model predicting the velocity field. In practice, we approximate the time integral by uniformly sampling $\tau ~\mathcal{U}\left[0,1\right]$ at each training step.

Inference

At inference, we follow a straightforward two-step generation process. First, we sample latent codes from our learned diffusion model by solving the ODE using a standard ODE solver (e.g., Runge-Kutta method):

$$\frac{d\mathbf{z}\left(t\right)}{dt}=\mathbf{g}\left(\mathbf{z}\left(t\right),t\right),t\in \left[0,1\right],$$

starting from a sample of standard Gaussian noise ${\mathbf{z}}_1~\mathcal{N}\left(0,\mathbf{I}\right)$ and integrating backward to obtain z₀. This process is deterministic and typically requires 20–50 function evaluations for high-quality samples. Second, we decode the sampled latent codes through the FAE decoder to obtain the final continuous function. Importantly, since our decoder operates on arbitrary query points, we can evaluate the generated function at any desired resolution or location in the domain.

Implementation details

We maintain consistent architectural choices and training procedures across all experiments unless explicitly stated otherwise. Our FAE architecture comprises an 8-layer encoder and a 4-layer decoder, both utilizing an embedding dimension of 256 and MLP width of 512. The encoder’s patch size is adaptively selected based on input resolution. The DiT employs 8 Transformer blocks with an embedding dimension of 256 and MLP width of 512. All attention operations use 8 heads.

Training protocol

For the FAE training, we employ the AdamW optimizer⁷⁷ with a weight decay of 10⁻⁵. The learning rate follows a two-phase schedule: an initial linear warm-up from zero to 10⁻³ over 2000 steps, followed by exponential decay with a factor of 0.9 every 2000 steps. Each training batch contains 16 samples, with 4096 query coordinates and their corresponding outputs randomly sampled from the grid. To enhance the model’s resolution invariance, we randomly downsample inputs by different factors (e.g., 1, 2, 4, 8) during training while optimizing the model to reconstruct the original high-resolution input function by minimizing the MSE loss.

The DiT follows a similar optimization strategy but operates with larger batches of 128 samples over 10⁵ iterations. For training the conditioned diffusion models, we generate conditions by randomly downsampling input functions and encoding them through the pretrained FAE encoder. Detailed hyperparameters for all experiments are provided in Supplementary Section D.1.

Author contributions

S.W., Z.D., and L.L. conceived the project and developed the main methodology. S.W. implemented the core algorithms and conducted the experiments. Z.D. contributed to the theoretical analysis. T.L. and S.S. assisted with numerical simulation. L.L. supervised the project and provided critical revisions. All authors discussed the results and contributed to writing the manuscript.

Peer review

Peer review information

Nature Communications thanks the anonymous reviewers for their contribution to the peer review of this work. A peer review file is available.

Data availability

The datasets generated in this study have been deposited in the Zenodo database 10.5281/zenodo.18911353. These data are openly available without restriction.

Code availability

All code used to train the models, run the numerical experiments, and reproduce the results in this study is publicly available at https://github.com/sifanexisted/fundiff.

Competing interests

The authors declare no competing interests.

Supplementary information

The online version contains supplementary material available at 10.1038/s41467-026-72292-0.