Generating Physical Dynamics under Priors

## Abstract

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 this 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, ), climate modeling (Rasp et al., ), and materials science (Choudhary et al., ). Accurately generating such dynamics is crucial for advancing our understanding and predictive capabilities in these fields (Bzdok & Ioannidis, ). Recently, generative models have revolutionized the study of physics by providing powerful tools to simulate and predict complex systems.

Generative v.s. discriminative models. Even when high-performing discriminative models for dynamics are available such as finite elements (Zhang et al., ; Uriarte et al., ), finite difference (Lu et al., ; Salman et al., ), finite volume (Ranade et al., ) or physics-informed neural networks (PINNs) (Raissi et al., ), 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., ), anomaly detection (Finke et al., ), and semi-supervised learning (Ma et al., ). They enhance robustness and interpretability by modeling the joint distribution of data and labels, offering insights into unseen scenarios (Takeishi & Kalousis, ). Generative models are also pivotal in creative domains, such as drug discovery (Lavecchia, ), 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., ), 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.

Figure: 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.

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 predictive tasks, where the objective is to estimate a *specific ground-truth value* $\bm{x}_{0}$, diffusion generative models aim to characterize an *full ground-truth distribution* $\nabla{\bm{x}}\log q_{t}(\bm{x}_{t})$ or $\mathbb{E}[\bm{x}_{0}\mid\bm{x}_{t}]$ (notations 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 into a diffusion-based generation framework, our approach 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.

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 (Martínez-Aranda et al., ). 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 $\bm{x}_{0}\in\mathbb{R}^{n}$ drawn from an unknown distribution $q_{0}$. Denoising diffusion probabilistic models (Song & Ermon, ; Song et al., ; Ho et al., ) describe a forward process ${\bm{x}_{t},t\in[0,T]}$ governed by an Ito stochastic differential equation (SDE)

$$ $\mathrm{d}\bm{x}_{t}=f(t)\bm{x}_{t}\mathrm{d}t+g(t)\mathrm{d}\mathbf{w}_{t}, \quad\bm{x}_{0}\sim q_{0},\quad 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}\left(\bm{x}_{t}\mid\bm{x}_{0}\right)=\mathcal{N}\left(\bm{x}_{t}\mid
\alpha_{t}\bm{x}_{0},\sigma_{t}^{2}\mathbf{I}\right)$ and has a corresponding reverse process of with the probability flow ordinary differential equation (ODE) (Song et al., ), running from time $T$ to 0 0, defined as

$$ $\frac{\mathrm{d}\bm{x}_{t}}{\mathrm{d}t}=f(t)\bm{x}_{t}-\frac{1}{2}g^{2}(t) \nabla{\bm{x}}\log q_{t}(\bm{x}_{t}),\quad\bm{x}_{T}\sim q_{T}(\bm{x}_{T}\mid \bm{x}_{0})\approx q_{T}(\bm{x}_{T}).$ (2) $$

