M2PDE: Compositional Generative Multiphysics and Multi-component PDE Simulation
Authors: Tao Zhang
Zhenhai Liu
Feipeng Qi
Yongjun Jiao
Tailin Wu
Venue: ICML 2025 poster (2025)
Keywords: multiphysics
multi-component
PDE simulation
physical simulation
generative
Primary area: applications->chemistry_physics_and_earth_sciences
TL;DR: We introduce a compositional multiphysics and multi-component simulation method by structurally composing the learned energy function of each conditional simulations

Abstract: Multiphysics simulation, which models the interactions between multiple physical processes, and multi-component simulation of complex structures are critical in fields like nuclear and aerospace engineering. Previous studies use numerical solvers or ML-based surrogate models for these simulations. However, multiphysics simulations typically require integrating multiple specialized solvers-each for a specific physical process-into a coupled program, which introduces significant development challenges. Furthermore, existing numerical algorithms struggle with highly complex large-scale structures in multi-component simulations. Here we propose compositional Multiphysics and Multi-component PDE Simulation with Diffusion models (M2PDE) to overcome these challenges. During diffusion-based training, M2PDE learns energy functions modeling the conditional probability of one physical process/component conditioned on other processes/components. In inference, M2PDE generates coupled multiphysics and multi-component solutions by sampling from the joint probability distribution. We evaluate M2PDE on two multiphysics tasks-reaction-diffusion and nuclear thermal coupling--where it achieves more accurate predictions than surrogate models in challenging scenarios. We then apply it to a multi-component prismatic fuel element problem, demonstrating that M2PDE scales from single-component training to a 64-component structure and outperforms existing domain-decomposition and graph-based approaches. The code is available at github.com/AI4Science-WestlakeU/M2PDE.

M2PDE: Compositional Generative Multiphysics and Multi-component PDE Simulation
Tao Zhang' Zhenhai Liu
Feipeng Qi!
Yongjun Jiao
Tailin Wu?
Abstract
1. Introduction
Multiphysies simulation, which models the in- teractions between multiple physical processes, and multi-component simulation of complex structures are critical in fields like nuclear and aerospace engineering. Previous studies use nu- merical solvers or ML-based surrogate models for these simulations. However, multiphysics simu- lations typically require integrating multiple spe- cialized solvers- each for specific physical pro- cess-into coupled program, which introduces significant development challenges. Further- more, existing numerical algorithms struggle with highly complex large-scale structures in multi- component simulations. Here we propose compo- sitional Multiphysics and Multi-component PDE Simulation with Diffusion models (M2PDE) to overcome these challenges. During diffusion- based training, M2PDE learns energy funetions modeling the conditional probability of one phys- ical process/component conditioned on other pro- cesses/components. In inference, M2PDE gener- ates coupled multiphysics and multi-component solutions by sampling from the joint probabil- ity distribution. We evaluate M2PDE on two multiphysics tekks-reaction-diiftsiiis and nu- clear thermal coupling-where it achieves more accurate predictions than surrogate models in challenging scenarios. We then apply it to a multi-cone coltiicomppnent prismatic fuel element problem, demonstrating that M2PDE scales from single- component training a 64-component structure and outperforms existing domain decompoicoooooiiiioon and graph- based approaches. The code is github.com/AAl4Sieenn otthooccom//AASSSceeeeeeaalalll WestlakeU/M2PDD..
Multiphysics simulation involves the concurrent modeling of multiple physical processes -uchh as heat conduction, fluid flow, and structural mechanics- ulation framework to accurately capture the coupling ef- fects between different physical processes. Similarly, multi- component simulation focuses on simulating complex struc- tures composed of multiple similar components. Compo- nent is defined as: repeatable basic unit that makes up complete structure. For example, the reactor core typi- cally consists of hundreds thousands of fuel elements arranged in square or hexagonal pattern. These simula- tions are essential across various scientific and engineering disciplines, including nuclear engineering (Ma et al., 2022; Chen et al., 2021). aerospace engineering (Candeo et al., 2011; Wang al., 2023a), civil engineering (Sun al., 2017: Meyer et al., 2022), and automotive industry (Ragone et al., 2021). Despite their significance, both multiphysies and multi- multi-component partial differential equation (PDE) simulations share common challenge: while simulating individual components or physical processes is relatively straightforward, modeling the entire system with all its in- teractions is vastly more complex.
single sim-
Numerous numerical algorithms have been developed for multiphysics simulation, which are broadly categorized into loose coupling and tight coupling (Hales et al., 2015). Loose coupling solves each physical process independently. itera- tively transferring solutions until convergence, using meth- ods like operator splitting (MacNamara & Strang, 2016) and Picard iteration (Terlizzi Kotlyar, 2022). Tight cou- pling solves all processes simultaneously in large system (Knoll & Keyes, 2004), which can be more accurate but faces challenges such as high computational costs, varying spatial and temporal resolutions, and differing numerical methods across physical processes. Thus, loose coupling is more commonly used engineering applications. In multi- component simulation, directly simulating the overall struc- ture requires high computational cost and may encounter difficulties convergence due to the increase in degrees of freedom. Substructure methods have been used in fields like nuclear engineering (Chen al., 2021) and civil engineer- ing (Sun et al.. 2017) to reduce modeling and computational
'Equal contribution 'State Key Laboratory of Advanced Nu- clear Energy Technology. Nuclear Power Institute of China, China DDpartmeent of Artificial Intelligence, Westlake University, China. Correspondence Tailin Wu <wutailin westlake. westlake.ccu.cn>> Yongjun Jiao <jjaoyongiun@npit acc.
Proceedings of the nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

M2PDE: Compositional Generative Multiphysics and Multi-component PDE Simulation
costs for repetitive components.
across engineering. First, we demonstrate its capability for multiphysics simulation by applying it coupled reaction- diffusion equations and nuclear thermal coupling combined with conjugate heat transfer. Second, we verify its capa- bility in componeet simulation through thermal and mechanical analysis of prismatic fuel elements.
Despite advances in numerical algorithms, several signif- icant challenges remain in engineering applications. In multiphysics simulations, considerable time and effort are required to develop programs that couple different special- ized solvers. Furthermore, the complexity of the system increases due to coupling, requiring more computing re- sources. While some studies employ machine learning- based surrogate models to accelerate simulations (Sobes et al., 2021; Park et al., 2021), these models still depend on coupled data for training, which necessitates the prior devel- opment of coupled numerical solution programs. In the case of multi nulti-comppnent simulations, the substructure method has primarily been applied to mechanical problems, with no widely applicable general method for multi-component sys- tems. Consequently, current methods often rely selecting representative units for detailed analysis or introducing sim- plifications to the overall structure, which may compromise the accuracy and scope of the simulations.
Concretely, our contributions are threefold: (1) We intro- duce a novel approach, M2PDE, for multiphysies and multi- component PDE simulations, framing the problem in terms of joint probabilistic modeling. By training decoupled (small structure) training data, it can generate coupled (large structure) solutions. (2) We create and open-source bench- mark datasets for both multiphysics and multi nulti-com PDE simulations, providing a valuable resource for future research. (3) Our method demonstrates success in both domains. For multiphysics simulation, M2PDE accurately predicts coupled solutions in complex problems where surro- gate models fail. In multi-component simulations, M2PDE, trained on single components, accurately predicts larger structures with up to 64 components.
To address these challenges in engineering applications, we propose compositional Multiphysics and coltoeco PDE Simulation with Diffusion models (M2PDE). The core innovation of M2PDE is its treatment of multiphysies and multi-component PDE simulations as generative probabilis- tic modeling, where interactions between multiple physical processes or components are captured through composing learned energy functions conditioned on others in a struc- tured way. In multiphysics simulation, M2PDE generates coupled solutions (accounting for interactions between dif- ferent physical processes) from decoupled data (assuming other fields are known and focused on solving a single field) by modeling the solutions of physical processes as joint probability distribution. The solution for each individual process is treated as a conditional probability distribution, based on Bayes' theorem. By training diffusion models (Ho et al., 2020) on decoupled data, we capture these condi- tional distributions. During inference, the model combines these distributions and performs reverse diffusion to produce the coupled solution. For multi-component simulations, M2PDE models each component's 's solution as a conditional probability distribution using the local Markov property, conditioned on neighboring components. By training dif- fusion models on small structures, we create conditional models for individual components. During inference, re- verse diffusion is applied iteratively across all components, yielding the solution for the entire structure. We have math- ematically derived the principles why M2PDE can obtain coupled solutions and large structure solutions in Sections 3.1 and 3.2. A schematic of M2PDE is provided in Fig. 1.
2. Related Work
Multiphysies simulation. Most existing studies develop unified surrogate models for all physical processes by cou- pling solutions (Tang et al., 2024; Ren et al., 2020; Park et al., 2021; Wang et al., 2023b). For complex problems, programs for each physical process are typically indepen- dent. It is often feasible to establish a surrogate model for one specific physical process and then integrate it with other numerical programs (EI Haber et al., 2022; Han al., 2019). Alternatively, surrogate models can be constructed sepa- rately for each physical processes and iteratively converged through an iterative process (Sobes et al., 2021). Because the purpose of our algorithm is to infer coupled solutions through models trained with decoupled data, and establish- ing the surrogate model for all physical processes requires coupling solution training models, we adopt the method of establishing surrogate models for each physical process separately as the baseline to validate the proposed algorithm.
Multi-component simulation. To our knowledge, there do not exist utilized machine learning methods specifically de- signed for multi-component simulation. A relevant study the CoAB-MLSim algorithm (Ranade et al., 2021). This al- gorithm combines neural networks with numerical iteration. It first partitions the computational domain into multiple sub- domains, and then trains a neural network to learn the flux conversation between subdomains, similar to the domain decomposition algorithm (Chan & Mathew, 1994). During inference, the neural network with flux conservation is ap- plied sequentially in each subdomain, looping until conver- gence. We further extend this algorithm to multi-component simulation and use it as baseline. Besides, graph neural
We illustrate the promise of this approach through simplified problems in practical engineering, both of which are impor- tant in nuclear engineering and has widespread implications
2

---

