M2PDE: Compositional Generative Multiphysics and Multi-component PDE Simulation

Re1: Relation to Gibbs Sampling.

Gibbs 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: $\mathbf{z}^{(0)} = \bigl(z_1^{(0)}, z_2^{(0)}, \ldots, z_d^{(0)}\bigr)$, then 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.

Re2: Open-Source Dataset/Code.

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

Re3: Why the validation dataset consists of decoupled data and not coupled data when it is known that the test data is coupled data...

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

Re4: When is it appropriate to split up a problem into learning conditional distributions? Can all the problems be split into the conditional distributions?

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

Re5: Why simulation-based inference (SBI) is missing.

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

Re6: When should a multiphysics approach be split into conditional models, and when should you model the full joint distribution?

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

Re7: Does the Sampling Order Matter?

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

Re8: If the approach is like EM, does that mean the overall log likelihood always increases?

Theoretically, 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.

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.

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

Re10: Purpose and Cost of f in Alg. 2?

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

Re1: What is the novelty for your proposed diffusion framework?

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

Re2: Why use diffusion models for deterministic PDE systems?

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

[1] Mardani M, et al. Residual corrective diffusion modeling for km-scale atmospheric downscaling. Commun. Earth Environ. 2025.

[2] Li Z, et al. Learning spatiotemporal dynamics with a pretrained generative model. Nat. Mach. Intell. 2024.

Re3: Could you compare with more baselines like domain-decomposition and GNN-based network?

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

Re4: What are the error distributions respective to spatial and temporal dimensions.

Currently, 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.

Re5: Convergence criteria and accuracy at the interface.

We 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:

• 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).
• 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.

Compared to the program’s default settings, we have actually tried to make the simulation faster wherever possible.

Regarding accuracy, the spatial errors at the interfaces are somewhat higher, as analyzed in Re4.

Re1: the line 145, 197, right column, "z=(z1,z2,...,zn)" and "V=v1∪v2∪...∪vn", should use capital "N" instead of "n" ?

Answer: Thank you for pointing this out. We wll make the correction.

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.

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

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.

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

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.

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.

Re1: The end of the abstract (3 last sentences) appears a bit narrow and verbose.

Thank you for the suggestion! We will revise the ending of the abstract as follows:

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

Re2: Would be nice to mention better ... sampled from $p_{\neq z_i}$ is not straightforward.

I gather that this section focuses on two issues: (1) Under which scenarios does our algorithm exhibit an advantage? The key factor to train it. (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.

For (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.

For (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.

Re3: I am a bit confused by the baselines you used... Similar remarks for CoAE-MLSim.

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

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.

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

Re5: Figure 1. is pixelized and not very clear (or at least not informative compared to the space it takes), clean...

We have replaced Fig.1 with a clearer version at: https://anonymous.4open.science/r/MultiSimDiff-D5A3/rebuttul/schematic.pdf

Re6: By definition you algorithm will output a distribution over plausible solutions, ... , decreases as a function of the exact setup considered.

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