Paper: Scaling physics-informed hard constraints with mixture-of-experts (ICLR-2024)

Reviewer comment (Experiments):
Reason A for utilizing Mixture-of-Experts aligns with the phenomenon observed in language tasks [1] where data from different domains can interfere with each other and hinder training. However, the authors did not provide evidence for the occurrence of this phenomenon in dynamic systems, nor did they demonstrate that the improvement of the PI-HC-MoE over PI-HC is due to avoiding such interference. One possible experiment to demonstrate the latter is to randomly assign data points to experts and compare it to the allocation based on spatio-temporal coordinates.

Author response:
**Question: Data from different domains can interfere with each other and hinder training. Occurrence in dynamical systems, and reason why PI-HC-MoE is better than PI-HC?**

We don’t think it is necessarily interference in the NLP sense. Our general motivation for this work is that solving smaller localized hard constraints is easier than imposing a single global hard constraint. Beyond dynamical systems, the reduced size of the non-linear least squares solve makes it an easier system to solve. To some extent, global information (e.g., initial and boundary conditions, which we provide to each expert) is important to finding a unique weighting of the learned basis functions.

Reviewer comment (Experiments):
Challenge B, which concerns the scalability of the non-linear least squares solver limiting the number of sample points (m) and thereby affecting prediction accuracy, can be overcome by utilizing MoE. However, a more straightforward solution is to employ scalable non-linear least squares solvers, e.g., [2][3]. The authors did not comment on this line of work. If there are suitable solvers available, a comparison between PI-HC and PI-HC-MoE with the same number of data points would be desirable.

Author response:
**Question: Using scalable non-linear least squares solvers [1][2].**

We thank the reviewer for bringing up scalable non-linear least squares solvers. [1] is not immediately amenable to our problem setting as they evaluate and target CPU operations. It is non-trivial to port their implementation to a GPU-based ML workflow. DeepLM [2], while promising, performs a series of specific optimizations (including stochastic domain decomposition) that we are not sure are compatible with the implicit function theorem. DeepLM also introduces a custom backward Jacobian network for computing gradients instead of using the implicit function theorem. While we think a comparison using a scalable NLLS would be interesting and a direction to explore in the future, there is, to our knowledge, no suitable solver to provide a comparison for the review period.

We would also like to note that PI-HC-MoE is not mutually exclusive with a different non-linear least squares solver, provided the solver satisfies the implicit function theorem assumptions. In the future, we believe that combining both directions is a promising approach for tackling extremely large and complex dynamical systems (e..g, long time scale 3D Navier-Stokes).

Reviewer comment (Evaluation):
The performance of the baseline appears to be inconsistent with the results reported in [4], for example in the 1D non-linear diffusion-sorption equation results (see [4], Appendix E).

Author response:
**Question: The performance of the baseline appears to be inconsistent with the results reported in PDEBench [3].**

The baselines provided by PDEBench are trained on available numerical solver data. The only baseline which doesn’t do this is PINNs, where they train and test on a single problem (averaged across 10 training/test runs). Their base architecture and training procedure differ from ours. PDEBench uses a DeepONet (akin to an MLP), and auto-differentiation to compute PDE residuals, while we use finite differences (because of the complexity of differentiating through an FNO architecture). They are also not in a neural operator setting (where many PDE parameters to solution mappings are predicted, rather than a single one). We train and evaluate all 3 models in the neural operator setting, and so our problem setting is harder than the PINNs baseline that PDEBench provides.

Reviewer comment (Writing):
There are several errors and confusions in the use of mathematical symbols in the text, including but not limited to:  
(1) in section 3, first paragraph, line 1, $0$ should be in bold.  
(2) in section 3, first paragraph, line 4, the definition of $\mathcal{G}$ and in the second paragraph, line 4, the definition of $f_{\theta}$ — other notation should be used to denote the spaces of $\phi$, $u_\theta$ and $\mathbf{b}$.  
(3) In section 3, second paragraph, line 5, $f^i: \phi \rightarrow \mathbb{R}$ should be $f^i: \Omega \rightarrow \mathbb{R}$.  
(4) In section 3, second paragraph, third to last line, $f_\theta(\phi(x))$ should be $f_\theta(\phi)(x)$.  
(5) On page 4, two lines above Property (1), the definition of $S_{\mathbf{b}}$ and $S_\omega$. Similar errors are repeated throughout the paper and should be corrected.

