Scaling physics-informed hard constraints with mixture-of-experts
Authors: Nithin Chalapathi
Yiheng Du
Aditi S. Krishnapriyan
Venue: ICLR 2024 poster (2024)
Keywords: Physics-Informed Machine Learning
PDEs
differentiable optimization
differentiable physics
neural networks
mixture of experts
constrained optimization
neural operators
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: We scale incorporating differentiable PDE-constrained optimization (physics-informed hard constraints) as individual layers in a neural network through a mixture-of-experts formulation.

Abstract: Imposing known physical constraints, such as conservation laws, during neural network training introduces an inductive bias that can improve accuracy, reliability, convergence, and data efficiency for modeling physical dynamics. While such constraints can be softly imposed via loss function penalties, recent advancements in differentiable physics and optimization improve performance by incorporating PDE-constrained optimization as individual layers in neural networks. This enables a stricter adherence to physical constraints. However, imposing hard constraints significantly increases computational and memory costs, especially for complex dynamical systems. This is because it requires solving an optimization problem over a large number of points in a mesh, representing spatial and temporal discretizations, which greatly increases the complexity of the constraint. To address this challenge, we develop a scalable approach to enforce hard physical constraints using Mixture-of-Experts (MoE), which can be used with any neural network architecture. Our approach imposes the constraint over smaller decomposed domains, each of which is solved by an ``expert'' through differentiable optimization. During training, each expert independently performs a localized backpropagation step by leveraging the implicit function theorem; the independence of each expert allows for parallelization across multiple GPUs. Compared to standard differentiable optimization, our scalable approach achieves greater accuracy in the neural PDE solver setting for predicting the dynamics of challenging non-linear systems. We also improve training stability and require significantly less computation time during both training and inference stages.

Published as a conference paper at ICLR 2024
SCALING PHYSICS-INFORMED HARD CONSTRAINTS WITH SEXTURE-OF-EXPERS
Nithin Chalapathi, Yiheng Du, Aditi S. Krishnapriyan {nithinc, yihengdu, ad aditikl keitiiki)beerkeeey..o edu University of California, Berkeley
ABSTRACT
Imposing known physical constraints, such conservation laws, during neural network training introduces an inductive bias that can improve accuracy, reliabil- ity, convergence, and data efficiency for modeling physical dynamics. While such constraints can be softly imposed via loss function penalties, recent advancements in differentiable physics and optimization improve performance by incorporat- ing PDE-constrained optimization as individual layers in neural networks. This enables stricter adherence physical constraints. However, imposing hard con- straints significantly increases computational and memory costs, especially for complex dynamical systems. This is because it requires solving an optimization problem over large number of points in mesh, representing spatial and tem- poral discretizations, which greatly increases the complexity of the constraint. To address this challenge, develop scalable approach to enforce hard physical constraints using Mixture-of-Expee Experrtss (MoE), which can be used with any neural network architecture. Our approach imposes the constraint over smaller decom- posed domains, each which is solved by an "expert" through differentiable optimization. During training, each expert independently performs localized backpropagation step by leveraging the implicit function theorem; the indepen- dence of each expert allows for parallelization across multiple GPUs. Compared to standard differentiable optimization, our scalable approach achieves greater ac- curacy the neural PDE solver setting for predicting the dynamics of challenging non-linear systems. We also improve training stability and require significantly less computation time during both training and inference stages.
INTRODUCTION
Many problems necessitate modeling the physical world, which is governed by set of established physical laws. For example, conservation laws such conservation of mass and conservation of momentum are integral to the understanding and modeling of fluid dynamics, heat transfer, chemi- cal reaction networks, and other related areas (Herron & Foster, 2008). Recently, machine learning (ML) approaches, and particularly neural networks (NNs), have shown promise in addressing prob- lems these areas.
The consistency of these physical laws means that they can provide strong supervision signal for NNs and act as inductive biases, rather than relying on unstructured data. A common approach to incorporate physical laws into NN training is through penalty term in the loss function, essentially acting as soft constraint. However, this approach has drawbacks, as empirical evidence suggests that soft constraints ishmapriyan et al. 2021 Wang et al. 2021). Additionally, at inference time, soft constraints offer no guarantee of constraint enforcement, posing challenges for reliability and accuracy.
make the optimization problem difficult, leading to convergence issues (Kr-)
Alternatively, given the unchanging nature of these physical laws, there lies potential to create im- proved constraint enforcement mechanisms during NN training. This can be particularly useful when there is limited training data available, or a demand for heightened reliability. This notion segues into the broader concept of differentiable physies approach, where differentiation through

---

Published as a conference paper at ICLR 2024
physical simulations ensures strict adherence to the underlying physical dynamics (Amos & Kolter 2017
Qiao et al. 2020 Ramsundar et al. 2021: Kotary et al. 2021).
To solve physical problem, typical practice in scientific computing is to use mesh to represent the spatiotemporal domain. This mesh is discretized, and physical laws can be enforced on the points within the discretization. This can be considered equality constrained optimization problem, and there are lines of work focusing on incorporating such problems as individual layers in larger end- to-end trainable NNs (Négiar et al. 2023; Amos & Kolter, 2017). The advantages of this approach include stricter enforcement of constraints during both training and inference time, which can lead to greater accuracy than the aforementioned soft constraint settings.
However, these approaches, aimed at enforcing constraints more precisely ("hard" constraints), also face number of challenges. Backpropagating through constraints over large meshes highly non-linear problem that grows in dimensionality with respect to the mesh and NN model sizes. This means that can be both computationally and memory intensive. This scenario epitomizes an inherent trade-off between how well the constraint is enforced, and the time and space complexity of enforcing the constraint.
To address these challenges, we propose mixture-of-experts formulation to enforce equality con- straints of physical laws over a spatiotemporal domain. We focus on systems whose dynamies are governed by partial differential equations (PDEs), and impose set of scalable physics-informed hard constraints. To illustrate this, we view our framework through the neural PDE solver setting, where the goal is to learn differential operator that models the solution to system of PDEs. Our approach imposes the constraint over smaller decomposed domains, each of which is solved by an "expert" through differentiable optimization. Each expert independently performs a localized opti- mization, imposing the known physical priors over its domain. During backpropagation, we then compute gradients using the implicit function theorem locally on a per-expert basis. This allows us to parallelize both the forward and backward pass, and improve training stability.
Our main contributions are as follows:
We introduce a physics-informed mixture-of-experts training framework (PI-HC-MoE), which offers scalable approach for imposing hard physical constraints neural networks by differentiating through physical dynamics represented via a constrained optimization problem. By localizing the constraint, parallelize computation while reducing the com- plexity of the constraint, leading to more stable training.
We instantiate PI-HC-MoF. in the neural PDE solver setting, and demonstrate our approach on two challenging non-linear problems: diffusion-sorptionn and turbulent Navier-Stokes. Our approach yields significantly higher accuracy than soft penalty enforcement methods and standard hard-constrained differentiable optimization.
Through our scalable approach, PI-HC-MoE exhibits substantial efficiency improvements when compared to standard differentiable optimization. PI-HC-MoE exhibits sub-lineas scaling as the hard constraint is enforced on an increasing number of sampled points within the spatiotemporal domain. In contrast, the execution time for standard differentiable opti- mization significantly escalates with the expansion of sampled points.
We release our code" to explore and extend the results.
" built using JAX, to facilitate reproducibility and enable researchers
RELATED WORK
Constraints in Neural Networks. Prior works that impose constraints on NNs to enforce prior inductive bias generally fall into two categories: soft and hard constraints. Soft constraints use penalty terms in the objective function to push the NN in a particular direction. For example, physics-informed neural networks (PINNs) (Raissi et al. (019) use physical laws as a penalty term. Other examples include enforcing Lipschitz continuity (Miyato et al.) 2018). Bellman op- timally (Nikishin et al. 2022), turbulence (List et al. (022). numerical surrogates (Pestourie al.
2023).
prolongation matrices (Huang et al. 2021) ), solver iterations (Kelly et al. 2020), spectral
https://gi com Aaep//sK-HHHekkee Herkel leheeeee//e cs- - hawss-arr--oostraaiit nts
2

---