The marginal probability densities $\{q_{t}(\bm{x}_{t})\}_{t=0}^{T}$ of the forward SDE align with the reverse ODE (Song et al., ). This indicates that if we can sample from $q_{T}(\bm{x}_{T})$ and solve Equation 2, then the resulting $\bm{x}_{0}$ will follow the distribution $q_{0}$. By choosing $\alpha_{t}\rightarrow 0$ and $\sigma_{t}\rightarrow 1$, the distribution $q_{T}(\bm{x}_{T})$ can be approximated as a normal distribution. The score $\nabla{\bm{x}}\log q_{t}(\bm{x}_{t})$ can be approximated by a deep learning model. The quality of the generated samples is contingent upon the models’ ability to accurately approximate the score functions (Kwon et al., ; Gao & Zhu, ). A more precise approximation results in a distribution that closely aligns with the distribution of the training set. To enhance model fit, incorporating priors of the distributions and 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 $\mathbf{G}\in\mathcal{G}$, we have $q(\mathbf{G}(\bm{x}))=q(\bm{x})$. 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 ($\operatorname{SE}(3)$ group)—is essential for spatial understanding (Zhou et al., ). Equivariant models are usually required to embed invariance. A function (or model) $\mathbf{f}:\mathbb{R}^{n}\rightarrow\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 $\mathbf{G}\in\mathcal{G}$, $\mathbf{f}(\mathbf{G}(\bm{x}))=\mathcal{L}(\mathbf{G})(\mathbf{f}(\bm{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 model’s output of a distribution related to the ground-truth samples (specifically, $\nabla{\bm{x}}\log q_{t}(\bm{x}_{t})$ or $\mathbb{E}[\bm{x}_{0}\mid\bm{x}_{t}]$). In this section, we consider the forward diffusion process given by Equation 1, where $\bm{x}_{t}=\alpha_{t}\bm{x}_{0}+\sigma_{t}\bm{\epsilon}$, with $\bm{\epsilon}\sim\mathcal{N}(\mathbf{0},\mathbf{I})$.

### 3.1 Incorporating distributional priors

In this section, we study the score function $\nabla_{\bm{x}}\log q_{t}(\bm{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_{0}$ to imply the invariance of $q_{t}$ ) .

Let $q_{0}$ be a $\mathcal{G}$-invariant distribution. If for all $\mathbf{G}\in\mathcal{G}$, $\mathbf{G}$ is volume-preserving diffeomorphism and isometry, and for all $0<a<1$, there exists $\mathbf{H}\in\mathcal{G}$ such that $\mathbf{H}(a\bm{x})=a\mathbf{G}(\bm{x})$, then $q_{t}$ is also $\mathcal{G}$-invariant.

###### Theorem 1 (Sufficient conditions for the invariance of q 0 subscript 𝑞 0 q_{0} italic_q start_POSTSUBSCRIPT 0 end_POSTSUBSCRIPT to imply the invariance of q t subscript 𝑞 𝑡 q_{t} italic_q start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) .

##### Property of score functions.

Let $q_{t}$ be a $\mathcal{G}$-invariant distribution. By the chain rule, we have $\nabla_{\bm{x}}\log q_{t}(\bm{x})=\nabla_{\bm{x}}\log q_{t}(\mathbf{G}(\bm{x})
)=\frac{\partial\mathbf{G}(\bm{x})}{\partial\bm{x}}\nabla_{\mathbf{G}(\bm{x})}
\log q_{t}(\mathbf{G}(\bm{x}))$, for all $\mathbf{G}\in\mathcal{G}$. Hence,

$$ $\nabla_{\mathbf{G}(\bm{x})}\log q_{t}(\mathbf{G}(\bm{x}))=\left(\frac{\partial \mathbf{G}(\bm{x})}{\partial\bm{x}}\right)^{-1}\nabla_{\bm{x}}\log q_{t}(\bm{x }).$ (3) $$

This implies that the score function of $\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}_{\textsubscript{score}}(\bm{\theta})=\mathbb{E}_{t,\bm{x}_{0},\bm{ \epsilon}}\left[w(t)\left\|\mathbf{s}_{\bm{\theta}}\left(\bm{x}_{t},t\right)- \nabla{\bm{x}}\log q_{t}(\bm{x}_{t})\right\|^{2}\right],$ (4) $$

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

In the context of simulating physical dynamics, two distributional priors are commonly considered: $\operatorname{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 .

($\operatorname{SE(n)}$-invariant distribution) If $q_{0}$ is an $\operatorname{SE(n)}$-invariant distribution, then $q_{t}$ is also $\operatorname{SE(n)}$-invariant. The score function of an $\operatorname{SE(n)}$-invariant distribution is $\operatorname{SO(n)}$-equivariant and translational-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.

###### Example 1 .

###### Example 2 .

##### Equivalence class manifold for invariant distributions.

An *equivalence class manifold* (ECM) refers to the minimum subset of samples where all the rest elements are considered equivalent to one of the samples in this manifold (informal). For example, in three-dimensional space, coordinates that have undergone rotation and translation maintain their intrinsic structures of 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 $\bm{x}$ follows an $\mathcal{G}$-invariant distribution $q_{t}$. Let $\varphi$ map $\bm{x}$ to the corresponding point having the same intrinsic structure in ECM. Then there exists $\mathbf{G}\in\mathcal{G}$ such that $\mathbf{G}(\varphi(\bm{x}))=\bm{x}$ . Since $q_{t}$ is $\mathcal{G}$-invariant, we have $q_{t}(\bm{x})=q_{\textsubscript{ECM}}(\varphi(\bm{x}))\cdot p_{\operatorname{
Uniform}(\mathcal{G})}(\mathbf{G})$, where $q_{\textsubscript{ECM}}$ is defined on the domain of ECM. Taking the logarithm and derivative, we have $\nabla_{\varphi(\bm{x})}\log q_{t}(\bm{x})=\nabla_{\varphi(\bm{x})}\log q_{
\textsubscript{ECM}}(\varphi(\bm{x}))$. Note that $\nabla_{\bm{x}}\log q_{t}(\bm{x})=\frac{\partial\varphi(\bm{x})}{\partial\bm{x
}}\nabla_{\varphi(\bm{x})}\log q_{t}(\bm{x})$. Hence,

$$ $\nabla_{\bm{x}}\log q_{t}(\bm{x})=\frac{\partial\varphi(\bm{x})}{\partial\bm{x }}\nabla_{\varphi(\bm{x})}\log q_{\textsubscript{ECM}}(\varphi(\bm{x})).$ (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}_{\bm{\theta}}(\bm{x})=\nabla_{\bm{x}}\log q_{\textsubscript{ECM}}(
\bm{x})$ almost surely on $\bm{x}\in\operatorname{ECM}$, then we have $\mathbf{s}_{\bm{\theta}}(\bm{x})=\nabla_{\bm{x}}\log q_{t}(\bm{x})$ almost surely.

###### Theorem 2 (Equivalence class manifold representation) .

##### Objective for fitting the score function.

The score function $\nabla{\bm{x}}\log q_{t}(\bm{x}_{t})$ is generally intractable and we consider the objective for *noise matching* and *data matching* (Vincent, ; Song et al., ; Zheng et al., ), where objectives and optimal values are given by

$$ $\displaystyle\mathcal{J}_{\textsubscript{noise}}(\bm{\theta})$ $\displaystyle=\mathbb{E}_{t,\bm{x}_{0},\bm{\epsilon}}\left[w(t)\left\|\bm{ \epsilon}_{\bm{\theta}}\left(\bm{x}_{t},t\right)-\bm{\epsilon}\right\|^{2} \right],\quad$ $\displaystyle\bm{\epsilon}_{\bm{\theta}}^{*}\left(\bm{x}_{t},t\right)$ $\displaystyle=-\sigma_{t}\nabla_{\bm{x}}\log q_{t}\left(\bm{x}_{t}\right);$ (6a) $\displaystyle\mathcal{J}_{\textsubscript{data}}(\bm{\theta})$ $\displaystyle=\mathbb{E}_{t,\bm{x}_{0},\bm{\epsilon}}\left[w(t)\left\|\bm{x}_{ \bm{\theta}}\left(\bm{x}_{t},t\right)-\bm{x}_{0}\right\|^{2}\right],\quad$ $\displaystyle\bm{x}_{\bm{\theta}}^{*}\left(\bm{x}_{t},t\right)$ $\displaystyle=\frac{1}{\alpha_{t}}\bm{x}_{t}+\frac{\sigma_{t}^{2}}{\alpha_{t}} \nabla_{\bm{x}}\log q_{t}\left(\bm{x}_{t}\right).$ (6b) $$

The diffusion objectives for both the noise predictor $\bm{\epsilon}_{\theta}$ and the data predictor $\bm{x}_{\bm{\theta}}$ are intrinsically linked to the score function, thereby inheriting its characteristics and properties. However, the data predictor incorporates a term, $\frac{1}{\alpha_{t}}\bm{x}_{t}$, whose numerical range exhibits instability. This instability complicates 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 $\bm{\epsilon}_{\bm{\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.

###### Remark 1 .

### 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 Tweedie’s formula (Efron, ; Kim & Ye, ; Chung et al., ), we have $\mathbb{E}[\bm{x}_{0}\mid\bm{x}_{t}]=\frac{1}{\alpha_{t}}\left(\bm{x}_{t}+
\sigma^{2}_{t}\nabla_{\bm{x}}\log q_{t}\left(\bm{x}_{t}\right)\right)$. Hence,

$$ $\mathbb{E}[\bm{x}_{0}\mid\bm{x}_{t}]=\frac{1}{\alpha_{t}}\left(\bm{x}_{t}- \sigma_{t}\bm{\epsilon}_{\bm{\theta}}^{*}\left(\bm{x}_{t},t\right)\right), \quad\mathbb{E}[\bm{x}_{0}\mid\bm{x}_{t}]=\bm{x}_{\bm{\theta}}^{*}\left(\bm{x} _{t},t\right).$ (7) $$

For both noise and data matching objectives, we are essentially training a model to approximate $\mathbb{E}[\bm{x}_{0}\mid\bm{x}_{t}]$. A purely data-driven approach is often insufficient to capture the underlying physical constraints accurately. Therefore, similar to PINNs (Leiteritz & Pflüger, ), we incorporate an additional penalty loss $\mathcal{J}_{\mathcal{R}}$ into the objective function to enforce physical feasibility priors $\mathcal{R}\left(\bm{x}_{0}\right)=\mathbf{0}$ and set the loss objective to be $\mathcal{J}(\bm{\theta})=\mathcal{J}_{\textsubscript{score}}(\bm{\theta})+
\lambda\mathcal{J}_{\mathcal{R}}(\bm{\theta})$, where $\mathcal{J}_{\textsubscript{score}}$ is the data matching or noise matching objectives and $\lambda$ is a hyperparameter to balance the diffusion loss and physical feasibility loss. We consider the data matching objective where $\bm{x}_{\bm{\theta}}\left(\bm{x}_{t},t\right)$ approximates $\mathbb{E}[\bm{x}_{0}\mid\bm{x}_{t}]$. For noise matching models, we can transform the model’s output by Equation 7. For general cases, we cannot directly add the constraints $\mathcal{R}\left(\bm{x}_{0}\right)=\mathbf{0}$ to the output of the diffusion model $\mathbb{E}[\bm{x}_{0}\mid\bm{x}_{t}]$ due to the presence of Jensen’s gap (Bastek et al., ), i.e., $\mathcal{R}\left(\mathbb{E}[\bm{x}_{0}\mid\bm{x}_{t}]\right)\neq\mathbb{E}[
\mathcal{R}\left(\bm{x}_{0}\right)\mid\bm{x}_{t}]=\mathbf{0}$. However, in some special cases, we can avoid dealing with this gap.

##### Linear cases.

When the constraints are linear/affine functions, Jensen’s gap equals 0. Hence, we can directly add the constraints to $\bm{x}_{\bm{\theta}}\left(\bm{x}_{t},t\right)$. We have $\mathcal{J}_{\mathcal{R}}\left(\bm{\theta}\right)=\mathbb{E}_{t,\bm{x}_{0},\bm
{\epsilon}}\left[w(t)\left\|\mathcal{R}\left(\bm{x}_{\bm{\theta}}\left(\bm{x}_
{t},t\right)\right)\right\|^{2}\right]$.

##### Multilinear cases.

A function is called multilinear if it is linear in several arguments when the other arguments are fixed. Denote $\bm{x}_{0}=\begin{bmatrix}\mathbf{u}_{0}\\
\mathbf{v}_{0}\end{bmatrix}\in\mathbb{R}^{m+n},\mathbf{u}_{0}\in\mathbb{R}^{m}
,\mathbf{v}_{0}\in\mathbb{R}^{n}$. When the constraints function is multilinear w.r.t. $\mathbf{u}_{0}$, we can write the constraints in the form of $\mathcal{R}\left(\bm{x}_{0}\right)=\mathbf{W}_{0}\mathbf{u}_{0}+\mathbf{b}_{0}
=\mathbf{0}$, where $\mathbf{W}_{0}$ and $\mathbf{b}_{0}$ are functions of $\mathbf{v}_{0}$. In this case, we can use the penalty loss as $\mathcal{J}_{\mathcal{R}}\left(\bm{\theta}\right)=\mathbb{E}_{t,\bm{x}_{0},\bm
{\epsilon}}[w(t)\|\mathbf{W}_{0}\mathbf{u}_{\bm{\theta}}\left(\bm{x}_{t},t
\right)+\mathbf{b}_{0}\|^{2}]$. Such a design is supported by the following theorem whose proof can be found in Appendx F.4.

###### Theorem 3 (Multilinear Jensen’s gap) .

The optimizer for $\mathbb{E}_{t,\bm{x}_{0},\bm{\epsilon}}[w(t)\|\mathbf{u}_{\bm{\theta}_{1}}
\left(\bm{x}_{t},t\right)-\mathbf{u}_{0}\|^{2}]$ is the reweighted optimizer of $\mathbb{E}_{t,\bm{x}_{0},\bm{\epsilon}}[w(t)\|\mathbf{W}_{0}\mathbf{u}_{\bm{
\theta}_{2}}\left(\bm{x}_{t},t\right)+\mathbf{b}_{0}\|^{2}]$ with reweighted variable $\mathbf{W}_{0}^{\top}\mathbf{W}_{0}$.

###### Theorem 3 (Multilinear Jensen’s gap) .

##### Convex cases.

If the constraints $\mathcal{R}$ is convex, by Jensen’s inequality, $\mathcal{R}\left(\mathbb{E}[\bm{x}_{0}\mid\bm{x}_{t}]\right)\leq\mathbb{E}[
\mathcal{R}\left(\bm{x}_{0}\right)\mid\bm{x}_{t}]=\mathbf{0}$. Hence, $0=\|\mathbb{E}[\mathcal{R}\left(\bm{x}_{0}\right)\mid\bm{x}_{t}]\|^{2}\leq\|
\mathcal{R}\left(\mathbb{E}[\bm{x}_{0}\mid\bm{x}_{t}]\right)\|^{2}$. When a data matching model is approximately optimized, directly applying constraints to the model’s output minimizes the upper bound of the constraints on $\bm{x}_{0}$. The upper bound of the Jensen’s gap is related the absolute centered moment $\sigma_{p}=\sqrt[p]{\mathbb{E}[|\bm{x}_{0}-\mu|\bm{x}_{t}|^{p}]}$, where $\mu=\mathbb{E}[\bm{x}_{0}|\bm{x}_{t}]$. If the constraints function $\mathcal{R}$ approach $\mathcal{R}(\mu)$ no slower than $|\bm{x}_{0}-\mu|^{\eta}$ and grow as $\bm{x}_{0}\rightarrow\pm\infty$ no faster than $\pm|\bm{x}_{0}|^{n}$ for $n\geq\eta$, then the Jensen’s gap $\mathbb{E}[\mathcal{R}\left(\bm{x}_{0}\right)\mid\bm{x}_{t}]-\mathcal{R}\left(
\mathbb{E}[\bm{x}_{0}\mid\bm{x}_{t}]\right)$ approaches to 0 no slower than $\sigma_{n}^{\eta}$ as $\sigma_{n}\rightarrow 0$ (Gao et al., ). Usually, in the reverse diffusion process, $\sigma_{n}\rightarrow 0$ as $t\rightarrow 0$ since the generated noisy samples converge to a clean one. In this case, we use the penalty loss of $\mathcal{J}_{\mathcal{R}}\left(\bm{\theta}\right)=\mathbb{E}_{t,\bm{x}_{0},\bm
{\epsilon}}\left[w(t)\left\|\mathcal{R}\left(\bm{x}_{\bm{\theta}}\left(\bm{x}_
{t},t\right)\right)\right\|^{2}\right]$.

In the aforementioned three scenarios, at the implementation level, the model’s output may be directly considered as the ground-truth sample $\bm{x}_{0}$ itself, rather than the conditional expectation $\mathbb{E}[\bm{x}_{0}\mid\bm{x}_{t}]$. These scenarios are referred to as “elementary cases”. In the following, we will discuss how to deal with nonlinear cases using the above elementary cases.

##### Reducible nonlinear cases.

For nonlinear constraints, mathematically speaking, we cannot directly apply the constraints to $\mathbb{E}[\bm{x}_{0}\mid\bm{x}_{t}]$. However, we may recursively use multilinear functions to decompose the nonlinear constraints into elementary ones as: $\mathcal{R}\left(\bm{x}_{0}\right)=\mathbf{g}_{1}\circ\cdots\mathbf{g}_{m}
\left(\bm{x}_{0}\right)=\mathbf{0}$, where all $\mathbf{g}_{i}$ are elementary. Using elementary functions for decomposition, we may 1) reduce nonlinear constraints into elementary ones by treating terms causing nonlinearity as constants, and 2) reduce the complex constraints into several simpler ones. In this case, the penalty loss is set to $\mathcal{J}_{\mathcal{R}}\left(\bm{\theta}\right)=\mathbb{E}_{t,\bm{x}_{0},\bm
{\epsilon}}\left[w(t)\left\|\textbf{g}_{1}\circ\cdots\textbf{g}_{m}\left(\bm{x
}_{\bm{\theta}}\left(\bm{x}_{t},t\right)\right)\right\|^{2}\right]$. See Sec. 4.2 for concrete examples of nonlinear formulas for the conservation of energy.

##### General nonlinear cases.

For general nonlinear cases, if it is not feasible to decompose the nonlinear constraints into their elementary components, it may be necessary to consider alternative approaches where we may reparameterize the constraints variable into elementary cases. Given the nonlinear constraints, we reparameterize it as $\mathcal{R}\left(\bm{x}_{0}\right)=\mathbf{g}\left(\mathbf{h}(\bm{x}_{0})
\right)=\mathbf{0}$, where $\mathbf{g}$ is elementary and $\mathbf{h}$ is non-necessarily elementary functions. Subsequently, another diffusion model, denoted as $\tilde{\bm{x}}_{\bm{\theta}}\left(\bm{x}_{t},t\right)$, is trained to predict $\mathbf{h}(\bm{x}_{0})$, utilizing the same hidden states as model $\bm{x}_{\bm{\theta}}\left(\bm{x}_{t},t\right)$. This training process employs the methods applicable to elementary cases. The objective is for model $\tilde{\bm{x}}_{\bm{\theta}}$ to learn the underlying physical constraints and encode these constraints into its hidden states. Consequently, when model $\bm{x}_{\bm{\theta}}$ predicts, it inherently incorporates the learned physical constraints $\mathbf{g}$ parameterized by $\mathbf{h}(\bm{x}_{0})$. To train model $\tilde{\bm{x}}_{\bm{\theta}}$, we set the penalty loss to be $\mathcal{J}_{\mathcal{R}}\left(\bm{\theta}\right)=\mathbb{E}_{t,\bm{x}_{0},\bm
{\epsilon}}\left[w(t)\left\|\tilde{\bm{x}}_{\bm{\theta}}\left(\bm{x}_{t},t
\right)-\mathbf{h}(\bm{x}_{0})\right\|^{2}\right]$. See Appendix E.1 for implementation details.

Notably, in our proposed methods for integrating constraints, the explicit form of prior knowledge, such as the physics constants required for energy calculations, is not necessary. Instead, it suffices to determine whether the model’s output parameters are elementary w.r.t. the constraints. This approach enhances the applicability of our methods to a broader spectrum of constraints.

###### Remark 2 .

In conclusion, incorporating the physics constraints can be achieved in different ways depending on their complexity. For elementary constraints, one can directly omit Jensen’s gap and impose the penalty loss on the model’s output. In the case of nonlinear constraints, decomposition or reparameterization techniques are utilized to transform constraints into elementary ones.

###### Remark 2 .

## 4 Experiments

In this section, we assess the enhancement achieved by incorporating physics constraints into the fundamental diffusion model across various synthetic physics datasets. We conduct a grid search to identify an equivalent set of suitable hyperparameters for the network to perform the data/noise matching, ensuring a fair comparison between the baseline method (diffusion objectives without penalty loss) and our proposed approach of incorporating physics constraints. Appendix E provides 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 omitting Jensen’s gap, finding that nonlinear constraints can hinder performance if not properly handled. However, appropriately managing these priors using our proposed methods can lead to significant performance improvements.

### 4.1 PDE datasets

PDE datasets, including advection (Zang, ), Darcy flow (Li et al., ), Burgers (Rudy et al., ), and shallow water (Klöwer et al., ), are fundamental resources for studying and modeling various 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: $\displaystyle\quad\partial_{t}u(t,x)+\beta\partial_{x}u(t,x)$ $\displaystyle=$ $\displaystyle\;0,$ (8a) Darcy flow: $\displaystyle\quad\partial_{t}u(x,t)-\nabla(a(x)\nabla u(x,t))$ $\displaystyle=$ $\displaystyle\;f(x),$ (8b) Burger: $\displaystyle\quad\partial_{t}u(x,t)+u(x,t)\partial_{x}u(x,t)$ $\displaystyle=$ $\displaystyle\;0,$ (8c) Shallow water: $\displaystyle\quad\partial_{t}u=-\partial_{x}h,\quad\partial v_{t}=-\partial h _{y},$ $\displaystyle\partial_{t}h=-c^{2}\left(\partial_{x}u+\partial_{y}v\right).$ (8d) $$

A detailed introduction and visualization of the datasets can be found in Appendix C.1. In this study, we investigate the predictive capabilities of generative models applied to advection and Darcy flow datasets. Our experiments focus on evaluating the models’ accuracy in forecasting future states given initial conditions. Additionally, we examine the models’ ability to generate physically feasible samples that align with the distribution of the training set on advection, Burger, and shallow water datasets. The evaluation metrics are designed to assess to what extent the solutions adhere to the physical feasibility constraints imposed by the corresponding PDEs.

##### Injecting physical feasibility priors.

We train the models that apply the data matching objective as suggested in Remark 1. We employ finite difference methods to approximate the differential equations. This approach renders the PDE constraints linear for the advection, Darcy flow, and shallow water datasets. However, PDE constraints become multilinear for the Burgers’ equation dataset (see Appendix C.2 for the proof). Thus, the first set of datasets: advection, Darcy flow, and shallow water—correspond to the linear case, while the Burgers’ equation dataset corresponds to the multilinear case. We can directly apply the physical feasibility constraints on the model’s output.

##### Experimental results.

Results can be seen in Tab. 1, 2. In Tab. 1, we analyze the performance of diffusion models in predicting physical dynamics, given initial conditions, within a generative framework that produces a Dirac distribution. The accuracy of these models is evaluated using the RMSE metric. The observed loss magnitude is comparable to the prediction loss using with FNO, U-Net, and PINN models (Takamoto et al., ) (refer to Appendix E.3 for further details). Our results indicate that the incorporation of constraints consistently enhances the accuracy of the prediction. In Tab. 2, the feasibility of the generated samples is evaluated by calculating the RMSE of the PDE constraints, which determine the impact of incorporating physical feasibility priors on diffusion models. We also provide visualization of the generated samples in Fig. 10, 11, 12.

**Table 1: Performance comparison of diffusion models with/without priors for predicting physical dynamics. The models’ accuracy is measured using the RMSE metric, highlighting the impact of incorporating constraints on improving prediction accuracy.**
| Method | Advection ($\times 10^{-2}$) | Darcy flow ($\times 10^{-2}$) |
| --- | --- | --- |
| w/o prior | 1.716 | 2.261 |
| w/ prior | 1.621 | 2.174 |

### 4.2 Particle dynamics datasets

We train diffusion models to simulate the dynamics of chaotic three-body systems in 3D (Zhou & Yu, ) and five-spring systems in 2D (Kuramoto, ; Kipf et al., ) (see Appenidx. D.1 for visualizations of datasets). In the case of the three-body, 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}^{(0)}=[\mathbf{C}^{(0)}\;\mathbf{V}^{(0)}]\in\mathbb{R}^{L\times K
\times 2D}$, where $\mathbf{C}^{(0)},\mathbf{V}^{(0)}\in\mathbb{R}^{L\times K\times D}$. Here, the matrix $\mathbf{C}_{0}$ encapsulates the coordinate features, while $\mathbf{V}_{0}$ encapsulates the velocity features. The superscript denotes the time for the diffusion process and the subscripts denote the matrix index. $L$ represents the temporal length of the physical dynamics, $K$ denotes the number of particles, and $D$ corresponds to the spatial dimensionality. We use the subscript $l$ to indicate time, while the subscripts $i,j$, and $k$ are used to denote the indices of particles. The subscript $d$ represents the index corresponding to the spatial axis. We also use the subscript of $\bm{\theta}$ to denote the corresponding values of the model’s prediction of $\mathbb{E}[\mathbf{X}^{(0)}\mid\mathbf{X}^{(t)}]$ with inputs $\mathbf{X}^{(t)}$ and $t$, and $\mathbf{X}_{\bm{\theta}}=[\mathbf{C}_{\bm{\theta}}\;\mathbf{V}_{\bm{\theta}}]$.

##### Injecting SE ⁢ ( n ) SE n \operatorname{SE(n)} roman_SE ( roman_n ) -invariance and permutation invariance.

Two physical dynamic systems are governed by the interactions between each pair of particles, resulting in a distribution that is $\operatorname{SE(n)}$ and permutation invariant. Our objective is to develop models that are SO(n)-equivariant, translation invariant, and permutation equivariant. We intend to apply a noise matching objective to achieve the desired invariant distribution.
However, to the best of our knowledge, no such architecture satisfying the above properties has been established within the context of diffusion generative models. Therefore, we opt to utilize a data augmentation method to ensure the model’s equivariance and invariance properties (Chen et al., ; Botev et al., ), i.e. we apply these group operations in the training process, which enforces models to be equivariant and invariant.

##### Conservation of momentum.

For both datasets, the conservation of momentum can be expressed as follows:

$$ $\sum_{k=1}^{K}m_{k}\mathbf{V}_{l,k,d}^{(0)}=\operatorname{constant}_{d},\quad \forall l=1,\ldots,L,d=1,\ldots,D.$ (9) $$

Here, $m_{k}$ represents the mass of the $k$-th particle, and $\mathbf{V}_{l,k,d}^{(0)}$ denotes the velocity along axis $d$ of the $k$-th particle at time $l$. The total momentum in each axis remains constant, as indicated by the equality. This constraint is linear w.r.t. $\mathbf{V}_{l,k,d}^{(0)}$, corresponding to the linear case. Hence, let $\mathbf{f}:\mathbb{R}^{L\times K\times D}\times\mathbb{R}^{K}\rightarrow
\mathbb{R}^{D}$ calculate the mean of the total momentum over time and we set the penalty loss as

$$ $\mathcal{J}_{\mathcal{R}}\left(\bm{\theta}\right)=\mathbb{E}_{t,\bm{x}_{0},\bm {\epsilon}}\left[w(t)\sum_{l=1}^{L}\sum_{d=1}^{D}\left\|\left(\sum_{k=1}^{K}m_ {k}\left(\mathbf{V}_{\bm{\theta}}\right)_{l,k,d}\right)-\mathbf{f}_{d}(\mathbf {V}_{\bm{\theta}},\{m_{k}\}_{k=1}^{K})\right\|^{2}\right].$ (10) $$

##### Conservation of energy for the three-body dataset.

The total of gravitational potential energy and kinetic energy remains constant over time. The energy conservation equation is given by:

$$ $-\sum_{i\neq j}^{K}\frac{Gm_{i}m_{j}}{\mathbf{R}_{l,ij}^{(0)}}+\sum_{k=1}^{K} \sum_{d=1}^{D}\frac{1}{2}m_{k}(\mathbf{V}_{l,k,d}^{(0)})^{2}=\operatorname{ constant},\quad\forall l=1,\ldots,L,$ (11) $$

where $G$ denotes the gravitational constant. $\mathbf{R}_{l,ij}^{(0)}=\|\mathbf{C}^{(0)}_{l,i}-\mathbf{C}^{(0)}_{l,j}\|$ denotes the Euclidean distance between the $i$-th and $j$-th particle at time $l$. This constraint is nonlinear with $\mathbf{X}^{(0)}$ but can be decomposed into elementary cases. Note that the constraint is multilinear w.r.t. $1/\mathbf{R}_{l,ij}^{(0)}$ and $(\mathbf{V}_{l,k,d}^{(0)})^{2}$. Hence, from the results of the general nonlinear cases, we can train another model sharing the same hidden size as the model for noise matching to predict these variables related to the conservation of energy. Furthermore, since these variables are convex w.r.t. $\mathbf{X}^{(0)}$, by the results of the convex case, we can directly apply the penalty loss as:

$$ $\displaystyle\mathcal{J}_{\mathcal{R}}\left(\bm{\theta}\right)=\mathbb{E}_{t, \bm{x}_{0},\bm{\epsilon}}\left[w_{1}(t)\sum_{l;i\neq j}\left\|\frac{1}{\left( \mathbf{R}_{\bm{\theta}}\right)_{l,ij}}-\frac{1}{\mathbf{R}_{l,ij}^{(0)}} \right\|^{2}+w_{2}(t)\sum_{l,k,d}\left\|\left(\mathbf{V}_{\bm{\theta}}\right)_ {l,k,d}^{2}-\left(\mathbf{V}_{l,k,d}^{(0)}\right)^{2}\right\|^{2}\right]$ , (12) $$

where $\left(\mathbf{R}_{\bm{\theta}}\right)_{l,ij}=\|\mathbf{C}_{\bm{\theta}}\left(
\mathbf{X}^{(t)},t\right)_{l,i}-\mathbf{C}_{\bm{\theta}}\left(\mathbf{X}^{(t)}
,t\right)_{l,j}\|$, i.e. model’s prediction of the Euclidean distance between two particles calculated from its prediction of coordinates. This penalty loss can also be derived from the reducible case. The detailed derivation can be found in Appendix D.2.

##### Conservation of energy for the five-spring dataset.

The combined elastic potential energy and the kinetic energy are conserved throughout time. The equation for the conservation of energy is represented by:

$$ $\sum_{(i,j)\in\mathcal{E}}\frac{1}{2}\kappa\left(\mathbf{R}_{l,ij}^{(0)}\right )^{2}+\sum_{k=1}^{K}\sum_{d=1}^{D}\frac{1}{2}m_{k}\left(\mathbf{V}_{l,k,d}^{(0 )}\right)^{2}=\operatorname{constant},\quad\forall l=1,\ldots,L,$ (13) $$

where $\kappa$ denotes the elastic constant, $\mathbf{R}_{l,ij}^{(0)}=\|\mathbf{C}^{(0)}_{l,i}-\mathbf{C}^{(0)}_{l,j}\|$ denotes the distance between the $i$-th and $j$-th particle at time $l$, and $\mathcal{E}$ denotes the edge set of springs connecting particles. $m_{k}$ represents the mass of the $k$-th particle. Analogue to the conservation of energy for the three-body dataset, we can reduce the nonlinear constraints into elementary cases.

Figure: 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, and 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 stronger adherence to physical feasibility.
Refer to caption: x1.png

**Table 3: Sample quality of 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 MSE error between the ground-truth motion and the generated ones. We also calculate the error of physical feasibility such as conservation of the energy and momentum, which should remain unchanged along the evolution of the systems.**
| Method | Three-body | Five-spring |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| Traj error | Vel error | Energy error | Dynamic error | Momentum error | Energy error |  |
| w/o prior | 2.5613 | 2.6555 | 3.8941 | 5.1929 | 5.3511 | 1.0891 |
| w/ prior | 1.6072 | 0.7307 | 0.5062 | 5.0919 | 0.3687 | 0.7448 |

##### Experimental results.

The results can be seen in Tab. 3 and Fig. 2, 13, 14, and we refer readers to Appendix E.2 for a detailed account of the experimental settings, as well as a more extensive comparison of the effects of hyperparameters and various methods for injecting constraints across the discussed cases. Our analysis indicates that for the three-body dataset, the incorporation of the conservation of energy prior, via the reducible case method, substantially enhances the model’s performance across all evaluated metrics. Similarly, applying the conservation of momentum prior to the five-spring dataset significantly reduces the momentum error in the generated samples. This also contributes to a reduction in the errors associated with dynamics and energies. Fig. 2 demonstrates that the total momentum and energy of samples generated with the incorporation of priors exhibit greater stability compared to those without priors. We also provide sampling results using the DPM-solvers (Lu et al., ) in Appendix E.5, which significantly lower computational expenses.

### 4.3 Ablation studies

Figure: Table 4: Results of an ablation study comparing the effects of data matching and noise matching techniques. The findings show that incorporating a distributional prior improves model performance.

##### Data matching vs noise matching.

We employ data matching and noise matching techniques for the PDE and particle dynamics datasets, respectively. An ablation study is conducted to investigate the effects of applying the alternative matching objective on the particle dynamics and PDE datasets, both without physical feasibility priors. The results, presented in Tab. 4, demonstrate that incorporating a distributional prior can significantly improve the model’s performance.

##### Omitting Jensen’s gap.

In the three-body dataset, we employ a multilinear function to simplify constraints into convex scenarios. We now conduct an ablation study in which the output of a diffusion model is considered the ground-truth, and the constraint of energy conservation is imposed similarly to the injection of constraints by penalty loss in the prediction tasks of PINNs. This configuration is referred to as “prior by PINN”. We define a penalty loss based on the variation of energy over time, analogous to the penalty loss used to enforce momentum conservation constraints. However, unlike the conservation of momentum, which is governed by a linear constraint and can thus be applied directly, the conservation of energy involves a nonlinear constraint. This introduces Jensen’s gap, preventing the direct application of the constraint. The results, presented in Tab. 5, indicate that directly applying nonlinear constraints can degrade the model’s performance. However, appropriately handling these constraints can significantly improve the sample quality.

## 5 Conclusion

In conclusion, this paper presents a groundbreaking and principled method for generating physically feasible dynamics using diffusion-based generative models by integrating two types of priors: distributional priors and physical feasibility priors. We inject distributional priors by choosing the proper equivariant models and applying the noising matching objective. We incorporate physical feasibility priors by decomposing nonlinear constraints into elementary cases. Empirical results demonstrate the robustness and high quality of the dynamics produced across a variety of physical phenomena, underscoring the significant promise of our method for data-driven AI4Physics research. This work emphasizes the importance of embedding domain-specific knowledge into learning systems, setting a precedent for future research bridging physics and machine learning through innovative use of physical priors.

## Appendix A Related work

### A.1 Generative methods for physics

Numerous studies have been conducted on the development of surrogate models to supplant numerical solvers for physics dynamics with GANs (Farimani et al., ; de Oliveira et al., ; Wu et al., ; Yang et al., ; Bode et al., ) and VAEs (Cang et al., ). Nevertheless, to generate realistic physics dynamics, one must accurately learn the data distribution or inject physics prior (Cuomo et al., ). Recent advancements in diffusion models (Song et al., ) have sparked increased interest in their direct application to the generation and prediction of physical dynamics, circumventing the need for specific physics-based formulations (Shu et al., ; Lienen et al., ; Yang & Sommer, ; Apte et al., ; Jadhav et al., ; Bastek et al., ). However, these approaches, which do not incorporate prior physical knowledge, may exhibit limited performance, potentially leading to suboptimal results.

### A.2 Score-based diffusion models

Score-based diffusion models are a class of generative models that create high-quality data samples by progressively refining noise into detailed data through a series of steps (Song et al., ). These models estimate the score function, the gradient of the log-probability density of the data, using a neural network (Song & Ermon, ). By applying this score function iteratively to noisy samples, the model reverses the diffusion process, effectively denoising the data incrementally, and generates samples following the same distribution as the training set. Although numerous studies on diffusion models have focused on generating SE(3)-invariant distributions (Xu et al., ; Yim et al., ; Zhou et al., ), there remains a lack of comprehensive research on the generation of general invariant distributions under group operations. Meanwhile, in contrast to GANs, the outputs produced by diffusion models represent the distributional properties of the data samples. This fundamental difference means that physical feasibility priors cannot be added directly to the model output due to the presence of a Jensen gap (Chung et al., ), i.e. $\mathcal{R}\left(\mathbb{E}[\bm{x}_{0}\mid\bm{x}_{t}]\right)\neq\mathbb{E}[
\mathcal{R}\left(\bm{x}_{0}\right)\mid\bm{x}_{t}]$. A potential solution to this problem involves iterating and drawing samples during the training process and subsequently incorporating the loss of physics feasibility on the generated samples (Bastek et al., ). However, this approach necessitates numerous iterations, often in the hundreds, rendering the training process inefficient.

### A.3 Physics-informed neural networks

Physics-Informed Neural Networks (PINNs) are a class of deep learning models that incorporate physical laws and constraints into their training process (Lawal et al., ). Unlike traditional training processes, which learn patterns solely from data, PINNs leverage priors including PDEs that describe physical phenomena to guide the learning process. By incorporating these physical feasibility equations as part of the penalty loss, alongside the data prediction loss, PINNs enhance their ability to model complex systems. This integration allows PINNs to be applied across various fields, including fluid dynamics (Cai et al., ), electromagnetism (Khan & Lowther, ), and climate modeling (Hwang et al., ). Their ability to integrate domain knowledge with data-driven learning makes them a powerful tool for tackling complex scientific and engineering challenges.

## Appendix B Extension on equivalence class manifold

### B.1 Formal definition

Let $X$ be a set and $\sim$ be an equivalence relation on $X$. The equivalence class manifold $\mathcal{M}$ is defined as the set of equivalence classes under the relation $\sim$. Formally,

$$ $\mathcal{M}=\{x\mid x\in X,y\in[x]\Rightarrow y=x\},$ (14) $$

where $\mathcal{M}$ is a Riemannian manifold and $[x]$ denotes the equivalence class of $x$, defined as:

$$ $[x]=\{y\in X\mid y\sim x\}.$ (15) $$

### B.2 Equivalence class manifold of SE(3)-invariant distribution

The following theorem provides a method to use a set of coordinates to represent all other coordinates having the same pairwise distances.

###### Theorem 4 (Equivalence class manifold of SE(3)-invariant distribution (Dokmanic et al., 2015 ; Hoffmann & Noé, 2019 ; Zhou et al., 2024 ) ) .

Given any pairwise distance matrix $D\in\mathbb{R}_{+}^{n\times n}$, there exists a corresponding Gram matrix $M\in\mathbb{R}^{n\times n}$ defined by

$$ $M_{ij}=\frac{1}{2}(D_{1j}+D_{i1}-D_{ij})$ (16) $$

and conversely

$$ $D_{ij}=M_{ii}+M_{jj}-2M_{ij}.$ (17) $$

By performing the singular value decomposition (SVD) on the Gram matrix $M\in\mathbb{R}^{n\times n}$ (associated with $D$), we obtain exactly three positive eigenvalues $\lambda_{1},\lambda_{2},\lambda_{3}$ and their respective eigenvectors $\mathbf{v}_{1},\mathbf{v}_{2},\mathbf{v}_{3}$, where $\lambda_{1}\geq\lambda_{2}\geq\lambda_{3}>0$. The set of coordinates

$$ $\mathcal{C}=[\mathbf{v}_{1},\mathbf{v}_{2},\mathbf{v}_{3}]\left[\begin{array}[ ]{ccc}\sqrt{\lambda_{1}}&0&0\\ 0&\sqrt{\lambda_{2}}&0\\ 0&0&\sqrt{\lambda_{3}}\end{array}\right]$ (18) $$

satisfies has the same pairwise distance matrix as $D$.

Define the above mapping from the pairwise distances to coordinates as $f$. Then, the equivalence class manifold of SE(3)-invariant distribution can be given by

$$ $\mathcal{M}=\{f(D)\mid D\text{ is a pairwise distance matrix}\}.$ (19) $$

$\mathcal{M}$ satisfies the property of being a Riemannian manifold (Zhou et al., ).

###### Theorem 4 (Equivalence class manifold of SE(3)-invariant distribution (Dokmanic et al., 2015 ; Hoffmann & Noé, 2019 ; Zhou et al., 2024 ) ) .

## Appendix C PDE datasets

A summary of the important properties of datasets can be found in Tab. 6.

**Table 6: Comparative summary of datasets. The table highlights key aspects such as the type of generation (conditional or unconditional), the matching objective, the distributional priors, and the nature of the constraint cases.**
| Datasets | Cond/Uncond generation | Matching objective | Distributional priors | Constraint cases |  |
| --- | --- | --- | --- | --- | --- |
| PDE | advection | both | data (Equation 6b) | PDE constraints | linear |
| Darcy flow | conditional | linear |  |  |  |
| Burger | unconditional | multilinear |  |  |  |
| shallow water | conditional | linear |  |  |  |
| particle dynamics | three-body | unconditional | noise (Equation 6a) | SE(3) + permutation invariant | all cases |
| five-spring | conditional | SE(2) + permutation invariant | all cases |  |  |

### C.1 Dataset settings

##### Advection.

The advection equation is a fundamental model in fluid dynamics, representing the transport of a scalar quantity by a velocity field. The dataset presented herein consists of numerical solutions to the linear advection equation, characterized by

$$ $\partial_{t}u(t,x)+\beta\partial_{x}u(t,x)=0,\quad x\in(0,1),t\in(0,2]$ (20) $$

where $u$ denotes the scalar field and $\beta=0.1$ is a constant advection speed. The visualization of training samples can be seen in Fig. 3. Based on the initial conditions provided for the advection equation, our model utilizes a generative framework to predict the subsequent dynamics, with the specific aim of forecasting the next 40 frames. We then compare these predictions with the ground-truth to assess performance. Additionally, we evaluate the model’s capability to generate samples unconditionally, without initial conditions, and measure performance using the physical feasibility implied by the PDE constraints.

Figure: Figure 3: Samples from the advection dataset with varying initial conditions. The horizontal axis represents the spatial coordinate $x$, while the vertical axis represents the parameter $t$.
Refer to caption: figs/advection/func_value_%C4%B1.pdf

##### Darcy flow.

Darcy’s law describes the flow of fluid through a porous medium, which is a fundamental principle in hydrogeology, petroleum engineering, and other fields involving subsurface flow. The mathematical formulation of the Darcy flow PDE is given by:

$$ $\partial_{t}u(x,t)-\nabla(a(x)\nabla u(x,t))=f(x),\quad x\in(0,1)^{2},$ (21) $$

where $u(x,t)$ represents the fluid pressure at location $x$ and time $t$, $a(x)$ is the permeability or hydraulic conductivity, and $f(x)$ denotes sources or sinks within the medium. Given the initial state at $t=0$, we use the generative scheme to forecast the state at $t=1$. Fig. 4 provides a visualization of training samples. The accuracy of these predictions is evaluated by comparing them with the ground-truth values.

Figure: Figure 4: The figure illustrates representative training samples from the Darcy flow dataset. The first row displays the values for the function $a(x)$, while the second row shows the values of $u(x,t)$ at time $t=1$.
Refer to caption: figs/darcy/condition_%C4%B1.pdf

##### Burger.

The Burgers’ equation is a fundamental PDE that appears in various fields such as fluid mechanics, nonlinear acoustics, and traffic flow. It is a simplified model that captures essential features of convection and diffusion processes. The equation is given by:

$$ $\partial_{t}u(x,t)+u(x,t)\partial_{x}u(x,t)=0,$ (22) $$

where $u(x,t)$ represents the velocity field, $x$ and $t$ denote spatial and temporal coordinates, respectively. We unconditionally generate samples following the distribution of the training set and evaluate feasibility within the realm of physics as dictated by the constraints of PDE.

Figure: Figure 5: The figure depicts representative training samples from the Burger dataset. The samples within this dataset exhibit minimal variability. The horizontal axis denotes the spatial coordinate $x$, whereas the vertical axis represents the parameter $t$.
Refer to caption: figs/burger/func_value_%C4%B1.pdf

##### Shallow water.

The linearized 2D shallow water equations describe the dynamics of fluid flows under the assumption that the horizontal scale is significantly larger than the vertical depth. These equations are instrumental in fields such as oceanography, meteorology, and hydrology for modeling wave and current phenomena in shallow water regions. Let $u$ and $v$ denote the components of the velocity field in the $x$- and $y$-directions, respectively. The variable $h$ represents the perturbation in the free surface height of the fluid from a mean reference level. The parameter $c$ denotes the phase speed of shallow water waves, which is a function of the gravitational acceleration and the mean water depth. The equations are expressed as follows:

$$ $\frac{\partial u}{\partial t}=-\frac{\partial h}{\partial x},\quad\frac{ \partial v}{\partial t}=-\frac{\partial h}{\partial y},\quad\frac{\partial h}{ \partial t}=-c^{2}\left(\frac{\partial u}{\partial x}+\frac{\partial v}{ \partial y}\right).$ (23) $$

We conditionally generate the dynamics of shallow water expressed by $h,u,v$ conditioned on the given $c$. We provide a visualization of one sample in Fig. 6.

Figure: Figure 6: In the figure, the sequence from left to right represents a sample of the dynamics of shallow water. Each sample consists of 50 frames, from which 10 frames have been uniformly selected for visualization.
Refer to caption: figs/sw/train/3/%C4%B1.png

### C.2 Converting to elementary cases by finite difference approximation

##### Advection and shallow water.

The original constraint of the advection equation is given by

$$ $\partial_{t}u(t,x)+\beta\partial_{x}u(t,x)=0.$ (24) $$

If we use the finite difference method to approximate the derivative, assume a grid with time steps $t_{n}$ and spatial points $x_{i}$. Let $u_{n}^{i}$ denote the approximation to $u(t_{n},x_{i})$. For the time derivative, use a forward difference approximation: $\partial_{t}u\approx\frac{u_{n+1}^{i}-u_{n}^{i}}{\Delta t}$. For the spatial derivative, use a central difference approximation: $\partial_{x}u\approx\frac{u^{i+1}_{n}-u^{i}_{n}}{\Delta x}$. Substituting these approximations into the PDE, we have

$$ $\frac{u_{n+1}^{i}-u_{n}^{i}}{\Delta t}+\beta\frac{u^{i+1}_{n}-u^{i}_{n}}{ \Delta x}=0.$ (25) $$

Rearrange to obtain an equation that involves only $u$ values and constants:

$$ $u_{n+1}^{i}-(1+\beta\frac{\Delta t}{\Delta x})u_{n}^{i}+\beta\frac{\Delta t}{ \Delta x}u^{i+1}_{n}=0.$ (26) $$

In this form, the constraint is a linear equation involving $u_{n+1}^{i},u_{n}^{i},u^{i+1}_{n}$. The linearization of the shallow water constraints can be performed in an analogous manner.

##### Darcy flow.

The given Darcy flow equation is:

$$ $\partial_{t}u(x,t)-\nabla(a(x)\nabla u(x,t))=f(x),$ (27) $$

Using the finite difference method, we discretize the domain into a grid with grid spacing $\Delta x=\Delta y$ and $\Delta t$. Let $u_{i,j}$ represent $u(x_{i},y_{j},t_{n})$ and $u^{n}$ represent $u(x_{i},y_{j},t_{n})$. The finite difference approximations for the gradients and divergence are:

$$ $\displaystyle\partial_{t}u$ $\displaystyle\approx\frac{u^{n+1}-u^{n-1}}{2\Delta t},$ (28a) $\displaystyle\partial_{x}u$ $\displaystyle\approx\frac{u_{i+1,j}-u_{i-1,j}}{2\Delta x},$ (28b) $\displaystyle\partial_{y}u$ $\displaystyle\approx\frac{u_{i,j+1}-u_{i,j-1}}{2\Delta y}.$ (28c) $$

The Hessian matrix of $\nabla^{2}u(x,t)$ is also linear w.r.t. $u(x,t)$ when approximated by the finite difference method. Hence, the left-hand side of the Darcy flow equation is the sum of terms linear w.r.t. $u(x,t)$ and thus the constraint is linear.

##### Burger.

The partial differential equation

$$ $\partial_{t}u(x,t)+u(x,t)\partial_{x}u(x,t)=0$ (29) $$

can be approximated using finite differences as follows: time derivative (forward difference): $\partial_{t}u(x,t)\approx\frac{u(x,t+\Delta t)-u(x,t)}{\Delta t}$, spatial derivative (central difference): $\partial_{x}u(x,t)\approx\frac{u(x+\Delta x,t)-u(x-\Delta x,t)}{2\Delta x}$. Substituting these into the PDE gives:

$$ $\frac{u(x,t+\Delta t)-u(x,t)}{\Delta t}+u(x,t)\cdot\frac{u(x+\Delta x,t)-u(x- \Delta x,t)}{2\Delta x}=0$ (30) $$

Hence, the constraint is multilinear w.r.t. $u(x,t)$ if we consider values of $u$ at other points as constants.

## Appendix D Particle dynamics datasets

### D.1 Dataset introduction

The dataset features are structured as $\mathbf{X}^{(0)}=[\mathbf{C}^{(0)}\;\mathbf{V}^{(0)}]$, where $\mathbf{C}^{(0)}$ and $\mathbf{V}^{(0)}$ are both elements of $\mathbb{R}^{L\times K\times D}$. In this context, $L$ refers to the temporal length of the physical dynamics, $K$ represents the number of particles, and $D$ denotes the spatial dimensionality. Specifically, $\mathbf{C}^{(0)}$ captures the coordinate features, and $\mathbf{V}^{(0)}$ captures the velocity features. For the three-body dataset, the parameters are set as $L=10$, $K=3$, and $D=3$, indicating a temporal length of 10, with 3 particles in a 3-dimensional space. Similarly, for the five-spring dataset, $L=50$, $K=5$, and $D=2$, corresponding to a temporal length of 50, 5 particles, and a 2-dimensional space. For both datasets, we generated 50k samples for training. We aim to generate samples following the same distribution as in the training dataset.

Fig. 7 provides visual representations of two samples from the particle dynamics dataset, showcasing the behavior of systems within the three-body and five-spring datasets.

Figure: Figure 7: The presented figures illustrate samples from the particle dynamics dataset. The first two rows depict a sample from the three-body dataset, while the subsequent two rows represent a sample from the five-spring dataset.
Refer to caption: figs/3body/2/%C4%B1.pdf

### D.2 Details of injecting the conservation of energy for the three-body dataset.

The total of gravitational potential energy (GPE) and kinetic energy (KE) remains constant over time. The formula of the energy conservation equation is given by:

$$ $-\sum_{i\neq j}^{K}\frac{Gm_{i}m_{j}}{\mathbf{R}_{l,ij}^{(0)}}+\sum_{k=1}^{K} \sum_{d=1}^{D}\frac{1}{2}m_{k}(\mathbf{V}_{l,k,d}^{(0)})^{2}=\operatorname{ constant},\quad\forall l=1,\ldots,L,$ (31) $$

where $G$ denotes the gravitational constant, and all three bodies have the same mass $m_{k}$. $\mathbf{R}_{l,ij}^{(0)}=\|\mathbf{C}^{(0)}_{l,i}-\mathbf{C}^{(0)}_{l,j}\|$ denotes the Euclidean distance between the $i$-th and $j$-th mass at time $l$. This constraint is nonlinear with $\mathbf{X}^{(0)}$ but can be decomposed into elementary cases. Note that the constraint is multilinear w.r.t. $1/\mathbf{R}_{l,ij}^{(0)}$ and $(\mathbf{V}_{l,k,d}^{(0)})^{2}$. Hence, from the results of the general nonlinear cases, we can apply the penalty loss $\mathcal{J}_{\mathcal{R}}\left(\bm{\theta}\right)=\mathcal{J}_{\textsubscript{
GPE}}\left(\bm{\theta}\right)+\mathcal{J}_{\textsubscript{KE}}\left(\bm{\theta
}\right)$, and

$$ $\displaystyle\mathcal{J}_{\textsubscript{GPE}}\left(\bm{\theta}\right)$ $\displaystyle=\mathbb{E}_{t,\bm{x}_{0},\bm{\epsilon}}\left[w_{1}(t)\sum_{l;i \neq j}\left\|\left(\tilde{\mathbf{R}}_{\bm{\theta}}\right)_{l,ij}-\frac{1}{ \mathbf{R}_{l,ij}^{(0)}}\right\|^{2}\right],$ (32a) $\displaystyle\mathcal{J}_{\textsubscript{KE}}\left(\bm{\theta}\right)$ $\displaystyle=\mathbb{E}_{t,\bm{x}_{0},\bm{\epsilon}}\left[w_{2}(t)\sum_{l,k,d }\left\|\left(\tilde{\mathbf{V}}_{\bm{\theta}}\right)_{l,k,d}-\left(\mathbf{V} _{l,k,d}^{(0)}\right)^{2}\right\|^{2}\right],$ (32b) $$

where $\tilde{\mathbf{R}}_{\bm{\theta}}$ and $\tilde{\mathbf{V}}_{\bm{\theta}}$ share the same hidden state as the model $\mathbf{X}_{\bm{\theta}}$ and $\mathbf{X}_{\bm{\theta}}$ predicts $\mathbb{E}[\mathbf{X}^{(0)}\mid\mathbf{X}^{(t)}]$. The setting details of these models are introduced in Appendix E.1. We refer such a setting as “*noise matching + conservation of energy (general nonlinear)*”. Meanwhile, note that $1/\mathbf{R}_{l,ij}^{(0)}$ and $(\mathbf{V}_{l,k,d}^{(0)})^{2}$ are convex w.r.t. model’s output. From the results in the convex case, we can directly apply the penalty loss to the output of $\mathbf{X}_{\bm{\theta}}$ and set the penalty loss to be

$$ $\displaystyle\mathcal{J}_{\textsubscript{GPE}}\left(\bm{\theta}\right)$ $\displaystyle=\mathbb{E}_{t,\bm{x}_{0},\bm{\epsilon}}\left[w_{1}(t)\sum_{l;i \neq j}\left\|\frac{1}{\left(\mathbf{R}_{\bm{\theta}}\right)_{l,ij}}-\frac{1}{ \mathbf{R}_{l,ij}^{(0)}}\right\|^{2}\right],$ (33a) $\displaystyle\mathcal{J}_{\textsubscript{KE}}\left(\bm{\theta}\right)$ $\displaystyle=\mathbb{E}_{t,\bm{x}_{0},\bm{\epsilon}}\left[w_{2}(t)\sum_{l,k,d }\left\|\left(\mathbf{V}_{\bm{\theta}}\right)_{l,k,d}^{2}-\left(\mathbf{V}_{l, k,d}^{(0)}\right)^{2}\right\|^{2}\right],$ (33b) $$

where $\left(\mathbf{R}_{\bm{\theta}}\right)_{l,ij}=\|\mathbf{C}_{\bm{\theta}}\left(
\mathbf{X}^{(t)},t\right)_{l,i}-\mathbf{C}_{\bm{\theta}}\left(\mathbf{X}^{(t)}
,t\right)_{l,j}\|$, i.e. model’s prediction of the Euclidean distance between two masses. We refer such a setting as “*noise matching + conservation of energy (reducible nonlinear)*”, since this penalty loss function can be derived using a multilinear function composed with convex functions. In comparison to the penalty loss described in Equation 32, the penalty loss presented in Equation 33 is applied directly to the output of $\mathbf{X}_{\bm{\theta}}$. This direct application imposes a stronger constraint, thereby more effectively ensuring that the model adheres to the specified physical constraints.

## Appendix E Experiments Details

##### Configuration.

We conduct experiments of advection, Darcy flow, three-body, and five-spring datasets on NVIDIA GeForce RTX 3090 GPUs and Intel(R) Xeon(R) Silver 4210R CPU @ 2.40GHz CPU. For the rest of the datasets, we conduct experiments on NVIDIA A100-SXM4-80GB GPUs and Intel(R) Xeon(R) Platinum 8358P CPU @ 2.60GHz CPU.

##### Training details.

We use the Adam optimizer for training, with a maximum of 1000 epochs. We set the learning rate to 1e-3 and betas to 0.95 and 0.999. The learning rate scheduler is ReduceLROnPlateau with factor=0.6 and patience=10. When the learning rate is less than 5e-7, we stop the training.

##### Diffusion details.

As for diffusion configuration, we set the steps of the forward diffusion process to 1000, the noise scheduler to $\sigma_{t}=\operatorname{sigmoid}(\operatorname{linspace}(-5,5,1000))$ and $\alpha_{t}=\sqrt{1-\sigma_{t}^{2}}$. The loss weight $w(t)$ is set to $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}$ (Song et al., ). We generate samples using the DPM-Solver-1 (Lu et al., ).

##### Experiment summary.

We summarize the choice for backbones and properties of datasets in Tab. 6. We conducted an equivalent search for the hyperparameters of both the baseline methods and the proposed methods. The specific search ranges for each dataset and the corresponding hyperparameters are summarized in Tab. 7.

**Table 7: Summary of the model hyperparameters.**
| Datasets | Backbone | Model hyperparameters | Batch size |  |
| --- | --- | --- | --- | --- |
| PDE | advection | GRU | hidden size: [128, 256, 512], layers: [3, 4, 5] | 128 |
| Darcy flow | Karras Unet | dim: [128] | 8 |  |
| Burger | Karras Unet | dim: [32] | 128 |  |
| shallow water | 3D Karras Unet | dim: [16] | 64 |  |
| particle dynamics | three-body | NN+GRU | RNN hidden size: [64, 128, 256, 512, 1024], layers: [3, 4, 5] | 64 |
| five-spring | EGNN+GRU | RNN hidden size: [256, 512, 1024] | 64 |  |

### E.1 Training general nonlinear cases

##### Three-body dataset.

To reduce the nonlinear conservation of the energy by general nonlinear cases, we apply the penalty loss $\mathcal{J}_{\mathcal{R}}\left(\bm{\theta}\right)=\mathcal{J}_{\textsubscript{
GPE}}\left(\bm{\theta}\right)+\mathcal{J}_{\textsubscript{KE}}\left(\bm{\theta
}\right)$, and

$$ $\displaystyle\mathcal{J}_{\textsubscript{GPE}}\left(\bm{\theta}\right)$ $\displaystyle=\mathbb{E}_{t,\bm{x}_{0},\bm{\epsilon}}\left[w_{1}(t)\sum_{l;i \neq j}\left\|\left(\tilde{\mathbf{R}}_{\bm{\theta}}\right)_{l,ij}-\frac{1}{ \mathbf{R}_{l,ij}^{(0)}}\right\|^{2}\right],$ (34a) $\displaystyle\mathcal{J}_{\textsubscript{KE}}\left(\bm{\theta}\right)$ $\displaystyle=\mathbb{E}_{t,\bm{x}_{0},\bm{\epsilon}}\left[w_{2}(t)\sum_{l,k,d }\left\|\left(\tilde{\mathbf{V}}_{\bm{\theta}}\right)_{l,k,d}-\left(\mathbf{V} _{l,k,d}^{(0)}\right)^{2}\right\|^{2}\right].$ (34b) $$

The models $\tilde{\mathbf{R}}_{\bm{\theta}}$ and $\tilde{\mathbf{V}}_{\bm{\theta}}$ share the same hidden state as the model $\mathbf{X}_{\bm{\theta}}$, where $\mathbf{X}_{\bm{\theta}}$ is tasked with predicting $\mathbb{E}[\mathbf{X}^{(0)}\mid\mathbf{X}^{(t)}]$. The GRU architecture serves as the backbone for $\mathbf{X}_{\bm{\theta}}$. Consequently, we have designed the outputs of the models $\tilde{\mathbf{R}}_{\bm{\theta}}$ and $\tilde{\mathbf{V}}_{\bm{\theta}}$ to be generated by an additional linear layer that takes the hidden state of the GRU within $\mathbf{X}_{\bm{\theta}}$ as input.

##### Five-spring dataset.

For the five-spring dataset, we apply the penalty loss $\mathcal{J}_{\mathcal{R}}\left(\bm{\theta}\right)=\mathcal{J}_{\textsubscript{
PE}}\left(\bm{\theta}\right)+\mathcal{J}_{\textsubscript{KE}}\left(\bm{\theta}\right)$, and

$$ $\displaystyle\mathcal{J}_{\textsubscript{PE}}\left(\bm{\theta}\right)$ $\displaystyle=\mathbb{E}_{t,\bm{x}_{0},\bm{\epsilon}}\left[w_{1}(t)\sum_{(i,j) \in\mathcal{E},l}\left\|\left(\tilde{\mathbf{R}}_{\bm{\theta}}\right)_{l,ij}- \left(\mathbf{R}_{l,ij}^{(0)}\right)^{2}\right\|^{2}\right],$ (35a) $\displaystyle\mathcal{J}_{\textsubscript{KE}}\left(\bm{\theta}\right)$ $\displaystyle=\mathbb{E}_{t,\bm{x}_{0},\bm{\epsilon}}\left[w_{2}(t)\sum_{l,k,d }\left\|\left(\tilde{\mathbf{V}}_{\bm{\theta}}\right)_{l,k,d}-\left(\mathbf{V} _{l,k,d}^{(0)}\right)^{2}\right\|^{2}\right].$ (35b) $$

The models $\tilde{\mathbf{R}}_{\bm{\theta}}$ and $\tilde{\mathbf{V}}_{\bm{\theta}}$ utilize the same hidden state as the model $\mathbf{X}_{\bm{\theta}}$, with $\mathbf{X}_{\bm{\theta}}$ responsible for predicting $\mathbb{E}[\mathbf{X}^{(0)}\mid\mathbf{X}^{(t)}]$. The underlying structure of $\mathbf{X}_{\bm{\theta}}$ is based on EGNN for extracting node and edge features and a GRU network for dealing with time series. As a result, the outputs of $\tilde{\mathbf{R}}_{\bm{\theta}}$ are produced by an additional linear layer that processes the edge features generated by EGNN within $\mathbf{X}_{\bm{\theta}}$, and the outputs of $\tilde{\mathbf{V}}_{\bm{\theta}}$ are produced by an additional linear layer that takes the hidden state of the GRU within $\mathbf{X}_{\bm{\theta}}$ as input.

### E.2 Details of experiment results

Tab. 8 and Tab. 9 present the outcomes of the grid search conducted on both the three-body and five-spring datasets. For the three-body datasets, the top three combinations of hyperparameters—hidden size and the number of layers—are highlighted for each training method. For the five-spring datasets, the top three hidden size hyperparameters identified for each training method are provided.

**Table 8: The outcomes of the grid search conducted on the three-body datasets are summarized. For each training method, we highlight the top three combinations of hyperparameters, focusing on hidden size and the number of layers. “linear” refers to the settings in 10, and “reducible nonlinear” and “general nonlinear” refer to the settings in Equation 33 Equation 34, respectively.**
| Method | Hyperparameter | Trajectory error | Velocity error | Energy error |
| --- | --- | --- | --- | --- |
| data matching | 256, 4 | 5.2455 | 4.2028 | 12.758 |
| 512, 5 | 5.7765 | 3.8985 | 13.636 |  |
| 256, 5 | 5.5098 | 4.4144 | 11.643 |  |
| noise matching | 256, 4 | 2.5613 | 2.6555 | 3.8941 |
| 245, 5 | 2.5695 | 2.6713 | 3.8944 |  |
| 512, 3 | 2.6368 | 2.7192 | 3.5427 |  |
| noise matching<br>+<br>conservation of momentum (linear) | 512, 5 | 2.1409 | 2.2529 | 4.1116 |
| 1024, 4 | 2.4179 | 2.5261 | 3.9003 |  |
| 512, 4 | 2.4188 | 2.5264 | 6.8971 |  |
| noise matching<br>+<br>conservation of energy (reducible nonlinear) | 128, 3 | 1.6072 | 0.7307 | 0.5062 |
| 128, 4 | 1.6659 | 0.7605 | 0.5198 |  |
| 128, 5 | 1.7821 | 0.8030 | 0.4532 |  |
| noise matching<br>+<br>conservation of energy (general nonlinear) | 512, 4 | 2.2745 | 2.4238 | 4.0223 |
| 512, 3 | 2.5335 | 2.6234 | 3.8091 |  |
| 1024, 3 | 2.5068 | 2.6737 | 5.2131 |  |

**Table 9: Grid search results for the five-spring datasets are provided. We present the top two hyperparameter hidden size identified for each training method. The results of our experiment indicate that the output of the GRU model is inadequate to accurately predict the distance between particles. This limitation renders the methods designed for reducible nonlinear cases inapplicable, and consequently, their results have been excluded from our study. In contrast, the convoluted edge features generated by the EGNN model are sufficiently informative for predicting particle distances. Moreover, the application of methods suitable for general nonlinear cases improves performance. “linear” refers to the settings in Equation 10 and “general nonlinear” refers to the settings in Equation 35.**
| Method | Hyperparameter | Dynamic error | Momentum error | Energy error |
| --- | --- | --- | --- | --- |
| data matching | 1024 | 5.3120 | 5.2320 | 1.1204 |
| 256 | 5.3872 | 5.1448 | 1.1030 |  |
| noise matching | 512 | 5.1929 | 5.3511 | 1.0891 |
| 256 | 5.1950 | 5.3468 | 1.0805 |  |
| noise matching<br>+<br>conservation of momentum (linear) | 256 | 5.0919 | 0.3687 | 0.7448 |
| 512 | 5.0990 | 0.4335 | 0.7652 |  |
| noise matching<br>+<br>conservation of energy (general nonlinear) | 256 | 5.1615 | 5.3032 | 1.0548 |
| 1024 | 5.1809 | 5.3902 | 1.0879 |  |

### E.3 Comparsion with prediction methods

We conduct a comparison between the performance of the generation and prediction methods. In Tab. 10, we present a comparative analysis of generative models and prediction models for predicting physical dynamics, specifically advection and Darcy flow. The results of the prediction methods are taken from Takamoto et al. (), which performs a comprehensive comparison of FNO (Li et al., ), Unet (Ronneberger et al., ), and PINN (Raissi et al., ) (using DeepXDE (Lu et al., )). Generative models that use diffusion techniques, both with and without prior information, exhibit comparable performance in both tasks. The diffusion model with priors shows an improvement over the one without priors. In this work, we do not conduct the procedure of super-resolution or denoising (Wang et al., ; Saharia et al., ), which are critical in practical applications to produce high-quality, clean, and detailed images from diffusion models. Hence, the performance of diffusion models can be further enhanced by the introduction of super-resolution and denoising.

**Table 10: Performance comparison of diffusion generative models with prediction models. The results of the prediction methods are brought from Takamoto et al. ().**
| Method | Backbone | Advection | Darcy flow |  |
| --- | --- | --- | --- | --- |
| Generation | diffusion w/o prior | Karras Unet<br>(Ho et al., ) | $1.716\times 10^{-2}$ | $2.261\times 10^{-2}$ |
| diffusion w/ prior | $1.621\times 10^{-2}$ | $2.174\times 10^{-2}$ |  |  |
| Prediction | forward propagator approximation | FNO | $\mathbf{4.9\times 10^{-3}}$ | $1.2\times 10^{-2}$ |
| autoregressive method | Unet | $3.8\times 10^{-2}$ | $\mathbf{6.4\times 10^{-3}}$ |  |
| PINN | DeepXDE | $7.8\times 10^{-1}$ | - |  |

### E.4 Why not transformer?

We attempted to implement a transformer architecture as the backbone for sequential data in particle dynamics datasets. However, our results indicate that the transformer-based model does not achieve performance comparable to that of recurrent structure backbones. This discrepancy is likely due to the nature of physical dynamics, where the next state is strongly dependent on the current state. The attention mechanism employed by transformers may reduce performance in this context, as it does not inherently account for the temporal evolution of states.

### E.5 Sampling in fewer steps using DPM-solvers

We conduct experiments of using the DPM-solvers (Lu et al., ) to sample in fewer steps. By reducing the number of diffusion steps required, DPM solvers significantly lower computational expenses in generating physics dynamics. This efficiency is achieved with minimal degradation in performance, ensuring that the resulting dynamics remain closely aligned with the underlying physical principles. We apply the DPM-Solver-3 (Algorithm 2 in Lu et al. ()) and the results can be seen in Fig. 9 and 9.

Figure: Figure 8: Results of sampling through DPM-Solver-3 on the three-body dataset. The x-axis denotes the values for $M$ in Algorithm 2 in Lu et al. ().
Refer to caption: x7.png

## Appendix F Proofs

### F.1 Sufficient conditions for the invariance of marginal distribution

###### Definition 5 (volume-preserving) .

A function whose derivative has a determinant equal to 1 is known as a volume-preserving function.

###### Definition 6 (isomorphism) .

An isomorphism is a structure-preserving mapping between two structures of the same type that can be reversed by an inverse mapping.

###### Definition 7 (diffeomorphism) .

A diffeomorphism is an isomorphism of differentiable manifolds. It is an invertible function that maps one differentiable manifold to another such that both the function and its inverse are continuously differentiable.

###### Definition 8 (isometry) .

Let $\mathbf{X}$ and $\mathbf{Y}$ be metric spaces with metrics (e.g., distances) $d_{\mathbf{X}}$ and $d_{\mathbf{Y}}$. A map $\mathbf{f}:\mathbf{X}\rightarrow\mathbf{Y}$ is called an isometry if for any $\mathbf{a},\mathbf{b}\in\mathbf{X}$, $d_{\mathbf{X}}(\mathbf{a},\mathbf{b})=d_{\mathbf{Y}}(\mathbf{f}(\mathbf{a}),
\mathbf{f}(\mathbf{b}))$.

###### Definition 9 (homothety) .

If for all $\mathbf{G}\in\mathcal{G}$ and for all scalar $0<a<1$, there exists $\mathbf{H}\in\mathcal{G}$ such that $\mathbf{H}(a\bm{x})=a\mathbf{G}(\bm{x})$. $\mathbf{H}$ would be a homothety (a transformation that scales distances by a constant factor but does not necessarily preserve angles). The group $\mathcal{G}$ formed by all such transformations is called the homothety group.

See 1

###### Proof.

For VE diffusion (defined in Sec. 3.4 in Song et al. ()), $q_{t}\left(\bm{x}_{t}\mid\bm{x}_{0}\right)=\mathcal{N}\left(\bm{x}_{t}\mid\bm{
x}_{0},\sigma^{2}_{t}\mathbf{I}\right)$. For any $\mathcal{G}$-invariant distribution $q_{0}$ and $\mathbf{G}\in\mathcal{G}$, we have

$$ $\displaystyle q_{t}(\mathbf{G}(\bm{x}_{t}))$ $\displaystyle=\int q_{t}(\mathbf{G}(\bm{x}_{t})\mid\bm{x}_{0})q_{0}(\bm{x}_{0} )\mathrm{d}\bm{x}_{0}\quad$ probability chain rule (36a) $\displaystyle=\int q_{t}(\mathbf{G}(\bm{x}_{t})\mid\mathbf{G}(\bm{x}_{0}))q_{0 }(\mathbf{G}(\bm{x}_{0}))\mathrm{d}\mathbf{G}(\bm{x}_{0})\quad$ change of variables (36b) $\displaystyle=\int\mathcal{N}\left(\mathbf{G}(\bm{x}_{t})\mid\mathbf{G}(\bm{x} _{0}),\sigma^{2}_{t}\mathbf{I}\right)q_{0}(\bm{x}_{0})\mathrm{d}\bm{x}_{0}$ volume-preserving diffeomorphism (36c) $\displaystyle=\int\mathcal{N}\left(\bm{x}_{t}\mid\bm{x}_{0},\sigma^{2}_{t} \mathbf{I}\right)q_{0}(\bm{x}_{0})\mathrm{d}\bm{x}_{0}$ isotropic Gaussian $\mathcal{N}$ and isometry $\mathbf{G}$ (36d) $\displaystyle=\int q_{t}(\bm{x}_{t}\mid\bm{x}_{0})q_{0}(\bm{x}_{t})\mathrm{d} \bm{x}_{0}$ (36e) $\displaystyle=q_{t}(\bm{x}_{t})$ (36f) $$

Hence, the marginal distribution $q_{t}$ at any time $t$ is an $\mathcal{G}$-invariant distribution.

For VP diffusion (defined in Sec. 3.4 in Song et al. ()), assume $\alpha_{t}>0$ at any time $t$. $q_{t}\left(\bm{x}_{t}\mid\bm{x}_{0}\right)=\mathcal{N}\left(\bm{x}_{t}\mid
\alpha_{t}\bm{x}_{0},(1-\alpha_{t})\mathbf{I}\right)$. Define $\hat{q}_{t}(\bm{x}_{t})=q_{t}(\frac{1}{\alpha_{t}}\bm{x}_{t})$. Note that $\frac{1}{\alpha_{t}}\bm{x}_{t}=\bm{x}_{0}+\frac{\sqrt{1-\alpha_{t}}}{\alpha_{t
}}\bm{\epsilon},\bm{\epsilon}\sim\mathcal{N}(\mathbf{0},\mathbf{I})$, is a random variable generated by some VE diffusion process. Hence, its marginal distribution at any time $t$ is $\mathcal{G}$-invariant. For any $\mathcal{G}$-invariant distribution $q_{0}$, we have

$$ $\displaystyle q_{t}(\mathbf{G}(\bm{x}_{t}))$ $\displaystyle=\hat{q}_{t}(\alpha_{t}\mathbf{G}(\bm{x}_{t}))\quad\quad\quad\quad\quad$ by definition (37a) $\displaystyle=\hat{q}_{t}(\mathbf{H}(\alpha_{t}\bm{x}_{t}))$ by sufficient conditions (37b) $\displaystyle=\hat{q}_{t}(\alpha_{t}\bm{x}_{t})$ $\mathcal{G}$ -invariance (37c) $\displaystyle=q_{t}(\bm{x}_{t})$ (37d) $$

∎

###### Definition 5 (volume-preserving) .

###### Definition 6 (isomorphism) .

###### Definition 7 (diffeomorphism) .

###### Definition 8 (isometry) .

###### Definition 9 (homothety) .

###### Proof.

##### Discussion of related theorems.

###### Theorem 10 (Proposition 1 in Xu et al. ( 2022 ) ) .

Let $q\left(\bm{x}_{T}\right)$ be an SE(3)-invariant density function, i.e., $q\left(\bm{x}_{T}\right)=q\left(\mathbf{G}\left(\bm{x}_{T}\right)\right)$ for $\mathbf{G}\in\operatorname{SE}(3)$. If Markov transitions $q\left(\bm{x}_{t-1}\mid\bm{x}_{t}\right)$ are SE(3)-equivariant, i.e., $q\left(\bm{x}_{t-1}\mid\bm{x}_{t}\right)=q\left(\mathbf{G}\left(\bm{x}_{t-1}
\right)\mid\mathbf{G}\left(\bm{x}_{t}\right)\right)$, then we have that the density $q_{\theta}\left(\bm{x}_{0}\right)=\int q\left(\bm{x}_{T}\right)q_{\theta}\left
(\bm{x}_{0:T-1}\mid\bm{x}_{T}\right)\mathrm{d}\bm{x}_{1:T}$ is also $\operatorname{SE}(3)$-invariant.

Xu et al. () explore the integration of invariance during the sampling process while disregarding it during the forward process. Xu et al. () propose that sampling through equivalent translational kernel results in invariant distributions. In contrast, our Theorem 1 demonstrates that even when the transition probabilities of the Markov chain are not $\operatorname{SE(n)}$-equivariant, the resulting composed distribution can still be $\operatorname{SE(n)}$-invariant. This result offers a stronger conclusion than that presented by Xu et al. ().

###### Theorem 11 (Proposition 3.6 in Yim et al. ( 2023 ) ) .

Let $\mathcal{G}$ be a Lie group and $\mathcal{H}$ a subgroup of $\mathcal{G}$. Let $\mathbf{X}^{(0)}\sim q_{0}$ for an $\mathcal{H}$ invariant distribution $q_{0}$. If $\mathrm{d}\mathbf{X}^{(t)}=b\left(t,\mathbf{X}^{(t)}\right)\mathrm{d}t+$ $\Sigma^{1/2}\mathrm{~{}d}\mathbf{B}^{(t)}$ for bounded, $\mathcal{H}$-equivariant coefficients $b$ and $\Sigma$ satisfying $b\circ L_{h}=\mathrm{d}L_{h}(b)$ and $\Sigma\mathrm{d}L_{h}(\cdot)=\mathrm{d}L_{h}(\Sigma\cdot)$, and where $\mathbf{B}^{(t)}$ is a Brownian motion associated with a left-invariant metric. Then the distribution $q_{t}$ of $\mathbf{X}^{(t)}$ is $\mathcal{H}$-invariant.

In contrast to our Theorem 1, Theorem 11 in Yim et al. () imposes specific conditions on the relationship between the forward diffusion scheduler and the group operators, whereas our theorem does not. However, Yim et al. () impose fewer constraints on the properties of the group operations.

###### Theorem 10 (Proposition 1 in Xu et al. ( 2022 ) ) .

###### Theorem 11 (Proposition 3.6 in Yim et al. ( 2023 ) ) .

### F.2 Invariant distribution examples

#### F.2.1 SE ⁢ ( n ) SE n \operatorname{SE(n)} roman_SE ( roman_n ) -invariant distribution

If $q_{0}$ is an $\operatorname{SE(n)}$-invariant distribution, then $q_{t}$ is also $\operatorname{SE(n)}$-invariant.

###### Proof.

Given any $\mathbf{G}\in\operatorname{SE(n)}$, let $\mathbf{G}(\bm{x})=\mathbf{R}\bm{x}+\mathbf{b}$, where $\mathbf{R}\in\operatorname{SO}(n),\mathbf{b}\in\mathbb{R}^{n}$.

- •
volume-preserving: $\det\left(\frac{\mathrm{d}\mathbf{G}(\bm{x})}{\mathrm{d}\bm{x}}\right)=\det
\left(\mathbf{R}^{\top}\right)=1$.
- •
diffeomorphism: *smoothness*: The transformation $\mathbf{G}(\bm{x})=\mathbf{R}\bm{x}+\mathbf{b}$ is smooth because it involves linear operations (rotation and translation) that are smooth. Specifically, the rotation $\mathbf{R}$ and the translation $\mathbf{b}$ are smooth functions of their parameters; *bijectivity*: The function $\mathbf{G}(\bm{x})$ is bijective. For any $\bm{x}\in\mathbb{R}^{n}$, the function $\mathbf{G}(\bm{x})$ is one-to-one and onto. The inverse function is given by: $\mathbf{G}^{-1}(\bm{y})=\mathbf{R}^{-1}\left(\bm{y}-\mathbf{b}\right)$, where $\bm{y}\in\mathbb{R}^{n}$. Since $\mathbf{R}$ is a rotation matrix, it is invertible, and its inverse $\mathbf{R}^{-1}$ is also smooth. Therefore, the inverse function is smooth.
- •
isometry: for all $\bm{x},\bm{y}\in\mathbb{R}^{n},\|\mathbf{G}(\bm{x})-\mathbf{G}(\bm{y})\|^{2}=
\|\mathbf{R}\bm{x}+\mathbf{b}-\left(\mathbf{R}\bm{y}+\mathbf{b}\right)\|^{2}=
\left(\bm{x}-\bm{y}\right)^{\top}\mathbf{R}^{\top}\mathbf{R}\left(\bm{x}-\bm{y
}\right)=\|\bm{x}-\bm{y}\|^{2}$. Hence, $\|\mathbf{G}(\bm{x})-\mathbf{G}(\bm{y})\|=\|\bm{x}-\bm{y}\|$.
- •
homothety: Given any $\mathbf{G}\in\operatorname{SE(n)}$ and $0<a<1$, let $\mathbf{H}(\bm{x})=\mathbf{R}\bm{x}+a\mathbf{b}\in\operatorname{SE(n)}$. Then, $\mathbf{H}(a\bm{x})=\mathbf{R}\left(a\bm{x}\right)+a\mathbf{b}=a\left(\mathbf{
R}\bm{x}+\mathbf{b}\right)=a\mathbf{G}(\bm{x})$.

Hence, sufficient conditions are satisfied and $q_{t}$ is also $\operatorname{SE(n)}$-invariant.
∎

Let $q$ be an $\operatorname{SE(n)}$-invariant distribution. Given a set of $m$ points $\mathcal{C}\in\mathbb{R}^{n\times m}$, we write it in the vector form $\bm{x}\vcentcolon=\operatorname{vec}(\mathcal{C})\in\mathbb{R}^{mn}$. For any $\mathbf{G}(\mathcal{C})=\mathbf{R}\mathcal{C}+\mathbf{b},\mathbf{R}\in
\operatorname{SO}(n)$, let $\hat{\mathbf{R}}\in\mathbb{R}^{mn\times mn}$ be a block matrix with $\mathbf{R}$ on its diagonal block and $\hat{\mathbf{b}}=[\mathbf{b}^{\top}\cdots,\mathbf{b}^{\top}]^{\top}\in\mathbb{
R}^{mn}$. We have $\hat{\mathbf{G}}(\bm{x})=\hat{\mathbf{R}}\bm{x}+\hat{\mathbf{b}}$. Then, $(\nabla_{\bm{x}}\hat{\mathbf{G}}(\bm{x}))^{-1}=(\hat{\mathbf{R}}^{\top})^{-1}=
\hat{\mathbf{R}}$. Hence, $\nabla_{\hat{\mathbf{G}}(\bm{x})}\log q(\hat{\mathbf{G}}(\bm{x}))=\hat{\mathbf
{R}}\nabla_{\bm{x}}\log q(\bm{x})$, which imples $\nabla_{\mathbf{G}(\mathcal{C})}\log q(\mathbf{G}(\mathcal{C}))=\mathbf{R}
\nabla_{\mathcal{C}}\log q(\mathcal{C})$. Thus, the score function of an $\operatorname{SE(n)}$-invariant distribution is $\operatorname{SO(n)}$-equivariant and translational-invariant.

###### Proof.

#### F.2.2 Permutation-invariant distribution

We first list some useful properties of the Kronecker product:

- •
$\operatorname{vec}(\mathbf{A}\mathbf{X}\mathbf{B})=\left(\mathbf{B}^{T}\otimes
\mathbf{A}\right)\operatorname{vec}(\mathbf{X})$.
- •
$\det\left(\mathbf{A}\otimes\mathbf{B}\right)=\det\left(\mathbf{A}\right)^{m}
\det\left(\mathbf{B}\right)^{n}$, where $\mathbf{A}\in\mathbb{R}^{m\times m},\mathbf{B}\in\mathbb{R}^{n\times n}$.
- •
$(\mathbf{A}\otimes\mathbf{B})^{T}=\mathbf{A}^{T}\otimes\mathbf{B}^{T}$.
- •
$(\mathbf{A}\otimes\mathbf{B})(\mathbf{C}\otimes\mathbf{D})=(\mathbf{A}\mathbf{
C})\otimes(\mathbf{B}\mathbf{D})$.
- •
For square nonsingular matrices $\mathbf{A}$ and $\mathbf{B}$: $(\mathbf{A}\otimes\mathbf{B})^{-1}=\mathbf{A}^{-1}\otimes\mathbf{B}^{-1}$.

If $q_{0}$ is a permutation-invariant, then $q_{t}$ is also permutation-invariant.

###### Proof.

Given any $\mathbf{G}\in\mathcal{G}$ with $\mathbf{G}(\mathbf{X})=\mathbf{P}\mathbf{X}\mathbf{P}^{\top}$, we consider its vector form of $\operatorname{vec}\left(\mathbf{G}(\mathbf{X})\right)=\operatorname{vec}\left(
\mathbf{P}\mathbf{X}\mathbf{P}^{\top}\right)=\left(\mathbf{P}\otimes\mathbf{P}
\right)\operatorname{vec}\left(\mathbf{X}\right)$.

- •
volume-preserving: $\det\left(\frac{\mathrm{d}\operatorname{vec}\left(\mathbf{G}(\mathbf{X})\right
)}{\mathrm{d}\operatorname{vec}\left(\mathbf{X}\right)}\right)=\det\left(\left
(\mathbf{P}\otimes\mathbf{P}\right)^{\top}\right)=\det\left(\mathbf{P}\otimes
\mathbf{P}\right)=\det\left(\mathbf{P}\right)^{n}\det\left(\mathbf{P}\right)^{
n}=\det\left(\mathbf{P}\right)^{2n}=\left(\pm 1\right)^{2n}=1$.
- •
diffeomorphism: *smoothness*: $\mathbf{G}$ is smooth because it involves matrix multiplication, which is a smooth operation in $\mathbf{X}$. Since $\mathbf{P}$ and $\mathbf{P}^{\top}$ are constant matrices (not functions of $\mathbf{X}$), $\mathbf{G}$ inherits the smoothness from the matrix operations; *bijectivity*: Suppose $\mathbf{Y}=\mathbf{G}(\mathbf{X})=\mathbf{P}\mathbf{X}\mathbf{P}^{\top}$. To recover $\mathbf{X}$ from $\mathbf{Y}$, we compute: $\mathbf{X}=\mathbf{P}^{\top}\mathbf{Y}\mathbf{P}$, which is a valid operation because $\mathbf{P}$ is invertible.
- •
isometry: note that $\left(\mathbf{P}\otimes\mathbf{P}\right)^{\top}\left(\mathbf{P}\otimes\mathbf{
P}\right)=\left(\mathbf{P}^{\top}\otimes\mathbf{P}^{\top}\right)\left(\mathbf{
P}\otimes\mathbf{P}\right)=\left(\mathbf{P}^{\top}\mathbf{P}\right)\otimes
\left(\mathbf{P}^{\top}\mathbf{P}\right)=\mathbf{I}\otimes\mathbf{I}=\mathbf{I}$. For all $\mathbf{X},\mathbf{Y}\in\mathbb{R}^{n\times n}$,
$\displaystyle\quad\|\operatorname{vec}\left(\mathbf{G}(\mathbf{X})\right)-
\operatorname{vec}\left(\mathbf{G}(\mathbf{Y})\right)\|^{2}$
(38a)
$\displaystyle=\|\left(\mathbf{P}\otimes\mathbf{P}\right)\operatorname{vec}
\left(\mathbf{X}\right)-\left(\mathbf{P}\otimes\mathbf{P}\right)\operatorname{
vec}\left(\mathbf{Y}\right)\|^{2}$
(38b)
$\displaystyle=\left(\operatorname{vec}\left(\mathbf{X}\right)-\operatorname{
vec}\left(\mathbf{Y}\right)\right)^{\top}\left(\mathbf{P}\otimes\mathbf{P}
\right)^{\top}\left(\mathbf{P}\otimes\mathbf{P}\right)\left(\operatorname{vec}
\left(\mathbf{X}\right)-\operatorname{vec}\left(\mathbf{Y}\right)\right)$
(38c)
$\displaystyle=\left(\operatorname{vec}\left(\mathbf{X}\right)-\operatorname{
vec}\left(\mathbf{Y}\right)\right)^{\top}\mathbf{I}\left(\operatorname{vec}
\left(\mathbf{X}\right)-\operatorname{vec}\left(\mathbf{Y}\right)\right)$
(38d)
$\displaystyle=\|\operatorname{vec}\left(\mathbf{X}\right)-\operatorname{vec}
\left(\mathbf{Y}\right)\|^{2}$
(38e)
Hence, $\|\mathbf{G}(\mathbf{X})-\mathbf{G}(\mathbf{Y})\|_{\textsubscript{F}}=\|
\mathbf{X}-\mathbf{Y}\|_{\textsubscript{F}}$.
- •
homothety: Given any $\mathbf{G}\in\mathcal{G}$ and $0<a<1$, let $\mathbf{H}=\mathbf{G}\in\mathcal{G}$. Then, $\mathbf{H}(a\mathbf{X})=\mathbf{P}\left(a\mathbf{X}\right)\mathbf{P}^{\top}=a
\mathbf{P}\mathbf{X}\mathbf{P}^{\top}=a\mathbf{G}(\mathbf{X})$.

Hence, sufficient conditions are satisfied and $q_{t}$ is also permutation-invariant.
∎

Let $q$ be a permutation-invariant distribution of feature $\mathbf{X}\in\mathbb{R}^{n\times n}$ such as affinity/connectivity matrices representing relationships or connections between pairs of entities (e.g., nodes in a graph) or Gram matrices in kernel methods representing similarities between a set of vectors. Let $q(\mathbf{\mathbf{P}}\mathbf{X}\mathbf{P}^{\top})=q(\mathbf{X})$ for any permutational matrix $\mathbf{P}$. Consider the vectorization of $\mathbf{X}$. Let $\hat{q}(\operatorname{vec}(\mathbf{X}))\vcentcolon=q(\mathbf{X})$. Note that $\operatorname{vec}(\mathbf{P}\mathbf{X}\mathbf{P}^{\top})=(\mathbf{P}\otimes
\mathbf{P})\operatorname{vec}(\mathbf{X})$. Hence, $\nabla_{\operatorname{vec}(\mathbf{P}\mathbf{X}\mathbf{P}^{\top})}\log\hat{q}(
\operatorname{vec}(\mathbf{P}\mathbf{X}\mathbf{P}^{\top}))=(\mathbf{P}\otimes
\mathbf{P})^{-T}\nabla_{\operatorname{vec}(\mathbf{X})}\log\hat{q}(
\operatorname{vec}(\mathbf{X}))=(\mathbf{P}\otimes\mathbf{P})\nabla_{
\operatorname{vec}(\mathbf{X})}\log\hat{q}(\operatorname{vec}(\mathbf{X}))$. This implies that $\nabla_{\mathbf{P}\mathbf{X}\mathbf{P}^{\top}}\log q(\mathbf{P}\mathbf{X}
\mathbf{P}^{\top})=\mathbf{P}\nabla_{\mathbf{X}}\log q(\mathbf{X})\mathbf{P}^{\top}$. Thus, the score function of a permutation-invariant distribution is permutation-equivariant.

###### Proof.

### F.3 ECM equivalence

See 2

###### Proof.

Suppose we have a $(\mathcal{G},\nabla^{-1})$-equivariant model $\mathbf{s}_{\bm{\theta}}$ and $\mathbf{s}_{\bm{\theta}}(\varphi(\bm{x}))=\nabla_{\varphi(\bm{x})}\log q_{
\textsubscript{ECM}}(\varphi(\bm{x}))$ almost surely. Then, we have

$$ $\displaystyle\mathbf{s}_{\bm{\theta}}(\bm{x})$ $\displaystyle=\frac{\partial\varphi(\bm{x})}{\partial\bm{x}}\mathbf{s}_{\bm{ \theta}}(\varphi(\bm{x}))$ by equivariance of the model (39a) $\displaystyle=\frac{\partial\varphi(\bm{x})}{\partial\bm{x}}\nabla_{\varphi( \bm{x})}\log q_{\textsubscript{ECM}}(\varphi(\bm{x}))$ (39b) $\displaystyle=\nabla_{\bm{x}}\log q_{t}(\bm{x})$ by Equation 5 (39c) $$

∎

###### Proof.

### F.4 Multilinear Jensen’s gap

The following lemma is directly from the results of the optimal values of noise matching, which will be used in proving Theorem. 3.

###### Lemma 12 .

The gradient of $\mathbb{E}_{t,\bm{x}_{0},\bm{\epsilon}}[w(t)\|\bm{\epsilon}_{\bm{\theta}}\left
(\bm{x}_{t},t\right)-\bm{\epsilon}\|^{2}]$ w.r.t. $\bm{\theta}$ at $\bm{\epsilon}_{\bm{\theta}}\left(\bm{x}_{t},t\right)=-\sigma_{t}\nabla_{\bm{x}
}\log q_{t}\left(\bm{x}_{t}\right)$ equals $\mathbf{0}$.

See 3

###### Proof.

Without loss of generality, suppose that the optimizer of $\mathbf{u}_{\bm{\theta}_{2}}$ is given by $\mathbf{u}_{\bm{\theta}_{1}^{*}}+\mathbf{u}_{\Delta\bm{\theta}}$. The loss optimizer of data matching is given by $\mathbf{u}_{\bm{\theta}_{1}^{*}}\left(\bm{x}_{t},t\right)=\frac{1}{\alpha_{t}}
\left(\bm{u}_{t}+\sigma^{2}_{t}\nabla_{\bm{u}}\log q_{t}\left(\bm{x}_{t}\right
)\right)$. Substituting into the PDE loss term, we have

$$ $\displaystyle\quad\mathbb{E}_{t,\bm{x}_{0},\bm{\epsilon}}[w(t)\|\mathbf{W}_{0} \mathbf{u}_{\bm{\theta}_{2}}\left(\bm{x}_{t},t\right)+\mathbf{b}_{0}\|^{2}]$ (40a) $\displaystyle=\mathbb{E}_{t,\bm{x}_{0},\bm{\epsilon}}[w(t)\|\frac{1}{\alpha_{t }}\mathbf{W}_{0}\left(\bm{u}_{t}+\sigma^{2}_{t}\nabla_{\bm{u}}\log q_{t}\left( \bm{x}_{t}\right)\right)+\mathbf{W}_{0}\mathbf{u}_{\Delta\bm{\theta}}+\mathbf{ b}_{0}\|^{2}]$ (40b) $\displaystyle=\mathbb{E}_{t,\bm{x}_{0},\bm{\epsilon}}[w(t)\|\frac{1}{\alpha_{t }}\mathbf{W}_{0}\left(\alpha_{t}\bm{u}_{0}+\sigma_{t}\bm{\epsilon}+\sigma^{2}_ {t}\nabla_{\bm{u}}\log q_{t}\left(\bm{x}_{t}\right)\right)+\mathbf{W}_{0} \mathbf{u}_{\Delta\bm{\theta}}+\mathbf{b}_{0}\|^{2}]$ (40c) $\displaystyle=\mathbb{E}_{t,\bm{x}_{0},\bm{\epsilon}}[w(t)\|\mathbf{W}_{0}\bm{ u}_{0}+\mathbf{b}_{0}+\frac{\sigma_{t}}{\alpha_{t}}\mathbf{W}_{0}\bm{\epsilon} +\frac{\sigma^{2}_{t}}{\alpha_{t}}\mathbf{W}_{0}\nabla_{\bm{u}}\log q_{t}\left (\bm{x}_{t}\right)+\mathbf{W}_{0}\mathbf{u}_{\Delta\bm{\theta}}\|^{2}]$ (40d) $\displaystyle=\mathbb{E}_{t,\bm{x}_{0},\bm{\epsilon}}[w(t)\|\frac{\sigma_{t}}{ \alpha_{t}}\mathbf{W}_{0}\bm{\epsilon}+\frac{\sigma^{2}_{t}}{\alpha_{t}} \mathbf{W}_{0}\nabla_{\bm{u}}\log q_{t}\left(\bm{x}_{t}\right)+\mathbf{W}_{0} \mathbf{u}_{\Delta\bm{\theta}}\|^{2}]$ (40e) $\displaystyle=\mathbb{E}_{t,\bm{x}_{0},\bm{\epsilon}}[w(t)\frac{\sigma_{t}^{2} }{\alpha_{t}^{2}}\|\mathbf{W}_{0}\left(\bm{\epsilon}+\sigma_{t}\nabla_{\bm{u}} \log q_{t}\left(\bm{x}_{t}\right)+\frac{\alpha_{t}}{\sigma_{t}}\mathbf{u}_{ \Delta\bm{\theta}}\right)\|^{2}]$ (40f) $$

Dropping the reweighting term $\sigma_{t}^{2}/\alpha_{t}^{2}$ does not change the optimal solution. When $\mathbf{u}_{\Delta\bm{\theta}}\equiv\mathbf{0}$, observing the above objective and the noise matching objective, the above objective is the reweighted objective of noise matching by replacing $\mathbf{I}$ with $\mathbf{W}_{0}^{\top}\mathbf{W}_{0}$.
∎

###### Lemma 12 .

###### Proof.

## Appendix G Visualization of generated samples

In this study, we refrain from applying super-resolution and denoising procedures, which are essential in practical applications for generating high-quality, clear, and detailed images from diffusion models. Consequently, the generated samples contain noise. For the three-body dataset, since we only generate 10 frames, we apply the cubic spline to visualize a smooth trajectory and this is not applied when evaluating the quality of samples.

Figure: Figure 10: Examples of Darcy flow samples generated by models trained with/without PDE constraints.
Refer to caption: x9.png

Figure: Figure 11: Examples of Burger samples generated by models trained with/without PDE constraints.
Refer to caption: x13.png

Figure: Figure 12: Examples of shallow water samples generated by models trained with PDE constraints. Each sample contains 50 frames and we uniformly visualize 10 of them.
Refer to caption: figs/sw/gen/%C8%B7/%C4%B1.png

Figure: Figure 13: The figure illustrates examples of three-body samples produced by models trained with and without conservation constraints. The left panel represents the baseline model, the middle panel corresponds to our model, and the right panel depicts the energy as a function of time.
Refer to caption: figs/3body/sample_baseline_%C4%B1.pdf

Figure: Figure 14: The figure illustrates examples of five-spring samples produced by models trained with and without conservation constraints. The left panel represents the baseline model, the middle panel corresponds to our model, and the right panel depicts the momentum as a function of time.
Refer to caption: figs/5spring/base_%C4%B1.pdf

## Peer Reviews

confidence: 2
contribution: 2
initial_rating: 6
presentation: 3
questions: - The three contributions listed in Section 1 appear similar and convey the same meaning. Accurately summarizing the manuscript's contributions will be more helpful for the readers.

- Why distribution priors (such as translational rotational invariance) are important for generating physical dynamics? As stated in [1], the invariant sampling and invariant loss functions by restricting architecture designs often sacrifice empirical performances. Are simple data augmentations that approximate probability equivariance also available in physical dynamics generation? Experiments with data augmentations that approximate probability equivariance might help analyze the effectiveness of equivariant models.

- How are the physical priors selected in this work? It appears that the roto-translational invariance and the priors based on energy and momentum conservation laws are chosen arbitrarily. There seems to be no systematic analysis that explores all possible categories of physical priors or how to incorporate each type of prior into diffusion models. I believe there are many more physical priors that generative models should be constrained by. Categorizing them and discussing each category systematically would provide more insight.

- In Line 498-508, the paragraph title is "Data matching vs noise matching". However, the ablation study is conducted for incorporating a distributional prior or not, which is expected to investigate the training objective instead.

---
[1] SwinGNN: Rethinking Permutation Invariance in Diffusion Models for Graph Generation. TMLR 2024.
soundness: 3
strengths: 1. The manuscript is well-written and easy to follow.
2. Thorough derivations are provided in the main content and appendix.
3. Sufficient experiments and ablation studies are conducted to validate the effectiveness of the proposed method.
summary: This work proposes a framework that incorporates physical priors, including distributional priors and physical feasibility priors, to generate physically realistic dynamics with diffusion-based models.
weaknesses: Please see the Question part.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: - The model includes various hyperparameters, especially in the weighting of physical feasibility losses. How sensitive is the model to these settings? Could the authors provide guidance on tuning these parameters, or propose default values based on their findings?

- What does "intrinsic structures of pairwise distance" refer to?

- Does the resolution matter for "reducible nonlinear cases"?

- As I noticed, all of the datasets in this work seem to be with multi-dimensional features. Yet in biology, we usually just have single-dimensional features possibly collected at a time-step. Will this affect the performance of the method?

- I think the five springs dataset is probably more like an ODE dataset instead of a PDE. The transition functions are second-order ODEs.

- Following last question, how about reporting the errors of NRI (Kipf et al., 2018) on the experiments in Section 4.2 as well? 

- I have a tiny suggestion on the setup of the paper. If it is hard to fit a whole section of related works in the main body, it would be a good practice to merge a part of the related works to the section of preliminaries. Readers need the clear message in the related works to recall the research gap.

- Please check the score function in line 146, it does not consistent with the one in line 115. Please be careful about the suffices.
soundness: 3
strengths: - The work is innovative in its focus on embedding physical feasibility directly into diffusion-based generative models, specifically through the use of both distributional and physical priors. By designing a process that enforces these constraints, the model is capable of generating realistic dynamics, distinguishing it from more traditional generative approaches that may ignore or only partially enforce physical laws.

- Moreover, the paper is well-written and is thorough with both the theoretical part and the empirical results for the proposed model. 

- The paper addresses a challenge in generating accurate physical dynamics with AI methods. It enables more reliable and scientifically sound simulations in areas like environmental, material science, and possibly the fluid simulations.
summary: The paper proposes a novel framework to generate physically feasible dynamics by incorporating physical priors into diffusion-based generative models. Unlike traditional generative approaches, which often disregard physical constraints, this model integrates two types of priors: distributional priors (such as invariance to roto-translation) and physical feasibility priors (e.g., conservation laws for energy and momentum, and partial differential equation constraints). By embedding these priors into the generative process, the method produces realistic dynamics across a variety of physical systems.
weaknesses: - The experiments focus on synthetic datasets or physics-inspired datasets, and while this is valuable, additional testing on real-world, noisy datasets would definitely help.

- If possible, comparing the model's performance with more baselines will make the empirical results more convincing. Yet it might be difficult to find the methods from the exactly same field, but some methods from, e.g. time-series forecasting, can be taken into consideration.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: 1. **Related to W1** - What is the element of novelty in your exposition of incorporating distributional priors? If there is not any, this is not a problem, but then the relevant related works should be cited and the paper should indicate this clearly.
2. **Related to W3** - Can you also provide error bars on all results?
3. **Related to W6** - Could you comment on the extra cost of introducing the physical constraints using finite difference methods?
4. **Alternative objective** - For the distributional prior, have you ever tried a different objective that is a mix between the noise and data objective (for example, v-prediction [Salimans et al. [4]])? 
5. For the PDE datasets, why don’t you provide the RMSE of the PDE constraints on all datasets?
6. From my understanding, you only impose either the momentum conservation constrain or the energy conservation constraint. What if in the particle dynamics datasets you imposed both a conservation of momentum and a conservation of energy penalty? Would this lead to instabilities?
7. When employing finite difference methods to approximate the differential equations, how do you make sure that the spatio-temporal discretisation of the datasets is fine enough to yield significant results? Is the discretisation the reason why the RMSE of the PDE constraints in Table 2 is not closer to 0?
8. Could you please include in the captions of Table 4 and 5 the quantity that you are reporting (I am assuming Traj error)?
9. Isn’t the result in Table 4 for Shallow Water with distributional prior the same as the result in Table 2 Shallow water? If so, shouldn’t it be 8.151 instead of 8.150? And for the three-body dataset, shouldn’t it be 2.5613 as in Table 3 instead of 2.6084, or is there a difference between those settings?

[4] Salimans, T., & Ho, J. (2022). Progressive distillation for fast sampling of diffusion models. arXiv preprint arXiv:2202.00512.
soundness: 2
strengths: 1. **Relevant topic.** There is increased interest in the dynamical system modelling community to incorporate relevant priors into models. Not only would this help with more efficient training, but perhaps more importantly, it should lead to better stability and capabilities to generalise out-of-distribution (OOD - when the test data comes from a different distribution than the training data).
2. **Wide range of datasets.** The experiments are performed on a wide range of datasets - four PDE ones, and two particle dynamics datasets. This helps at showcasing how the techniques would be employed in various scenarios, studying both linear and multilinear PDE constraints in the PDE datasets, as well as (reducible and general) nonlinear ones in the particle dynamics datasets.
summary: This work explores ways of incorporating physics-informed priors in diffusion models for dynamical systems. In particular, they investigate **distributional priors**, aiming to incorporate appropriate symmetries (such as roto-translational invariance, permutation invariance, etc.), and **physical feasibility priors** as a way to encourage adherence to physical laws (e.g. momentum and energy conservation, PDE constraints, etc.). The manner in which the latter are imposed depends on the nature of the constraints, and the authors analyse several such cases separately (e.g. linear, multilinear, nonlinear, etc.).

The work provides empirical evidence on i) four PDE datasets, where PDE constraints are used, and ii) two particle dynamics datasets, where the diffusion model leverages the relevant physical laws, and training involves data augmentation to incorporate appropriate symmetries.
weaknesses: 1. **Contribution in distributional prior incorporation.** I believe there is some related work that is not being mentioned in section 3.1, as there are now plenty of works studying the incorporation of relevant priors in diffusion models. For example, Mathieu et al. [1] incorporate a series of geometric priors in infinite-dimensional modelling, also focusing on conditions for the diffusion process to be G-invariant assuming that the score network is G-equivariant.
2. **Lack of empirical evidence from the first methods subsection.** I do not think that the theory behind the first Methods subsection (Incorporating distributional priors) is sufficiently evidenced empirically.
- In the PDE datasets, the only result used empirically is that the data matching objective should be used, but this is rather weakly supported theoretically with the smoothness argument (although evidenced empirically in Table 4).
- In the particle dynamics datasets, I understand that finding the right architecture that incorporates all the relevant inductive biases can be hard, but if in the end data augmentation was used, how is the theory supported empirically in this case? I think you should include a discussion on this, as well as more details on how data augmentation was performed. 
- Also related to the particle dynamics dataset, it seems that for the five-spring dataset you are using the EGNN architecture, (but you do not cite it - I am assuming you are referring to Satorras et al. [2]). I think this should be stressed in the main text with a brief discussion on the choice of architecture (referring to the inductive biases it incorporates). It's an important component given the chosen topic and I do not understand why it wasn't more clearly emphasised. And why are you using it in the five-spring dataset, but not in the three-body one?
- Finally, the section on ECM from Methods is not referenced later on in the paper. From my understanding, its purpose is to justify the data augmentation approach (in which case this connects to the previous point). I think a clearer connection between methods and how they are applied should be made.
3. **Lack of error bars.** The results do not contain any sort of error bars, which actually makes it very difficult to assess the effectiveness of the proposed methods. Especially in the PDE datasets results, the differences between the mean RMSEs are fairly small, and, depending on the errors, the distinction between the method w/o and w/ prior might actually turn out not to be significant. Could you please include them?
4. **Inadequate sample quality.** For the Darcy flow samples, the quality is very poor. I understand that no super-resolution or denoising procedures have been employed, but it seems like there is too much noise left in the samples. Do you take a final denoising step after the end of the diffusion process (by applying Tweedie’s formula)? Do you normalise the data? I can see that the data values for Darcy flow (with poor quality) are much lower (0.025-0.150) vs. Burgers (0.0-1.0), but the last sigma value is probably the same, so the proportional influence of the noise is higher in the Darcy flow dataset, potentially leading to poorer quality. 
5. **Effectiveness of general nonlinear constraints.** The results on general nonlinear constraints are not that promising. In the three-body dataset, the results for noise matching + conservation of energy (general nonlinear) are very much comparable (especially for certain hyperparameters) to plain noise matching (although this is hard to assess because of the lack of error bars) - see Table 8. In the five-spring dataset (Table 9), the improvement is once again marginal for general nonlinear. Therefore, although I liked the idea of utilising the same hidden states as the model, I would argue that it did not prove to be effective in these experiments. If the authors agree, I think this limitation should be highlighted properly in the paper.
6. **No discussion on extra computational cost.** The paper does not discuss anything about the extra computational costs of the proposed method.

