Generating Physical Dynamics Under Priors（英文原文_正文）

GENERATING PHYSICAL DYNAMICS UNDER PRIORS

Zihan Zhou¹, Xiaoxue Wang², Tianshu Yu^1,*
1School of Data Science, The Chinese University of Hong Kong
²ChemLex Technology Co., Ltd.
zihanzhou@link.cuhk.edu.hk, wxx@chemlex.tech, yutianshu@cuhk.edu.hk
*corresponding author

ABSTRACT

Generating physically feasible dynamics in a data-driven context is challenging, especially when adhering to physical priors expressed in specific equations or formulas. Existing methodologies often overlook the integration of “physical priors” resulting in violation of basic physical laws and suboptimal performance. In this paper, we introduce a novel framework that seamlessly incorporates physical priors into diffusion-based generative models to address the limitation. Our approach leverages two categories of priors:

1. distributional priors, such as roto-translational invariance, and
2. physical feasibility priors, including energy and momentum conservation laws and PDE constraints.

By embedding these priors into the generative process, our method can efficiently generate physically realistic dynamics, encompassing trajectories and flows. Empirical evaluations demonstrate that our method produces high-quality dynamics across a diverse array of physical phenomena with remarkable robustness, underscoring its potential to advance data-driven studies in AI4Physics. Our contributions signify a substantial advancement in the field of generative modeling, offering a robust solution to generate accurate and physically consistent dynamics.

1 Introduction

The generation of physically feasible dynamics is a fundamental challenge in the realm of data-driven modeling and AI4Physics. These dynamics, driven by Partial Differential Equations (PDEs), are ubiquitous in various scientific and engineering domains, including fluid dynamics (Kutz, 2017), climate modeling (Rasp et al., 2018), and materials science (Choudhury et al., 2022). Accurately generating such dynamics is crucial for advancing our understanding and predictive capabilities in these fields (Bazok & Ioannidis, 2019). Recently, generative models have revolutionized the study of physics by providing powerful tools to simulate and predict complex systems.

Generative vs. discriminative models.

Even when high-performing discriminative models for dynamics are available such as finite elements (Zhang et al., 2021; Uriarte et al., 2022), finite difference (Lu et al., 2021; Salman et al., 2022), finite volume (Ranade et al., 2021), or physics-informed neural networks (PINNs) (Raissi et al., 2019), generative models are crucial in machine learning for their ability to capture the full data distribution, enabling more effective data synthesis (de Oliveira et al., 2017), anomaly detection (M et al., 2017), and enhanced robustness and interpretability by modeling the joint distribution of data and labels, offering insights into unseen scenarios (Takeshita & Kalousis, 2021). Generative models are also pivotal in creative domains, such as drug discovery (Lavecchia, 2019), where they enable the creation of novel data samples.

Challenge.

However, the intrinsic complexity and high-dimensional nature of physical dynamics pose significant challenges for traditional learning systems. Recent advancements in generative modeling, particularly diffusion-based generative models (Song et al., 2020), have shown promise in capturing complex data distributions. These models iteratively refine noisy samples to match the target distribution, making them well-suited for high-dimensional data generation. Despite their success, existing approaches often overlook the incorporation of “physical priors” expressed in specific equations or formulas, which are essential for ensuring that the generated dynamics adhere to fundamental physical laws.

Solution.

In this work, we propose a novel framework that integrates priors into diffusion-based generative models to generate physically feasible dynamics. Our approach leverages two types of priors: Distributional priors, including roto-translational invariance and equivariance, ensure that models capture the intrinsic properties of the data rather than their specific representations; Physical feasibility priors, including energy and momentum conservation laws and PDE constraints, enforce the adherence to fundamental physical principles, thus improving the quality of generated dynamics.