Author response:
**Minor text typos**

We thank the reviewer for pointing these out. We have updated the main text accordingly.

Reviewer comment (Experiments):
The paper lacks empirical evaluation of the approach on a wider range of physical problems. While the authors demonstrate the effectiveness of their approach on several physical problems, including fluid dynamics, structural mechanics, and quantum mechanics, it would have been beneficial to have more empirical evaluation of the approach on a wider range of physical problems. This would have helped to better understand the generalizability and limitations of the approach.

Author response:
**While the authors demonstrate effectiveness of their approach on several physical problems, including fluid mechanics, structural mechanics, and quantum mechanics, it would have been beneficial to have more empirical evaluation of the approach on a wider range of physical problems.**

We note that the experiments in the paper represent challenging non-linear problems, and demonstrate both the improved accuracy and efficiency of our method. Nevertheless, we have added another experiment on another challenging non-linear system: the 2D reaction-diffusion system (similar to Navier-Stokes, it is technically a 3D, with 2 spatial dimensions and 1 temporal dimension) to Appendix G. Our results demonstrate a similar pattern as our other experiments, that PI-HC-MoE has the lowest error when compared to PI-HC and PI-SC.

Reviewer comment (Theory):
The paper does not provide a detailed theoretical analysis of the approach. While the authors provide a detailed description of the approach, they do not provide a detailed theoretical analysis of the method. This could have helped to better understand the underlying principles and assumptions of the method, and to provide a more solid foundation for future research.

Author response:
**Theoretical analysis of approach.**

As our method incorporates iterative solvers within NNs, much of the relevant theoretical analysis here is on error analysis and convergence for non-linear least squares iterative solvers. Since we only use Levenberg-Marquadt (LM), we cite a few relevant references.

LM finds $x\in \mathbb{R}^n$ such that $F(x): \mathbb{R}^n \mapsto \mathbb{R}^m = \mathbf{0}$. Let $\hat{x}$ be such an optimal solution. If $n = m$, the Jacobian at $F(\hat{x})$ is nonsingular, and the initial guess for $x$ is sufficiently close to $\hat{x}$, LM has a quadratic rate of convergence [1]. There are a few works which establish superlinear convergence [2] and error bounds [3]. This is by no means exhaustive, but a sample of some existing literature to the best of our knowledge.

Reviewer comment (Theory):
The paper does not discuss the limitations of the approach. While the authors discuss the advantages of their approach, they do not discuss the limitations of the method, such as the assumptions made about the physical system or the potential impact of the choice of experts on the overall performance. This could limit the ability of readers to fully understand the applicability and limitations of the approach.

Author response:
**The paper does not discuss the limitations of the approach.**

We have included a limitations section in Appendix H.

Reviewer comment (Theory):
Unless I am mistaken, there is no clear formulation of the output function $u(x, t)$ at inference except in Figure 1, and if I understand the figure correctly, then $u(x, t) = \sum_k b(x,t)^T w_k = b(x,t)^T( \sum_k w_k)$. In this case, the sum of $w_k$ is the weights used to query the function over the domain $\Omega$, and as $w \neq w_k$ a priori, we do not know if the constraints are hardly imposed on any sampled points. Therefore, at inference I do not think that the PDE can be constrained in a hard setting with this architecture.

Author response:
**Formulation of the output function $u(x,t)$ at inference.**

We are summing over the domains after applying the constraint individually to each expert’s subdomain. This is different from $\Sigma_k \omega_k$. For any point in the domain, there is only one $\omega_k$ applicable to any spatiotemporal point. Tangibly, we are concatenating the predictions by each expert to produce the final $u_\theta$. To clarify this idea, we have added an example inference step for the 1d diffusion-sorption in Appendix B that details the exact steps taken. Any sampled points are guaranteed to be constrained, since for any given spatiotemporal point, only one expert computes the final prediction.