**Minor**

7. **Lack of thorough details/references on architecture.** In Table 7, I believe the Karras Unet should be referenced (from the EDM [3] paper). The paper also mentions that in the five-spring dataset the EGNN architecture is used but no reference is given (Satorras et al. [2]) and maybe the acronym should be defined for completeness.
8. **Normalised metrics.** For easier comparison between datasets, it might be better to use normalised metrics rather than absolute ones.
9. **Small legend in graphs and typos.** The font of the legends in the plots is too small (see Figure 2, Figure 8, Figure 9, etc.). There are also several typos throughout the paper (L107 extra “of”, L114 should be $\alpha_T$ and $\sigma_T$?, L271 related to the absolute, etc.). 
10. **Inconsistent notation for score.** Sometimes you specify the score as $\nabla_{\mathbf{x}} \log q_t(\mathbf{x})$, sometimes as $\nabla_{\mathbf{x}} \log q_t(\mathbf{x}_t)$.

Overall, I think that the paper contains some nice ideas, but fails to clearly connect them to the empirical evidence, the experiments are poorly presented (with lack of error bars) and the effectiveness of some cases seems overstated.

[1] Mathieu, E., Dutordoir, V., Hutchinson, M.J., Bortoli, V.D., Teh, Y.W., & Turner, R.E. (2023). Geometric Neural Diffusion Processes. ArXiv, abs/2307.05431.