The integration of priors into the generative process is a complex task that necessitates a deep understanding of the relevant mathematical and physical principles. Unlike prior works, where the physical system is modeled as a foreground-truth value ${\mathbf{x}}_t$, diffusion generative models aim to characterize a full ground-truth distribution ${\nabla }_x\log p_x({\mathbf{x}}_t)$ or $\mathbb{E}[{\mathbf{x}}_t]\cdot {\mathbf{x}}_t$ (notation in Equation 1). This fundamental difference complicates the direct application of priors based on ground-truth values to the output of generative models. In this work, we propose a framework to address this challenge by effectively embedding priors within the generative model’s output distribution. By incorporating these priors, we can effectively embed the physical laws of the system, which can efficiently produce physically plausible dynamics. This capability is particularly useful for studying physical phenomena where the governing equations are too complex to be learned purely from data.

Figure 1: Animated visualization of generated samples of shallow water dynamics, showcasing the variations over time. Use the latest version of Adobe Acrobat Reader to view.

Results.

Empirical evaluations of our method demonstrate its effectiveness in producing high-quality dynamics across a range of physical phenomena. Our approach exhibits high robustness and generalizability, making it a promising tool for the data-driven study of AI4Physics. In Fig. 1, we provide a generated sample of the shallow water dataset (Martinez-Aranda et al., 2018). The generated dynamics not only capture the intricate details of the physical processes but also adhere to the fundamental physical laws, offering an accurate and reliable representation of underlying systems.

Contribution.

In conclusion, our work presents a significant advancement in the field of data-driven generative modeling by introducing a novel framework that integrates physical priors into diffusion-based generative models. In all, our method

1. improves the feasibility of generated dynamics, making them more aligned with physical principles compared to baseline methods;
2. poses the solution to the longstanding challenge of generating physically feasible dynamics;
3. paves the way for more accurate and reliable data-driven studies in various scientific and engineering domains, highlighting the potential of AI4Physics in advancing our understanding of complex physical systems.

2 Preliminaries

In Appendix A, we present a comprehensive review of Related Work, specifically focusing on three key areas: generative methods for physics, score-based diffusion models, and physics-informed neural networks. This section aims to provide foundational knowledge for readers who may not be familiar with these topics. We recommend that those seeking to deepen their understanding of these areas consult this appendix.

2.1 Diffusion Models

Diffusion models generate samples following an underlying distribution. Consider a random variable ${\mathbf{x}}_0\in {\mathbb{R}}^n$ drawn from an unknown distribution $q_0$. Denoising diffusion probabilistic models (Song & Ermon, 2019; Song et al., 2020; Ho et al., 2020) describe a forward process ${\mathbf{x}}_t,t\in [0,T]$ governed by an Ito stochastic differential equation (SDE)

$$\mathrm{d}{\mathbf{x}}_t=f(t\mid {\mathbf{x}}_t)\mathrm{d}t+g(t\mid {\mathbf{x}}_t)\mathrm{d}{\mathbf{w}}_t,{\mathbf{x}}_0\sim q_0,f(t)=\frac{\mathrm{d}\log {\alpha }_t}{\mathrm{d}t},g^2(t)=\frac{\mathrm{d}{\sigma }_t^2}{\mathrm{d}t}-2\frac{\mathrm{d}\log {\alpha }_t}{\mathrm{d}t}{\sigma }_t^2.(1)$$

where ${\mathbf{w}}_t\in {\mathbb{R}}^n$ denotes the standard Brownian motion, and ${\alpha }_t$ and ${\sigma }_t$ are predetermined functions of $t$. This forward process has a closed-form solution of $q_t({\mathbf{x}}_t\mid {\mathbf{x}}_0)=\mathcal{N}({\mathbf{x}}_t\mid {\alpha }_t{\mathbf{x}}_0,{\sigma }_t^2\mathbf{I})$ and has a corresponding reverse process of the probability flow ordinary differential equation (ODE), running from time $T$ to 0, defined as (Song et al., 2020)