Compositional models. Recent research has extensively ex- plored the compositional combination of generative models for various applications, including image video synthe- sis (Du al., 2020; 2023; Po & Wetzstein, 2024; Yang et al., 2023), multimodal perception (Li et al., 2022), trajec- tory planning (Du et al., 2019; Urain et al., 2023), inverse design (Wu et al., 2024b), and data assimilation (Rozet & Louppe, 2023). A particularly effective approach for combining predictive distributions from local experts the product of experts framework (Hinton, 2002; Cohen al., 2020). Their focus is on how single object is influenced by multiple factors, such as generating images that meet various requirements in image generation (Du et al., 2023) or enhancing the lift-to-drag ratio under the influence of two wings in inverse design (Wu et al., 2024b). However, our problem involves multiple objects, such as multiple physical processes and components, requiring the capture of interac- tions between these fields or components. Existing research not applicable to multiphysics and multi-component PDE simulation. To the best of our knowledge, we are the first to introduce compositional generative approach to multi- physics and multi multi-component PDE simulations, demonstrat- ing how this framework enables generalization to far more M2PDE: Compositional Generative Multiphysics and Multi-component PDE Simulation
Multiphysies Simulation
Pegiocal compling inles faue suupling
decouple
22 8
22
Valle, 2 assumm))
M2PDE
Learn ()() from decoupled data
Compose to predict coupled solution
Multi-component Simulation
Eddvidaal sumounding componeees
V'
Naretial salve
(Zep )
Viy
B(Zu Jaze A Vz,, E Faree
7,09
Zv
Ragars
+ Volle, Eo(E S, Rovell
2g1
Zav:
M2PDE
Learn Fa(2,, (aa,) from smaller structure V
Compose to predict larger structure
Figure I. M2PDE schematic. Our proposed algorithm can use models trained with decoupled data to predict coupled solutions (top) and use models trained with small structure to predict large structures (here 64 components/(bottom))
network (GNN) (Wu et al., 2020; Fan et al., 2024) can learn on small graphs and inference on larger graphs (Xu al., 2019); Graph Transformer (Kreuzer al., 2021) employs the Laplacian matrix of graph characterize its structure, and by leveraging the Transformer architecture, it achieves learning on graphs. We also compare M2PDE with GNN and Graph Transformer.
complex simulation tasks than those encountered during training.
3. Method
In this section, we introduce the principle of M2PDE solv- ing multiphysics and multi-component PDE simulation in section 3.1 and section 3.2, respectively.
3.1. Multiphysics simulation
Consider a complex multiphysics simulation that consists of multiple physical processes: 2 = (#1, 22, ), each z; represents the trajectory of one or more physical fields that belong to that specific i-th physical process. For example, the mechanics might contain both the stress and strain fields in three directions over given time interval [0, Each process 2; has its own governing equation which depends on other processes and solving equations for other pro- cesses also requires that process. Therefore, all equations must be solved simultaneously to achieve the most accurate representation of the physical system.
-..,,,
Simulating all the processes together can be challeng- ing, while it will be simple if simulate single process 24. By specifying the other processes zại (Zi, Zi Zi+1 EN) and the given outer inputs'
4; f(ztir C)
(I)
In this paper, "outer inputs" refers the inputs of the physical system.
3

---

M2PDE: Compositional Generative Multiphysics and Multi-component PDE Simulation
where f is numerical solver. Omitting the given condition C, then: 2; = ((zz1). Now we consider the results of mul- tiple physical processes as a joint probability distribution:
where the denoising network eg(+) corresponds to the gra- dient of the energy function V-Eo(+) (Du et al., 2023), e~ (0, I). During inference, we sample from the joint distribution p(21, 22....N)) via (Ho et al., 2020):
(21, (2, ..., 2N) ~ p(21, 22, ...: 2N)
(2)
1
1- 022
For each process, we consider as a conditional distribution: 2: ~ p( p(2il2 244), which relates to the joint distribution via:
2 Vas
Fi,s - v1 ag
c (2(2 =3(00))
(8)
+ U,W, w (),
p(z1, 22...., ...... NN) p p(==\2)))p((()))
(3)
Writing the probability distribution in the form of (learnable) energy functions Ez)) (Du et al., 2023; LeCun et al., 2006), the energy functions relates to the joint probability of z, the
zi =
(Zi,a - v1 - -ā,td(Zia ==/):) 8))
(9)
1....1 and i = 1. Here, 2f rep- for = 5,8 resents the estimated value for the ith field 2, 24z = -f.-.-=f---=--....-..) and 0₂ is the noise level, the su- perscript e represents estimation.
conditional probability of zi, and the marginal distribution of 243 respectively by:
p(2)= zo -E(=)
This iterative method is similar to the Expectation- Maximization (EM) algorithm (Moon, 1996), refining each (4) variable's 's estimation based current estimates of oth- ers. An external loop can be added to repeat the diffusion model's ' inference, using the previous step's physical fields to improve the initial estimate. The ablation study about hyperparameters K, À and the estimation method are dis- cussed in Appendix F. The algorithm is shown in Algorithm 1. Line 2 is the external loop, while lines 6 11 represent the denoising cycle of the diffusion model. In each diffusion step, the physical physical processes are updated sequen- tially, with 2मl using the estimated physical processes that have already been updated at this diffusion step.
p(2s| 2x1)= = (12zz))
1
E(xllzz))
1 _e- -((ges)
p(zz1)= =
Zypi
where Z, Zzs are normalization coefficients (constants). since
Note that for p(zs|zzi), since 244 is the condition, the nor- malization (2z42) depends on 2 Substituting Eq. into Eq. 3, then taking logarithms of both sides, we have:
E(z) + log = [E(z; 2411 + (241)
(5)
[E(z)) + 10g2
Taking the derivative w.r.t. 2 on both sides, we have:
3.2. Multi-component simulation
E2, 22, ................................ 2N) = Vz, E(ii244) which uses the fact that log log lzgi( and E(zzz) are all independent of 2,.
(6)
Consider a complex structure that is composed of many com- ponents: V = 01 U U2 U...
... U vu. and the solution in each component eú is Z It should be noted that v; represents an entity here, and if there are multiple physical processes on this entity, it is also a multiphysics problem. Each compo- nent shares similarities and is arranged in specific patterm, like an array, to compose this complex structure. Simulating the entire structure V can be challenging while simulating an individual component is easier. By specifying the boundary condition 20,,, the given outer inputs C, and the geometry v; of component U;, we can compute Zn,:
Eq. 6 is the foundation of our compositional multiphysies simulation method. We see that when sampling the joint distribution p(21, Z2, 22., 2N), ). we can simply use the learned conditional diffusion model to sample each 2, while us- ing the estimated =2, of other physical processes as condi- tions. This means that to learn the multiphysics simulation of multiple physical processes 21, 22,,..... ..., 2N. we no longer need to develop a coupled algorithm that simultaneously solves all physical processes. Instead, we can simply use decoupled solvers (each physical process is solved inde- pendently while treating the other physical processes as known) to generate data, learn the conditional distributions p(=/(=z)) x e (( and in the inference time, sample from the joint distribution via Eq. 6, achieving multiphysics simulation. During training, the energy E(zi (zzs) is implic- itly learned via the diffusion objective below, which learns the gradient of the energy:
Zez = f C,
(10)
where f is a numerical solver. The outer inputs C' and geometry i are given conditions, zōe, is boundary condi- tions. Omitting the given condition, then: ze; = f(zv.) Then we divide the whole geometry V to three parts: V = i; U dr; U Vio. where Vio represents other parts of V except e, U dv; The solution of the whole geometry V can be written as the following probability distribution:
LASEE = ||cc col colvi- 3,2 + BB,2244 1)
(Zuzy 20o7: Znc.)) ~ pl ZZzz; Z0 2neo)
(11)
4

---

M2PDE: Compositional Generative Multiphysics and Multi-component PDE Simulation
Algorithm 1 Algorithm for multiphysics simulation by M2PDE.
the estimated zšo, as conditions. The multi-component simulation can be achieved using algorithm similar to multiphysics simulation. Since each zu; is inferred with the same model, it can be processed together, improving inference efficiency by eliminating the need for additional loops for each physical process. We provide Alg. 2 in Appendix A for multi-component simulation. Additionally, we use the assumption of Markov random fields in the derivation. The rationality of this assumption and the application scenarios of the algorithm are discussed in Appendix J.
input Compositional set of diffusion model ( C, 1,2,..., N, outer inputs C, dif- fusion step S, number of external loops K, number of physical processes N.
1: 2f ~ I) // initialize estimated fields 2f // An external loop to improve the estimated fields of: 2: for k= = -I,..,, K do
3: // ''initiaiize each physical fields 2, for i in 1,... N
4:
2f f zf // previous estimated fields 2f
5: of (O, I/ current estimated fields of
Our proposed framework of multiphysics simulation in Sec- tion 3.1 and multi compoiet simulation in Section 3.2, con- stitute our full method of compositional Multiphysies and Multi-component PDE Simulation with Diffusion models (M2PDE). It circumvents the development of coupled pro- grams that requires huge development efforts, and achieves multiphysics and multi-component PDE simulation by com- posing the learned conditional energy functions according to the variable dependencies. Below, we test our method's 's capability challenging engineering problems.
6: 2i.5 I) // intial physical fields z;
// denoising cycle of diffusion model:
7:
for == S.... 1 do
8:
==1- if k= > 1 else 1 // define the weights of Ef and of
// llooss for each physical process:
9:
for i = 1, N do
10:
w (0,, // use weighted estimated fields as conditions for single step denoising:
11:
21. = (2,, - 1-22 et
(2,,
I AaG +
VI-a.
4. Experiments
(1 2) 0) + o.w
// update the estimation of current field:
In the experiments, we aim to answer the following ques- tions: (1) Can M2PDE predict coupled solutions (account- ing for interactions between different physical processes) from models trained in decoupled data (assuming other pro- cesses are known and focus on solving a single process)? (2) Can M2PDE predict large structure solutions from a model trained in small structure data? (3) Can M2PDE outperform surrogate model? in both tasks? To answer these questions, we conduct experiments to assess our algorithm's 's perfor- mance on two problems across three scenarios. In Section 4. we solve the reaction-diffusion equation. While it's not a classic multiphysics coupling issue since both quantities are part of concentration fields, we consider them as sepa- rate physical processes to validate the capability of M2PDE for multiphysics simulation. Section 4.2 explores complex nuclear engineering scenario, involving various types of cou- pling: region, interface, strong, weak, unidirectional, and bidirectional to further test the algorithm's 's capacity to han- dle multiphysics simulation. Section 4.3 simplifies actual nuclear engineering problems evaluate the algorithm's effectiveness with multi-component simulation. Each ex-
12:
of = _(2,, - vI- (2,, +
1) s)s)
13:
end for
14: end for
15: end for
output 21 22.0. 2N,0
Consider the complex structure as an undirected graph G = (V, E), and the random variable že, is the property of component uy. The graph G satisfies the local Markov prop- erty: A variable is conditionally independent of all other variables given its neighbors. Thus, 2e, satisfies
(zv, 1 2v)N((:)) 20v.
(12)
Here dv; is the set of neighbors of Us Nuu vy Ud and V V. [[c]] Uio- By using this property of Markov random field, p(z., 20 2oco) can be written as:
p(z,, 20 Zoia) = p(zv, I20v. (pzz.. zav. )p(z),,) (13)
"The surrogate model in this article refers to a neural network that directly maps system inputs to outputs in one step, unlike diffu- sion models that generate outputs dep-by-step through denoising. For multiphysies problems, surrogate models are trained separately for each physical field using decoupled data, and obtain coupled solutions through iteration.. For multi-component problems, sur- rogate models use domain decomposition concepts, predicting each component iteratively based on surrounding components to achieve convergence.
Writing the probability distribution in the form of energy, and through the same derivation as in Section 3.1, we obtain:
V sug (2 Zde, 2e...) = V... E(2,, 1200.)
(14)
Therefore, when sampling the joint distribution p(Zer+ 2 Z..)) we can simply use the learned conditional diffusion model to sample each 2n while using
5

---

M2PDE: Compositional Generative Multiphysics and Multi-component PDE Simulation
Table 1. Relative 1.2 norm of error reaction-diffusion equation for multiphysics simulation.
method
decoupled
coupled
decoupled
coupled
surrogate + FNO
0.0669
0.0600
0.0080
0.0320
M2PDE (ours) + FNO
0.0270
0.0290
0.0102
0.0264
surrogate + U-Net
0.0152
0.0184
0.0039
0.0174
M2PDE (ours) + U-Net
0.0119
0.0141
0.0046
0.0174
periment uses two network architectures, training both with their respective diffusion and surrogate models for compari- son, employing consistent hyperparameters and settings to ensure fairness. The computational domains of experiment I and experiment 2 are on regular meshes, using Fourier neural operator (FNO) (Li et al., 2021; Lim et al., 2023) and U-Net (Ronneberger et al., 2015) as network architec- tures. The computational domain of experiment is on a finite element mesh, using Geo-FNO (Li et al., 2023) and Transolver (Wu et al., 2024a) the network architecture. Additionally, in experiment 3, we also compare M2PDE with graph neural networks GIN (Xu et al., 2019) Graph transformer SAN (Kreuzer et al., 2021) and MeshGraphNee (Pfaff et al., 2020).
network architecture, and training process are provided in Appendix B.
Ou
Mu An + n- u² - or, E [0, [0,5]
ōt
Óv uo + (u v) E [0, 11. t e [0,5]
Ot
[u, v) = [uo, vol E
[0, 11, t
(15)
Table presents the relative L.2 norm (1.2 norm of prediction error divided by L2 norm of the ground-truth) in predictions made by surrogate model and M2PDE on validation set of decoupled data and a test set of coupled data. For FNO, the prediction error for u is comparable between M2PDE and surrogate models on decoupled data; however, M2PDE shows a significantly larger error in predicting v, which is four times that of the surrogate model. As a result, the error in predicting the coupled solution for e greater than that of the surrogate model, while the error for u is lower. For U-Net, the prediction errors for and are similar between M2PDE and surrogate models on decoupled data, but M2PDE achieves lower error for the coupled solution.
To highlight the difference between training and testing data, Appendix G calculates the Wasserstein distances between decoupled and coupled data in multiphysics problems, as well as between small and large structural data in multi- component problem. and visualizes them. In Appendix K, we conduct a comparison of our dataset with existing scientific datasets (Takamoto et al., 2022). As another con- tribution the community, we also open-source the data at here to facilitate future method development of multiphysics and multi component PDE simulations.
This straightforward experiment tests the correctness of M2PDE but shows no significant advantages over the surro- gate model. However, the surrogate model fails in solving more complex problems, which will be discussed in the next section.
4.1. Reaction-diffusion
Reaction-diffusion (RD) equations have found wide applica- tions in the analysis of pattern formation, including chemical reactions. This experiment uses the ID -aaguuog-Naagmoo reaction-di diffusion equation (Rao et al., 2023). it has two concentration fields: u, v, the governing equations are pre- sented in Eq. 15. The objective is to predict the system's 's evolution under different initial conditions. The training data consists of decoupled data, where other physical pro- cesses are assumed and treated as inputs to solve the equa- tions governing the current physical process Gaussian random field (Bardeen et al., 1986) is employed to gener- ate the other physical processes and initial conditions, and numerical algorithms are used to compute the solution of the current physical process. The validation data similarly consists of decoupled data not used during training. The test data consists of coupled solutions obtained using fully cou- pled algorithms. Further details on the datasets, equation,
4.2. Nuclear thermal coupling
This experiment evaluates M2PDE's performance in han- dling various physical processes and coupling modes, in- cluding regional and interface coupling, strong and weak coupling, and unidirectional and bidirectional coupling. The focus is on nuclear-therma coupling in transient conditions for plate fuel elements, using a simplified pin cell analysis with a transient disturbance modeled as change in neutron flux density at the boundary. This physical system requires solving neutron physics, heat conduction, and flow heat transfer equations, taking into account negative feedback between neutron physics and temperature, unidirectional coupling from fluid to neutron field, and strong interface coupling between solid and fluid. The goal is to predict the system's 's evolution under different neutron boundary
6

---

M2PDE: Compositional Generative Multiphysics and Multi-component PDE Simulation
v,t DAØ (0E, Ea(T)),, 10, Lyby [0, Ly], [0,
Ot (.), t(y.t)
(16)
o(L. L1..811) = 6(2, Ly,
pcp To v.t) -Vig VT. + A E [0, L,). E [0, Ly (0,5)
Øt
ÖT, (z. 0,t)
OT, Ly,
(17)
dy
dy
Ts (is: V, (T.(tss))11)
V 17 = 0, E |Los Lig + Lig, y E (0,
+ u Vū = -Vp μ u + f.x E (L,, L, + Lyl. E [0, Lyl. E [0.5]
Ot
OT,
+ ū. VT;) = k VeT,I [Tss Tus + Lsl. E Tyy, [0,5]
(18)
pcp
Øt
0T, (Ts, U.
ØT (Ts.
ky
-is
dr
dx
conditions. The geometric and coupling relationships are illustrated in Fig. 2. The governing equations for the pled evolution of neutron flux (ø), solid temperature (T,), and fluid temperature, pressure, and velocity (T' ū)) are presented in Eq. 16, Eq. 17, and Eq. 18, respectively.
(a) Solid temperature.
(b) Fluid temperature.
Sald
Figure 3. Comparison of prediction results between M2PDE and surrogate model. The surrogate model fails on the test set of the coupled scenario.
That
hast Bax
serperstare
Neuiren
and test set of coupled data. In single physical process prediction (decoupled data), surrogate models outperform M2PDE. However, in predicting the coupled solution, all surrogate models fail except for the neutron physics field, with the predicted solid and fluid temperature fields shown in Fig. 3 (for more visualizations, physics field remains relatively accurate because the feed- back from solid and fluid temperatures is weak and primarily driven by external input boundary conditions. In contrast, solid temperature and fluid fields are significantly influenced by other physical processes, leading to non-physical predic- tions due to the lack of iterative process data during train- ing. In comparison, M2PDE more accurately captures the morphology of coupled solutions and demonstrates higher accuracy. In addition, we further use DDIM (Song et al., 2021) to accelerate sampling and compare the operational efficiency of different methods. M2PDE achieves an accel- eration of up to 29-fold speedup, with detailed information in Appendices and
Figure 2. Geometric structure (left); Coupling relationship be- tween different physical processes (right).
Fig. 6). The neutron
Generating estimated physical fields in this two-dimensional time series problem with three physical processes is chal- lenging using Gaussian random fields. To address this, we employ a pre-iteration method for data generation. The val- idation dataset consists of decoupled data not used during training, while the test dataset comprises coupled data. Cou- pled data is computed using the operator splitting iterative algorithm (MacNamara & Strang, 2016), which exchanges information between physical processes at each time step. Additional details on the datasets, network architecture, and training process can be found in Appendix C.
Table 2 displays the relative prediction errors of surrogate models and M2PDE on a validation set of decoupled data

---

M2PDE: Compositional Generative Multiphysics and Multi-component PDE Simulation
Heat pipe boundary
boundary
Matrix
(a) The structure used generate training data.
(h) Larger structure for testing.
Figure 4. In (a), the left figure shows the entire reactor. The right figure illustrates portion of the reactor core. This structure composed 16 fucl elements used to generate training data. (b) is a large structure composed 64 elements used for testing.
Table Relative 1.2 norm of prediction error on nuclear thermal coupling for multiphysics simulation. The unit
10
neutron
solid
fluid
method
decoupled
coupled
decoupled
coupled
decoupled
coupled
surrogate + FNO
0.251
22.1
0.0445
31.8
0.106
10.2
M2PDE (ours) FNO
0.738
8.42
0.175
9.72
0.615
7.31
surrogate + U-Net
0.181
4.45
0.0800
18.2
0.0927
8.03
M2PDE (ours) + U-Net
0.741
1.38
0.140
2.75
0.228
3.86
4.3. Prismatic fuel element
ture of 64 fuel elements. The average relative error strain Exy Ey denoted as e. GIN, SAN, and MeshGraphNet learn small graphs with 16 components and on large graphs with 64 components. Due the uniformity of graph struc- tures in all training data and the fact that SAN learns global relationship. SAN fails to predict larger structures. GIN and MeshGraphNet are capable of learning local relation- ships, thereby successfully handling larger-scale structures. Specifically, MeshGraphNet achieves the lowest error when learning structures with 16 components. However, when ex- trapolating
This experiment tests the ability of M2PDE to solve multi- component simulation problems, focusing on the thermal and mechanical analysis of prismatic fuel elements for new consists
type of reactor (Ma al., 2022). The reactor of three components: fuel, matrix, and heat pipe. Since engineering focuses mainly on the matrix, we consider the fuel and fluid as boundary conditions for analysis. Different heat fluxes will be assigned to the fuel boundary to simulate various heat release behaviors of the fuel rods. The aim is train model that predicts its temperature T and strain ea, by based on the solutions of its three neighbors and its heat flux.
larger structures, M2PDE exhibits lower error than MeshGraphNet.
Subsequently, a comparative analysis between the surro- gate model and M2PDE has been conducted. The surrogate model performs better in predicting a single component, but for medium structure, M2PDE outperforms it. It's im- portant to note that the surrogate model's predictions occa- sionally diverge, necessitating adjustments to the relaxation factor to maintain stability. For the large structure, U-Net in the surrogate model demonstrates better stability, while the FNO model continues to diverge even after relaxation factor adjustments. M2PDE is very stable and accurate, and no divergence phenomenon has been observed. term of strain, the relative error of M2PDE has been reduced by 41% Compared with the surrogate model. In addition,
The training data originate from medium structure simu- lation that includes 16 fuel elements; hence, single sim- ulation data point can generate 16 training data. The well- trained model will be tested on two structures: one is the medium structure used for data generation, and the other is large structure containing 64 fuel elements, as shown in Fig. 4. Further details the datasets, network architecture, and training process are provided in Appendix D.
Table 3 presents the prediction relative errors of surrogate model and M2PDE across three tasks: a single fuel element, medium structure of 16 fuel elements, and large struc-
8

---

M2PDE: Compositional Generative Multiphysics and Multi-component PDE Simulation
Table Relative 1.2 of prediction error on prismatic fuel element experiment, for single-componett and 16-component (medium)
and 64-component (large) simulations. The unit is 1 x 10-2
single
16-component
64-component
method
T
T
T
GIN
1.96
3.18
4.63
7.02
SAN
0.114
16.5
1.00x 102
1.18
MeshGraphNet
0.0583
0.427
1.28
2.40
surrogate Geo-FNO
0.0883
0.195
0.337
2.59
divergent 0.950
divergent 3,52
M2PDE (ours) Geo-FNO
0.139
0.459
0.338
2.42
surrogate Transolver
0.0764
0.251
0.314
1.13
1.25
3.31
M2PDE (ours) + Transolver
0.144
0.339
0.288
1.01
0.793
2.07
we further use DDIM to accelerate sampling and compare the operational efficiency of different methods. M2PDE achieves an accelcration of up to 41-fold speedup, with detailed information in Appendices H and I.
multiphysics simulations where surrogate models fail, and exhibits greater accuracy in predicting larger structures in multi-component simulations. We believe this approach provides a new approach to address multiphysics and multi- component PDE simulations, important across science and engineering.
5. Limitation and future work
There are also several limitations of our proposed M2PDE that provide exciting opportunities for future work. Firstly, in multiphysics simulation, although the M2PDE trained on decoupled data can predict coupled solutions more accu- rately than baseline surrogate models, the prediction errors are still higher compared to single physical processes pre- dictions. In addition, there is certain gap in accuracy com- pared to models trained through coupled data, as shown in Appendix E. Future efforts can focus improving dataset generation, training methods, and incorporating physical information to boost accuracy. Secondly, we plan to explore additional accelerated sampling algorithms, aiming to sig- nificantly improve efficiency while maintaining prediction accuracy. Lastly, while the experiments in this paper have simplified certain aspects of woal-world engineering prob- lems to focus on key issues and improve research efficiency, future work will aim to apply this method to more com- plex practical engineering problems to further validate its effectiveness and practicality.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences our work, none which we feel must be specifically highlighted here.
References
Bardeen, M., Bond, J., Kaiser, N., and Szalay, A. The statistics of peaks of gaussian random fields. Astrophysi- cal Journal, Part I (ISSN 0004-637X), vol. 304, May 1, 1986, p. 15-61, SERC-supported research., 304:15-61, 1986.
Candeo, A., Ducassy, C., Bocher, P., and Dughiero, F. Multi- physics modeling of induction hardening of ring gears for the aerospace industry. IEEE Transactions on Magnetics, 47(5):918-95) (11:)911-21,, 2011.
Chan, T. F. and Mathew, T. P. Domain decomposition algo- rithms. Acta numerica, 3:61-143, 1994.
6. Conclusion
This work presents M2PDE novel method for multi- physics and multi-component PDE simulations, driven by the needs of real-world engineering applications. In mul- tiphysics scenarios, models trained on decoupled data can predict coupled solutions, while in multi-componnnt simula- tions, models trained on small structures can extrapolate to larger ones. We develop three datasets validate M2PDE and compare it to the surrogate model method. Results show that M2PDE effectively predicts coupled solutions in
Chen, W., Wang, D., and Zhang, Y. Seismic fragility analysis of nuclear power plants based on the sub- structure method. Nuclear Engineering and De- sign, 382:1 11389, 2021. ISSN 0029-5493. doi:
com/ URL scceenee/rrrcaee//////2222222222
Cohen, S., Mbuvha, R., Marwala, T., and Deisenroth. M. Healing products of gaussian process experts. Interna-

---

M2PDE: Compositional Generative Multiphysics and Multi-component PDE Simulation
tional Conference on Machine Learning, pp. 2068-2077. PMLR, 2020.
Hinton, G. E. Training products of experts by minimizing contrastive divergence. Neural computation, 14(8):1771- 1800, 2002.
Du, Y.. Lin, T., and Mordatch. I. Model based planning with energy based models. CORL, 2019.
Ho, J., Jain, A., and Abbeel, P. Denoising diffusion probabilisticc models. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H. (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 6840-6851. Curran Associates, Inc. 2020. URL https::/prooceednggg.oomm neurips cc/paper (iwaaperffieeeepppee/////////// 4c5icfee3554ar000871111100000000
Du, Y., Li, S., and Mordatch, I. Compositional visual gener- ation with energy based models. In Advances in Neural Information Processing Systems, 2020.
Du, Y., Durkan, C., Strudel, R., Tenenbaum, J. B., Diele- man, S., Fergus, R., Sohl-Dickstein, J., Doucet, A., and Grathwohl, W. S. Reduce, reuse, recycle: Composi- tional generation with energy-based diffusion models and meme. In International conference on machine learning, pp. 8489-8518 PMLR, 2023.
4 pdf.
Icenhour, C., Keniley, S., DeChant, C.,, Permann, C., Lind- say, A., Martineau, R., Curreli, D., and Shannon, S. Multi- physics object oriented simulation environment (moose), Technical report, Idaho National Lab, Idaho Falls, ID (United States), 2018
El Haber, G., Viquerat, J., Larcher, A., Ryckelynck, D., Alves, J., Patil, A., and Hachem, E. Deep learning model to assist multiphysics conjugate problems. Physics of Fluids, 34(1). 2022.
Knoll, D. and Keyes, D. Jacobian-free newton-krylov meth- ods: a survey of approaches and applications. Journal of Computational Physics, 193(2):377 397, 2004. ISSN 0021-9991. doi: https://bbi.or///...011.......pp
Fan, S., Hitt, A. L.,, Tang, M., Sadigh, B., and Zhou, F. Ac- celerate microstructure evolution simulation using graph neural networks with adaptive spatiotemporal resolution. Machine Learning: Science and Technology, 5(2):025027, 2024.
010. URL cateree/aariicel//a////111999919
Kreuzer, D., Beaini, D., Hamilton, W., Létoumeau, V., and Tossou, P. Rethinking graph transformers with spectral attention. Advances in Neural Information Processing Systems, 34:21618 34:21618-21629, 2021.
Feydy, J., Séjouméé T., Vialard, F.-X., Amari, S.-i., Trouve, A., and Peyré, G. Interpolating between optimal trans- port and mmd using sinkhorn divergences. In The 22nd International Conference Artificial Intelligence and Statistics, pp. 2681-26900 2019.
LeCun, Y., Chopra, S., Hadsell, R., Ranzato, M.. Huang, F. et al. A tutorial on energy-based learning. Predicting structured data, 1(0), 2006.
Gong, D., Huang, S., Wang, G., and Wang, K. Heat trans- fer calculation on plate-type fuel assembly of high flux research reactor. Science and Technology of Nuclear Installations, 2015(1):198654, 2015.
Li, S., Puig, X., Paxton, C., Du, Y., Wang, C., Fan, L., Chen, T., Huang, D.-A., Akyürek, E., Anandkumar, A., et al. Pre-trained language models for interactive decision- making. Advances in Neural Information Processing Systems, 19:31199-31222 2022.
Gupta. J. K. and Brandstetter, J. Towards multi- spatiotemporal-scale generalized PDE modeling. Trans- actions on Machine Learning Research, 2023. ISSN 2835-8856. URL nttps://openrevitt forum? wtw.iiddddSSSDbETTT
Li, Z., Kovachki, N. B., Azizzadenesheli, K., liu, B., Bhat- tacharya, K., Stuart, A., and Anandkumar, A. Fourier neu- ral operator for parametric partial differential equations. In International Conference on Learning Representations, 2021. URL https: /opnreviewww net forum? i id=c 9NQVtmnO.
Hales, J., Tonks, M., Gleicher, F., Spencer, B., Novascone, S., Williamson, R., Pastore, G., and Perez, D. Advanced multiphysics coupling for Iwr fuel performance analysis. Annals of Nuclear Energy, 84:98-110, 2015.
Li, Z.,, Huang, D. Z., Liu, B., and Anandkumar, A. Fourier neural operator with learned deformations for pdes on general geometries. Journal of Machine Learning Re- search, 24(388): 2023.
Han, Z., Rahul, and De. S. A deep baaring--bsee hybrid approach for the solution of multiphysics problems in electrosurgery. Computer Methods in Applied Mechanics and Engineering, 357:1,2603, 2019. ISSN 0045-
Lim, J. H., Kovachki, N. B., Baptista, R., Beckham, C., Arizzadenesheli, K., Kossaifi, J., Voleti, V., Song, J., Kreis, K., Kautz, J., et al. Score-based based diffusion models in function space. arXiv preprint arXi::2302.074000 2023.
7825. doi: http://aa..ogy1...006........... URL httpe://ww..eeienneddeeeet
10

---

M2PDE: Compositional Generative Multiphysics and Multi-component PDE Simulation
Ma, Y., Liu, J., Yu, H., Tian, C., Huang. S., Deng, J., Chai, X., Liu, Y., and He, X. Coupled nnndiation-thermall mechanical analysis of the solid-state core in heat pipe cooled reactor. Nuclear Engineering and Technology, 54 (6):2094-2106 2022.
Ren, T., Wang. L.,, Chang, C., and Li, X. Machine learning- assisted multiphysics coupling performance optimization in a photocatalytic hydrogen production system. Energy Conversion and Management, 216:1 12935, 2020.
Ronneberger, O., Fischer, P., and Brox, T. U-net: Con- volutional networks for biomedical image segmenta- tion. In Medical image computing and computer-asiisted intervention 2015: 18th international confer- ence, Munich, Germany, October 5-9, 2015, proceedings, part III 18, pp. 234 241. Springer, 2015.
MacNamara, S. and Strang, G. Operator Splitting, Pp- 95-114. Springer International Publishing, Cham, 2016. ISBN 978-3-319-4158955 978- 978-3-319-41589--53, URL /ttt::s://doo.rooo// 1110//9883-319--1----
doi: 10.1007/
Rozet, F. and Louppe. G. Score-based data assimilation. Advances in Neural Information Processing Systems, 36: 40521-40541, 2023.
Meyer, P. P., Bonatti, C., Tancogne- Dejean, T., and Mohr, D. Graph-based metamaterials: Deep learn- ing of structure-proped proppery relations. Materials &
Design, 223:111175, 2022. ISSN 0264-1275. doi:
Sobes, V., Hiscox, B., Popov, E., Archibald, R., Hauck, C., Betzler, B., and Terrani, K. Ai-based design of nuclear reactor core. Scientific reports, 1 11(1):19644) 2021.
URL httpp::/www..ccieneeeieeee. ssience/arrice///////22212777272
Song, J., Meng, C., and Ermon, Denoising diffusion implicit models. In International Conference on Learning Representations, 2021. URL https: openrevie.. net/ / wtwpmmmmdddtttiidddeeeeae
Moon, T. K. The expectaaionn-maxxiizzaaooo algorithm. IEEE Signal processing magazine, 13(6):47-60, 1996.
Park, T.-R., Park, H., Kim, K. Im, C.-N., and Cho, J.- H. Heat and weight optimization methodology of ther- mal batteries by using deep learning method with multi- physics simulation. Energy Conversion and Management, 236: 44033, 2021.
Sun, B., Gu, Q. Zhang, P., and Ou, J. A practical numerical substructure method for seismic nonlinear analysis of tall building structures. The Structural Design of Tall and 2017.
Special Buildings,
Takamoto, M., Praditia, T., Leiteritz, R., MacKinlay, D., Alesiani. F., Pflüger, D., and Niepert, M. Pdebench: An extensive benchmark for scientific machine learning. In Advances in Neural Information Processing Systems, volume 35, pp. 1596-1611, 2022.
Pfaff, T., Fortunato, M., Sanchez-Goozallez A., and Battaglia, P. Learning mesh-based simulation with graph networks. In International conference on learning repre- sentations, 2020.
Po, R. and Wetzstein, G. Compositional 3d scene genera- tion using locally conditioned diffusion. In 2024 Inter- national Conference on 3D Vision (3DV), pp. 651-663. IEEE, 2024.
Tang, Z., Zhao, S., Li, J., Zuo, Y., Tian, J., Tang. H., Fan, J., and Zhang, G. Optimizing the chemical vapor deposition process of 4h-sic epitaxial layer growth with machine- learning-assisted multiphysics simulations. Case Studies in Thermal Engineering, 59: 19:107, 2024.
Ragone, M., Yurkiv, V., Ramasubramanian, A., Kashir, B., and Mashayek, F. Data driven estimation of electric vehi- cle battery state-of-charge informed by automotive sim- ulations and multi phitii-phycsss modeling. Journal of Power Sources, 483:229108, 2021.
Terlizzi, and Kotlyar, D. A perturbation based accelera- tion for monte carlo-thermal hydraulics picard iterations. part ii: Application to 3d pwr-based problems. Annals of Nuclear Energy, 166: 108713 13, 2022.
Urain, J., Li, A., Liu, P., D D'Eramo, C., and Peters, Com- posable energy policies for reactive motion generation and reinforeement learning. The International Journal of 2023.
Ranade, R., Hill, C., He, H., Maleki, A., Chang, N., and Pathak, J. A composable autoencoder-based iterative algorithm for accelerating numerical simulations. CoRR, abs/2110.03788 2021. URL https:////// https://axxiv..oor //tp:/xxxxv.////////.///////../t abs/2110.033 03780.
Robotics Research,
Van der Maaten, and Hinton, G. Visualizing data using t-sne, Journal of machine learning research, 9(11), 2008.
Rao, C., Ren, P., Wang, Q., Buyukozturk, O., Sun, H., and Liu, Y. Encoding physics to learn reaction-diffusion processes. Nature Machine Intelligence, 5(7):765-779, 2023.
Wang, C., han Song, L., Yuan, Z., and sheng Fan, J. State-of-the-art ai-based computational anal- ysis in civil engineering. Journal of Industrial

---

M2PDE: Compositional Generative Multiphysics and Multi-component PDE Simulation
Information Integration, 33:10047)) 2023a. ISSN 2452-
414X. doi: http://bbb.boom/10000001.1.1.22.
URL httpss//www.cceeneeddre...cc. cccene/s/arttcee////////44411300
Wang, J. Jiang, H., Chen, G., Wang, H., Lu, L., Liu, J., and Xing. I.. Integration of multi-physics and machine dearning-based surrogate modelling approaches for multi- objective optimization of deformed gdl of pem fuel cells. Energy and AI, 14: 10:261, 2023b.
Wu, H., Luo, H., Wang, H., Wang, J., and Long, M. Tran- solver: A fast transformer solver for pdes general geometries. In International Conference on Machine Learning, 2024a.
Wu, T, Maruyama, T., Wei, L., Zhang, T., Du, Y., Iac- carino, G., and Leskovec, J. Compositional generative inverse design. In The Twelfth International Confer- ence on Learming Representations, 2024b. URL https: /ppenreview net forum? id=wmXCC@FSd1.
Wu, Z., Pan, S.,, Chen, F., Long, G., Zhang, C., and Philip, S. Y. A comprehensive survey on graph neural networks. IEEE transactions on neural networks and learning sys- tems, 32(1):4422 2020.
Xu, K., Hu, W., Leskovec, J., and Jegelka, S. How powerful are graph neural networks? In International Conference on Learning Representationss 2019. URL https:// openreview. / forum?id=ryGa 61A5Km.
Yang, M., Du, Y., Dai, B., Schuurmans, D., Tenenbaum, J. B., and Abbeel, P. Probabilistic adaptation of text-to- video models. arXiv preprint 306.:2206.01172,
12

---

M2PDE: Compositional Generative Multiphysics and Multi-component PDE Simulation
A. Algorithm for multi-component simulation.
Multi-componett simulation first requires training a diffusion model to predict the solution of the current component based on the solutions of its neighboring components. Additionally, it is necessary to define the connectivity of all components and the functionn f to update the surrounding components' solutions for each component. The multi-component simulation algorithm is presented in Algorithm 2. Lines 6 to 11 are the denoising cycle of the diffusion model, in each diffusion step, the solutions of each component are updated together.
Algorithm 2 Algorithm for multi-component simulation by M2PDE
input diffusion model co (z C, outer inputs C, diffusion step S, number of external loops K, number of component N, connectivity of all components adj, update function [( -.. Zuy adj) of zae.
1: 26, ~ (0, // initialize estimated solution for each component 02 I/ Add an external loop to improve 261-
2: for k == -,,..,, K do
3: // l'initiaiize each component Fe., for i in 1,...,N N
4: 2T, + z5, // update previous estimated solutions for each component &
5: 20, -O( // initialize current estimated solutions for each component 2f
6: 2 ~ I), // initialize solutions for each component Au;
I/ denoising cycle of diffusion model:
7:
for S,... do
8:
A ==1- if k >> 1 else 1 // define the weights of 2 and of
9:
=0v- = -( -.... z adj) // update the solutions of surrounding components for each Zng
10:
w ~ 10
// use weighted estimated solutions as conditions for single step denoising,
// update all components together:
11:
20,,8-1 vas (Zeg,s VI-a, (o A=SS, + (1 A) C)) + a.w // update the estimated solutions of all components together
12:
2 Va,, -Z, vI- (zee Az + (1 A) C?)
13:
end for
end for output
14:
00 2v2,0. Zuy.0
B. Additional details for reaction-diffusion
This section provides additional details for Section 4.1.
Dataset. We employed the solve ivp function in Python to solve the reaction-diffusion equations. The coefficients Pau Pa, O,, and 3 of Eq. 18 are set to 0.01, 0.05, 0.1, and 0.25, respectively. The spatial mesh consisted of n₂ = 20 points, the time step adaptively controlled by the algorithm, but only outputs the results of 10 time steps. To train the data for a single physical process, it was necessary assume the initial conditions of the other physical processes and the current field. For instance, training u required assumptions about 24 and v. The dimension of uo is [ng]] which was generated using a one-dimensional Gaussian random field, and v has dimensions [Re, Rz,], and was generated by sampling a one-dimensional Gaussian random field n, times. Similarly, training v also requires preparation of on and u.
Model structure. The 2D U-Net and 2D FNO serve as both the surrogate and M2PDE. U-Net consists of modules: a downsampling encoder, a middle module, and an upsampling decoder. The encoder and decoder comprise four layers, each with three residual modules and downsumpiing/upsampiin convolutions, with the third module incorporating attention mechanisms. The middle module also contains three residual modules, with attention mechanisms included in the second module. The input data is encoded into hidden dimension before undergoing sequential downsampling and upsampling. FNO consists of three modules: a lift-up encoder, n FNO layers, and a projector decoder. Each FNO layer includes spectral convolution, a spatial convolution, and a layer normalization. The surrogate model predicts the evolution of the current physical process using its initial conditions and those of other physical processes. Its input dimension is [b, 1, 10, 20] and output dimension is [b, 1, 10, 20J. The diffusion model has an input dimension of [b. 2, 10, 20] and an output dimension of [b, 1, 10, 20], with representing the batch size. The shape of initial condition of [b, 1, 202 and will repeat to align the
13

---

M2PDE: Compositional Generative Multiphysics and Multi-component PDE Simulation
Algorithm 3 Surrogate model combination algorithm for multiphysics simulation.
input Compositional set of surrogate model (2(4, C), 1,2, ... N, outer inputs C. maximum number of iterations M, tolerance (max. relaxation factor a. Initialize constant fields 24, i I,... N. m
2: while < M and 6maz do + 1
3:
4: for = 1,.., N do
5:
-==+, =
6:
2, = or ( 24, C) + (1 - 2
7: end for
8:
6= (2 2,) N.... N
9: end while
output 21: 22. ... NN
required shape. The diffusion step of the diffusion model is set to 250. More details are shown in Table 4. Training, The surrogate model and M2PDE are trained similarly, with further details in Table 5.
Inference. The hyperparameter K set to 2. The surrogate models' combination algorithm in experiments and identical, demonstrated in Alg 3. The relaxation factor a is 0.5.
Table Hyperparameterss of model architecture for reaction-diffusion task.
Hyperparameter name
Hyperparameters for U-Net architecture:
Channel expansion factor
(1,2)
(1,2)
Number of downsampling layers
2
2
Number of upsampling layers
2
2
Number residual blocks for each layer
3
Hidden dimension
24
24
Hyperparameters for FNO architecture:
FNO width number of FNO layer FNO mode padding
24
24
4
4
[6,12]
[6.12]
[8,8]
[8.8]
Table 5, Hyperparameters of training for reaction- diffusion Hyperparameters for U-Net and FNO training Loss function
and
MSE
Number of examples for training dataset
10a
Total number of training steps (surrogate; diffusion)
10a;2 x 105
Gradient accumulate every per epoch
2
learning rate
10-1
Batch size
256
C. Additional details for nuclear thermal coupling
This section provides additional details for Section 4.2.
Problem description. The goal of this problem to predict the performance of plate-type fuel assembly under transient conditions. A typical pin cell in JRR-3M fuel assembly (Gong et al., 2015) is adopted as the computational domain, as shown in Fig. For simplicity, the cladding in the fuel plate omitted here without losing the representativeness of its
14

---

M2PDE: Compositional Generative Multiphysics and Multi-component PDE Simulation
multiphysics coupling feature. U-Zr alloy and lead-bismuth fluid are adopted as fuel and coolant materials, respectively. Their physical property parameters can be found in the repository. We consider single-group diffusion equation for the neutron physics process and employ an incompressible fluid model for coolant modeling. Temperature fields in solid and fluid can influence the macroscopic absorption cross-section in the neutron physics equation, while neutron flux affects the heat source in the fuel domain. Conjugate heat transfer occurs at the interface between the fluid and solid domains. While the feedback of temperature on neutrons is inherently complex, a linear negative feedback is assumed for simplicity. The governing equations are presented in Eq. 16, Eq. 17, and Eq. 18.
Here is neutrons/ per fission, D is the diffusion coefficient of the neutron, Ei, E. are the fission and absorption cross- section, respectively, and we only consider the feedback of temperature on the absorption cross-section Eg here. k,s ky are the conductivity of solid and fluid, respectively, both being functions of T.
Dataset, We utilize the open-source finite element software MOOSE (Multiphysics Object-Orienter Simulation Environ- ment) (Icenhour et al., 2018) to tackle the nuclear thermal coupling problem. The solid temperature field uses a mesh of [64, 8), the fluid fields have a mesh of [64, 12], and the neutron physics field employs a mesh of [64, 20J. The neutron physies and solid temperature fields are calculated using the finite element method at mesh points, while the fluid domain uses the finite volume method at mesh centers. Interpolation applied align the neutron physics and solid temperature values with the fluid fields. The time step is adaptively controlled by the algorithm, but only outputs the results of 16 time steps. So the input dimensions for the surrogate models of neutron physics field, solid temperature field, and fluid fields are [b, 2, 16, 64, 20J. [b, 2, 16, 64, 8], and [b. 1, 16, 64, 12]. respectively. The input dimensions for the diffusion model of the three fields are [b. 3, 16, 64, 20). [b, 3. 16, 64, 8]. and [b, 5, 16, 64, 12], respectively. The output dimensions of the three fields are [b, 1, 16, 64, 20), [b, 1, 16, 64, 8], and [b, 4. 16, 64, 12], respectively.
As noted, assuming the distribution of physical field data in high-dimensional problems is challenging. We recommend a pre-iteration method for data generation. Initially, we assume constant values for all other physical fields and calculate the current field. This process repeats until all fields are computed. If there are n physical fields, pre-iteration requires n I calculations plus one iteration for data generation, totaling 2n- 1 calculations. To accelerate data generation, the most iime-consuming field can be excluded from pre-iteration. In this problem, the fluid fields' computation time is approximately three times that of the other fields, so it is excluded from pre-iteration. The process begins by assuming constant fluid fields and solid temperatures to calculate the neutron physics field, followed by using the resulting neutron physics field and assumed fluid fields' temperature to calculate the solid temperature. The data generation proceeds sequentially with calculations for the fluid fields, neutron physics field, and solid temperature field. Through this process, the input datasets for all physical fields are not constant fields, but rather fields with distribution information after one iteration.
Model structure. The 3D U-Net and 3D FNO serve as both the surrogate model and M2PDE, using a layer design identical to the 2D. For regional coupling, concatenation is directly applied to the channel dimension using the concat function. In contrast, for interface coupling, dimensions must be replicated to align spatially before concatenation. The conditioning of the diffusion step for FNO is operating in spectrum space (Gupta & Brandstetter, 2023), which is better than in the original space for this problem. The diffusion step of the diffusion model is set to 250. More details are shown in Table 6.
Training. The surrogate model and M2PDE are trained similarly, but training neutron physics fields using M2PDE requires more time to converge, with further details in Table 7.
Inference. The hyperparameter K is set to 2. The relaxation factor for surrogate model a is set to 0.5.
Detailed results. Fig. 5 Ilustrates how the relative errors are distributed along the x, y, t direction. Overall, the neutron field exhibits a relatively uniform distribution of errors along all three directions. In x direction, the error in both the solid and the fluid slightly increases near their interface, likely due to coupling effects in this region. Along the y direction, the solid temperature error gradually increases with distance, whereas the fluid error shows the opposite trend, which may be related to the distinct thermal and physical properties of the two media. Fig. 6 presents the results of predicting various physical fields using the last time step surrogate model and M2PDE + U-Net on the final test data. The neutron physics field and solid temperature field are represented by é and T., respectively. The fluid fields include four physical quantities: Tr, P, Ua, My. totaling six quantities. Since the neutron physics field and the tyy component of the fluid fields are less influenced by other physical processes, the surrogate model can still make predictions, but the accuracy lower than that of the M2PDE. Besides, the surrogate model has failed predict the other physical processes. In contrast, M2PDE continues provide relatively accurate predictions, although some distortions are observed in certain regions.
15

---

M2PDE: Compositional Generative Multiphysics and Multi-component PDE Simulation
Table Hyperparameters of model architecture for nuclear thermal coupling task.
Hyperparameter name
neutron
solid
fluid
Hyperparameters for U-Net architecture
Channel Expansion Factor
(1,2,4)
(1,2,4)
(1,2,4)
Number of downsampling layers
3
3
3
Number of upsampling layers
3
3
3
Number of residual blocks for each layer
3
3
3
Hidden dimension
8
8
16
Hyperparameters for FNO architecture
FNO width number of FNO layer FNO mode padding
8
8
16
3
3
3
[6.16,8]
[6,16,4]
[6,16,6]
[8,8,8]
[8,8,]]
[8,8,8]
Table 7. Hyperparameters of training for nuclear thermal coupling task. neutron&&solid&ffuu
Hyperparameters for U-Net and FNO training Loss function
MSE
Number of examples for training dataset
5 10
Total number of training steps (surrogate; diffusion)
10 10
Gradient accumulate every per epoch
2
learning rate
10
Batch size
32
solld stid
solld fuidd
reaaive cIror
relatve CTTOr
Blid
632-
126
time step
direction
direction
(a) t direction.
(b) x direction.
(C) y direction.
Figure 5. The distribution of errors along the x, y, and directions.
16

---

M2PDE: Compositional Generative Multiphysics and Multi-component PDE Simulation
Ground truth
Surrogate
M2PDE
Ground truth
Surrogate
M2PDE
500
34
BE
0G
730
04
01
780
06-
CA-
INS
92
62
350
800
1.0
82
0.
0.
1.
60
63
10
66
(a) d.
(b) T..
Ground truth
Surrogate
MZPDDE
Ground truth
Surrogate
MZPDE
14.
10
300
02
0.8
392
120
04-
133
000
403
360
022
02
100
05
1-
4:1
0.
30
05
1.0
64
05
L4
65
10
(c) Ty-
(d) P.
Ground truth
Surrogate
M2PDE
Ground truth
Surrogate
M2PDE
14"
0.0066
000000
0.1
00004
0.5
04
040
14
60022
0000
6.7
035
0.000
6.22
3:0
425
000000
COOOO
320
620
c000
355
0.5
0.5
0.3
0.5
1.3
040
6.5
2:00
0.5
35
1.0
34
05
1.0
(c) Mar.
(f) Ugy.
Figure 6. Comparison of surrogate model and M2PDE for predicting all physical fields.
D. Additional details for prismatic fuel element
This section provides additional details for Section 4.3.
Problem description. This problem aims to predict the thermal and mechanical performance of prismatic fuel elements in heat pipe reactor (Ma al., 2022) at different source power. The reactor core is stacked up using hexagonal prism SiC matrix, with multiple holes dispersed in the matrix for containing fuel elements and heat pipes as shown in Fig. 7. The SiC matrix plays role in locating the fuel and heat pipes at expected positions in the core. The entire structure consists of two basic components, one oriented upwards and the other downwards, as illustrated in Fig. 8. Fission energy released in fuel elements is dissipated using heat pipes. Both the fuel elements and heat pipes are considered as boundaries here, and only the more concerned matrix behavior is analyzed in the demonstration. Only strain is predicted here since stress can be
17

---

M2PDE: Compositional Generative Multiphysics and Multi-component PDE Simulation
derived from the mechanical constitutive equation, and displacement is obtained through strain integration. The analysis uses the plane strain assumption (ez 0) and excludes irradiation effects, simplifying it to steady-state problem.
Fuel element
Heat pipe
Matrix
Figure Schematic heat pipe core structure. The left figure shows partial structure of the entire reactor, with multiple holes dispersed in the matrix for containing fuel elements and heat pipes. The right figure shows how to select a medium structure for analysis from overall structure.
Figure 8. Two basic components: facing upwards (left) and the other facing downwards (right).
Dateset. We MOOSE to calculate the thermal and mechanical problems. The training data comes from medium structure simulation with 16 fuel elements, allowing each simulation to generate 16 training data, as shown in Fig. 4. This structure is chosen because fundamental component, along with its neighboring components, is entirely contained within the interior, which where most components that need to be predicted in large structures are located. When generating data, the heat flux density is uniformly sampled from the range [10$, 10 W/m. free boundary condition is randomly assigned to one edge, while symmetric boundary conditions are applied to the remaining two edges. Each fundamental component is uniformly meshed with 804 points, each requiring the prediction of three physical quantities. To predict the central component, the heat flux density of this component and the coordinates of each mesh point are concatenated with data from its three neighboring components, yielding an input dimension of [b, 804, 15] for the diffusion model and [b, 804, 12] for the surrogate model, where b is the batch size. The output dimension is [b, 804, 3]. The sequence of neighboring elements is consistent, with the downward-facd faccad center element being the upward-facing center element rotated by 180 degrees. This arrangement is illustrated in Fig. 8. Boundary conditions are considered only for symmetric and free types, represented as [o, 1, and [O, O]. respectively, and are replicated dimension of [804, 3].
18

---

M2PDE: Compositional Generative Multiphysics and Multi-component PDE Simulation
Algorithm 4 Surrogate model combination algorithm for multi-component simulation.
input surrogate model c (2av. C). outer inputs C, maximum number of iterations M. tolerance Emars relaxation factor
1: Initialize constant fields Zuts = I,...... N. 0
2: while < M and > 6maz do
3: m = m + 1
4: for = -,,.., N do
5: i.; = Ew;
6: Z1: =cco(200.: C) + (1 2
7: end for
8: 6 L1 (zov. 200,); 1, N.... NN
9:
end while
output 2001: 7002' ..., 20vN
Model structure. The Geo-FNO and Transolver serve as both the surrogate model and M2PDE. Geo-FNO enhances FNO for irregular meshes using three modules: geometry encoder that converts physical fields from irregular to latent uniform meshes, FNO functioning in latent space, and a geometry decoder that transforms physical fields from the uniform mesh back to the original irregular mesh. We utilize a 2D Geo-FNO that transforms into a 2D uniform mesh. Transolver is designed to tackle complex structural simulation problems involving numerous mesh points by learning the intrinsic physical states of the discretized domain. Given mesh set with N points and C features per point, the network first assigns each mesh point to M potential slices, transforming the shape from N x C to M x N x C. It then applies spatially weighted aggregation, resulting in shape of M x C. Self-attention is used to capture intricate correlations among different slices, after which the data is transformed back to the mesh points. The conditioning of diffusion step for Geo-FNO is also operating in spectrum space (Gupta & Brandstetter, 2023). More details about the network can be found in (Wu et al., 2024a). The setting of hyperparameters is shown in Table 8. The diffusion step of the diffusion model is set 250.
Training. The surrogate model and M2PDE are trained similarly, but training M2PDE requires more time to converge, with further details in Table 9.
Inference. This problem uses the same neural network to predict the performance of all elements, allowing for simultancous updates of the physical fields and enhancing inference speed. This method applies to both diffusion and surrogate models. The surrogate models' combination algorithm for mult-component simulation demonstrated in Alg 4, the relaxation factor a is set to 0.5. The hyperparameter K set to 3.
Detailed results. Fig. 9 and Fig. 10 compares the results of predicting the large structure using the surrogate model and M2PDE Transolver. Because the surrogate model of FNO fails in predicting large structures, only the results of M2PDE + FNO are provided in Fig. 11. The strain is only displayed in the x-direction due to its similarity in both and y. The error graph indicates that M2PDE offers more accurate predictions.
Graph neural network and Graph Transformer configuration. For GIN, SAN and MeshGraphNet, each component treated as a node in the graph, with training conducted on a small graph of 16-component, and ultimately tested on a larger graph of 64-component. Compared with the surrogate model and M2PDE, they use only the system's 's input as input features. In contrast, the surrogate model and M2PDE enrich its input by incorporating the solutions from the surrounding component, thereby improving accuracy, demonstrated in Table 3. The input to the GIN, SAN and MeshGraphNet is the heat flux density and boundary conditions of each component, and the output is the physical quantities at all grid points on the component. GIN updates the nodes on the graph through the graph structure, whereas SAN captures graph structural information by inputting the eigenvalues and eigenvectors of the graph Laplacian matrix into a transformer. The training settings of GIN, SAN and MeshGraphNet are consistent with the M2PDE. We have adjusted the number of network layers and the size of hidden layers to obtain the model with optimal performance.
19

---

M2PDE: Compositional Generative Multiphysics and Multi-component PDE Simulation
Ground buuth
Ground buih
1250
1340
1350
1100
0.23
1100
1544
1000
-1.075
1000
-0.55
s/s
0.p
0.0
b/n
s1n
0.15
Predition [K]
Predition (K)
1300
0375
1300
0.32
0350
0025
0325
1700
0.02
1130
0.50
1150
+0.02
11.00
-1.035
1050
00.05
0.00
306
010
2.15
00.0
000
0.5
510
0.5
(K)
Error (K)
0.11%
150
0.30
0.50
0023
0.23
0.00
0.00
+4.05
-10
-0.055
8.000
100
0/33
100
0073
-0.10
-4.00
(a) Surrogate model when predicting T.
(b) M2PDE when predicting T..
Figure Comparison of surrogate models and M2PDE Transolver predicting the temperature large structures.
20

---

M2PDE: Compositional Generative Multiphysics and Multi-component PDE Simulation
Ground bruth
Ground Iruth
60033
CAS3C
00220
0.20-
0.23
00340
-1.075
00330
-0.55
s/s
0:p
000
b/n
sin
0.15
Predition
Predition
10033
0375
0.32
0350
00250
0025
0.25
0.02
0.50
+0.02
-0.035
00240
00035
00032
00.05
0.00
3/6
3.0
2.15
00.0
0.0
0.5
510
013
Errrr
0..1%
0.30
0.50
00025
0023
0.23
0.00
0.00
+4.05
-0.35
8.000
-200005
0/35
-0,000
0035
-29219
-00030
-0.10
-4.00
(a) Surrogate model when predicting Ex.
(b) M2PDE when predicting Ex--
Figure Comparison of surrogate models and M2PDE Transolver for predicting the strain large structures.
21

---

M2PDE: Compositional Generative Multiphysics and Multi-component PDE Simulation
Ground buuth
Ground Iruth
000255
1350
1300
1200
10000
1330
60345
0.23
00040
1250
-1.075
1000
-0.55
s/s
0:p
0.0
b/n
sin
0.55
Predition [K]
Predition
60325
0375
00240
0.32
0350
1750
0025
0325
1200
0.02
0.50
+0.02
-1.035
1100
00340
1250
00335
1000
00220
00.05
0.00
306
310
2.55
00.0
0.0
0.5
510
015
(K)
0.11%
0.30
0.50
00025
0023
0.23
0.00
0.00
+4.05
-10
-0.055
8.900
100
0/35
-0,000
0073
-00030
-0.10
-4.00
(a) M2PDE when predicting T'.
(b) M2PDE when predicting Ex--
Figure II. The results M2PDE FNO for predicting large structure.
22

---

M2PDE: Compositional Generative Multiphysics and Multi-component PDE Simulation
Table 8. Hyperparameters of model architecture for prismatic fuel element task.
Hyperparameter name
Hypperparaeeeess for Transolver
Number of layers
5
Number of head
8
Number of slice
16
Hidden dim
64
Hyperparameters for Geo-FNO
Uniform grid size
[64, 64]
FNO width
FNO mode
[8,8]
Number of FNO layer
3
Hidden dim
64
Table 9. Hyperparameters of training for prismatic fuel element task.
Hyperparameters for Transolver and Geo-FNO training
Loss function
MSE
Number of examples for training dataset
16000
Total number of training steps(
10 x 105
Gradient accumulate every per epoch
2
learning rate
10-4
Batch size
256
E. Comparison of models trained using coupled and decoupled data.
To further investigate the model's boundaries for multiphysics simulation, utilize coupled data train diffusion models and compare them to models trained on decoupled data in experiments and 2. The input of the diffusion model is the external input of the physical system, while the output is the solution of the coupled physical fields. In experiment the input consists of the initial conditions of u and with the output being their trajectories. Since and v are defined on the same grid, a single network can be employed to predict u and together. In experiment 2, the input is the variation of neutron boundaries over time, and the output is the trajectories of the neutron field, solid temperature and fluid fields. Since that the three fields are defined in different computational domains, three separate networks are trained. Aside from the differences input and output dimensions, all other parameters remained consistent with those used in the decoupled scenario. The coupled datasets for experiments 1 and 2 consist of 10,000 and 5,000 samples, respectively, which is consistent with decoupled datasets. The model is evaluated using unseen coupling data during training.
The result is shown in Table 10, the accuracy of the model trained with decoupled data decreased by about 1 order of magnitude.
Table 10. Comparison of models trained coupled and decoupled data. Coupled data model Decoupled data model
Reaction-diffusion
u
0.00151
0.0141
v
0.00185
0.0174
Nuclear thermal coupling
neutron
0.00512
0.0197
solid
0.00098
0.0287
fluid
0.00302
0.0391
23

---

M2PDE: Compositional Generative Multiphysics and Multi-component PDE Simulation
F. ABLATION STUDY
F.I Method for calculating the estimated physical fields
We compare two methods for estimating physical fields: one using of from Eq. and the other using the current physical field ži.s with noise. As shown in Table 11, 2f provides significantly better results than 2i,s, indicating that the estimate from 2j accurate.
Table 11. Comparison methods for estimating physical fields.
2i,8
2 (Eq.9)
Reaction-diflusion
0.0525 0.0355
0.0141
0.0174
Nuclear thermal coupling
neutron
0.0184
0.0197
solid
0.0913
0.0287
fluid
0.1000
0.0391
Prismatic fuel element
T
0.0289
0.0076
E
0.0083
0.0194
F22. Selection of Hyperparameter K
This section examines how hyperparameter K affects the predictive performance of multiphysics and multi-component problems in experiments 2 and 3. As shown in Tables 12 and 13, setting K to 2 for multiphysics problems and K to 3 for multi-con colti--ommpoent problems is adequate. The multiphysics algorithm updates physical fields at each diffusion time step, leading to faster convergence. In contrast, the multi-component problem relies on the field estimated in the previous time step for each diffusion iteration, resulting in slower convergence. Additionally, increasing further has a negligible effect on model performance.
Table 12. Hyperparameters of K for multiphysics simulation.
K
neutron
solid
fluid
0.0199
0.0304
0.0524
0.0206
0.0287
0.0391
0.0203
0.0288
0.0395
Table 13. Hyperparameters of K for multi-component simulation.
K
T
E
0.00907
0,0236
0.00833
0.0222
0.00785
0,0206
0.00772
0.0207
0.00750
0.0203
F.3. Selection of Hyperparameter of X
The hyperparameters À determine the weight of the current physical field, theoretically requiring a reliable estimate of at the beginning and gradually increasing to as diffusion progresses to provide better estimates of the current results the later stage. We demonstrate this with experiment 3, which involves solving 64 components with slow convergence, making sensitive to hyperparameter A. We set the values to 1, 0.5, and linear increase. When K is too large, result differences
24

---

M2PDE: Compositional Generative Multiphysics and Multi-component PDE Simulation
minor, except when À is I; thus, K is set 2. As shown in Table 14, employing linearly increasing setting yields superior performance, which is consistent with the analysis.
Table 14. Hyperparameters of for multi-component simulation.
À
0
0.00878
0.0228
1
0.00913
0.0237
0.5
0.00895
0.0233
linear increase
0.00816
0.0217
G. The difference between training dataset and testing dataset.
For multiphysics simulation, we train models for each physical process using decoupled data and combine them during testing to predict coupled solutions; for multi-component simulation, we train model to predict individual component, then combine it during testing to predict the large structure composed of multiple components. To quantify the difference between the model's training and testing data, we calculate the Wasserstein distance (Feydy et al., 2019) between the training and validation data, as well as between the training and testing data, with the training and validation data originating from the same distribution. In addition, we also used the t-SNE (Van der Maaten & Hinton, 2008) algorithm to visualize this difference.
The results are presented in Table 15. In experiment 1, there is significant difference between the training and testing data, as can be seen from Fig. 12, where only small fraction of decoupled data points fall within the range of coupled data. In experiment 2, the difference between the training and testing data for the neutron physics field is relatively small, likely due to the weak coupling effect of other physical processes on the neutron physics field. For the solid temperature field and fluid field, the difference between the training and testing data also very pronounced, with almost no overlapping points in the Fig. 13. In experiment since the range of training data has been expanded during data generation cover as many potential scenarios of large structures as possible, the difference between the training and testing data is not significant as in the multiphysics problem, and the testing data are also within the range of the training data, shown in the Fig. 14.
Table 15. Wasserstein distance of datasets. Training and validation Training and testing
Reaction-diffusion
0.343
52.7
0.0435
20.3
Nuclear thermal coupling
neutron
42.4
31.3
solid
1.35
56.3
fluid
1.22
986
Prismatic fuel element
T
0.233
9,05
0.625
12.5
25

---

M2PDE: Compositional Generative Multiphysics and Multi-component PDE Simulation
(a) M.
(b)
Figure 12. Visualization of experiment 1 Dataset.
(a) Neutron.
(b) Solid.
(c) Fluid.
Figure 13. Visualization experiment Dataset.
(a) T.
(b) E.
Figure 14. Visualization of experiment 3 Dataset.
H. Sampling acceleration.
By employing the DDIM algorithm expedite the sampling process of diffusion models, we have successfully enhanced the efficiency of model inference. The DDIM algorithm encompasses two parameters: the number of time steps S and the parameter which controls the noise (Song et al., 2021):
Ctt
B
(19)
where 17 [0, 1]. We conduct tests across various parameter combinations, including S 10, 25, 50, and 17 = 0, 0.5, 1, with particular focus on the model's performance in coupled and large structure prediction. These three experiments all use the most accurate model, which is: U-Net, U-Net, Transolver. Table 16 indicates that in experiment 1, the setting of S = 25 closely mirrors the results of S = 50, with 1) having a relatively minor impact. Table 17 indicates that in experiment 2, the setting of S = 25 also approximates the outcome of S = 50 but is more sensitive to 7), with 1 = 1 yielding the best performance. Table 18 indicates that in experiment 3, S 50 is the optimal setting, and 7) = 0 provides the best results.
During the training of diffusion models, we uniformly set the number of time steps to 250. By employing accelerated
26

---

M2PDE: Compositional Generative Multiphysics and Multi-component PDE Simulation
sampling techniques, achieved a 10-fold acceleration for multiphysies problems and 5-fold acceleration for multi- component problem while ensuring the maintenance of predictive accuracy.
Table 16. Relative norm of error on reaction -iiftusion equation DDIM sampling.
method
decoupled
coupled
decoupled
coupled
Original DDPM
0.0119
0.0141
0.0046
0.0174
S ==10, 1 0
0.0143
0.0170
0.0117
0.0215
S 25, 17 = 0
0,0123
0,0151
0,0082
0,0190
S 50, 17 = 0
0.0123
0.0147
0.0059
0.0179
25, 1 = 0.5
0.0123
0.0152
0.0082
0.0191
S 25, 17
0.0119
0,0151
0.0081
0.0192
Table 17. Relative L2 norm of prediction nuclear thermal coupling for DDIM sampling.
The unit is x 10-
neutron decoupled
solid decoupled
fluid
method
coupled
coupled
decoupled
coupled
Original DDPM
0.487
1.97
0.108
2.87
0.303
3.91
S = 10,
0.638
1.89
0.261
4.45
0.478
4.42
= 25, =
0.552
2.03
0.142
3.64
0.343
4.08
= 50, =
0.533
1.96
0.138
3.21
0.346
4.02
25, 0.5
2.82
2.78
0.793
5.28
0.970
4.70
S 25,
10.9
10.3
2.99
14.4
1.82
8.20
Table Relative norm of prediction error on prismatic fuel element experiment DDIM sampling. The unit x 10"2.
single
16-component
64-component
method
T
E
Original DDPM
0.107
0.303
0.213
1.03
0.759
1.94
S 10,
0.207
0.425
1.69
3.81
1.89
4.13
S = 25, 17 =
0.166
0.353
0.952
2.55
1.30
3.26
S 50, 17
0.158
0.337
0.669
1.87
0.865
2.31
S 50, / 0.5
0.150
0.352
0.586
1.69
0.954
2.61
S 50, 17 1
0.130
0.322
0.553
1.62
1.05
2.80
I. Efficiency analysis.
This section compares the computational efficiency M2PDE, surrogate model, and numerical programs. The time unit for each experiment defined as the time required for single neural network inference. These three experiments all use the accurate model. Since the surrogate model and M2PDE both use the same network architecture and have consistent network parameters, it is assumed that the time for single inference using these two methods is equal. The numerical programs are run on the CPU and have all been optimized to the best parallel count.
Let the number of physical processes be denoted by N, the number of iterations for the surrogate model by M, the number of diffusion steps by 5, and the number outer loop iterations for the diffusion model by K. The computation time for the surrogate model is M x N, while the diffusion model is x S x N. The specific choices of N, M, S, K for each experiment are presented in Table 19.
27

---

M2PDE: Compositional Generative Multiphysics and Multi-component PDE Simulation
The results are presented in Table 20. In experiment 1, the problem is relatively simple, and the numerical algorithm achieves efficient solutions through explicit time stepping, while the introduction of M2PDE actually reduces efficiency. However, in experiment 2, which addresses more complex problems, M2PDE achieves 29-fold acceleration compared to numerical programs. In experiment 3, comparing the results of 16 components with 64 components. it is observed that as the computational scale increases, the acceleration effect of M2PDE becomes increasingly significant. Furthermone, when dealing with mulli-component problems, the surrogate model requires iteration to ensure the convergence of solutions across all components, Due to the large number of components, the number of iterations needed significantly increases compared to multiphysics problems, resulting in higher efficiency for M2PDE. In addition, we have only compared the efficiency of single computations for all experiments. When dealing with multiple problems simultaneously, the acceleration provided by M2PDE will be even more pronounced due the parallel nature of GPU computing.
In general, the more complex the problem, the more pronounced the acceleration effect of M2PDE becomes. In fact, the problems in experiment 2 and experiment 3 have been simplified to a certain extent, and the actual situations are even more complex. Therefore, M2PDE holds significant value in solving real-world complex engineering problems.
Table 19, Values K, N, M, and for the three experiments.
experiment
N
M
S
K
cractionddiffusion
2
27
25
2
Nuclear thermal coupling
3
21
25
Prismatic fuel element (16-component)
1
309
50
Prismatic fuel element (64-component)
1
324
50
Table 20. Comparison of running time.
Experiment
Unit (s)
Numerical program
Surrogate model
M2PDE
Speedup
Reaction-diffusion
0.0115
54
100
0.064
Nuclear thermal coupling
0.0242
4368
63
150
29
Prismatic fuel element (16-component)
0.0067
834
309
150
5.6
Prismatic fuel element (64-ccomooent)
0.0256
6170
324
150
41
J. Application Scenarios for Multi-Commpoent simulation.
In this section, we discuss the application scenarios of M2PDE for multi-component simulation from both theoretical and practical perspectives.
From a theoretical perspective, in the derivation of Section 3.2, we make an assumption: the solution on a multi-component structure is undirected graph that satisfies the local Markov property, meaning that any two non-adjacent variables conditionally independent given all other variables. Using this property, we derived Eq. 14. We believe this assumption is applicable to most problems because physical fields are continuous in space, and the information exchange between any two points must be transmitted through the points in between. However, there is class of problems to which current methods cannot be directly applied, which is the PDE that requires determining eigenvalues:
Mό Xd
(20)
Here M is the operator, À is the eigenvalue, ø is the physical field to be solved. The À varies with different systems, and the relationships we learn small structures may not be applicable to large structures. Solutions to these problems may be similar to numerical algorithms, requiring the addition of an eigenvalue search process, which will be undertaken in future work.
From a practical implementation perspective, for a complex structure, it is necessary to clearly determine its basic components and the relationships between these components and their surrounding components, so that we can understand how the
28

---

M2PDE: Compositional Generative Multiphysics and Multi-component PDE Simulation
components affected by their surrounding components. In addition, training data must encompass all possible scenarios that each component in large structure might encounter, such all possible boundary conditions and the relationships with surrounding components.
K. Datasets Description.
This section provides a concise description of the datasets utilized the three experiments, with their detailed backgrounds introduced in Appendix standard scientific datasets PDEbench (Takamoto al., 2022). Comparison is shown in Table 21, where Na is the spatial dimension, N is the number of physical processes, and Ne the number of components. Table 21 only lists some of the datasets in PDEBench, but all of its datasets have N, and N₂ values of 1. The dataset of Experiment 1 in this paper exists in the benchmark, but Experiments 2 and 3 are completely new datasets.
D. We outline the principal characteristics of these datasets and compare them with the
Table 21. Datasets Description.
PDE
Na
Time
Computational domain
Ne
Burgers'
yes
Line
compressible Navier-Stokes
3
yes
Cube
incompressible Navier-Stokes
2
yes
Rectangle
1
shallow-water
2
yes
Rectangle
reaction-diffusion (Expl)
yes
Line
heat conduction + neutron diffusion
incompressible Navier-Stokes (Exp2)
yes
3 Rectangle
3
heat conduction mechanics (Exp3)
no
Irregular domain
16, 64
29


Reviews:
[{"review_id": "RVi69C3bNA", "reviewer": "Reviewer_WGGT", "summary": "This paper introduces MultiSimDiff a new method for solving multi-physics/multi-compnennts simulations efficiently by learning the conditional score of each component's solution given its parameters and solutions of other components. Experiments demonstrate that MultiSimDiff outperforms largely a simpler surrogate that model each component's solution independently on a class of coupled cases.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> Re1: The end of the abstract (3 last sentences) appears a bit narrow and verbose.\n\nThank you for the suggestion! We will revise the ending of the abstract as follows: \n\nWe demonstrate the effectiveness of MultiSimDiff through two multiphysics tasks—reaction-diffusion and nuclear thermal coupling—where it achieves more accurate predictions than surrogate models in challenging scenarios. We then apply it to prismatic fuel element simulation, an exemplary multi-component problem, where MultiSimDiff successfully extrapolates from single-component training to a 64-component structure and outperforms existing domain-decomposition and graph-based approaches.\n\n> Re2: Would be nice to mention better ... sampled from $p_{\\neq z_i}$ is not straightforward.\n\nI gather that this section focuses on two issues:\n(1) Under which scenarios does our algorithm exhibit an advantage? The key factor to train it.\n(2) For multiphysics problems, it is not easy to find a suitable $p_{\\neq i}$, or in other words, you believe that finding a suitable $p_{\\neq z_i}$ is almost equivalent to solving the problem itself.\n\nFor (1), our algorithm aims to solve multiphysics and multi-component simulation problems that can be numerically challenging. In the more complex experiments 2 and 3, we observed up to 29× and 41× speedups, with higher accuracy than other methods. In real engineering applications with large-scale coupling, these gains can be substantial. The key factor is to view multi physics/component simulation problems from the perspective of probability, and replace complex joint probabilities with easily obtainable conditional probability.\n\nFor (2), indeed, for multiphysics simulation, find $p_{\\neq z_i}$ that perfectly respects all coupling is nearly as hard as solving the entire problem. But **we do not solve the coupled physical field equations**, we currently use a pre-iteration technique (lines 762–772) to approximate it; however, Figure 12 shows a notable gap from the true coupled data distribution. Hence, the accuracy of predicting the coupled solutions of physical fields still needs improvement, we discuss this limitation in Section 5 and plan further improvements.\n\n> Re3: I am a bit confused by the baselines you used... Similar remarks for CoAE-MLSim.\n\nFor multi-component simulation, we use CoAE-MLSim as a baseline (implemented per the paper’s description) since it is not open-sourced. This corresponds to the “surrogate model” in Table 3. The Related Work section also discusses GNN (GIN) and Graph Transformer (SAN) approaches, which also can be seen in Table 3. We have also added MeshGraphNet for comparison, as noted in our response to Review QAdZ (Re3). However, these methods are not applicable to multiphysics simulations. The baseline for multiphysics simulations is shown in alg.3 in manuscript.\n\n> Re4: It seems that what 1 discusses in 3.1 regarding Markov Blanket and using diffusion model to jointly sample from dependent components is very close to your work from a technical standpoint.\n\nTheir method replaces the log-likelihood gradient at one time step in a long sequence with a smaller subsequence. We instead replace the gradient of a complex joint distribution (for each physical field or component) with a tractable conditional distribution. For multiphysics problems, some fields typically depend on all other fields,  so their subsequence substitution is not valid. For multi-component problems, the posterior estimation step (their Algorithm 3) introduces further challenges. In their scenario, the known information is observational data; in ours, it is the entire physical system’s inputs, which would constrain the model to the same scale as the training configuration. Overall, the core difference is that we must be able to **generalize to tasks beyond the training distribution**, while their method focuses on the same domain for training and inference.\n\n> Re5: Figure 1. is pixelized and not very clear (or at least not informative compared to the space it takes), clean...\n\nWe have replaced Fig.1 with a clearer version at: https://anonymous.4open.science/r/MultiSimDiff-D5A3/rebuttul/schematic.pdf\n\n> Re6: By definition you algorithm will output a distribution over plausible solutions, ... , decreases as a function of the exact setup considered.\n\nWe think that you are referring to uncertainty quantification. In general, there is model uncertainty and data uncertainty. In our current experiments, we learn from simulation data, which do not contain noise. Ideally, if a physical process is deterministic, the predictive uncertainty should be close to zero. We conducted uncertainty quantification experiments in both Experiment 2 and 3, and found that the standard deviation of the model predictions is extremely small, indicating that the diffusion model can approximate a deterministic physical process. We have drawn a table of the results: https://anonymous.4open.science/r/MultiSimDiff-D5A3/rebuttul/std.md.", "claims_and_evidence": "Overall the paper is well-written and claims are supported by evidence.", "methods_and_evaluation_criteria": "Yes, the idea of building a global solutions by representing each solution as a distribution conditioned on other components' solution appears original and well motivated.", "theoretical_claims": "There is not strict theoretical claims.", "experimental_designs_or_analyses": "Experiments are sound although I would have appreciated if other alternative strategies such as the ones discussed in the related works section would have been used as additional baselines.", "relation_to_broader_scientific_literature": "In my opinion, the paper provides an insightful related work section which position this work well into the literature.", "essential_references_not_discussed": "-", "other_strengths_and_weaknesses": "Overall I very much enjoyed reading the paper and I believe the problem statement is very sound as well as the proposed solution. \nThat said, I am not very knowledgeable in the field of compositional mulitphysics/multi-component solvers and related works using ML for that sake. For that reason and for the few remarks below, I will only weakly support the acceptance of the paper but might increase my score in light of the other reviews and response from the author.\nHere is a list of remarks/questions (ordered as noted when reading the manuscript):\n- The end of the abstract (3 last sentences) appears a bit narrow and verbose. I am not sure it really helps the reader position your work.\n- The intro is very clear\n- Would be nice to mention better the gain in computation time you may have and exactly when using your method will be beneficial, what are the key ingredients to train it. For instance, it is not totally clear to me as whether you need a large datasets of coupled simulations or not for the method to work well. I understand this consideration is problem dependent. For some problems it may be easy and meaningful to simply solve independently each component conditioned on a plausible state for the other components that are part of the Markov blanket but I imagine that in many cases finding these \"plausible\" sets for the Markov blanket really requires to have jointly solved this Markov blanket... For Multi-components I understand this may still make your method a nice to scale to a larger number of components, however for multi-physics it seems that you cannot to very much.\n- To emphasise my point above, I think you should discuss more clearly the issue of training the conditional scores and in particular that for tightly coupled systems having these sampled from $p_{\\neq i}$ is not straightforward.\n- I am a bit confused by the baselines you used. You discussed in the relate work that you compare to GNN and Graph Transformer but I did not see these results in the tables. Similar remarks for CoAE-MLSim.\n- It seems that what [1] discusses in 3.1 regarding Markov Blanket and using diffusion model to jointly sample from dependent components is very close to your work from a technical standpoint.\n- Figure 1. is pixelized and not very clear (or at least not informative compared to the space it takes), clean... \n- By definition you algorithm will output a distribution over plausible solutions, it would interesting to discuss a bit better what exactly is modelled by that uncertainty and how you get rid of it in your numerical experiments. Also, how much uncertainty increases, decreases as a function of the exact setup considered.\n\n[1]:https://proceedings.neurips.cc/paper_files/paper/2023/file/7f7fa581cc8a1970a4332920cdf87395-Paper-Conference.pdf", "other_comments_or_suggestions": "-", "questions_for_authors": "See my comments above.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2ywVo8pvGH", "reviewer": "Reviewer_itfA", "summary": "This paper proposes MultiSimDiff, a novel compositional generative model for multiphysics and multi-component simulations. The core idea is to use diffusion models to learn energy functions representing the conditional probability distributions of different physical processes or components. During inference, MultiSimDiff reconstructs the joint solution by sampling from these learned distributions. This approach circumvents the need for coupled numerical solvers and allows generalization from decoupled (small structure) training data to coupled (large structure) predictions. The authors validate the proposed method in several tasks.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> Re1: the line 145, 197, right column, \"z=(z1,z2,...,zn)\" and \"V=v1∪v2∪...∪vn\", should use capital \"N\" instead of \"n\" ?\n\n**Answer**: Thank you for pointing this out. We wll make the correction.\n\n> Re2: The statement of z\\_{i}^{e} in eq.9 is not clear enough. For readers who are familiar with diffusion, it's smooth to get the intuition of \" estimate z\\_{i}^{0} when we at still at z\\_{i}^{s} \" there. However, for readers who are not very familiar with diffusion, the mixture usage of z\\_{i}^{s} and z\\_{i}^{e} could lead to confusion.\n\n**Answer**: Thank you for your suggestion. We will add an explanation of the superscript “e” in the paper to indicate that it is an estimated value, distinguishing it from the subscripts $i$ and $s$.\n\n> Re3: In the Algorithm1 table, the inner loop (step 8-11) is set over the \"i\", but before that, the notation of \"i\" is already used in step 1-7. My understanding is the operations involved \"i\" in step 1-7 should be applied over i=1,2..N. Authors should make it clear.\n\n**Answer**: Thank you for pointing this out. Indeed, that is the meaning we intended to convey. We will add a statement clarifying that the operations in steps 1–7 are applied to all $i=1,\\ldots,N$. You can see in: https://anonymous.4open.science/r/MultiSimDiff-D5A3/rebuttul/ALG1.png\n\n> Re4: While the compositional approach is compelling, the choice of conditioning structure is manually designed and attached to single task. I'm curious whether this method can self-adapt to different partitioning schemes for new physical systems, and whether the learned decoupled component (represent. by diffusion model) has the generalizability to other task.\n\n**Answer**: Your understanding is correct. For multi-physics problems, the choice of conditional structure does indeed depend on the specific task, because the conditional diffusion model needs the other physical fields in order to predict the current physical field. As such, the set of physical fields in the system is already fixed. If a new physical field is added or an existing one is removed, any physical field model that is coupled to it must be retrained. This is indeed a very interesting open question, and we will look into it in future work. For multi-component simulation, the current method can be extended to various combinations of components.", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "No  theoretical claims.", "experimental_designs_or_analyses": "Yes, the paper provides ablation studies on diffusion parameters and crucial hyper-paprameters, but a more systematic hyperparameter study (e.g., different architectures, sampling steps) would be encouraged.", "relation_to_broader_scientific_literature": "This work provides valuable insights and new ways of using generating models to handle real-world physical problems.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "Strength: \n\n- The paper frames multiphysics and multi-component simulation as a generative modeling problem, leveraging diffusion models to compose energy-based representations. This is a pretty fresh and novel perspective on simulation-based learning. \n\n- The setting of experiment, especially the prismatic fuel element part, is solid and impressive, which  demonstrates strong generalization from small to large structures, which is crucial for real-world engineering applications.\n\nWeaknesses:\n\n- The presentation of the method still could be improved. There are some typos and inconsistent of notations, and more preliminary on the training part could help readers to better understand the work. For example: \n\n  - the line 145, 197, right column,  \"z=(z1,z2,...,zn)\" and \"V=v1∪v2∪...∪vn\", should use capital \"N\" instead of  \"n\" ?   \n\n   - The statement of z_{i}^{e} in eq.9 is not clear enough. For readers who are familiar with diffusion, it's smooth to get the ituition of \" estimate z_{i}^{0} when we at still at  z_{i}^{s} \" there. However, for readers who are not very familiar with diffusion, the mixure usage of z_{i}^{s} and z_{i}^{e} could lead to confusion.\n\n  - In the Algorithm1 table, the inner loop (step 8-11) is set over the \"i\", but before that, the notation of \"i\" is already used in step 1-7. My understanding is the operations invloved \"i\" in step 1-7 should be applied over i=1,2..N. Authors should make it clear.", "other_comments_or_suggestions": "While the compositional approach is compelling, the choice of conditioning structure is manually designed and attached to single task. \nI'm curious whether this method can self-adapt to different partitioning schemes for new physical systems, and whether the learned decoupled component (represnt. by diffusion model) has the generalizability to other task.", "questions_for_authors": "see other parts", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SOfHny0HiA", "reviewer": "Reviewer_QAdZ", "summary": "The paper proposed a compositional diffusion model framework to handle multi-physics multicomponent surrogate model for physics systems. It leverages the existing diffusion backbone and demonstrated it effects on three multi-physics/multicomponent PDE systems. The paper claims the contribution are: 1. Introducing the compositional diffusion model for multi-physics and 2. Creating novel benchmark dataset for multi-physics/multi-components surrogate models research. The metrics comparison also indicates the effectiveness of the proposed model compared with baselines.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> Re1: What is the novelty for your proposed diffusion framework?\n\nOur study is application-driven rather than aiming to improve diffusion models directly. Our contribution is not diffusion model itself, but the higher-level algorithm on top of diffusion models for multiphysics and multi-component simulation tasks. Specifically, our framework can learn from single-field or small-scale data and generate complete, coupled solutions or full-structure predictions at inference. This cuts development costs for coupled solvers, simplifies large-scale simulations, and extends readily to more complex scenarios.\n\n> Re2: Why use diffusion models for deterministic PDE systems?\n\nWe believe diffusion models can be used to simulate deterministic PDE systems because a deterministic structure can be considered a Gaussian distribution with a fixed mean and a very small variance—effectively negligible. **We provide experimental evidence showing this variance is very small** (see [link](https://anonymous.4open.science/r/MultiSimDiff-D5A3/rebuttul/std.md) and Re6 in our response to Reviewer WGGT). Moreover, when physical systems involve noise, diffusion models can also help quantify predictive uncertainty. Several studies have applied diffusion models to simulate physical systems, such as in weather prediction [1] and spatiotemporal field forecasting [2]. In our experiments, our baseline is effectively a deterministic surrogate model, which generally performs better on validation sets (decoupled data / small structures). However, our main objective is to do well on the test sets (coupled / large structures), and our experiments confirm that the diffusion model outperforms this deterministic baseline in that scenario. In summary, we believe diffusion models are a viable approach for modeling PDE systems.\n\n[1] Mardani M, et al. Residual corrective diffusion modeling for km-scale atmospheric downscaling. Commun. Earth Environ. 2025.\n\n[2] Li Z, et al. Learning spatiotemporal dynamics with a pretrained generative model. Nat. Mach. Intell. 2024.\n\n> Re3: Could you compare with more baselines like domain-decomposition and GNN-based network?\n\nDomain decomposition in numerical simulation splits the computational domain rather than multiple physical fields. Therefore, machine learning approaches derived from domain decomposition are not used to solve multi-physics problems. Hence we use such methods only for multi-component baselines. For multi-component simulations, our manuscript already compares domain decomposition (surrogate), Graph Neural Networks (GIN), and Graph Transformer (SAN). Based on your suggestion, we further added MeshGraphNet for comparison and tuned certain hyperparameters (e.g., hidden layer size, number of message passing layers). The result is shown in: https://anonymous.4open.science/r/MultiSimDiff-D5A3/rebuttul/Table3.png. Indeed, MeshGraphNet is a very strong baseline. We find that on the 16-component dataset used to train MeshGraphNet, it achieves performance significantly better than any of the other models; in the 64-component dataset, our method still performs the best, demonstrating that our method’s strong multi-component generalization capability.\n\n\n> Re4: What are the error distributions respective to spatial and temporal dimensions.\n\nCurrently, our predictions are jointly conducted in both time and space. Therefore, theoretically, the errors in space and time should be similar. For multiphysics modeling, the results are shown in the figure at the following link: https://anonymous.4open.science/r/MultiSimDiff-D5A3/rebuttul/Fig5.png. Overall, due to the coupled physical fields, the spatial errors at the interfaces are slightly higher, while the temporal errors remain close. The flow field initially has larger errors, which decrease as the flow stabilizes. For multi-component simulations (see Fig.9 in the original submission), larger errors typically occur at the component boundaries.\n\n> Re5: Convergence criteria and accuracy at the interface.\n\nWe did not deliberately loosen the convergence criteria to slow down the numerical simulation. We used MOOSE to create our dataset, and the code repository includes the input files used to generate the data. The convergence criteria include an outer nonlinear iteration loop and an inner linear iteration loop:\n- Experiment 2: Nonlinear iteration has a relative error of 1e-8 (default 1e-8) and absolute error of 1e-8 (default 1e-50), with a maximum of 20 iterations (default 50). The linear iteration has a relative tolerance of 1e-5 (default 1e-5) and a maximum of 100 steps (default 10000).\n- Experiment 3: The maximum number of linear iterations is set to 20, and the maximum number of nonlinear iterations is 5, with other settings left as default.\n\nCompared to the program’s default settings, we have actually tried to make the simulation faster wherever possible.\n\nRegarding accuracy, the spatial errors at the interfaces are somewhat higher, as analyzed in Re4.", "claims_and_evidence": "The claims that it introduces a effect multi-physics/multicomponents diffusion framework for physics application is supported by the result section. Moreover, some of the datasets are novel but not all of them. e.g, reaction-diffusion dataset is a standard questions researchers have used.", "methods_and_evaluation_criteria": "I don't think the proposed methods make sense for multiphysics/multicomponents surrogates. Firstly, the diffusion model itself doesn't have novelty. It just leverages the existing pipeline for a new dataset. Secondly, I am not convinced that the diffusion model is a top candidate for modeling multiphysics/multicomponents process. The underlying physics are all deterministic. And there exists extensive literatures based on domain decompositions or GNNS to build deterministic surrogates. Moreover, the success of the multicomponents model is more relied on the self-similar feature of the graphical data and the model that could capture these resolution invariants features, instead of the diffusion model itself. Moreover, the diffusion model is also not used to account for any uncertainty quantifications for the surrogate models in the current context to help explain the sample variety caused by diffusion models. Therefore, I am not convinced by the novelty contribution and the justification of using diffusion model for current applications.", "theoretical_claims": "I have checked the equations in the main manuscript and they are correct with minor typos to the best of my knowledge.", "experimental_designs_or_analyses": "I have checked the soundness and validity of experimental designs and have found issues.\n\nFor the multiphysics experiments. I think it is needed to compare with other domain-decomposition based surrogate models to justify the effectiveness of the papers. I also think it is beneficial to show spatial/temporal error plots to give better idea of how the coupling surface works more than a simple number.\n\nFor the multicomponents experiments, I think the GNN based benchmarks is not optimized for performance. I have reviewed research papers that leverage GNN to successfully spatially extrapolate to larger domains. I will list the reference below.", "relation_to_broader_scientific_literature": "I think the contribution lies in trying to extend the diffusion based framework to multi-physics/multi components surrogate modeling, altought I am not convinced by the current script. Domain decomposition and multi components (train small, test big) ideas have been explored in separate works and the authors tried to show that a diffusion model could be a unified work for these tasks.", "essential_references_not_discussed": "I think the several essential references about the benchmark needs to be included in the result comparison part. For example, [1] paper demonstrated that GNN itself could work very well on training on 64 by 64 domain and directly generalizes well to 1024 to 1024 domain for physics systems. However, the current paper reports that GNN-based model all performs very poor. Moreover, the author mentioned domain-decomposition based method like [2], but didn't compare with them in the result part.\n\n[1] Fan, Shaoxun, et al. \"Accelerate microstructure evolution simulation using graph neural networks with adaptive spatiotemporal resolution.\" Machine Learning: Science and Technology 5.2 (2024): 025027.\n\n[2] Ranade, R., Hill, C., He, H., Maleki, A., Chang, N., andPathak, J. A composable autoencoder-based iterativealgorithm for accelerating numerical simulations. CoRR,abs/2110.03780, 2021. URL https://arxiv.org/abs/2110.03780.", "other_strengths_and_weaknesses": "Strength is trying to building a unified surrogate model for multiphysics/multicomponents systesm. These applications are time consuming and there is a strong research need to get faster surrogate. Moreover, the paper writing is clear to follow the key idea. The weakness is it seems directly apply the existing pipeline on a new dataset, without fair comparison to the high performance benchmarks.", "other_comments_or_suggestions": "None", "questions_for_authors": "I have briefly mentioned my questions in the above sections but I would like to summarize it here.\n\n1. What is the novelty for your proposed diffusion framework?\n\n2. Why do you think it is justified to use diffusion model to model the deterministic PDE systems. The challenges for multi physics is how to model the interface accurately while not being too expensive. And the challenges of multi-components lies in learning a scale-invariant features leveraging GNN-like NNs. None of these are related to the strength of diffusion models. \n\n3. Could yo compare with more baselines. For multi physics, it is better to include more advanced domain-decomposition based surrogate. And for the multicomponents, it is better to compare with existing GNN-based framework that alright shows great potential on doing the spatial extrapolation tasks. It could be related but not limited to the reference I mentioned above.\n\n4. What are the error distributions respective to spatial and temporal dimensions. \n\n5. As for the speed up, the numerical simulation cost heavily depends on the convergence criteria set in the interface. Could you report what is your convergence or residual set for numerical simulation. And the mean time, how accurate your diffusion surrogate is in the interface?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9qbwWCZQAr", "reviewer": "Reviewer_hPBr", "summary": "The paper proposes a compositional multiphysics and multicomponent simulation model that uses diffusion. The approach, MultiSimDiff, consists of learning conditional distributions of individual components conditional on the other physical processes. At test-time, the approach samples from the conditional distributions to build a sample from the joint distribution. The paper tests their approach on three tasks: reaction-diffusion, nuclear thermal coupling, and prismatic fuel thermal and mechanical analysis. The paper’s reported contributions are: a novel approach to multi-physics / multi-component modeling by framing in terms of a joint distribution; open-source data sets; demonstrating success for new simulation models.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> Re1: Relation to Gibbs Sampling.\n\nGibbs sampling splits a problem into conditional distributions and samples them sequentially. In our setting, this is essentially our baseline (named **surrogate model** in manuscript), where we iteratively update each physical field using a surrogate model (See Algorithm 3 in the manuscript). Specifically, in a multiphysics problem, if you randomly initialize all physical fields as:\n$\\mathbf{z}^{(0)} = \\bigl(z_1^{(0)}, z_2^{(0)}, \\ldots, z_d^{(0)}\\bigr)$,\nthen at each iteration $k$, you loop over each physical field $z_i$ and update it with the conditional probability $p(z_i \\mid z_{\\neq i})$. If that conditional probability is replaced by our trained surrogate model, it essentially becomes Algorithm 3.\n\n> Re2: Open-Source Dataset/Code.\n\nWe verified the anonymous repository is accessible. It contains solver input files and download links for the data, mentioned at line 294. We also provide the url here: https://anonymous.4open.science/r/MultiSimDiff-D5A3/README.md.\n\n> Re3: Why the validation dataset consists of decoupled data and not coupled data when it is known that the test data is coupled data...\n\nWe believe this is fair for both the surrogate model and the diffusion model. For multiphysics problems, the goal is to train the model on decoupled data but eventually predict the coupled solution. Consequently, the training and validation sets are composed of decoupled data, and the test set is composed of coupled data. Both models use the same training, validation, and test sets.\n\n> Re4: When is it appropriate to split up a problem into learning conditional distributions? Can all the problems be split into the conditional distributions?\n\nOur tasks naturally decompose into conditional fields—each physical field or component is determined by the others. In multi-physics, each field typically has its own solver requiring boundary information from other fields. Similarly, in multi-component setups, a component’s solution depends on its neighbors. But there is a type of eigenvalue problem that cannot be solved by current methods, which can be seen in Appendix J.\n\n> Re5: Why simulation-based inference (SBI) is missing.\n\nWe understand SBI as inferring model parameters from given data. Gloecker et al. trained a single diffusion model for $p(x, \\theta)$, thus achieving their “all in one” model to obtain both the likelihood and posterior. But for our problem, due to the difficulty in modeling the joint distribution of multiple physical fields or components, our approach is to combine individual models to achieve similar functionality. Although these two methods look very similar, their directions of operation seem to be opposite.\n\n> Re6: When should a multiphysics approach be split into conditional models, and when should you model the full joint distribution?\n\nModeling the full joint distribution (simultaneously solving all coupled equations) in a multiphysics system is difficult, especially in fields like nuclear engineering, where each physical process (e.g., neutron, fluid, mechanic) often has its own independently developed solver. Building a fully coupled solver to unify these modules is both development-intensive and computationally expensive. Instead, our approach trains simpler conditional distributions – each field conditioned on the others (decoupled solution) – and then composes these models during inference to approximate the fully coupled solution.\n\n> Re7: Does the Sampling Order Matter?\n\nWe have found that in multiphysics, the order of updating each field has negligible influence on final results, as shown in https://anonymous.4open.science/r/MultiSimDiff-D5A3/order.md. For multi-component simulation, we update all components simultaneously, making order irrelevant.\n\n> Re8: If the approach is like EM, does that mean the overall log likelihood always increases?\n\nTheoretically, yes. The log-likelihood can be written as $-E(z) - \\log Z$. Ignoring constants, it is $-E(z)$. Moving in the direction of $\\nabla E$ decreases $E(z)$ and thereby increases the log-likelihood.\n\n> Re9: How much compute is needed to train each individual conditional diffusion model? It looks like there is a different diffusion model for each component.\n\nFor multiphysics with d fields, we train d diffusion models. In Experiment 2 we used three (solid, neutron, fluid). For multi-component tasks, we only need one model because each component is analogous. Even in practical engineering, the number of \"d\" rarely exceeds four. So the training cost is not particularly high.\n\n> Re10: Purpose and Cost of f in Alg. 2?\n\nIn multi-component setups, we update all components concurrently and need function f to gather neighboring solutions for each component, costing O(n) for n components but no extra neural inference, which is very fast. In Algorithm 1 (multiphysics), each field is updated sequentially, and there is no unified function f to update their inputs.", "claims_and_evidence": "Claim 1: “To the best of our knowledge we are the first to introduce a compositional generative approach to Multiphysics and multi-component simulations”\n* Splitting up a problem into its conditionals and sampling from them sequentially has been known for a long time. This is what Gibbs sampling is. The other relevant literature not covered is simulation-based inference. A paper that seems to be tackling a somewhat similar goal in that literature is Gloecker at al. (https://arxiv.org/pdf/2404.09636v1). I would say that potentially the claim could be either grounded in the specific architecture, or in the specific applications shown in the paper. \n\nClaim 2: “We create and open-source benchmark datasets … ”. “The code is available at the anonymous repository”.\n* I might have missed it, but I can’t see a link anywhere to the code or the benchmark datasets. Also, it would potentially be more valuable if it were the simulators as well as the data that were being released. Is that something the authors would consider releasing?\n\nClaim 3: “MultiSimDiff, trained on single components, accurately predicts larger structures with up to 64 components.”\n* I agree with this claim.", "methods_and_evaluation_criteria": "* The evaluation across three novel application datasets makes sense for the work. \n* One question is why the validation dataset consists of decoupled data and not coupled data when it is known that the test data is coupled data. This would seem to put the surrogate models at a disadvantage, since they are being optimized on decoupled data but tested on coupled data.", "theoretical_claims": "* Appendix E shows that there is a gap between models trained on coupled data vs. decoupled. My main question related to the theoretical claims of the paper is when is it appropriate to split up a problem into learning conditional distributions? For example, for Gibbs sampling, it may not be that easy to make this assumption. The Prismatic fuel element seems to make sense as it enables generalization. The paper seems to imply that it is taken for granted that all the problems can be split into the conditional distributions.", "experimental_designs_or_analyses": "* No addition to what has previously been mentioned.", "relation_to_broader_scientific_literature": "* The paper is focused at multiphysics/multi-component simulations. There has not been a large amount of work focused on these multi-component scenarios. This is in part due to a lack of available simulators/datasets.", "essential_references_not_discussed": "* As mentioned earlier, it seems like the area of simulation-based inference (SBI) is missing.", "other_strengths_and_weaknesses": "### Strengths:\n* Paper is well-written, and presentation is good.\n* The application domain is a strength of the paper. The three proposed applications are nice to see and it is not easy to introduce three new examples in a single paper without relying on previous benchmarks.\n### Weaknesses:\n* The novelty of the approach does not seem especially strong given that splitting into conditional distributions is a common procedure for generative modeling and sampling.", "other_comments_or_suggestions": "* See below", "questions_for_authors": "In addition to the above comments:\n* When should a multiphysics approach be split into conditional models, and when should you model the full joint distribution?\n* Is there any advantage to the order of the conditional distributions that are sampled from? For example, in Gibbs sampling, the order of sampling from the distributions can sometimes be varied.\n* If the approach is like EM, does that mean the overall log likelihood always increases?\n* How much compute is needed to train each individual conditional diffusion model? It looks like there is a different diffusion model for each component.\n* Line 8 of Alg. 2 runs an additional update step using the function $f(\\cdot)$. What is the purpose of this function, how costly is it, and why is it not needed for Alg. 1?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper introduces a diffusion-based approach for multiphysics and multi-component simulations. The reviewers found this work valuable for several reasons. A lack of available datasets has so far
limited research into data-driven approaches to such simulations, which previously largely required expensive numerical simulations with tight coupling. 
The authors introduce a set of new datasets and demonstrate that their approach which is trained on decoupled data shows strong generalization from small to large structures.
The reviewers overall found the paper well-written except for some confusing notation which was addressed by the authors during the rebuttal.
Decision: Accept (poster)