[2] Satorras, V.G., Hoogeboom, E., & Welling, M. (2021). E(n) Equivariant Graph Neural Networks. ArXiv, abs/2102.09844.

[3] Karras, T., Aittala, M., Aila, T., & Laine, S. (2022). Elucidating the Design Space of Diffusion-Based Generative Models. ArXiv, abs/2206.00364.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: * Why does performance improve only incrementally after adding priors?
* In Equation (4) and line 205, are there typos? \nabla x -> \nabla_x?
soundness: 3
strengths: * The motivation is clear - this paper is a clear contribution to the AI4Science (esp AI4Physics) community where data-driven models should be combined with scientific inductive biases. Although physics-informed learning has been mainstream in scientific machine learning, physics-informed learning in the context of generative modeling is relatively new. 
* This paper is both theoretically (supported by theorems) and practically sound (supported by experiments). The idea of decomposing a complex constraint into elementary cases is a good strategy.
summary: This paper proposes two ways to incorporate physical priors into generative models: (1) they inject distributional priors by choosing the proper equivariant models and (2) they incorporate physical feasibility priors by decomposing nonlinear constraints into elementary cases. The key is to identify elementary cases where Jensen's gap can be omitted and decompose complicated constraints into elementary cases. In general, this paper is a good contribution to the AI4Science community, where many scientific priors are available and should be incorporated into learning systems.
weaknesses: * Empirical improvement is marginal (Table 1, 2, 3).
* This framework is useful when (1) we don't understand the system fully but (2) we know some partial information (existence of energy etc). However, all the examples in the paper are synthetic, so we know both the underlying equation and energy but "pretend" that we don't know about the equation. I understand this is only a proof of concept, but a more realistic example would greatly strengthen the paper.

Meta-review: The submission proposes methods for incorporating physical priors into diffusion-based generative models when used for solving physical problems (e.g., generating physical fields), which is a topic not yet heavily explored. The presented formulations seem sufficiently general and sound. The reviewers also raised a few insufficiencies, e.g., the usage of equivariant architectures and data augmentation (which may mix with the effect of the proposed loss), insignificant improvements in some cases, lack of necessary citations and discussions with relevant works, and not sufficiently diverse physical laws and more realistic demonstrations. During rebuttal, the authors have provided abundant further results and explanations, which have fixed some of the insufficiencies, and made reasonable arguments on some other points. Overall, the reviewers tend to believe the submission makes a helpful exploration on methods to build physical priors into diffusion models.

Additional personal comments: Do the formulations in Sec. 3.1 only apply to measure-preserving group operations (it is indeed the case for the two examples in the paper, though)? And the authors may be interested in a discussion on fulfilling the physical constraints by building a diffusion model directly on ECM.

Decision: Accept (Poster)