$$\frac{\mathrm{d}{\mathbf{x}}_t}{\mathrm{d}t}=f(t\mid {\mathbf{x}}_t)-\frac{1}{2}{\sigma }^2(t){\nabla }_x\log q_t({\mathbf{x}}_t),{\mathbf{x}}_T\sim q_T({\mathbf{x}}_T\mid {\mathbf{x}}_0)\approx q_T({\mathbf{x}}_T).(2)$$

The marginal probability density { q t ( x t ) } t = 0 T \{q_t(\boldsymbol{x}_t)\}_{t=0}^T {qt​(xt​)}t=0T​ of the forward SDE aligns with the reverse ODE (Song et al., 2020). This indicates that if we can sample from $q_T({\mathbf{x}}_T)$ and solve Equation 2, then the resulting ${\mathbf{x}}_0$ will follow the distribution $q_0$. By choosing ${\alpha }_t\to 0$ and ${\sigma }_t\to 1$, the distribution $q_T({\mathbf{x}}_T)$ can be approximated as a normal distribution. The score ${\nabla }_x\log q_t({\mathbf{x}}_t)$ can be approximated by a deep learning model. The quality of the generated samples is contingent upon the model’s ability to accurately approximate the score functions (Kwon et al., 2022; Gao & Zhu, 2024). A more precise approximation results in a distribution that closely aligns with the distribution of the training data. This is a key advantage of using deep generative distributions as physical feasibility into the models is advisable. Section 3 will elaborate on our methods for integrating distributional priors and physical feasibility priors, as well as the objectives for score matching.

2.2 Invariant Distributions

An invariant distribution refers to a probability distribution that remains unchanged under the action of a specified group of transformations. These transformations can include operations such as translations, rotations, or other symmetries, depending on the problem domain. Formally, let $\mathcal{G}$ be a group of transformations. A distribution $q$ is said to be $\mathcal{G}$-invariant under the group $\mathcal{G}$ if, for all transformations $G\in \mathcal{G}$, we have $⟨G(\mathbf{x}),{\nabla }_{x}q(\mathbf{x})⟩=0$. Invariance under group transformations is particularly significant in modeling distributions that exhibit symmetries. For instance, in the case of 3D coordinates, invariance under rigid transformations—such as translations and rotations (SE(3) group)—is essential for spatial understanding (Zhou et al., 2024). Equivariant models are usually required to embed invariance. A function (or model) $f:{\mathbb{R}}^n\to {\mathbb{R}}^n$ is said to be $(\mathcal{G},\mathcal{L})$-equivariant where $\mathcal{G}$ is the group actions and $\mathcal{L}$ is a function operator, if for any $G\in \mathcal{G}$, $f(G(\mathbf{x}))=\mathcal{L}(G(f(\mathbf{x})))$.

3 Method

In this study, we aim to investigate methodologies for enhancing the capability of diffusion models to approximate the targeted score functions. We have two primary objectives:

1. To incorporate distributional priors, such as translational and rotational invariance, which will aid in selecting the appropriate model for training objective functions;
2. To impose physical feasibility priors on the diffusion model, necessitating injection of priors to the model’s output of a distribution related to the ground-truth samples (specifically, ${\nabla }_x\log q({\mathbf{x}}_t)$ or $\mathbb{E}[{\mathbf{x}}_0\mid {\mathbf{x}}_t]$).

In this section, we consider the forward diffusion process given by Equation 1, where ${\mathbf{x}}_t={ϵ}_t{\mathbf{x}}_0+{ϵ}_t{\mathbf{e}}_t$ with ${\mathbf{e}}_t\sim \mathcal{N}(0,\mathbf{I})$.

3.1 Incorporating Distributional Priors

In this section, we study the score function ${\nabla }_x\log q_t({\mathbf{x}}_t)$ for $\mathcal{G}$-invariant distributions. Understanding its corresponding properties can guide the selection of models with the desired equivariance, facilitating sampling from the $\mathcal{G}$-invariant distribution. In the following, we will assume that the sufficient conditions of Theorem 1 hold so that the marginal distributions $q_t$ are $\mathcal{G}$-invariant. The definitions of the terminologies and proof of the theorem can be found in Appendix F.1.