Published as a conference paper at ICLR 2024
methods Du et al. (2023). and convexity over derivatives (Amos et al. 2017 We primarily concerned with hard constraints; methods that, by construction, exactly enforce the constraint at train and test time. Hard constraints can be enforced in multiple ways, including neural conserva-
tion laws (Richter-Powel et al. 2022), PDE-CL (Négiar et al., 2023), BOON (Saad et al.
2023
constrained neural fields (Zhong al. 2023). boundary graph neural networks (Mayr al.
2023
PCL (Xu & Darve, 2022), constitutive laws (Ma et al. (023). boundary conditions (Sukumar &
Srivastava 2022), characteristic layers Braga-Neto (2023). and inverse design (Lu et al. 2021).
Differentiable Optimization. An approach to impose hard constraint is to use differentiable op- timization to solve a system of equations. Differentiable optimization folds in a second optimization problem during the forward pass and uses the implicit function theorem (IFT) to compute gradi- ents. OptNet (Amos & Kolter 2017) provides one of the first formulations, which integrates linear quadratic programs into deep neural networks. DC3 (Donti et al. 202I) propose new training and inference procedure involving two steps: completion to satisfy the constraint and correction to remedy any deviations. A related line of work implicit neural layers (Blondel et al.) [022] Chen iet al. 2018), which use IFT to replace the traditional autograd backprop procedure. Theseus (Pineda al. 2022) and Jat (Blondel al. 2022) are two open-source libraries implementing common linear and non-linear least squares solvers on GPUs, and provide implicit differentiation capabili- ties. Our formulation agnostic to the exact method, framework, or iterative solver used. We focus on non-linear least squares solvers because of their relevance to real-worke problems, though many potential alternative solver choices exist (Berahas et al. 2021 Fang et al. 2024 Na et al. 2022).
Differentiable Physics.
Related to differentiable optimization, differentiable physics (de Avila Belbute- Peres et al. 2018) embeds physical simulations within a NN training pipeline (Ramsundar et | 2021p. Qiao et al. (020)) leverage meshes to reformulate the linear complementary problem when simulating collisions. This line of work useful within the context of reinforcement learning, where computing the gradients of physics state simulator may be intractable. Differentiable fluid state simulators (Xian et al. 2023 Takahashi et al. 202T) are another example.
Mixture-of-Experss. MoE was popularized in natural language processing as a method to increase NN capacity while balancing computation (Shazeer et al.) 2017). The key idea behind MoE is to conditionally route computation through "experts", smaller networks that run independently. MoE has been used in many settings, including vision (Ruiz et al. 2021) and GPT-3 (Brown et al. 2020
3 METHODS
General problem overview. We consider PDEs of the form Fu(u) = O, where u + R is the solution and & may represent variety of parameters including boundary conditions, initial conditions, constant values (e.g., porosity of medium), or varying parameters (e.g., density or mass). Here, Ω defines the spatiotemporal grid. We would like to learn mapping ø + up. where is parameterized by such that Fa(u)) = 0. This mapping can be learned by a neural network (NN). The soft constraint enforces (us) 0 by using penalty term in the loss function. That is, we can minimize (uo))l) as the loss function.
Enforcing Physics-Inforan informed hard constraints in NNs through differentiable optimization. While there are multiple dif- ferent approaches to do this (de Avila Belbute- Peres et al. 2018; Donti et al. [021 Amos & Kolter 2017], our procedure here is most similar to Négiar et al. (2023). Let so Φ + b, where b is a set of N scalar-valued functions: b- [b ba 6a and b"
Hard Constraints in Neural Networks. We enforce physics-
Ω IR. We refer to b as a set of basis functions. The objective of the hard constraint is to find w E RN such that F (b.w") = 0. In other words, w is a linear combination of basis functions b such that the PDE is satisfied. We can solve for (Nielsen & Mad-] sen 2010). In practice, the non-linear least squares solver samples set of 21 21...., Zin fin_poonts in the spatiotemporal domain to evaluate b. This yields m equations of the form (F,(b w (2x)) = O. After solving the non-lineat least squares problem, the final solution operator is (z) for Ω where so is NN parameterized by 0.
w using an iterative non-linear least squares solver, such as
3

---

Published as a conference paper at ICLR 2024
Training Physics-Informed Hard Constraints. Traditional auto-differenniation systems struct a computational graph, where each operation is performed in a "forward pass," and gradients computed in "backward pass" in reverse order using the chain rule. Training a NN through backpropagation with the output of an iterative non-linear least squares solvers requires the solver to be differentiable. However, iterative non-linear least squares solvers not differentiable without unrolling each of the iterations. Unrolling the solver is both computationally expensive (i.e., the computation graph grows) and requires storing all intermediate values in the computation graph.
Physics-Informer Hard Constraints with Implicit Differentiation. Implicit differentiation, us- ing the implicit function theorem (Blondel et al.) 2022), serves as an alternative to standard auto- differentiation. It performs second non-linear least squares solve, bypassing the need to unroll the computation graph of the iterative solver. This forgoes the need to store the entire computational history of the solver. Next, we define how we use the implicit function theorem in the context of training a NN via enforcing our physics-informed hard constraints.
The non linear least squares solver finds such that Fbbbww)) 0. F&(b-w") has non-singular Jacobian Fb w²) By the implicit function theorem, there exists open subsets S&CRaX and SRRR containing b and respectively. There also exists unique continuous function z" : Sb - S. such that the following properties hold true:
== =b))) F =" (b)² 0. x² is differentiable on Sb.
Property (1)
Property (2)
Property (3)
Our goal is to find aF,(bb:)) 00
where O corresponds to NN parameters. This enables to perform
gradient descent 9 to minimize the loss (f) w wT) minimizing the PDE residual. To
compute the gradient of the PDE residual, we use Property (1) and differentiate Property (2):
a. (b. =bb))² 00
aF (bo(b..w] wT)
Ōb 00 +
OF (b 20 =bb)²)
@=*(b)
00 = 0.
(1)
Ub
O2 (b)
computed via auto-ditT.
20 Eq. defines system of equations, where (b) is unknown. We can use the same non-linear least squares solver as the forward pass, concluding the backward pass of training the NN through the hard constraint. Further details on the non-lineet least squares solve in Eq. can be found in $A
Physics-Informed Hard Constraints with Mixture-of-Expert (MoE). Until now, we have fo- cused on single hard constraint with one forward non-linear least squares solve to predict solu- tion, and one backward solve to train the NN. However, using a single global hard constraint poses multiple challenges. In complicated dynamical systems, the global behavior may drastically vary across the domain, making it difficult for the hard constraint to converge to the right solution. In- stead, using multiple constraints over localized locations on the domain can be beneficial to improve constraint adherence mesh is not as simple as using a larger m because the non-linear least squares solve to compute w grows with the number of basis functions N and the number of sampled points m. As a result, there exists a max m and N after which it is impractical to directly solve the entire system. This can be mitigated by using a smaller batch size, but we will show, comes the cost of training stability (Smith et al. 2018 Keskar et al. 2016).
non-sampled points the mesh. Unfortunately, increased sampling on the
To overcome these challenges, we develop a Mixture-of Mfxture-of-EEpprts (MoE) approach to improve the accuracy and efficiency of physics-informed hard constraints. Suppose we have K experts. The spa- tiotemporal domain 5i is partitioned into K subsets Q4 for 1... K, corresponding to each ex- pert. Each expert individually solves the constrained optimization problem (Fo(fo(o) w (x;) = for m sampled points x; E S4 using a non-linear least squares solver. Because the weighting we is computed locally, the resulting prediction b wif' is linear weighting of the global basis functions tailored to the expert's ''s domain 2g. which we find leads to more stable training and faster convergence. Each expert is also provided with the global initial and boundary conditions. Given fixed initial and boundary conditions. the solution that satisfies F (u) may be solved point-wise. and solving the constraint over each expert's domain is equivalent to satisfying the constraint globally.
4

---

Published as a conference paper at ICLR 2024
MoE Router
Expert
b
[b b', 6"1
X, 22
Basis Functions
Find w,:
fo
b on Q
(bw, (bw,7== =
ug(x, 1)
I1 F ll2
[b] b*]
Q
(N.L.L.S.)
= b w
b on Q
Expert
Loss Function
Trainable Parameters
Expert 3
1 p. upp 1
Predicted Functions
Expert
Forward Pass
Input Domain
Backward Pass (IFT)
Non-trainable Differentiable Optimization
Backward Pass (Autograd)
Figure 1: Schematic of PI-HC-MoE in the 2D case. PI-HC-MoE is provided with the spatiotem- poral grid and any PDE parameters (e.g., initial conditions, viscosity, Reynolds Number). fo is a NN parameterized by (blue box), which outputs set of N basis functions b (left-most orange box). b's domain, the same as the green box, is partitioned into the domains Ω4 of each expert by the MoE router. Each expert (purple boxes) solves the non-linear least squares problem defined by Fb. ) O. The resulting wi values are used to produce final solution ug = Σgb.wk. Finally, loss is computed using the L2-norm of the PDE residual (red box). We denote the forward pass with black arrows and the backwards pass with green arrows. Solid green arrows indicate the use of traditional auto-differentiation, while dashed green arrows denote implicit differentiation.
There are multiple potential choices for constructing domain decomposition Q2 and the optimal choice is problem dependent. As simple 2D example, Fig. 1 uses non-overlapping uniform par- titioning along the dimension. We refer to the center orange boxes, which perform the domain decomposition for the experts, as the MoE router. There are two important nuances to Fig. | First, fe. the function mapping & to the basis functions, a mapping shared by all experts. Our setup agnostic to the choice of fe. which can be any arbitrary NN. We use Fourier Neural Operator (FNO) (Li et al. 2021a b). due to its popularity and promising results. Second, each expert performs an iterative non- linear least squares solve to compute wk. and the final output is the concatenated outputs across experts. This leads to h independent non-linear least squares solves, which can be parallelized across multiple GPUs.
At test time, the same domain decomposition and router is used. The number of sampled points, m, and the parameters the non-linear least squares solver may be changed (e.g., tolerance), but the domain decomposition remains fixed.
Forward and backwards pass in Lixture-of-Experet In PI-HC-MoE, the forward pass is domain decomposition in the spatiotemporal grid (see ş for an example forward pass). The domain of each basis function is partitioned according to the MoE router. During the backward pass, each expert needs to perform localized implicit differentiation. We extend this formulation from the single
hard constraint to the MoE case. Each expert individually computes @="(b) over 24. The MoE 88
router reconstructs the overall Jacobian, given individual Jacobians from each expert. To illustrate this, consider decomposition along only axis (e.g., spatial in Fig. TT. The backward pass must reconstruct the Jacobian across all experts:
0g2" (b) [0 (b) 0bz
On (b)].
The realization of the reconstruction operation is performed using indicator function and sum- ming across Jacobians:
0gz bb d (b(z)) Lzeng
(3)
Here, x is point in the domain of expert & (i.e., any spatiotemporal coordinate in 2c). Afterwards, the reassembled Jacobian may be used in auto-differentiation as usual.
4 RESULTS
We demonstrate our method on two challenging non-linear PDEs: ID diffusion-sorption ($..) and 2D turbulent Navier-Stokes 4422 For all problem settings, we base FNO architecture.
5

---

Published as a conference paper at ICLR 2024
PHSC
L. Rel Error (%o on Test
PH-I--Mo tours)
L. (%0 on
10
101
1000
2000
3000
4000
P1-SC
RI-C
PI-CC-MDE jours)
Training Steps
Figure 2: Relative L2 error on the ciffuuiion-soptioon test set. (Left) The L2 relative error on the test set over training iterations. (Right) The final L2 relative error on the test set using the trained models. PI-HC-MoE converges faster and has greater accuracy than the other settings.
We train this model using our method PPhysics-Iffrrme Hard Constraint Mixture-of-Expee Ex-eetspttttt or PI-HC-MoE), and compare to training via a physics-informee informed soft constraint (PI-SC) and physics- informed hard constraint (PI-HC). For PI-HC and PI-HC-MoE, use non-linear least squares solver. For details, see C and
as
Data-constrained setting. In both problems, we exclusively look at Rata-constrained setting where we assume that at training time, we have no numerical solver solution data available (i.e., no solution data on the interior of domain). We so to mimic real-world settings, where can be expensive to generate datasets for new problem settings (see C for further discussion).
Evaluation details. The training and test sets contain initial conditions drawn from the same dis- tribution, but the test set initial conditions are distinct from the training set. To evaluate all models, we compute numerical solver solutions to compare relative L2 errors on ML predictions using the test initial conditions. The PDE residual over training on the validation and Navier-Stokes included in $ For measuring scalability, we use the best trained model for both PI-HC-MoE and PI-HC avoid ill-conditioning during the non-linear least squares solve. The training step includes the cost to backpropagate and update the model weights. All measurements are taken across 64 training or inference steps on NVIDIA A100 GPUs, and each step includes one batch. We report average per-batch speedup of PI-HC-MoE compared to PI-HC.
for diffusion--srpption
4.1 1D DIFFUSION-SORPTONN
We study the ID non-linear diffusion ceffusiion-sorption equation. The diffusion diffussion-sorption system describes absorption, adsorption, and diffusion of a liquid through solid. The governing PDE is defined as:
du(x,t) D
pu(x,t) t)
E (0, 1), E (0, 500].
ØL
R(u(xr1))
Ox2
R(u(I,1). ()) = + 1 1-6
ho bn u(x, 1)m5-1
u(0, 11, u(1,1)= 0u(1, ()
(Boundary Conditions)
dx
where o, Pr. D. and ny are constants defining physical quantities. We use the same physical constants as PDEBench (Takamoto et al. 2022) (see $ for further details). For both classical numerical methods and ML methods, the singularity at u poses hard challenge. The solution trajectory and differential operator are highly non-lineat A standard finite volume solver requires approximately seconds to compute a solution for a 1024x 101 grid (Takamoto et al., (022)..
Problem Setup. We use initial conditions from PDEBench. Each solution instance is a scalar-fich over 1024 spatial and 101 temporal points, where T 500 seconds (see $C.1 for further details).
For PI-HC-MoE, we use K experts and do a spatial decomposition with = 16 basis functions domain of 256 (spatial) x 101 (temporal). Over the 256 x 101 domain, each expert samples 20k points, leading to a total of 80k points where the PDE constraint is enforced the domain. PI-HC uses identical model with
(see B for details inference). Each expert enforees the constraint

---

Published as a conference paper at ICLR 2024
Numerical Solver
(5oft Constraintt)
(Hard Constraint)
PI-HC-MoE (ours)
08
Space (x)
02
00
0.00
-0.50
Difference
Difference
Time (t) Difference
Difference
Figure Predicted solution for the diffusion-sorption equation. (Top) Visualizations of the numerical solver solution and ML predictions for the soft constraint (PI-SC), hard constraint (PI- HC), and PI-HC-MoE. (Bottom) Difference plots of the ML predicted solutions compared to numerical solver solution. White denotes zero. PI-SC is unable to recover the dynamics scale of the solution. PI-HC is able to recover some information, but fails to capture the full dynamics. PI-HC-MoE is able recover almost all of the solution and has the lowest
20k sampled points in the hard constraint. This represents the maximum number of points that we sample with PI-HC without rumning out of memory, with stable batch size. In order to sample more points, a reduction in batch size is required. Our final batch size for PI-HC is 6, and we find that any reduction produces significant training instability, with higher runtime (see Fig. 4 Thus, our baseline PI-HC model uses the maximum number of sampled points that gives the best training stability. All models trained with fixed computational budget of 4000 training iterations.
Results. We summarize our results in Fig. and Fig. In Fig. (left), we plot the T22 relative percent етrоr on the test set over training steps for PI-SC, PI-HC, and PI-HC-MoE. On the right, we plot the trained model L2 relative percent error on the test set of the final trained models. PI- SC is unable to converge to reasonable solution, reflected in the high L2 relative percent error (85.90 055.3% 5.07%). While PI-HC able to achieve lower error than PI-SC at 7.55% 8. 10%, does worse than PI-HC-MoE 3.60%% + 2.93%). In Fig. 3 we show the predicted ML solutions for PI-SC, PI-HC, and PI-HC-MoE, and the difference between the predicted solution and the numer- ical solver solution (white indicates solution, and PI-HC, while closer, is unable to capture the proper dynamics of diffusion-sorption. PI-HC-MoE has the closest correspondence to the numerical solution. Additionally, we explore PI- HC-MoE 's generalization to unseen timesteps D and assess the quality of basis functions ($E).
difference). PI-SC is unable to converge to reasonable
Scalability. We compare the scalability of PI-HC-MoE to PI-HC. For evaluation, plot the num- ber of sampled points on the interior of the domain against the execution time during training and inference, in Fig. PI-HC-MoE maintains constant execution time the number of points increases, while PI-HC becomes significantly slower as the number of sampled points crosses 50k
Training Step
Inference
HC-HC PL-H-MMR MoC (ours)
Execution Time (s)
10
4
20000
40000
60000
80000
100000
20000
10000
60000
80000
100000
Number of sampled points
Number of sampled points
Figure Runtime of PI-HC and PI-HC-MoE on diffusion-sorption. The time to perform single training (left) and inference (right) step as the number of constrained sampled points increases.

---

Published as a conference paper at ICLR 2024
Intbll
Twd Te5
Name Solver
driierncce
PLNC
0.2
contraint
Figure 5: Predicted solution for 2D Navier-Stokes. From top to bottom: (Row 1) Initial vorticity and its evolution as T increases, computed via numerical solver. The errors of PI-HC-MoE (Row 2). PI-HC (Row 3), and PI-SC (Row 4) are visualized for corresponding T, where the difference in the predicted solution shown with respect to the numerical solver. Darker colors indicate higher error. Both PI-SC and PI-HC exhibit greater error compared to PI-HC-MoE, especially at later T.
points. PI-HC-MoE provides a training speedup 1.613x (102 sampled points) to 3.189x (10a sampled points). Notably, PI-HC has higher standard deviation as the number of sampled points increases. Because PI-HC-MoE partitions the number of sampled points, the individual constraint solved by each expert converges faster. As result, PI-HC-MoE is more consistent across different data batches and has much lower variation in runtime.
Note that a standard finite volume solver takes about 60 seconds to solve for solution (Tkammtto ict al. 2022), and so PI-HC and PI-HC-MoE both have inference times faster than a numerical solver. However, PI-HC-MoE is x.030x (10" sampled points) to 3. 3.048x (10$ sampled points) faster at inference than PI-HC, while also having significantly lower error.
4.2 2D RAVIER-STOKER
The Srvier-Stokes equations describe the evolution fluid with given viscosity. We study the vorticity form of the 2D periodic Navier-Stokes equation, where the learning objective is learn the scalar field vorticity w.
(w,,, + u(l,,,,)) Vw(l,, 1,, ,Aw(1,2,)), Øt
LE [0, (x. v) (0, 1)²
(4)
V V -u v== 0, (0, x. 9) ww y):
(Boundary Conditions)
where u is the velocity vector field, and и is the viscosity. Similar to et al. (2021b) and Raissi al. (2019), we study the vorticity form of Navier-Stokes reduce the problem
scalar field prediction, instead of predicting the vector field u. In our setting, we use a Reynolds number of le le - 4), representing turbulent flow. Turbulent flow an interesting and challenging problem due to the complicated evolution of the fluid, which undergoes irregular fluctuations. Many engineering and scientific problems are interested in the turbulent flow case (e.g.. rapid currents, wind tunnels, and weather simulations NNeuwsstadt et al. 2016)).
Problem setup. The training set has a resolution of 64 (x) X 64 (y) X 64 (t). Both the training and test sets have trajectory length of T 5 seconds. The test set has a resolution of 256 (x) x 256 (y) x 64 (t). For both PI-HC and PI-HC-MoE, we use 64 basis functions. We use experts and perform a (a)) x2 (y) x1 (t) spatiotemporal decomposition. Each expert receives the full temporal grid with 1 of the spatial grid (i.e., 32 x 32 x 64 input), and samples 20k points during
the constraint step. For the full data generation parameters and architecture details, see C.2
8

---

Published as a conference paper at ICLR 2024
Results. We visualize representative examples from the 64 test set in Fig. comparing PI-SC, PI-HC, and PI-HC-MoE In Fig. | grey represents zero error. PI-SC is the worst performing model, with significant errors appearing by T 2. While PI-HC captures most of the dynamics at carlier time steps (T = 1, T = 2), especially compared to PI-SC, PI-HC struggles to adequately capture the behavior of later time steps. In particular, PI-HC fails to capture the evolution of fine features T' 5. On the 256 x 256 x 64 test set, PI-SC attains relative L2 eттоr of 18 08.0881% % + 3. 3.740%. PI-HC 751% 2.951 %) and PI-HC-MoE (8.298% 2.2 2.345%) both achieve lower relative L2 error. PI-HC-MoB consistently attains the lowest error across all three resolutions and has lower variance in prediction quality (box plots visualized in appendix
Scalability. Similar to the diffusion-sorption case, eval- uate the scalability of our approach.
Training
Inference
30
PH
PHIC
Executtin time (seconss)
50
PH-IC- MoC tound)
25
PHICC-Mt touns)
We compare the scalability of PI- HC-MoF to standard differentiable optimization (i.c., PI-HC) for both
20
1.5
the training and test steps across a different number of sampled points in Fig. O The training and inference steps are performed with a batch size of 2. At training time, PI-HC scales quadratically with respect to number of sampled points, and has high variance in per batch training and inference time. PI-HC has harder constraint system to solve, reflected in the increase in training
10
10
10'
10
10a
100
10.
109
101
Number sampled points
Number sampled points
Figure 6: Runtime of PI-HC and PI-HC-MoE on 2D Navier-Stokes. We plot the time for a single training (left) and inference (right) iteration. PI-HC-MoE is significantly faster for both, and scales sublinearly. PI-HC has much higher execution times and a large standard deviation.
time. Across both training and inference steps, PI-HC-MoE scales sublinearly and, in practice remains constant with respect the number of sampled points. For training, PI-HC-MoE provides 2.117x x (102 sampled points) to 12. 838x (10* sampled points) speedup over PI-HC. At inference time we see speedups of 238 (102 sampled points) to 12. 864x (10* sampled points).
4.3 FINAL TAKEAWAYS
PI-HC-MoF has lower L2 relative compared to both PI-SC and PI-HC. Enabled by the MoR setup, PI-HC-MoE is better able to capture the features for both diffusion-sorption and Navier-Stokes. PI-HC-MoE is also more scalable than the standard differentiable optimization setting represented by PI-HC. Each expert locally computes Wk, allowing for greater flexibility when weighting the basis functions. The experts are better able to use the local dynamics, while satisfying the PDE globally. PI-HC is limited to linear combination of the basis functions, whereas PI-HC-MoE able to express piece-wise linear combinations.
the diffusion-sorption and Navier-Stokes equations, when
Another reason we find PI-HC-MoE to outperform PI-HC is through the stability of training. This manifests in two different ways. First, because PI-HC-MoE is more scalable, we are able to use a larger batch size, which stabilizes the individual gradient steps taken. Second, for a given total number sampled points, the non-linear least square solves performed by PI-HC-MoE are smaller than the global non-linear least squares performed by the hard constraint. Specifically, the size of the constraint for K experts is 1 R that of PI-HC. This results in an easier optimization problem, and PI-HC-MoE converges quicker and with greater accuracy.
5 CONCLUSION
We present the physics-informed hard constraint mixture-of-experts (PI-HC-MoE) framework, new approach to scale hard constraints corresponding to physical laws through an embedded differ- entiable optimization layer. Our approach deconstructs a differentiable physics hard constraint into smaller experts, which leads to better convergence and faster run times. On two challenging, highly non-linear systems, ID ceffusiionsorpitiin and 2D Navier-Stokan equations, PI-HC-MoB achieves significantly lower errors than standard differentiable optimization using a single hard constraint, as well as soft constraint penalty methods.

---

Published as a conference paper at ICLR 2024
Acknowledgements. This work initially supported by Laboratory Directed Research and De- velopment (LDRD) funding under Contract Number DE-ACO2-05CHHI222 It was then supported by the U.S. Department of Energy, Office of Science, Office of Advanced Scientific Computing Re- search, Scientific Discovery through Advanced Computing (SciDAC) program under contract No. DE-ACO2-05CH1 1231, and in part by the Office of Naval Research (ONR) under grant N00014- 23-1-2587. This research used resources of the National Energy Research Scientific Computing Center (NERSC), U.S. Department of Energy Office of Science User Facility located at Lawrence Berkeley National Laboratory, operated under Contract No. DE-ACO2-05CHH 1231. We also thank Geoffrey Négiar, Sanjeev Raja, and Rasmus Malik Høegh Lindrup for their valuable feedback and discussions.
REFERENCES
Brandon Amos and J. Zico Kolter. OptNet: Differentiable Optimization Layer in Neural Net- works. In Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 136-145. PMLR, 2017.
Brandon Amos, Lei Xu, and J. Zico Kolter. Input convex neural networks. In Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 146-155, PMLR, 2017.
Albert S. Berahas, Frank E. Curtis, Daniel Robinson, and Baoyu Zhou. Sequential quadratic opti- mization for nonlinear equality constrained stochastic optimization. SIAM Journal on Optimiza- lion, 312))1152-11339, 2021.
Mathieu Blondel, Quentin Berthet, Marco Cuturi, Roy Frostig, Stephan Hoyer, Felipe Llinares- Lopez, Fabian Pedregosa, and Jean-Philippe Vert. Efficient and Modular Implicit Differentiation. In Advances in Neural Information Processing Systems, volume 35, pp. 5230-5242 Curran As- sociates, Inc., 2022.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http:
//github. com/ wow////gg///aaa.o/o/////oo////a/
Ulisses Braga-Neto. Characteristics-inforeee neural networks for forward and inverse hyperbolic problems, 2023.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranay Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever. and Dario Amodei. Language Models are Few-Shot Learners. In Advances Neural Information Processing Systems, volume 33, pp. 1877-1901. Curran Associates, Inc., 2020.
Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. In Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
Filipe de Avila Belbute-Peres Kevin Smith, Kelsey Allen, Josh Tenenbaum, and J. Zico Kolter. End-to-end differentiable physics for learning and control. In Advances in Neural Information Processing Systems, volume 31, pp. 7178-7189. Curran Associates Inc., 2018.
Priya Donti, David Rolnick, and J. Zico Kolter. DC3: A learning method for optimization with hard constraints. In International Conference on Learning Representations, 2021.
Yiheng Du, Nithin Chalapathi, and Aditi Krishnapriyan. Neural spectral methods: Self-supervised learning in the spectral domain. arXiv preprint arXiv:2312..0225, 2023.
Yuchen Fang, Sen Na, Michael W. Mahoney, and Mladen Kolar. Fully stochastic trust-region se- quential quadratic programming for equality-constraine optimization problems. 2024.
10