Reviewer comment (Theory):
I also assume that solving an equation on a domain $\Omega$ with boundaries $\delta \Omega$ is not always equivalent to solving the equation separately on different subdomains and that precautions must be taken. There is no mathematical derivation of the domain decomposition problem, and therefore we do not know if the solution found with the subdomains is close enough to that of the solution on the full problem.

Author response:
**Solving an equation on a domain $\Omega$ with boundaries $\delta \Omega$  is not always equivalent to solving the equation separately on different subdomains and that precautions must be taken.**

A key point in our setting is that each expert still receives the global boundary and initial conditions. While each subdomain is solving a localized problem, it still has information about the full problem, i.e., the PDE residual is enforced on each subdomain, and consistency between $\omega_k$ is enforced through the known global initial and boundary conditions. Given fixed initial and boundary conditions, the PDE residual can be satisfied point-wise, and solving individual sub-domains is equivalent to solving the global solution. We have also updated the manuscript to reflect that initial and boundary conditions are provided to the non-linear least squares solver.

Reviewer comment (Presentation):
I also understood from Figure 1 that the different domains were not overlapping and represented "chunks" of the domain, but I am puzzled by the grid artifacts of the method in Figure 3. The solution seems to showcase some strange aliasing which would suggest overlapping domains.

Author response:
**Question: Grid artifacts in figure 3.**

This is a good point, and we went back and looked at a number of different solution predictions. In general, our predictions don’t exhibit artifacting, except for a few cases. We’ve added 16 random example predictions from our test set to Appendix D (Figure 9) to show this. We are not quite sure about the exact reason for the occasional artifacting, but it seems to be relatively rare.

Reviewer comment (Experiments):
The experiments are relatively limited - the paper only tests on 2 cases, one for 1D and another for 2D. In each experiment, only one set of environment parameters (e.g., only Reynolds number = 1e4 is used for the Navier-Stokes case) are tested.

Author response:
We note that the two experiments in the main text represent challenging non-linear problems, and demonstrate both the improved accuracy and efficiency of our method. Nevertheless, we have added another experiment on another challenging non-linear system: the 2D reaction-diffusion system (similar to Navier-Stokes, it is technically a 3D, with 2 spatial dimensions and 1 temporal dimension) to Appendix G. Our results demonstrate a similar pattern as our other experiments, where PI-HC-MoE has the lowest $L_2$ relative error (23.940% $\pm$ 3.074%) when compared to PI-HC (29.031% $\pm$ 4.418%) and PI-SC (96.876% $\pm$ 0.088%). Note that given the limited time reviewer period, we did not do that many searches over hyperparameters and train longer (which we will do), but the trend remains the same that PI-HC-MoE is the most accurate, while being much more efficient than PI-HC.

Reviewer comment (Experiments):
What are the temporal intervals used for generating the training data and for testing the model? In the 2D case, “Both the training and test sets have a trajectory length of T = 5 seconds.“. But it’s not clear for the 1D case.

Author response:
1D diffusion-sorption uses a maximum of T=500 seconds [1]. We have updated the main text to make this more clear.

Reviewer comment (Experiments):
How does the model's performance extend to temporal domain outside the training dataset? For example, if we generate the training data within the time frame of [0 seconds, 10 seconds], how well does the trained model predict the system in the interval of [10s, 20] compared with predicting the interval of [0s, 10s]? And how about the spatial domain outside the training dataset?

Author response:
This is an interesting question, and we have included an additional experiment in Appendix C to see how well the models generalize to future timesteps for the diffusion-sorption case. Here, the models are trained on the temporal domain of $t$ = [0 s, 400 s], and we report the $L_2$ relative error for predictions on $t$ = [400 s, 500 s] (a temporal region unseen during training).