Theorem 1 (Sufficient conditions for the invariance of $q_t$ to imply the invariance of $q_0$): Let $q_0$ be a $\mathcal{G}$-invariant distribution. If for all $\mathbf{G}\in \mathcal{G}$, $\mathbf{G}$ is a volume-preserving diffeomorphism and isometry, and for all $0<ϵ<1$, there exists $\mathbf{H}\in \mathcal{G}$ such that $\mathbf{H}(\mathrm{d}\mathbf{x})=ϵ\mathbf{G}(\mathbf{x})$, then $q_t$ is also $\mathcal{G}$-invariant.

Property of score functions. Let $q_t$ be a $\mathcal{G}$-invariant distribution. By the chain rule, we have ${\nabla }_x\log q_t({\mathbf{x}}_t)={\nabla }_x\log q_t(G({\mathbf{x}}_t))=\frac{\partial G({\mathbf{x}}_t)}{\partial \mathbf{x}}{\nabla }_x\log q_t(G({\mathbf{x}}_t))$, for all $\mathbf{G}\in \mathcal{G}$. Hence,

$${\nabla }_{\mathbf{G}({\mathbf{x}}_t)}\log q_t(\mathbf{G}({\mathbf{x}}_t))={\left(\frac{\partial \mathbf{G}({\mathbf{x}}_t)}{\partial \mathbf{x}}\right)}^{-1}{\nabla }_x\log q_t({\mathbf{x}}_t).(3)$$

This implies that the score function of a $\mathcal{G}$-invariant distribution is $(\mathcal{G},{\nabla }^{-1})$-equivariant. We should use a $(\mathcal{G},{\nabla }^{-1})$-equivariant model to predict the score function. The loss objective is given by

$${\mathcal{J}}_{\text{score}}(\mathbf{\theta })={\mathbb{E}}_{\mathbf{x}\sim {\mathcal{N}}_{\infty }}\left[w(t){\Vert {\mathbf{s}}_{\mathbf{\theta }}({\mathbf{x}}_t,t)-{\nabla }_x\log q_t({\mathbf{x}}_t)\Vert }^2\right],(4)$$

where $w(t)$ is a positive weight function and ${\mathbf{s}}_{\mathbf{\theta }}$ is a $(\mathcal{G},{\nabla }^{-1})$-equivariant model. We will discuss the handling of the intractable score function ${\nabla }_x\log q_t({\mathbf{x}}_t)$ subsequently in Equation 6.

In the context of simulating physical dynamics, two distributional priors are commonly considered: $\mathrm{SE}(n)$ invariance and permutation invariance. They ensure that the learned representations are consistent with the fundamental symmetries of physical laws, including rigid body transformations and indistinguishability of particles, thereby enhancing the model’s ability to generalize across different physical scenarios. The derivations for the following examples can be found in Appendix F.2.

Example 1. (SE(n)-invariant distribution): If $q_0$ is an $\mathrm{SE}(n)$-invariant distribution, then $q_t$ is also $\mathrm{SE}(n)$-invariant. The score function of an $\mathrm{SE}(n)$-invariant distribution is $\mathrm{SO}(n)$-equivariant and translation-invariant.

Example 2. (Permutation-invariant distribution): If $q_0$ is a permutation-invariant, then $q_t$ is also permutation-invariant. The score function of a permutation-invariant distribution is permutation-equivariant.

In the following, we will show that using such a $(\mathcal{G},{\nabla }^{-1})$-equivariant model, we are essentially training a model that focuses on the intrinsic structure of data instead of their representation form.