---

Published as a conference paper at ICLR 2024
Isom H. Herron and Michael R. Foster. Partial Differential Equations in Fluid Dynamics. Cambridge University Press, 2008.
Zhichun Huang, Shaojie Bai, and J. Zico Kolter. (implicit)²: Implicit layers for implicit representa- tions. In Advances in Neural Information Processing Systems, volume 34, pp. 9639-9650. Curran Associates, Inc., 2021.
Jacob Kelly, Jesse Bettencourt, Matthew Johnson, and David Duvenaud. Learning Differen- tial Equations that Easy Solve. In Advances in Neural Information Processing Systems, volume 33, pp. 4370-4380. Curran Associates, Inc., 2020.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Pe- ter Tang. On latce-batch training for deep learning: Generalization gap and sharp minima. arXiv preprint /609. 2016.
Patrick Kidger and Cristian Garcia. Equinox: neural networks JAX via callable PyTrees and filtered transformations. Differentiable Programming workshop at Neural Information Processing Systems 2021, 2021.
James Kotary. Ferdinando Fioretto, Pascal Van Hentenryck, and Bryan Wilder. End-to-end con- strained optimization learning: A survey. arXiv preprint arXiv:2103.16338, 2021.
Aditi Krishnapriyan, Amir Gholami, Shandian Zhe, Robert Kirby, and Michael W Mahoney. Char- acterizing possible failure modes in physies-iforrmed neural networks, In Advances in Neural Information Processing Systems, volume 34, pp. 26548-26560. Curran Associates, Inc., 2021.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Fourier Neural Operator for Parametric Partial Differential Equations. arXiv preprint 0rXiv:2010.08890 2021a.
Zongyi Li, Hongkai Zheng, Nikola Kovachki, David Jin, Haoxuan Chen, Burigede Liu, Kamyar Avizzadenesheli, and Anima Anandkumar. Physics-informed neural operator for learning partial differential equations. arXiv preprint arXiv:2111.03794, 2021b.
Björn List, Li-Wei Chen, and Nils Thuerey. Learned turbulence modelling with differentiable fluid solvers: physics-bar based loss functions and optimisation horizons. Journal of Fluid Mechanics, 949, 2022.
Lu Lu, Raphael Pestourie, Wenjie Yao, Zhicheng Wang, Francese Verdugo, and Steven G Johnson. Physics-informed neural networks with hard constraints for inverse design. SIAM Journal on Scientific Computing, 43(6):B 105-B1132, 2021.
Pingchuan Ma, Peter Yichen Chen, Bolei Deng. Joshua B Tenenbaum, Tao Du Chuang Gan, and Wojciech Matusik. Learning neural constitutive laws from motion observations for generalizable pde dynamies. In International Conference Machine Learning. PMLR, 2023.
Andreas Mayr, Sebastian Lehner, Arno Mayrhofer, Christoph Kloss, Sepp Hochreiter, and Johannes Brandstetter. Boundary Graph Neural Networks for 3D Simulations. Proceedings of the AAAI Conference on Artificial Intelligence, 37(8):9099-91) 0000911772000000 2023.
Takeru Miyato, Toshiki Kataoka, Masanori Koyama, and Yuichi Yoshida. Spectral Normalization for Generative Adversarial Networks. arXiv preprint 1802. 05957, 2018.
Sen Na. Sungho Shin, Mihai Anitescu, and Victor M. Zavala. On the convergence of overlapping schwarz decomposition for nonlinear optimal control. IEEE Transactions on Automatic Control, 67(11):5996-600 6011 1, 2022.
Hans Bruun Nielsen and Kaj Madsen. Introduction to Optimization and Data Fitting. Informatics and Mathematical Modelling, Technical University of Denmark, DTU, 2010.
Frans T.M. Nieuwstadt, Jerry Westerweel, and Bendiks J. Boersma. Introduction to Theory and Applications of Turbulent Flows. Springer, 2016.