The average error across test samples for PI-SC is 805.58% $\pm$ 19.06%, PI-HC is 48.047% $\pm$ 1.23%, and PI-HC-MoE is 14.79% $\pm$ 3.00%. PI-SC is completely unable to generalize to later timesteps. While PI-HC performs better than PI-SC, PI-HC-MoE generalizes the best to future time steps outside the training set. We hypothesize that one reason for this is because PI-HC and PI-HC-MoE both enforce the PDE equation constraints at test-time and so, as expected, generalize better than PI-SC. Finally, we see that, similar to our other experiments, PI-HC-MoE is more scalable and expressive than standard differentiable optimization (represented by PI-HC). While one advantage of the hard constraint is that constraints can be enforced at inference time, and so should do better for later timestep predictions not in the training set, this comes at the cost of computational expense, which PI-HC-MoE is able to address.

Reviewer comment (Experiments):
The paper only gives examples in 1D and 2D, what are the potential risks of using the proposed model for predicting 3D systems? Alternatively, can it be seamlessly adapted for 3D applications?

Author response:
Our method is model agnostic and there is no inherent limitation to extending PI-HC-MoE to 3D scenarios. Note that technically, our 2D Navier-Stokes evaluation and 2D reaction-diffusion are 3 dimensional (2 spatial, 1 temporal). For 3D spatial settings (4D with time), one limiting factor for all current base NN architectures with generally good expressivity (such as FNO) is the memory requirement, even for a soft constraint. Since our paper is meant to show a proof-of-concept for employing MoE for hard physical constraints, a thorough search over base NN architectures that meet memory requirements (while being reasonably expressive) is something we leave to future work.

Reviewer comment (Writing):
The paper is mostly quite clear, but it would have been beneficial to have an ongoing example where it is made clear what the outputs of the neural networks are, what the constraints are, and how the problem is divided in that case.

Author response:
We have added an example in Appendix B that walks through the inference procedure of PI-HC-MoE and PI-HC in the diffusion-sorption setting. We hope this example is illustrative to understanding our method.

Reviewer comment (Experiments):
The authors compare the execution time of PI-HC vs PI-HC-MoE, but what about PI-SC?

Author response:
For our execution time experiments, we compared PI-HC and PI-HC-MoE to establish the improved runtime and scalability of our method (compared to standard differentiable optimization). PI-SC is less relevant here because the accuracy is much worse. However, PI-SC, due to the lack of any imposed hard constraint at both training and test time, is faster than PI-HC and PI-HC-MoE. 

On 1D diffusion-sorption, the training time per iteration for PI-SC is 0.0595 (s), compared to 0.294 - 0.990 (s) for PI-HC-MoE. The inference time for PI-SC is 0.0496 (s), compared to 0.034 - 0.729 (s) for PI-HC-MoE. On 2D Navier-Stokes, the training time per iteration for PI-SC is 0.010 (s), compared to 0.427 - 7.219 (s) for PI-HC-MoE. The inference time for PI-SC is 0.005 (s), compared to 0.155 - 5.016(s) for PI-HC-MoE. We note that PI-HC-MoE is much closer to PI-SC in run time (and much faster than PI-HC), while being significantly more accurate.

Reviewer comment (Presentation):
The scales in Figure 3 are not clearly understood. If the row below represents the difference between the predictions made by the NN and the prediction made by the numerical solver, how is it possible to have a scale between -3 and 3?

Author response:
PI-SC predicts values in the range of [-3, 3]. To prevent excessive colorbars and for consistency across the top row of predictions, we clip the PI-SC prediction to [0, 1]. We do not perform the same clipping when plotting the bottom difference plots, leading to a scale of [-3, 3] for the difference plot of PI-SC. We have also updated the figure caption to make this clear.

Reviewer comment (Evaluation):
The authors compare their method against PI-HC and PI-SC in terms of computation time, but what about space? Is the memory requirement of PI-SC-MoE also advantageous and/or comparable?

Author response:
There is indeed a difference between the memory requirements of PI-HC and PI-HC-MoE. In order for PI-HC to sample a sufficiently large number of points, we were forced to decrease the batch size. With too many sampled points, the corresponding batch size that fits in memory for PI-HC leads to training instability and does not converge. For the same batch size, PI-HC runs out of memory more quickly than PI-HC-MoE.