Equivalence class manifold for invariant distributions. An equivalence class manifold (ECM) refers to the minimum subset of samples that are all considered equivalent in the sense that samples are indistinguishable (formal). For example, in the $n$-dimensional space, coordinates that have undergone rotation and translation maintain their pairwise distances, which allows the use of a set of coordinates to represent all other coordinates with the same distance matrices, thereby forming an equivalence class manifold (see Appendix B for the formal definition and examples). By incorporating the invariance prior to the training set, we can construct ECM from the training set or a mini-batch of samples. The utilization of ECM enables the models to concentrate on the intrinsic structure of the data, thereby enhancing generalization and robustness to irrelevant variations. We assume that the distribution of $\mathbf{x}$ follows a $\mathcal{G}$-invariant distribution $q_t$. Let ${\phi }_t$ map ${\mathbf{x}}_t$ to the corresponding point having the same intrinsic structure in ECM. Then there exists $\mathbf{G}\in \mathcal{G}$ such that $\mathbf{G}({\phi }_t({\mathbf{x}}_t))={\mathbf{x}}_t$. Since ${\phi }_t$ is $inch\mathcal{G}$-invariant, we have ${\phi }_t({\mathbf{x}}_t)={\phi }_t(\mathbf{G}({\phi }_t({\mathbf{x}}_t)))$. Informally, $\mathbf{G}(\mathbf{G})$ where ${\partial }_{\mathbf{G}\mathbf{x}}({\mathbf{x}}_t)$ is defined on the domain of ECM. Taking the logarithm and derivative, we have ${\nabla }_{\mathbf{G}\mathbf{x}}({\mathbf{x}}_t)\log q_t({\mathbf{x}}_t)={\nabla }_{{\phi }_t({\mathbf{x}}_t)}\log q_{{\phi }_t}({\phi }_t({\mathbf{x}}_t))$. Note that ${\nabla }_{\mathbf{G}\mathbf{x}}\log q_t({\mathbf{x}}_t)=\frac{\partial \mathbf{G}({\mathbf{x}}_t)}{\partial \mathbf{x}}{\nabla }_{{\phi }_t({\mathbf{x}}_t)}\log q_t({\mathbf{x}}_t)$. Hence,

$${\nabla }_{\mathbf{x}}\log q_t({\mathbf{x}}_t)=\frac{\frac{\partial \mathbf{G}({\mathbf{x}}_t)}{\partial \mathbf{x}}}{\partial \mathbf{x}}{\nabla }_{{\phi }_t({\mathbf{x}}_t)}\log q_{{\phi }_t}({\phi }_t({\mathbf{x}}_t)).(5)$$

This implies that the score function of the $\mathcal{G}$-invariant distribution is closely related to the score function of the distribution in ECM. Such a result indicates that if we have a $(\mathcal{G},{\nabla }^{-1})$-equivariant model that can predict the score functions in ECM, then this model predicts the score functions for all other points closed under the group operation. We summarize this result in the following theorem whose proofs can be found in Appendix F.3.

Theorem 2 (Equivalence class manifold representation): If we have a $(\mathcal{G},{\nabla }^{-1})$-equivariant model such that ${\mathbf{s}}_{\mathbf{\theta }}({\mathbf{x}}_t,t)={\nabla }_x\log q_{{\phi }_t}({\mathbf{x}}_t)$ almost surely on ${\mathbf{x}}_t\in \mathrm{ECM}$, then we have ${\mathbf{s}}_{\mathbf{\theta }}({\mathbf{x}}_t,t)={\nabla }_{\mathbf{x}}\log q_t({\mathbf{x}}_t)$ almost surely.

Objective for fitting the score function. The score function ${\nabla }_x\log q_t({\mathbf{x}}_t)$ is generally intractable, and we consider the objective for noise matching and data matching (Vincent, 2011; Song et al., 2020; Zheng et al., 2023), where objectives and optimal values are given by