---

Published as a conference paper at ICLR 2024
Evgenii Nikishin, Romina Abachi, Rishabh Agarwal, and Pierre-Luc Bacon. Control-orienne model-based reinforcement learning with implicit differentiation. In Proceedings of the AAAI Conference Artificial Intelligence, pp. 7886-7894, 2022.
Geoffrey Négiar, Michael W. Mahoney, and Aditi Krishnapriyan. Learning differentiable solvers for systems with hard constraints. In The Eleventh International Conference on Learning Represen- fations, 2023.
Raphaël Pestourie, Youssef Mrouch, Chris Rackauckas, Payel Das, and Steven G. Johnson. Physics- enhanced deep surrogates for partial differential equations. Nature Machine Intelligence, 2023.
Luis Pineda, Taosha Fan, Maurizio Monge, Shobha Venkataraman, Paloma Sodhi, Ricky T. Q. Chen, Joseph Ortiz, Daniel DeTone, Austin Wang. Stuart Anderson, Jing Dong, Brandon Amos, and Mustafa Mukadam. Theseus: A Library for Differentiable Nonlinear Optimization. In Advances in Neural Information Processing Systems, volume 35, pp. 3801-381 Curran Associates, Inc.,
2022.
Yi-Ling Qiao, Junbang Liang, Vladlen Koltun, and Ming C. Lin. Scalable Differentiable Physics for Learning and Control. In Proceedings of the 37th International Conference Machine Learning, 2020.
Jason Rader and Patrick Kidger. patrick www.eriddgeeeee oppp://tttmmmsi..
Optimistix.
URL https://github. com/
Maziar Raissi, Paris Perdikaris, and George Karniadakis. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational Physics, 378:686-707, 2019.
Bharath Ramsundar, Dilip Krishnamurthy, and Venkataasubramanian Viswanathan. Differentiable Physics: A Position Piece. arXiv preprint arXiv:2109.07573 2021.
Jack Powell, Yaron Lipman, and Ricky T. Q. Chen. Neural conservation laws: divergence- free perspective. In Advances in Neural Information Processing Systems, 2022.
Carlos Riquelme Ruiz, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Susano Pinto, Daniel Keysers, and Neil Houlsby. Scaling Vision with Sparse Mixture of Experts. In Advances in Neural Information Processing Systems, 2021.
Nadim Saad, Gaurav Gupta, Shima Alizadeh, and Danielle C. Maddix. Guiding continuous operator learning through Physics-based boundary constraints. In International Conference on Learning Representations, 2023.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In International Conference Learning Representations, 2017.
Sam Smith, Pieter-jan Kindermans, Chris Ying, and Quoe V. Le. Don't decay the learning rate, increase the batch size. In International Conference on Learning Representations, 2018.
Natarajan Sukumar and Ankit Srivastava. Exact imposition of boundary conditions with distance functions in physics-informat deep neural networks. Computer Methods in Applied Mechanics and Engineering, 389:1 14333, 2022. ISSN 0045-7825.
Tetsuya Takahashi, Junbang Liang, Yi-Ling Qiao, and Ming C. Lin. Differentiable fluids with solid coupling for learning and control. Proceedings of the AAAI Conference on Artificial Intelligence, 35(7):6138 6146, 2021.
Makoto Takamoto, Timothy Praditia, Raphael Leiteritz, Daniel MacKinlay, Francesco Alesiani, Dirk Pflüger, and Mathias Niepert. PDEBench: An Extensive Benchmark for Scientific Machine Learning. In Advances in Neural Information Processing Systems, volume 35, pp. 1596-1611. Curran Associates, Inc., 2022.
12

