Scaling physics-informed hard constraints with mixture-of-experts

We thank you for your thoughtful comments and appreciate the acknowledgment of the improvements our work provides over the baselines. We answer your questions below and have made changes to the manuscript, noting your comments on the mathematical notation.

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.

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

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.

Minor text typos

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

[1] Cartis, C., Roberts, L. Scalable subspace methods for derivative-free nonlinear least-squares optimization. Math. Program. 199, 461–524 (2023). https://doi.org/10.1007/s10107-022-01836-1

[2] J. Huang, S. Huang and M. Sun, "DeepLM: Large-scale Nonlinear Least Squares on Deep Learning Frameworks using Stochastic Domain Decomposition," 2021 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), Nashville, TN, USA, 2021, pp. 10303-10312, doi: 10.1109/CVPR46437.2021.01017.

[3] Takamoto, M., Praditia, T., Leiteritz, R., MacKinlay, D., Alesiani, F., Pflüger, D., & Niepert, M. (2022). PDEBENCH: An Extensive Benchmark for Scientific Machine Learning. ArXiv, abs/2210.07182.

We are glad to see that you noted that our work offers a fresh and effective way to use MoE, the advantages of our novel approach, and importance for real-world applications. While most of your questions are already answered in our paper, we have additionally replied to these questions below with the relevant section. We have also added another experiment to our paper, and all updates are in blue.

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.

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.

The paper does not discuss the limitations of the approach.

We have included a limitations section in Appendix H.

Question: What are the benefits of imposing known physical constraints during neural network training?

As we discuss in our introduction (Section 1), many problems necessitate modeling the physical world, which is governed by a set of established physical laws. The consistency of these physical laws means that they can provide a strong supervision signal for NNs and act as inductive biases. As we demonstrate, adding these physical laws as hard constraints in a scalable manner greatly improves accuracy, training and inference speed, and data efficiency.

Question: How does the MoE approach differ from standard differentiable optimization in enforcing hard physical constraints?

As we discuss in our methods section (Section 3), our MoE approach imposes the constraint over smaller decomposed domains. Each domain is solved by an “expert” through differentiable optimization. During training, each expert independently performs a localized backpropagation step by leveraging the implicit function theorem, and the independence of each expert allows for parallelization across multiple GPUs. Compared to standard differentiable optimization, our scalable approach achieves greater accuracy on challenging non-linear systems. Compared to standard differentiable optimization, our approach improves both training stability and requires significantly less computation time during both training and inference stages.

Question: Can you provide an example of a complex dynamical system where this scalable approach would be particularly useful?

Our scalable approach is particularly useful for challenging, complex dynamical systems that are expensive to solve with numerical solvers. As we see in our experiments section (Section 4), at inference time, our method is much faster than a numerical solver.

[1] Yamashita, N., Fukushima, M. (2001). On the Rate of Convergence of the Levenberg-Marquardt Method. In: Alefeld, G., Chen, X. (eds) Topics in Numerical Analysis. Computing Supplemental, vol 15. Springer, Vienna. https://doi.org/10.1007/978-3-7091-6217-0_18

[2] Ipsen, I. C. F., C. T. Kelley, and S. R. Pope. “Rank-deficient nonlinear least squares problems and subset selection.” SIAM Journal On Numerical Analysis (2011). https://www.jstor.org/stable/23074331

[3] Christian Kanzow, Nobuo Yamashita, Masao Fukushima. “Levenberg–Marquardt methods with strong local convergence properties for solving nonlinear equations with convex constraints.” Journal of Computational and Applied Mathematics, Volume 172, Issue 2, 2004. https://doi.org/10.1016/j.cam.2004.02.013.

Thank you for your comments, and we are glad to see that you appreciated the performance and scalability of our method. We answer all questions below, and have also made updates to the manuscript (in blue), including additional solution visualizations and PDE residual plots.

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.

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.

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.

Question: Did you explore aspects of each basis function?

We did try to look at each basis function, but they are not very interpretable (a general problem with NNs), as they are akin to NN activations. However, we did an experiment where we looked at the performance improvement from our learned basis functions vs. doing an interpolation over the constrained points in Appendix E. The interpolated solutions have a higher relative $L_2$ error when compared to the predictions of PI-HC-MoE, indicating that the basis functions learned have information about the underlying dynamics of the system.

Question: Did you check the PDE residual of sampled points to see if the constraint was respected?

We did, and we have added plots of the PDE residuals on the validation set during training to the paper in Appendix F. As we see, in general, the PDE residual is quite low for PI-HC-MoE.

Question: Could you provide more explanations on the way to compute $\frac{\partial z^\*}{\partial \theta}$?