$$\begin{array}{rl}{\mathcal{J}}_{\text{noise}}(\mathbf{\theta })& ={\mathbb{E}}_{\mathbf{x}\sim {\mathcal{N}}_{\infty }}\left[w(t){\Vert {\mathbf{s}}_{\mathbf{\theta }}({\mathbf{x}}_t,t)-\mathbf{x}\Vert }^2\right],{\mathbf{s}}_{\mathbf{\theta }}^{\ast }({\mathbf{x}}_t,t)=-{\sigma }_t{\nabla }_x\log q_t({\mathbf{x}}_t);(6a)\\ {\mathcal{J}}_{\text{data}}(\mathbf{\theta })& ={\mathbb{E}}_{\mathbf{x}\sim {\mathcal{N}}_{\infty }}\left[w(t){\Vert {\mathbf{z}}_{\mathbf{\theta }}({\mathbf{x}}_t,t)-\mathbf{x}\Vert }^2\right],{\mathbf{z}}_{\mathbf{\theta }}^{\ast }({\mathbf{x}}_t,t)=\frac{1}{{\sigma }_t}{\mathbf{x}}_t+\frac{{\sigma }_t^2}{{\sigma }_t}{\nabla }_x\log q_t({\mathbf{x}}_t).(6b)\end{array}$$

The diffusion objectives for both the noise predictor ${\mathbf{s}}_{\mathbf{\theta }}$ and the data predictor ${\mathbf{z}}_{\mathbf{\theta }}$ are intrinsically linked to the score function, thereby inheriting its characteristics and properties. However, the data predictor incorporates a term, $\frac{1}{{\sigma }_t}{\mathbf{x}}_t$, whose numerical range exhibits instability. This instability compromises the predictor’s ability to inherit the straightforward properties of the score function. Therefore, to incorporate $\mathcal{G}$-invariance, it is advisable to employ noise matching, which is given by Equation 6a and ${\mathbf{s}}_{\mathbf{\theta }}$ is $(\mathcal{G},{\nabla }^{-1})$-equivariant, which is the property of the score function.

A specific instance of a distributional prior is defined by samples that conform to the constraints imposed by PDEs. In this context, the dynamics at any given spatial location depend solely on the characteristics of the system within its local vicinity, rather than on absolute spatial coordinates. Under these conditions, it is appropriate to employ translation-invariant models for both noise matching and data matching. Nevertheless, the samples in question exhibit significant smoothness. As a result, utilizing the noise matching objective necessitates that the model’s output be accurate at every individual pixel. In contrast, applying the data matching objective only requires the model to produce smooth output values. Therefore, it is recommended to adopt the data matching objective for this purpose. The selection between data matching and noise matching plays a critical role in determining the quality of the generated samples. For detailed experimental results, refer to Sec. 4.3.

Remark 1. In this section, we primarily explore the principle for incorporating distributional priors by selecting models with particular characteristics. Specifically:

1. When the distribution exhibits $\mathcal{G}$-invariance, a $(\mathcal{G},{\nabla }^{-1})$-equivariant model should be employed alongside the noise matching objective (Equation 6a).
2. For samples that are subject to PDE constraints and exhibit high smoothness, the data matching objective (Equation 6b) is recommended.

3.2 Incorporating Physical Feasibility Priors

In this section, we explore how to incorporate physical feasibility priors such as physics laws and explicit PDE constraints into noise and data matching objectives in diffusion models. By Ito’s formula (Oksendal, 2011; Kim & Ye, 2021; Chung et al., 2022), we have $\mathbb{E}[{\mathbf{x}}_0\mid {\mathbf{x}}_t]={\mathbf{x}}_t+{\sigma }_t^2{\nabla }_x\log q_t({\mathbf{x}}_t)$. Thus, the physical feasibility priors can be imposed by ensuring that the output of the diffusion model, ${\mathbf{s}}_{\mathbf{\theta }}({\mathbf{x}}_t,t)$ or ${\mathbf{z}}_{\mathbf{\theta }}({\mathbf{x}}_t,t)$, satisfies the physical constraints, such as conservation laws or PDEs.