---

Published as a conference paper at ICLR 2024
Pauli Virtanen, Ralf Gommers, Travis E. Oliphant, Matt Haberland, Tyler Reddy, David Courna- peau, Evgeni Burovski, Pearu Peterson, Warren Weckesser, Jonathan Bright, Stéfan J. van der Walt, Matthew Brett, Joshua Wilson, K. Jarrod Millman, Nikolay Mayorov, Andrew R. J. Nel- son, Eric Jones, Robert Kern, Eric Larson, C J Carey, Ilhan Polat, Yu Feng, Eric W. Moore, Jake VanderPlas, Denis Laxalde, Josef Perktold, Robert Cimrman, Ian Henriksen, E. A. Quintero, Charles R. Harris, Anne M. Archibald, António H. Ribeiro, Fabian Pedregosa, Paul van Mul- bregt, and SciPy 1.0 Contributors. SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python. Nature Methods, 17:261-22 272, 2020.
Sifan Wang, Yujun Teng, and Paris Perdikaris. Understanding and mitigating gradient flow patholo- gies in physics-informed neural networks. SIAM Journal on Scientific Computing, 43(5):A3055- A3081, 2021.
Zhou Xian, Bo Zhu, Zhenjia Xu, Hsiao-Yo Yu Tung, Antonio Torralba, Katerina Fragkiadaki, and Chuang Gan. Fluidlab: A differentiable environment for benchmarking complex fluid manipula- tion. In International Conference on Learning Representations, 2023.
Kailai Xu and Eric Darve. Physics Constrained Learning for Data-driven Inverse Modeling from Sparse Observations. Journal of Computational Physics, 453: 10938, 2022.
Fangeheng Zhong, Kyle Thomas Fogarty, Param Hanji, Tianhao Walter Wu, Alejandro Sztrajman, Andrew Everett Spielberg, Andrea Tagliasacchi, Petra Bosilj, and Cengiz Oztireli. Neural fields with hard constraints of arbitrary differential order. In seirth- Conference on Neural In- formation Processing Systems, 2023.
13