Equation 1 in the paper defines a set of non-linear equations where $\frac{\partial z^\*}{\partial \theta}$ is unknown. We have labeled which sections are computed via auto-differentiation in the updated manuscript. Because equation 1 defines a system of variables, $\frac{\partial z^\*}{\partial \theta}$ may be solved with another non-linear least squares solve during the backward pass (i.e., each training iteration performs two solves).

Additionally, we have updated section 3, incorporating the feedback from multiple reviewers, to make the notation more succinct and clear. We hope that this is more informative.

Thank you for your response. Indeed, we desire to solve for $\frac{\partial F_\phi(b \cdot \omega^T)}{\partial\theta}$ which requires $\frac{\partial z^*}{\partial \theta}$. We have added more details to Appendix I (explicitly listing the system that we solve), and also include this additional information here:

The implicit function theorem allows us to compute $\frac{\partial z^*}{\partial \theta}$ with a non-linear least squares system (Eqn. 1 in the main text). Specifically, using property (1) and differentiating property (2) yields:

$\frac{\partial F_\phi(\mathbf{b}\cdot\omega^T)}{\partial \theta} = \frac{\partial F_\phi(\mathbf{b}\cdot z^*(\mathbf{b})^T)}{\partial \theta} = \frac{\partial F_\phi(\mathbf{b} \cdot \omega^T)}{\partial \mathbf{b}} \cdot \frac{\partial \mathbf{b}}{\partial \theta} + \frac{\partial F_\phi(\mathbf{b}\cdot z^*(\mathbf{b})^T)}{\partial z^*(\mathbf{b})} \cdot \frac{\partial z^*(\mathbf{b})}{\partial \theta} = \mathbf{0}$.

Rearranging terms: $\frac{\partial z^*(\mathbf{b})}{\partial \theta} = \left[ \frac{\partial F_\phi(\mathbf{b}\cdot z^*(\mathbf{b})^T)}{\partial z^*(\mathbf{b})} \right]^{-1} \cdot - \frac{\partial F_\phi(\mathbf{b} \cdot \omega^T)}{\partial \mathbf{b}} \cdot \frac{\partial \mathbf{b}}{\partial \theta}$.

In the rearranged equation, all terms in the right-hand side (sans the matrix inverse) can be computed using auto-differentiation, and this defines a non-linear system of equations, where we solve for $\frac{\partial z^*(\mathbf{b})}{\partial \theta}$ using a non-linear least squares solver. Note that the right hand side involves a matrix inverse, which is computationally intractable to exactly compute, which is why we use a non-linear least squares solver. This allows us to get the final expression for $\frac{\partial z^*(\mathbf{b})}{\partial \theta}$.

As mentioned earlier, we have added both the above and further additional details to Appendix I in the paper. Please let us know if you have any additional questions.

Thank you for your comments, and for noticing the scalability and accuracy of our method. We respond to your questions below, and have also made updates to the manuscript (in blue) which include experiments on another system (2D reaction-diffusion), and an experiment testing predictions on future timesteps that the models were not trained on.

Experiments are relatively limited.

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.

Question: Temporal intervals for training data and for testing for diffusion-sorption.

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

Question: How do the models generalize to future time steps outside the training dataset?

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.

Question: Can PI-HC-MoE be extended to 3D scenarios?

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.

[1] Makoto Takamoto, Timothy Praditia, Raphael Leiteritz, Daniel MacKinlay, Francesco Alesiani, Dirk Pflüger, and Mathias Niepert. PDEBench: An Extensive Benchmark for Scientific Machine Learning. Advances in Neural Information Processing Systems, volume 35, pp. 1596–1611. 2022.

Thank you for your comments, we’re glad to see that you found our work novel and noted that it will be of high interest to a subset of the ML community. We respond to your comments below, and also made updates to the manuscript (in blue) to include an ongoing example.

Paper mostly clear, ongoing example: outputs of NN, constraints, how problem is divided.

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.

Question: Experimental analysis is comprehensive. Are PI-SC and PI-HC the only models against which you can compare?

We appreciate the recognition of our experimental analysis. PI-SC and PI-HC are the only settings we compare to because our method is primarily centered around the constraint optimization procedure, as PI-HC-MoE is agnostic to the base NN architecture. The soft constraint is the most common setting to compare against. As one of our contributions is to improve the scalability of the hard constraint setting represented by standard differentiable optimization, we also compared against this (PI-HC).

Question: The authors compare the execution time of PI-HC vs PI-HC-MoE, what about PI-SC?

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.

Question: Figure 3 scaling.

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.

Question: Memory requirements of PI-HC vs PI-HC-MoE.

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.

Dear reviewers,

As the discussion period comes to an end soon, we believe that we have addressed all your comments through a number of new experiments and updates to our manuscript. We would be happy to answer any additional questions. Thank you!