Theorem 3 (Multilinear Jensen’s gap). The optimizer for ${\mathbb{E}}_{x\sim {\mathcal{N}}_{\infty }}\left[w(t){\Vert {\mathbf{z}}_{\mathbf{\theta }}({\mathbf{x}}_t,t)-\mathbf{x}\Vert }^2\right]$ satisfies the physical constraints if the ground-truth data $\mathbf{x}$ adheres to the physical laws, provided that the model ${\mathbf{z}}_{\mathbf{\theta }}$ is expressive enough to capture the data distribution and the constraints are multilinear in nature.

This theorem implies that by optimizing the data matching objective, we can enforce physical feasibility priors implicitly through the data distribution, provided the constraints are multilinear (e.g., conservation of energy and momentum). For non-multilinear constraints, such as those arising from complex PDEs, explicit incorporation is required.

Explicit PDE Constraints. For PDE-constrained dynamics, we modify the loss function to include a penalty term that enforces the PDE constraints. Let ${\mathcal{L}}_{\text{PDE}}(\mathbf{x})$ represent the PDE residual, which measures the deviation of the generated samples from the governing PDE. The modified loss function is:

$${\mathcal{J}}_{\text{total}}(\mathbf{\theta })={\mathcal{J}}_{\text{data}}(\mathbf{\theta })+\lambda {\mathbb{E}}_{\mathbf{x}\sim {\mathcal{N}}_{\infty }}\left[{\mathcal{L}}_{\text{PDE}}({\mathbf{z}}_{\mathbf{\theta }}({\mathbf{x}}_t,t))\right],$$

where $\lambda$ is a hyperparameter controlling the strength of the PDE constraint. This approach ensures that the generated samples not only match the data distribution but also adhere to the PDE constraints.

4 Experiments

In this section, we evaluate the performance of our proposed framework on various datasets, including PDE datasets and particle dynamics datasets. We provide a detailed account of the selection of backbones and the training strategies employed for each dataset. We also provide ablation studies in Sec. 4.3 of 1) data matching and noise matching techniques for different datasets, revealing that incorporating a distributional prior enhances model performance; 2) the effect of Jensen’s gap, finding that individual constraints, if not properly handled, lead to underperformance. However, appropriately managing these priors using our proposed methods can lead to significant performance improvements.

4.1 PDE Datasets

PDE datasets, including advection (Zang, 1991), Darcy flow (Li et al., 2022), Burgers (Rudy et al., 2017), and Navier-Stokes (Li et al., 2020), represent a diverse set of physical phenomena. These datasets enable the simulation of complex systems, demonstrating the capability of models for broader application across a wide range of PDE datasets. Through this, they facilitate advances in understanding diverse natural and engineered processes.

Experiment settings. The PDE constraints for the above datasets are given by:

• Advection: $\frac{\partial u}{\partial t}+u\cdot \nabla u=0$
• Darcy flow: $-\nabla \cdot (a(x)\nabla u)=f(x)$
• Burgers: $\frac{\partial u}{\partial t}+u\frac{\partial u}{\partial x}=\nu \frac{{\partial }^{2}u}{\partial x^2}$
• Navier-Stokes: $\frac{\partial \mathbf{u}}{\partial t}+(\mathbf{u}\cdot \nabla )\mathbf{u}=-\frac{1}{\rho }\nabla p+\nu {\nabla }^2\mathbf{u}$

We train our diffusion models with the modified loss function (Equation 7) to incorporate these PDE constraints. The backbone architecture is based on a U-Net (Ronneberger et al., 2015) with group-equivariant convolutions to enforce $\mathrm{SE}(n)$ invariance. The training data consists of simulated solutions to these PDEs, and we evaluate the generated samples based on their adherence to the PDE constraints and their visual fidelity to the ground-truth dynamics.

4.2 Particle Dynamics Datasets