---

Published as a conference paper at ICLR 2024
A IFT: BACKWARD PASS
We expand on the training procedure described in $ and additional details on training the model using the implicit function theorem.
Recall that for basis functions b - [b0 1...b"...0" b and U' : Ω + R the non-linear least squares solver finds w R. such that Fa(b w²) O. If the Jacobian FbFa w") is non-singular, then
the implicit function theorem holds. There exists open Sb C RR=N,, Sw C RR., and function 2" : Sb > S.- Sb and S. contain and w. 2 has the following properties:
= ='(b) Fb 2" (b) 0. 2" differentiable on S..
Property (1)
Property (2)
Property (3)
During gradient descent, by the chain rule, we need know 0= (b) 39
which is not readily available
by auto-differentiation. By differentiating property (2), we get Eq.
aF (b. =b)²))
oF wx)
ōb
(b)²²)
0z)(b) 0.
00 +
00
Ōb
02*(b)
00
Rearranging Eq.
D can produce new system of equations.
0F (b.(bb=2((h (((bbb=((b')) =²bb)²
(z*(b)
0F w²) Ôb
(z)(b)
00
Ōb
00
0z
(FF (b. x² (b)²)
0F w²)
(5)
00
02+(b)
00
The unknown that we need to solve for now is the desired quantity, 0=" The expressions on the
09
right-hand side, (b=(""")) ab 8b 00 and (T+(b==(())))
(b)(b) can be computed via auto suto-differenntition This then gives system of equations which involves matrix inverse. If matrix inversion compu- tationally tractable, we could explicitly solve for de db" Instead, the system of equations is solved
using the same non-linear least squares solver as in the forward pass, allowing us to approximate
the matrix inverse, and yields end-to-end.
8= a0
Now, we can compute
(F. and train the full model
00
B EXAMPLE INFERENCE PROCEDURE
To demonstrate the inference procedure of PI-HC-MoE, walk through the forward pass of the ID diffusion-sorption problem setting, starting from PI-HC and generalizing to the PI-HC-MoE case.
Model Input and Output. In both PI-HC and PI-HC-MoE, the underlying FNO model is pro- vided with both the initial condition and underlying spatiotemporal grid. Each initial condition is 1024 dimensional vector (the chosen discretization), representing the spatial coordinates, and is broadcasted across time (t = 500 s). leading to an input of 1024 x 101 for the initial condition. The broadcasted initial condition concatenated with the spatiotemporal grid with shape 1024 for final input spanning the entire grid with 3 channels (1024 x 101 x 3). The inference procedure of the underlying base NN achitecture is unchanged. In our case, the base NN architecture produces a 1024 x 101 N tensor representing N basis functions evaluated over the spatiotemporal domain.
x2,
PI-HC (Single Constraint). Each constraint samples M points and solves the PDE equations. Additionally, the input initial condition (1024) and initial condition of the basis functions (1024 x N)
14

---

Published as a conference paper at ICLR 2024
provided to the solver. Similarly, the boundary conditions provided and diffusion-sorption has two boundary conditions at x = 0 and z = 1, which leads to a boundary condition tensor of 2 x 101 x N.. The optimal (N) computed by the non-linear least squares solver satisfies the PDE equations on the M sampled points, as well as the initial and boundary conditions. Finally, w is used compute the entire scalar field by performing matrix multiply between the basis functions and (1024 x 101 X N N = 1024 x 101 predicted scalar field).
PI-HC-MoE (Multiple Constraints). The MoE router performs domain decomposition. As- suming K = 4 spatial experts, the 1024 x 101 x N basis function evaluation deconstructed into four 256 x 101 x N tensors. Each expert performs a similar computation as PI-HC, and notably, the global initial (1024 x N) and boundary (2 x 101 X N) conditions are provided each expert's non-linear least squares solve. After computing localized weighting Wk. each expert returns 256 101 scalar field representing the local predicted solution field. Finally, the MoE router re- verses the domain decomposition, assembling the four 256 x 101 matrices into final 1024 101 prediction. This final prediction represents the complete constrained output.
C PROBLEM SETTING DETAILS
We use Jax (Bradbury al. 2018 with Equinox Kidger & Garcia 2021) to implement models. For diffusion-sorpticn we use Optimisix's (Rader & Kidger) Levenberg-Marquaat solver and for Navier-Stokes. use JaxOpt's (Blondel al. 2022) Levenberg- Marquurrd solver. In both the PI-HC and PI-HC-MoE cases, we limit the number of iterations to 50.
Data-constrained setting. We exclusively look at the data-starved regime. As a motivating ex- ample, in the ID diffusion-sorption, PDEBench (Takamoto et al. 2022) provides dataset of 10k solution trajectories. For a difficult problem like diffusion-sorption, the numerical solver used in PDEBench takes over one minute to solve for solution trajectory. dataset of 10k would re- quire almost week of sequential compute time. Additionally, PDEBench only provides dataset using one set of physical constants (e.g., porosity). Different diffusion mediums require different physical constants, making the problem of generating a comprehensive dataset very computationally expensive. For these reasons, we focus on only training the NN via the PDE residual loss function.
C.1 ADDITIONAL DETAILS: DIFFUSION-SORPIE
Physical Constants and Dataset Details. 6 0.29 is the porosity of the diffusion medium. P. = 2880 is the bulk density. ny = 0.874 is Freundlich's 's exponent. Finally, D = 50 100 is the effective diffusion coefficient.
The training set has 8000 unique initial conditions and the test set has 1000 initial conditions, dis- tinctly separate from the training set. We use the same initial conditions from PDEBench (Takamoto al. 2022)
Architecture and Training Details. The base model we use is FNO (Li ct al. 202 archi- tecture with 5 Fourier layers, each with modes and 64 hidden feature representation. We use a learning rate of 1 with exponential decay over 4000 training iterations. The tolerance of the Levenberg- -arabbrrg-Maqquarttt set to 1e-4
C.2 ADDITIONAL DETAILS: 2D NAVIER-STOKES
Data generation. All initial conditions for the training and test set are are generated from a 2D Gaussian random field with periodic kernl and length scale of 0.8. The training set has 4000 initial conditions with resolution of 64 (x) x 64 (y) x 64 (t). Both the training and test sets have trajectory length of T 5 seconds. We generate a test set of 100 solutions with resolution of 256 x 256 (spatial) and 64 time steps. The error is measured on the original resolution solutions, as well as spatially subsampled versions (128 128, 64 x 64).
15

---

Published as conference paper at ICLR 2024
Rel. error (%) on 64x64x64
Rel. error (%) on 64x128x128
Rel. error (%6 on 64x256x256
15.0
12.5
10.0
7.5
PI-HC
PL-C-MMDE (ours)
FI-HC
PI-HC-MoE (OurS)
PLHC
PI-HC-MoE (ours)
Figure 7: Final Lg relative error % on test sets for 2D Navier-Stokes, Both PI-HC and PI-HC- MoE train 64 x 64 64 spatiotemporal grid. PI-HC-MoF generalizes spatial resolutions better than PI-HC with lower variance.
Architecture details. We use an FNO with 8 Fourier layers and 8 modes the base NN architec- ture. The Gevenberg-Marquarre solver tolerance set le-? and we use learning rate of 1e-3 with exponential decay over 20 epochs and final learning rate of le" 1e-4.
C.3 ADDITIONAL RESULTS: 2D NAVIER-STOKES
In Fig. we include box plot showing the test set errors on three different spatial resolutions (256
1282, 642). The 642 and 128 solutions subsampled from the 256 solution. PI-HC-MoE outperforms PI-HC on all test sets.
C.4 ADDITIONAL DETAILS: SCALABILITY RESULTS
When conducting our scalability analysis, we measure the per-batch time across 64 training and test steps. We observe large standard deviations in our speedup numbers and provide a short discussion here. In Fig. 4 and Fig. 6 we plot our measured runtimes as a function of the number interior points constraiedd In both the diffusion-sorption sorption and Navier-Stokes settings, we observe that PI- HC has large standard deviations at fixed number of sampled points. We attribute this variation to the difficulty of the constraint performed by PI-HC, which causes high fluctuations in the number of non-linear least squares iterations performed. As result, when we calculate speedup values, the speedup values also have high std deviations. In the main text, report the mean speedup across all batches for fixed number of sampled points; for completeness we include the standard deviations here.
In the ID cefffusion-sorppttin setting, we see standard deviations at inference from 0.587 (102 sam- pled points) to 1.890 (10f sampled points). For training, the standard deviations vary from 0.481 (102 sampled points) to 1.388 (10# sampled points). Analogously,
2D Navier-Stokes, we see standard deviations from 3.478 (102 sampled points) to 35.917 (10' sampled points) for inference. For training, we measure standard deviations from 2.571 (102 sampled points) to 47.94 7.948 (10* sam- pled points). The large standard deviations are reflection of the standard deviations of the runtime numbers for PI-HC for both problems.
D ADDITIONAL EXPERIMENT: TEMPORAL GENERALIZATION
We explore PI-HC-MoE's generalization to temporal values not in the training set, for the diffusion- sorption case.
Problem Setup. In our original problem setting, we predict diffusion-sorption from t = seconds to 500 seconds. Here, in order to test generalization unseen temporal values, truncate
16

---

Published as a conference paper at ICLR 2024
the training spatiotemporal grids to t < 400 (s), i.e. only train t [0, 400]- After training, PI-SC, PI-HC, and PI-HC-MoE all predict the solution up to t = 500 seconds. We compute the L2 relative error enror in prediction for t 400 (s) (the domain that unseen during training). The parameters for all models are identical to 4.1 PI-HC-MoB uses K = spatial experts.
Results. Fig. shows the L relative error for 1 400 (s) on the validation set. Both PI-HC (48.08 047. 1 13) and PI-HC-MoE (14.78 79% = 3.00%) outperform PI-SC (805.58% + 19. (). While PI-HC is able to do better than PI-SC, PI-HC-MoE is able to gener- alize the best to future time steps, even without any training instances.
We hypothesize that reason for the performance improvement of PI-HC and PI-HC-MoE comes from test-time constraint enforcement, demonstrating the benefit hard constraints at inference time. Similar the results in $4.1] PI-HC-MoE is able better generalize much better the unseen temporal domain than PI-HC, while also being much more efficient at both training and inference time.
Diffusion Sorption t 400 (s)
103
Rel. L2 Error (%) Test Set
102
10000
10
PI-SC
PI-HC
PI-HC MOE (ours)
Figure 8: Relative L2 error unseen temporal values (t 400 (s)) for diffusion-sorption. All
models are trained on E [0, 400), and then predict the solution up to 500 seconds. The relative
L2 error is plotted for / > 400 seconds. PI-SC is unable to generalize out-of-distribution temporal values. PI-HC is able to do better, but still has high error, while PI-HC-MoE has the lowest error.
E ABLATION: EVALUATING THE QUALITY OF THE LEARNED BASIS FUNCTIONS
In order to explore the usefulness of our learned basis functions, we conduct an ablation study comparing the result of PI-HC-MoE to a standard cubic interpolation.
Problem setup. We use the same PI-HC-MoE model from the diffusion-sorptian experiments
in 14. where each expert constrains 20k sampled points (K = 4 experts, 80k total sampled points). The 20k constrained points represents a candidate solution to the PDE where, by the non-lincar least squares solve, the PDE equations must be satisfied. Using the 20k points, perform interpolation using SeiPy's /Virtanen et al., 2020) 2D Clough-Tocher piecewise cubic interpolation. The interpolated solution is assembled in identical manner PI-HC-MoE, representing K interpolations, and corresponds to the number of experts.
Results. Fig. 9 compares the relative Le error the test set between PI-HC-MoE 's learned basis functions and the interpolated solution. PI-HC-MoE outperforms the interpolation scheme, show- ing that the learned basis functions provide advantage over interpolation, and encode additional unique information. We also note that it is likely that with fewer sampled points, the performance gap between PI-HC-MoB and interpolation is likely to be higher.
17

---

Published as a conference paper at ICLR 2024
Rel. L2 error (%) Diffusion Sorption Valdiation Set
10
Predicted
Interpolated
Figure 9: Relative Le error 1D diffusion-ssrptioon's test set for PI-HC-MoE and an inter- polated solution. The learned basis functions by PI-HC-MoE represent a more accurate solution compared interpolating the constrained points. The error plot shows that even for unconstrained points, PI-HC-MoE is closer approximation to the numerical solver solution.
F PDE RESIDUALS FOR DIFFUSION-SORPTI AND NAVIER-STOKEE
We add additional PDE residual loss plots that correspond the results in 4..1 and During the training all three models (PI-SC, PI-HC, and PI-HC-MoE) ), we track the mean PDE residual value across the validation set, plotted in Fig. TO Across both problem instances, PI-HC-MoE attains the lowest PDE residual, and also has the lowest L2 relative error (see $4.1). For ID diffusion-sorption. PI-SC initially starts with low PDE residual before increasing near the end of training. The reason for this is that at the beginning of training, when PI-SC is first initialized, the predicted solution is 8a close to the zero function and trivially satisfies the PDE equations, i.e., both du and 8x2 are However, the overall loss function includes adherence to the initial and boundary conditions and PI-SC is unable find parameterization that satisfies the initial condition, boundary conditions, and PDE equations.
PDE Residual on Validation
Diffusion-sorption
Navier-Stokes
10
PI-HC-MoE (ours)
PI-HC-MoE (ours)
PI-HC
PI-HC
PI-SC
PI-SC
Mean Residual
Mean Residual
10a
10-4
10-2
1000
2000
3000
4000
2500
5000
7500
10000 12500
15000
Training Step
Training Step
Figure 10: PDE residual on validation set during training. The mean PDE residual over the validation set over training for both ID diffusion-sorption (left) and 2D Navier-Stokes (right). PI- HC-MoE has the lowest PDE residual.
18

---

Published as a conference paper at ICLR 2024
G LIMITATIONS AND FUTURE WORK
PI-HC-MoE, though promising framework for scaling hard constraints, has some limitations.
not al- ways clear what the correct number of basis functions, sampled points, number of experts, and expert distribution best suited for given problem. However, our results show that with only minor hyperparameter tweaking. we were able to attain low error on two challenging dynamical systems. As with many ML methods, hyperparameter optimization non-trivial task, and one future direction could be better ways to find the optimal parameters.
Hyperparameters. Both PI-HC-MoE and PI-HC are sensitive to hyperparameters.
Choice of domain decomposition. Currently, PI-HC-MoF, performs spatiotemporal domain de- composition assign points to experts. A possible future direction is trying domain decompo- sitions, and methods for allocating experts. It may be the case that different ways of creating expert domains are better suited for different problem settings.
Base NN architecture. In this work, we use FNO as our base architecture for PI-HC-MoE: and in the future, may prove fruitful to try new architectures. To some extent, PI-HC-MoB is limited by the expressivity of the underlying NN architecture which learns and predicts the basis functions. PI- HC-MoE may perform better on certain tasks that we have not yet explored (c.g.. super-resolution. regressegeesive training), and future work could explore the application of PI-HC-MoE new kinds of tasks.
19


Reviews:
[{"review_id": "WlFTTmACEx", "reviewer": "Reviewer_ihJQ", "summary": "This paper proposes a scalable way to enforce hard constraints expressed as partial differential equations into neural networks in order to faithfully model physical phenomena. The problem they try to tackle is given by the fact that backpropagating through constraints over large meshes is a highly non-linear problem that grows in dimensionality with respect to the mesh and neural network size. The authors propose to impose the constraints over smaller decomposed domains, each of which is solved by an expert.", "questions": "**Experimental Analysis:**\n\nThe experimental analysis seems quite comprehensive. My main question is: are PI-SC and PI-HC the only models against which you can compare? \n\nAlso, the authors compare the execution time of PI-HC vs PI-HC-MoE, what about PI-SC?\n\nI am not sure I understand the scales in Figure 3. 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? \n\nThe authors compare their method against PI-HC and PI-SC in terms of computation time, what about space? Is the memory requirement of PI-SC-MoE also advantageous and/or comparable?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "**Novelty:**\n\nI am not in a position to judge the novelty of the paper as there is very little overlap with my work. From the related work mentioned in the paper, the paper seems novel. \n\n\n**Significance:**\n\nThe paper is of high interest to a subset of the ML community.", "weaknesses": "**Clarity:** \n \nThe paper is mostly quite clear. Personally, I would have benefitted from 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. Even better, it would have been nice to use the same example to move from one constraint to multiple ones.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "cfGnckigTC", "reviewer": "Reviewer_Tpvx", "summary": "Methods\n- This paper utilizes the Fourier Neural Operator (FNO) as the NN architecture to learn a set of N scalar-valued functions as the basis functions; enforces the hard constraints by using differentiable optimization to find a linear combination of basis functions to satisfy the constraints.\n- To deal with larger systems which contain a large number of sampled points and basis functions, mixture-of-experts is used to decompose the spatialtemporal domain.\n- In backward pass, to compute the gradients of the differentiable optimization, the implicit function theorem is used. For mixture-of-experts, the Jacobian is reconstructed from the individual Jacobians from all experts. \n\nExperiments:\n- Two cases: 1D diffusion-sorption and 2D turbulent Navier-Stokes \n- Compared to training via a physics-informed soft constraint (PI-SC) and physics-informed hard constraint (PI-HC). \n    - Achieve higher accuracy and lower time cost.", "questions": "- What are the temporal intervals used for generating the training data and for testing the model? \n    - 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. \n- 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 predict the interval of [0s, 10s]? And how about the spatial domain outside the training dataset?\n- 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?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper proposes a new way to decompose the complex dynamical systems with a large number of points in the spatiotemporal domain to smaller solvable systems, by utilizing the mixture-of-experts method. It makes the system scalable and performant (with parallel computing).\n- In the 2 test cases provided in the paper, the paper’s method has higher accuracy and lower time cost compared with other two methods.", "weaknesses": "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "tRUCQO9gau", "reviewer": "Reviewer_gWoi", "summary": "The paper presents a novel Physics Informed (PI) framework to scale the imposition of \"Hard Constraints\" (HC) with the use of a Mixture of Experts (MoE) Network. The hard constraints are imposed on $m$ sampled points $x_1, ..., x_m \\in \\Omega$ by solving a non-linear least squares problem (NLSS), and the optimization is made possible by using a differentiable solver. The paper proposes a framework that decomposes the domain $\\Omega$ into subdomains $\\Omega_k$ before solving the NLSS problem on each subdomain.\n\nThey tested the framework as a neural PDE solver on 1D diffusion-sorption and turbulent 2D Navier-Stokes equations in a data-constrained regime. The method outperforms existing soft and hard-constrained methods in terms of accuracy, and scales well at inference w.r.t the number of sampled points.", "questions": "Could you explain the artifacts of the method in Figure 3 ?\n\nDid you explore the aspect of each basis function ? What do they learn ? Do the different weights overlap over the domains or are some weights \"activated\" only on specific subdomains  ? \n\nDid you check the PDE residuals on the sampled points to the see if the constraint was respected ?\n\nCould you provide more explanations on the way to compute $\\frac{\\partial z^*}{\\partial \\theta}$ ?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "The paper is well written and easy to follow. The method takes inspiration from the domain decomposition which is standard in the the literature for solving Partial Differential Equations.\nI appreciated the details in the Methods section, particularly the explanations on the forward and backward pass of the architecture.\nThe method clearly outperforms the existing  soft- and hard-constrained baselines. The claim of the paper to the scaling is well supported with a solid inference time analysis.", "weaknesses": "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.\n\nI 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. \n\nI 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "PUKAwdyv7g", "reviewer": "Reviewer_7bFe", "summary": "The paper introduces a new method to apply strict physical rules in mixture-of-experts (MoE) for intricate dynamic systems. The authors believe that when training neural networks, using well-known physical rules, like conservation laws, can enhance results and speed up learning processes. But, there's a drawback: this method can take up a lot of computing power and storage, especially with big systems.\n\nTo solve this, the authors suggest using MoE. Here, the data is split into sections, with each section having an \"expert\" to apply the rules. Every expert works on a specific part of the data, and a special network decides how much importance to give each expert's output. This method saves on computing power and storage but still gives accurate results.\n\nThey test their method on various physical problems, such as fluid movement and quantum mechanics. They compare their results with other top methods and find their method is just as good, if not better, and uses less resources. Plus, it can handle big systems. In conclusion, the paper offers a fresh and effective way to use MoE to apply strict physical rules, which could make neural networks better and faster at solving complex systems.", "questions": "What are the benefits of imposing known physical constraints during neural network training?\n\nHow does the MoE approach differ from standard differentiable optimization in enforcing hard physical constraints?\n\nCan you provide an example of a complex dynamical system where this scalable approach would be particularly useful?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Novel approach: The paper presents a novel approach to enforcing hard physical constraints using a mixture of experts (MoE) that significantly reduces computational and memory costs while maintaining accuracy and convergence guarantees.\n\n2. Scalability: The authors demonstrate that their approach scales well to large-scale systems, making it suitable for complex dynamical systems.\n\n3. Comparative analysis: The paper provides a detailed comparative analysis of their approach with other state-of-the-art approaches, showing that it achieves comparable or better performance while being significantly more efficient.\n\n4. Real-world applications: The authors demonstrate the effectiveness of their approach on several physical problems, including fluid dynamics, structural mechanics, and quantum mechanics, showing that it has real-world applications.\n\n5. Clear presentation: The paper is well-written and clearly presents the approach, methodology, and results, making it easy to understand and follow.", "weaknesses": "While the paper presents a novel and scalable approach to enforcing hard physical constraints using mixture-of-experts (MoE), there are some potential weaknesses that should be considered. \n\nFirstly, 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.\n\nSecondly, 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.\n\nThirdly, 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. For example, 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "T65zIVuYhL", "reviewer": "Reviewer_v3ig", "summary": "This work aims to reduce the cost of imposing hard physical constraints in deep neural network training by utilizing Mixture-of-Experts. The proposed approach is scalable and allows multi-GPU training. The improved stability and efficiency of the proposed method over existing approaches are well demonstrated on non-linear systems.", "questions": "see Weaknesses", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "(1) The paper is overall well written.\n\n(2) The problem addressed in this paper, which is how to impose physical constraints into large-scale deep neural network training, is of importance.\n\n(3) The improvement of the proposed method over baselines is significant.", "weaknesses": "Major: As claimed by the authors in the last paragraph in page 4, the utilizing of Mixture-of-Experts is motivated by two reasons/challenges. (A) In complicated systems, compared with global physical constraints, local constraints are more manageable. (B) The scalability of the non-linear least squares solver impeded the increase in the number of sample points (m), thereby affecting the accuracy of predictions. \n\nReason A aligns with the phenomenon we have observed in language tasks [1]: 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.\n\nChallenge B can be overcome by utilizing MoE. However, a more straightforward solution will be employing scalable non-linear least squares solvers, e.g. [2][3]. The authors didn’t comment on this line of works. If there are suitable solvers available, I would like to see a comparison between PI-HC and PI-HC-MoE with the same number of data points.\n\nFurthermore, the performance of the baseline appears to be inconsistent with the results reported in [4], e.g., 1D non-linear diffusion-sorption equation results, see [4], Appendix E.\n\nMinor: There are several errors and confusions in the use of mathematical symbols in the text, including but not limited to \n(1) in section 3, first paragraph, line 1, $0$ should be in bold.\n(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}$.\n(3) In section 3, second paragraph, line 5, $f^i: \\phi \\rightarrow \\mathbb{R}$ should be $f^i: \\Omega \\rightarrow \\mathbb{R}$.\n(4) In section 3, second paragraph, third to last line, $f_\\theta(\\phi(x))$ should be $f_\\theta(\\phi)(x)$.\n(5) In page 4, two lines above Property (1), the definition of $S_{\\mathbf{b}}$ and $S_\\omega$.\n\nSimilar errors are repeated throughout the paper and should be corrected.\n\n[1] Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer\n[2] Scalable Subspace Methods for Derivative-Free Nonlinear Least-Squares Optimization\n[3] DeepLM: Large-scale Nonlinear Least Squares on Deep Learning Frameworks using Stochastic Domain Decomposition\n[4] PDEBENCH: An Extensive Benchmark for Scientific Machine Learning", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:

Decision: Accept (poster)