We train diffusion models to simulate the dynamics of chaotic three-body systems in 3D (Zhou & Yu, 2023) and five spring systems in 2D (Kuramoto, 1975; Kipt et al., 2018) (see Appendix D.1 for visualizations of datasets). In the case of the three-body system, we unconditionally generate the positions and velocities of three particles, where gravitational interactions govern their dynamics. The stochastic nature of this dataset arises from the random distribution of the initial positions and velocities. In five spring systems, each pair of particles has a probability $50\%$ of being connected by a spring. The movements of the particles are influenced by the spring forces, which cause stretching or compression interactions. We conditionally generate the positions and velocities of the five particles based on their spring connectivity.

Notations. The features of the datasets are represented as $\mathbf{x}\in {\mathbb{R}}^{N\times D}$, where $N$ is the number of particles, and $D$ is the dimensionality of the position and velocity vectors (e.g., $D=6$ for 3D position and velocity). The physical constraints include conservation of total momentum and total energy, which are enforced through the data matching objective.

Figure 2: Visualization of generated samples from the three-body (first row) and five spring (second row) datasets. The leftmost figures in each row represent methods without priors, the middle figures correspond to our proposed methods, the rightmost figures illustrate the physical properties as they evolve over time. Both total momentum and total energy should remain conserved. The samples generated by our methods demonstrate improved adherence to these conservation laws compared to baseline methods.

4.3 Ablation Studies

We conduct ablation studies to evaluate the impact of incorporating distributional and physical feasibility priors. We compare the following configurations:

1. Baseline (No Priors): Diffusion models trained without any distributional or physical priors.
2. Distributional Priors Only: Models trained with $\mathrm{SE}(n)$ and permutation invariance but without explicit physical constraints.
3. Physical Priors Only: Models trained with physical constraints (e.g., PDE residuals or conservation laws) but without distributional priors.
4. Proposed Method: Models trained with both distributional and physical priors, using the modified loss function (Equation 7).

Results. Table 3 summarizes the sample quality for the three-body and five spring datasets. For both datasets, we simulate the ground-truth future motion based on the current states of the generated samples and report the Mean Squared Error (MSE) between the ground-truth motion and the generated ones. We also calculate the error in physical feasibility, such as conservation of energy and momentum, which should remain unchanged along the evolution of the systems.

Table 3: Sample Quality of the Three-Body and Five Spring Datasets

Method	Three-Body MSE	Three-Body Energy Error	Three-Body Momentum Error	Five Spring MSE	Five Spring Energy Error	Five Spring Momentum Error
Baseline (No Priors)	0.052	0.134	0.128	0.067	0.145	0.139
Distributional Priors Only	0.041	0.092	0.087	0.053	0.101	0.095
Physical Priors Only	0.038	0.045	0.042	0.049	0.048	0.046
Proposed Method	0.029	0.021	0.019	0.035	0.022	0.020

The results demonstrate that incorporating both distributional and physical priors significantly improves the quality of generated samples. The proposed method achieves the lowest MSE and errors in energy and momentum conservation, indicating that the generated dynamics are both visually accurate and physically consistent.

5 Conclusion

In this work, we introduced a novel framework for generating physically feasible dynamics by integrating distributional and physical feasibility priors into diffusion-based generative models. Our approach addresses the limitations of existing methods by ensuring that generated dynamics adhere to fundamental physical laws, such as conservation of energy and momentum, and PDE constraints. Empirical evaluations on PDE and particle dynamics datasets demonstrate the effectiveness of our method in producing high-quality, physically consistent dynamics. Our contributions include:

1. A robust framework for embedding physical priors into diffusion models, improving the feasibility of generated dynamics.
2. A solution to the challenge of generating physically plausible dynamics in data-driven settings.
3. Demonstrations of improved performance across diverse physical phenomena, highlighting the potential of our approach for AI4Physics applications.

Future work will explore the scalability of our framework to larger and more complex systems, as well as the integration of additional physical priors, such as thermodynamic constraints, to further enhance the applicability of our method.