Learning differentiable solvers for systems with hard constraints
Authors: Geoffrey Négiar
Michael W. Mahoney
Aditi Krishnapriyan
Venue: ICLR 2023 poster (2023)
Keywords: differentiable optimization
PDEs
physics
neural networks
differentiable constraints
dictionary learning

Abstract: We introduce a practical method to enforce partial differential equation (PDE) constraints for functions defined by neural networks (NNs), with a high degree of accuracy and up to a desired tolerance. We develop a differentiable PDE-constrained layer that can be incorporated into any NN architecture. Our method leverages differentiable optimization and the implicit function theorem to effectively enforce physical constraints. Inspired by dictionary learning, our model learns a family of functions, each of which defines a mapping from PDE parameters to PDE solutions. At inference time, the model finds an optimal linear combination of the functions in the learned family by solving a PDE-constrained optimization problem. Our method provides continuous solutions over the domain of interest that accurately satisfy desired physical constraints. Our results show that incorporating hard constraints directly into the NN architecture achieves much lower test error when compared to training on an unconstrained objective.

Published as a conference paper at ICLR 2023
LEARNING DIFFERENTIABLE SOLVERS FOR SYSTEMS WITH HARD CONSTRAINTS
Geoffrey Négiar
Michael W. Mahoney
Aditi S. Krishnapriyan' 1.2
Universityy of California, Berkeley
2 Lawrence Berkeley National Laboratory
International Computer Science Institute
ABSTRACT
We introduce a practical method to enforce partial differential equation (PDE) constraints for functions defined by neural networks (NNs), with high degree of accuracy and up to a desired tolerance. We develop a differentiable PDE- constrained layer that can be incorporated into any NN architecture. Our method leverages differentiable optimization and the implicit function theorem to effec- tively enforce physical constraints. Inspired by dictionary learning, our model learns a family of functions, each of which defines mapping from PDE parame- ters to PDE solutions. At inference time, the model finds an optimal linear com- bination of the functions in the learned family by solving PDE-conme eDE-srnstrained opti- mization problem. Our method provides continuous solutions over the domain of interest that accurately satisfy desired physical constraints. Our results show that incorporating hard constraints directly into the NN architecture achieves much lower test error when compared to training unconstrained objective.
INTRODUCTION
Methods based on neural networks (NNs) have shown promise in recent years for physics-based problems (Raissi et al. 2019: Li et al. 2020 Lu et al. 2021a Li et al. 202T). Consider pa- rameterized partial differential equation (PDE), Fu( PDE parameters ø and solution u are functions over a domain X. Let be a distribution of PDE- parameter functions áb The goal is to solve the following feasibility problem by training NN with parameters E RP, i.e., find such that, for all functions & sampled from $ the NN solves the feasibility problem,
O. F& differential operator, and the
Fo(uo(o)) = 0.
(1)
Training such model requires solving highly nonlinear feasibility problems in the NN parameter space, even when Fo describes a linear PDE.
Current NN methods use two main training approaches to solve Equation The first approach is strictly supervised learning, and the NN is trained on PDE solution data using a regression loss (Lu et al.] 2021a Li et al. 2020). In this case, the feasibility problem only appears through the data; does not appear explicitly in the training algorithm. The second approach (Raissi et al. 2019) aims solve the feasibility problem in Equation by considering the relaxation.
min Rewe
This second approach does not require access to any PDE solution data. These two approaches have also been combined by having both data fitting loss and the PDE residual loss (Li ct al. 202I).
However, both of these approaches come with major challenges. The first approach requires po- tentially large amounts of PDE solution data, which may need to be generated through expensive numerical simulations or experimental procedures. can also be challenging to generalize outside the training data, there is no guarantee that the NN model has leamned the relevant physics. For the second approach, recent work has highlighted that in the context of scientific modeling, the relaxed feasibility problem in Equation 2 is difficult optimization problem (Krishnapriyan al.

---

Published as a conference paper at ICLR 2023
2021 Wang et al. 2021 Edwards 2022 There are several reasons for this, including gradient imbalances in the loss terms (Wang et al. 2021) D and ill-conditioning (Krishnapriyan et al. 2021), as well as only approximate enforcement of physical laws. In numerous scientific domains includ- ing fluid mechanics, physics, and materials science, systems are described by well-known physical laws, and breaking them often lead to nonphysical solutions. Indeed, if physical law is only approximately constrained (in this case, "soft-constrained." as with popular penalty-based optimiza- tion methods), then the system solution may behave qualitatively differently or even fail reach an answer.
In this work, we develop a method overcome these challenges by solving the PDE-constrained problem in Equation [ directly. We only consider the data-starved regime, i... we do not assume that any solution data available on the interior of the domain (however, note that when solution data is available, we can easily add data fitting loss to improve training). To solve Equation [ we design PDE-constrained layer for NNs that maps PDE parameters to their solutions, such that the PDE constraints are enforced as "hard constraints." Once our model is trained, we can take new PDE parameters and solve for their corresponding solutions, while still enforcing the correct constraint.
In more detail, our main contributions are the following:
We propose method to enforce hard PDE constraints by creating a differentiable layer, which we call Sare--oootruind-aayee or PDE-CI. We make the PDE-CL. differentiable using implicit differentiation, thereby allowing to train our model with gradient- optimization methods. This layer allows us to find the optimal linear combination func- tions in a learned basis, given the PDE constraint.
At inference time, our model only requires finding the optimal linear combination of the fixed basis functions. After using a small number of sampled points to fit this linear com- bination, we can evaluate the model
much higher resolution grid.
We provide empirical validation of our method on three problems representing different types of PDEs. The 2D Darcy Flow problem is an elliptic PDE on stationary (steady-state) spatial domain, the ID Burger's problem a mon-linear PDE on spatiotemporal domain, and the ID convection problem is hyperbolic PDE on a spatiotemporal domain. We show that our approach has lower error than the soft constraint approach when predicting solutions for new, unseen test cases, without having access to any solution data during training. Compared to the soft constraint approach, our approach takes fewer iterations to converge to the correct solution, and also requires less training time.
2 BACKGROUND AND RELATED WORK
The layer we design solves a constrained optimization problem corresponding to a PDE constraint. We outline some relevant lines of work.
Dictionary learning. The problem we study can be seen as PDE-constrained dictionary learning. Dictionary learning (Mairal et al., 2009) aims to learn an over-complete basis that represents the data accurately. Each datapoint is then represented by combining sparse subset of the learned basis. Since dictionary learning is a discrete method, it is not directly compatible with learning solutions to PDEs, as we need to be able to compute partial derivatives for the underlying learned functions. NNs allow us to do exactly this, as we can learn a parametric over-complete functional basis, which is continuous and differentiable with regard to both its inputs and its parameters.
NNS and structural constraints. Using NNs to solve scientific modeling problems has gained interest in recent years (Willard al., 2020). NN architectures can also be designed such that they are tailored to specific problem structure, e.g. local correlations in features (LeCun et al.] 1998 Bronstein et al.) 2017; Hochreiter & Schmidhuber 1997), symmetries data (Cohen & Welling 2016), convexity (Amos et al. 2017), or monotonicity (Sill 1997) with regard to input. This reduces the class of models to ones that enforce the desired structure exactly. For scientific problems, NN generalization can be improved by incorporating domain constraints into the ML framework. in order to respect the relevant physics. Common approaches have included adding PDE terms as part of the optimization loss function (Raissi al. 2019), using NNs to learn differential operators in
2

---

Published as a conference paper at ICLR 2023
PDEs such that many PDEs can be solved inference time (Li et al., 2020: Lu et al. 2021a 22l. and incorporating numerical solvers within the framework of NNs (Um et al., 2020). It is sometimes possible to directly parameterize Gaussian processes LLane--Hegrraan or NNs (Hendriks et al. 2020) to satisfy PDEs, and fit some desired loss function. However, in the PDEs we study, we cannot have closed-form parameterization for solutions of the PDE. Previous work in PDE-solving has tried to enforce hard constraints by enforcing boundary conditions (Lu al. 2021bb.
2018 Didliing et al.) 2017)
We instead enforce the PDE constraint on the interior domain.
Implicit layers. A deep learning layer is a differentiable, parametric function defined as So : i + y. For most deep learning layers, the two Jacobians Øf and Øf are computed using the chain rule. For these explicit layers, fe is usually defined as the composition of elementary operations for which the Jacobians are known. On the other hand, implicit layers create an implicit relationship between the inputs and outputs by computing the Jacobian using the implicit function theorem (Krantz & Parks 2002), rather than the chain rule. Specifically, if the layer has imput r E I output R and parameters E R", we suppose that y solves the following nonlinear equation g(2, y. 6) = 0 for some g. Under mild assumptions, this defines an implicit function fo x + y. In our method, the forward function solves a constrained optimization problem. When computing the Jacobian of the layer, it is highly memory inefficient to differentiate through the optimization algorithm (i.e., all the steps of the iterative solver). Instead, by using the implicit function theorem, a set of linear systems can be solved to obtain the required Jacobians (see Amos & Kolter y) Barratt (2018 8); Blondel let al.] (202T): Agrawal et al. (2019); El Ghaoui et al. (2021) and the Deep Implicit Layer NeurlPS 2021 tutoria for more details). Implicit layers have been leveraged in many applications, including solving ordinary differential equations (ODEs) (Chen et al. 2018), optimal power flow Donti et al.
2s de
(202T. and rigid many-body physics (de Avila Belbute-
al. 2018).
Differentiable physics. In a different setting, recent work has aimed to make physics simulators differentiable. The adjoint method (Pontryagin al. (962) is classically used in PDE-constrained optimization, and it has been incorporated into NN training (Chen et al. 2018 Zhang et al., 2019 Krishnapriyan et al. 2022). In this case, the assumption is that discrete samples from a function satisfying an unknown ODE are available. The goal is to learn the system dynamics from data. A NN model is used to approximate the ODE, and traditional numerical integration methods are applied on the output of the NN to get the function evaluation at the next timestep. The adjoint method is used to compute gradients with regard to the NN parameters through the obtained solution. The adjoint
method also allows for differentiating through physics simulators (Degrave et al. 2019] [Scoenholz & Cubuk] 2020; Um et al., 2020). Our setup is different. In our case, the underlying physical law(s) are known and the NN is used to approximate the solutions, under the assumption that there is no observational data in the interior of the solution domain.
METHODS
We describe the details of our method for enforeing PDE constraints within a NN model.
3.1
PROBLEM SETUP
Our goal to learn a mapping between a PDE parameter function Φ : X - IR and the corresponding PDE solution u(Ø)= X + R, where the domain X' is an open subset of Rdd for some d. The PDE parameters could be parameter functions such as initial condition functions, boundary condition functions, forcing functions, and/or physical properties such as wavespeed, diffusivity, and viscosity. We consider well-posed PDEs. following previous work exploring NNs and PDEs (Raissi et al. 2019 Li et al. 2020 Wang al.) 2021). Let F& be functional operator such that for all PDE parameter functions @ sampled from 4. the solution u(Ø) satisfies F (u())) = O. The inputs to our NN vary depending on the domain of interest and the PDE parameters. In the simplest case, the input is pair (ær, Φ(x)), where x E X and Φ(x) is the value of the PDE parameter at æ. The output of the NN is the value of the corresponding approximated solution us (e)), for a given x. We want to learn the mapping,
G :
(3)
PDE parameters
PDE solutions
http://implicie http://mpppltttt---eeeetttooooo .org/
3

---

Published as a conference paper at ICLR 2023
0.14
0.012
0010
G
0.008
0.006
0.004
u(Φ)
Figure 1: Mapping PDE parameters ø to PDE solutions u(Ø). The goal of our model is to learn mapping G
u()). without access solution data. As an example, we study the Darcy Flow PDE, which describes the chemical engineering problem of fluid flow through porous medium (Darcy) 1856). The system composed of two materials in a given spatial domain X' 1) each with specific diffusion coefficients which depend the position. The left figure shows 6, which encodes the locations and diffusion properties the two materials. The right figure shows the corresponding solution u(d). The function is solution of the Darcy Flow PDE with diffusion coefficients ø if, for all (x, y) E (0, 1)2, it satisfies -V (ø(2, () Vu(z, y)) = 1. The boundary condition is u(x y) 0, (z, E 1)0
We show an example such a mapping Figure Importantly, we consider only the unsupervised learning setting, where solution data in the interior domain of the PDE is not available for training model. this setting, the training done by only enforcing the PDE, and the initial and/or boundary conditions.
3.2 A DIFFERENTIABLE CONSTRAINED LAYER FOR ENFORCING PDEs
There are two main components to our model. The first component is a NN parameterized by 0, denoted by fe- The NN So takes the inputs described in Section 3.1 and outputs vector in RR- The output dimension N the number of functions in our basis, and is hyperparameter.
The second component of our model, the PDE-constrained layer PDE-CL, is main design contribution. We implement a layer that performs a linear combination of the N outputs from the first component, such that the linear combination satisfies the PDE on all points 2; in the discretized domain. Specifically, let w be the weights in the linear combination given by the PDE-CL. The output of our system is Ug = Σ wiff where fø is the i-th coordinate output of fe. We now describe the forward and backward pass of the PDE-CL.
Forward pass of the differentiable constrained layer. Our layer is a differentiable root-finder for PDEs, and focus on both linear and non-linear PDEs. As an example, describe a system based on an affine PDE, where F& is an affine operator, which depends on ø. In our experiments, we study an inhomogeneous linear system Section 4.1 a homogeneous non-linear system Section 4.2 and a homogeneous linear system in Section 4.3 The operator Go is linear when, for any two functions u, v from to R. and any scalar À E K, we have that,
(4) The operator F& is affine when there exists function b such that the shifted operator Fg is linear. Let Go be the linear part of the operator: Go = Fe
Go(u + Au) = (v( + AGo(v).
We define the PDE-CL to find the optimal linear weighting w of the N ID functions encoded by the first NN component, over the set of sampled inputs ø1, En. The vector W E R. solves the linear equation system,
Vj = 1 n,
Go
Eu wifh
(x,) b(zz)
Σ (. b(z,).
(5)
This linear system is discretization the PDE Fa(ue) = O: we aim to enforce the PDE at the sampled points 21, - Tnn The linear system has constraints and N variables. These
4

---

Published as a conference paper at ICLR 2023
both hyperparameters, that can be chosen to maximize performance. Note that once N is fixed, it cannot be changed. On the other hand, n can be changed at any time. When the PDE non-linear, the linear system is replaced by non-linear least-squares system, for which efficient solvers are available (Levenberg 1944 Marquardt et al. 1963).
Backward pass of the differentiable constrained layer. To incorporate the PDE-CL into an end- to-end differentiable system that can be trained by first-order optimization methods, we need to compute the gradients of the full model using an autodiff system (Bradbury et al. 2018). To this, must compute the Jacobian of the layer.
The PDE-CL solves a linear system of the form g(w, A, b) = Aw bb = 0 in the forward pass, where A E R=*N, w E IR,, b R", Differentiating with respect to A and b using the chain rule gives the following linear systems, which must be satisfied by the Jacobians au and 24
du 26
0 Ow's
Wk A dw
Vi, N, VjEl. , n,
1
(6)
0Ajk
0Ajk
0g:
A. Ow
ViEL N.
0 = 0b,
1i=j
Dbg
where 1inj is 1 when i = j and 0 otherwise. Given the size and conditioning of our problems, we cannot directly solve the linear system. Thus, we use an indirect solver (such as conjugate gradient (Hestenes & Stiefel 1952) or GMRES (Sand & Schultz| 1986)) for both the forward system Aw = and the backward system given by Equation 6 and Equation 7 We use the JAX autodiff' framework (Bradbury et al. 2018 Blondel et al. 2021) to implement the full model. We include an
analysis and additional information on the enforcement of the hard constraints in D We also include an ablation study in E to evaluate the quality of functions in our basis.
Loss function. Our goal is to obtain NN parameterized function which verifies the PDE over the whole domain. The PDE-CL only guarantees that we verify the PDE over the sampled points. In the case where N > n, the residual over the sampled points ; is zero, up to numerical error of the linear solver used our layer. It is preferable to not use this residual for training the NN as it may not be meaningful, and an artifact of the chosen linear solver and tolerances. Instead, we sample new points rf....., Ty, and build the corresponding linear system A', B. Our loss value is ||'''w b'112, where w comes from the PDE-CL and depends on A and b, not A', b'. We compute gradients of this loss function using the Jacobian described above. Another possibility is to use > N. In this case, the residual A bll will be non-zero, and while the "hard constraints" will not be satisfied during training, we can minimize this residual loss directly. Let ((X) denote the uniform distribution over our (bounded) domain X. Formally, our goal is to solve the bilevel optimization problem.
min E E (a.................. (a ( ((()) (l (,) II. (...)........)) 6' 2i....,
0
s.t. w arg argmin min A 21 M 0yw 1.. 0. 0)
(8)
We approximate this problem by replacing the expectations by sums over finite samples. The ma- trices A, A', and vectors b, are built by applying the differential operator Fs to each function in our basis fg. using the sampled gridpoints. It is straightforward to extend this method in the case of non-linear PDEs by replacing the linear least-squares with the relevant non-linear least squares problem.
Inference procedure. At inference, when given new PDE parameter test point d. the weights A are fixed as our function basis is trained. In this paragraph, we discuss guarantees in the linear PDE case. Suppose that we want the values of up over the (new) points ar! ...., lest If test < N. we can fit w in the PDE-CL using all of the test points. This guarantees that our model satisfies the PDE on all test points: the linear system in the PDE-CL, is underdetermined. In practice, nes often larger than N. In this case, we can sample J C {1.... n SS,), [J] N. and fit the PDE-CL over these points. Over the subset (aten, .j E J}. the PDE will be satisfied. Over the other points, the residual may be non zero. non-zero. Another option is to fit the PDE-CL using all the points lest ...., ne in which case the residual may be non-zero for all points, but is minimized on average over the
lest
5

---

Published as a conference paper at ICLR 2023
4002
00005
2008
00000
51006
0.006
4.0006
00005
(a) Target
(b) Hard-constrained difference
(c) Soft-constrained difference
Figure 2: Heatmaps of Darcy Flow example test set predictions. We compare our hard- constrained model and the baseline soft-constrained model on a test set of new diffusion coefficients the same except for additional PDE-CL in the hard-constrained model. (a) Target solutions of a subset of PDEs in the test set. (b) Difference between the predictions of our constronnt PDE-CL model and the target solution.
The NN architectures
(c) Difference between the predic- tions of the baseline soft-constrained model and the target solution. Over the test dataset, our model achieves 1.82% 0.04% relative error and 0.0457 0.0021 interior domain test loss. In con- trast, the soft-constrained model only reaches 3.86% + 0.3% relative error and 1.1355 + 0.0433 interior domain loss. Our model achieves 71% less relative error than the soft-constrained model. While the heatmaps show subset of the full test set, the standard deviation across the test set for our model very low, as shown by the box plot in Appendix
sampled points by PDE-CL. Our method controls the trade-off between speed and accuracy: choosing a larger N results in larger linear systems, but also larger sets of points on which the PDE enforced. A smaller N allows for faster linear system solves. Once is fit, can query uo for any point in the domain. In practice, we can choose |J] to be much smaller than neest here, the linear system we need to solve much smaller than the linear system required by a numerical solver.
EXPERIMENTAL RESULTS AND IMPLEMENTATION
We test the performance of our model on three different scientific problems: 2D Darcy Flow (Sec- tion 4.1.. ID Burgers' equation (Section 4.2). and ID convection (Section 4.3). In each case, the model is trained without access to any solution data in the interior solution domain. The train- ing contains 1000 PDE parameters d. The model is then evaluated
separate test set with M = 50 PDE parameters & that are not seen during training. We compare model results on the test set using two metrics: relative T.2 error 1 M 2i=1 [ma( and the PDE residual loss Iu(,))l2
M DM [|F (u,) which measures how well the PDE is enforced on the interior domain. We demonstrate that our constrained NN architecture generalizes much better on the test set than the comparable unconstrained model for all three problems
4.1 2D DARCY FLOW
We look at the steady-state 2D Darcy Flow problem, which describes, for example, fluid flow through porous media. In this section, the PDE parameter (denoted by ø in the previous sections as a general variable) is E 0 (0, 1)2: R,,, )), diffusion coefficient. The problem formulated as follows:
(x(x)V((x)) f(x),
Ya E (0, 10,
=
Va E 0 1)2.
Here, is the forcing function (f ( 1)2; R). and u(x) = 0 is the boundary condition. The
differential operator then (u) = - Vu(z). Given set of variable coefficients v, the goal is to predict the correct solution u. The v(x) values are generated from a Gaussian and
We used single Titan RTX GPU for each in our experiments.

---

Published as a conference paper at ICLR 2023
10
Hard constraint (ours) constraint
Hard constraint (ours)
constraint
Lowest relative error for soft constraint
Intorior test loxs loxs
Relative test error
104
102
10-1
10
10 2
10
10
10
Batch
Batch
(a) PDE residual loss on test set
(b) Relative error on test set
Figure 3: 2D Darey Flow: Error test set during training, We train a NN architecture with the PDE residual loss function ("soft constraint" baseline), and the same NN architecture with our PDE-CL ("hard constraint"). During training,
track error the test set, which plot log-log scale. (a) PDE residual loss on the test set, during training. This loss measures how well (b) Relative error on the test set, during training. This metric measures the distance between the predicted solution and the target solution obtained via finite differences. Both measures show that our hard-constrained PDE-CL model starts at a much lower error (over an order of magnitude lower) on the test set at the very start of training, and continues to decrease as training proceeds. This is particularly visible when tracking the PDE residual test loss.
the PDE is enforced.
then mapped to two values, corresponding to the two different materials in the system (such as the fluid and the porous medium). We follow the data generation procedure from Li et al. (002)). We use the Fourier Neural Operator (FNO) (Li et al., 2020) architecture, trained using PDE residual loss as the baseline model "ooottcconttrainedd) Our model uses the FNO architecture and adds our PDE-CL ("hard-com contraannna.i."" The domain (0, 1)² is discretized over fy X ny points. For each point on this grid, the model takes input the coordinates, E (0, 1)2, and the corresponding values of the diffusion coefficients, ((r). The boundary condition is satisfied by using mollifier function (Li jet al., 2020), and so the only term in the loss function is the PDE residual. We use a constant forcing function I equal to 1. We provide more details on our setup and implementation in Appendix
We compare visually how close our hard-constrained model is to the target solution (Figure 2b). and how close the soft-constrained baseline model is to the target solution (Figure 2cp. Our hand-constrained model is much closer to the target solution, as indicated by the difference plots mostly being white (corresponding to zero difference).
Results. We plot example heatmaps from the test set in Figure 2
During the training procedure for both hard- and soft-constrained models, we track error on an unseen test set of PDE solutions with different PDE parameters from the training set. We show these error plots in Figure 3 In Figure 3a, our model starts at a PDE residual test loss value two orders of magnitude smaller than the soft constraint baseline. The PDE residual test loss continues decrease as training proceeds, remaining significantly lower than the baseline. Similarly, in Figure 3b] we show the curves corresponding to the relative eгror and the PDE residual loss metric on the test dataset. Our model starts at much smaller relative error immediately and continues to decrease, achieving final lower test relative error.
On the test set, our model achieves 1.82% + O. 0.04% relative error, versus 3. 3.86% + 0.3% for the soft-constrained baseline model. Our model also achieves 0.0457 + 0.0021 for the PDE residual test loss, versus 1.1355 + 0. 0.0433. Our model has almost two orders of magnitude lower PDE residual test loss, and it has significantly lower standard deviation. On the relative error metric, our model achieves 71% % improvement over the soft-constee conttrinst model. While the example heatmaps show a subset of the full test set, the standard deviation across the test set for our model is very low. This indicates that the results are consistent across test samples.
7

---

Published as a conference paper at ICLR 2023
635
601
60
-687
-44
(a) Target
(b) Hard-constrained difference
(c) Soft-constraneed difference
Figure 4: Heatmaps of ID Burgers' example test set predictions. We compare our hard- constrained model and the baseline soft-constrained model on a test set of new initial conditions un- Both architectures are the same, except for additional PDE-CL in the hard-constrained model. (a) Target solutions of a subset PDEs in the test set. (b) Difference between the predictions of hard-constrained model and the target solution.
(c) Difference between the predictions of the
baseline soft-constrained model and the target solution. Over the test dataset, our model achieves 1.11 + 0.11% relative error. The baseline coft-constrained model achieves only 4.34% + 0.33% relative error. We use the same base network architecture (MLPs) for both the soft-constrained and hard-constrained model. The errors in both models are concentrated around the "sharp" features in the solution, but these errors have 4x higher magnitude in the coft-constraieed model.
4.2 ID BURGERS' EQUATION
We study a non-linear ID PDE, Burgers' equation, which describes transport phenomena. The problem can be written as,
Ou(x, t)
1 (a, gu((zt)) Ox.
0t + 2 dr
E (0,1), (0,1), x (0,1), 1), E R, (0,
(10)
u(r,0) = un (z), u(I,t) u(x 1,t),
Here, ug is the initial condition, and the system has periodic boundary conditions. We aim to map the initial condition uo to the solution We consider problems with a fixed viscosity parameter of 0.01. We follow the data generation procedure from Li et al. (2020), which can be found here We use physics-ifformed DeepONet baseline model (Wang et al.) (02T) with regular multi-layer perceptrons as the base NN architecture, trained using the PDE residual loss. Our hard-constraing model composed of stacked dense layers and our PDE-CL, which allows for fair comparison. Because this PDE is non-linear, our PDE-CL solves non-linear least-squares problem, using the PDE residual loss and the initial and boundary condition losses.
Results. We plot example heatmaps from the test hard-constrained model is to the target solution (Figure 4b and similarly for the soft-constrained baseline model (Figure 4c). The solution found by our hard-constrained model is much closer to the target solution than the solution found by the baseline model, and our model captures "sharp" features in the solution visibly better than the baseline model. Our hard-constrained model achieves 1.11 + 0.1 1.11% % relative error after less than 5, 000 steps of training, using just dense layers and our PDE-CL. In contrast, the soft-constrained baseline model only achieves 4.34% 0.33% relative eπror after many more steps of training.
Figure We compare how close our
During the training procedure for both hard- and soft-constrained models, we track the relative error on validation set of PDE solutions with different PDE parameters from the training set. We show the error plot in Figure 5 The target solution was obtained using the Chebfun package (Driscoll et al. 2014), following Wang et al. 2021). Our model achieves lower error much earlier in training. The large fluctuations are due to the log-log plotting, and small batches used by our method for memory reasons.
8

---

Published as a conference paper at ICLR 2023
consiruint soft constraitt
139
Relsativ
Vunna wmwnmaman
10"4
10 30 10
Figure 5: 1D Burgers' equation: Error on validation set during training. We train NN with the PDE residual loss function ("soft constraint" baseline) and the same NN architecture with our PDE-CL ("hard constraint"). Both architectures are MI.P.. During training, we track relative error on the test set, which plot log-log scale. Our hard-constrained model Iearns low error predictions much carlier in training. The hard constrained model achieves lower relative error than the soft-constrained method.
4.3 1D CONVECTION
We study ID convection problem, describing transport phenomena. The problem can be formu- lated as follows:
Ou(x,t)
0, Ou(x,t) t)
+ ((x))
* E (0, 1),1 (),
Øt.
Ox
h(x) = sin(TT)) T g(t) = sin (
x E (0,
(11)
te (0, 1).
Here, f((z) is the initial condition (at t = 0), g(t) is the boundary condition (at x 0), and (xx) represents the variable coefficients (denoted by Φ in Section 3). Given of variable coefficients, 8(x), and spatiotemporal points (z,, f,). the goal is to predict the correct solution u(x,)) t). We provide results and more details in Appendix(A]
CONCLUSIONS
We have considered the problem of mapping PDEs to their corresponding solutions, in particular in the unsupervised setting, where no solution data is available on the interior of the domain during training. For this situation, we have developed method to enforce hard PDE constraints, when training NNs, by designing differentiable PDE-constrained layer (PDE-CL). We add our layer to any NN architecture to enforce PDE constraints accurately, and then train the whole system end- to-end. Our method provides two hyperparameters. We evaluate our proposed method on three problems representing different physical settings: 2D Darcy Flow problem, which describes fluid flow through a porous medium; 1D Burger's problem, which describes viscous fluids and dissipative system; and a ID convection problem, which describes transport phenomena. Compared the baseline soft-constrained model, our model can be trained in fewer iterations, achieves lower PDE residual error (measuring how well the PDE is enforced on the interior domain), and achieves lower relative error with respect to target solutions generated by numerical solvers.
to control the trade-off between speed and accuracy through
Acknowledgements.
The authors would like to thank Quentin Berthet, David Duvenaud, Romain Lopez. Dmitriy Morozov, Parth Nobel, Daniel Rothchild, Hector Roux de Bézieux, and Alice Schoe- nauer Sebag for helpful comments on previous drafts. MWM would like to acknowledge the DOE, NSF, and ONR for providing partial support of this work. This material is based in part upon work

---

Published as a conference paper at ICLR 2023
supported by the Intelligence Advanced Research Projects Agency (IARPA) and Army Research Of- fice (ARO) under Contract No. W911NF-20-C-0035. ASK and MWM would like to acknowledge 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. A00AC0000CCH1211.
10

---

Published as a conference paper at ICLR 2023
REFERENCES
Akshay Agrawal, Brandon Amos, Shane T. Barratt, Stephen P. Boyd, Steven Diamond, and J. Zico Kolter. Differentiable convex optimization layers. In NeurIPS. 2019.
Brandon Amos and J. Zico Kolter. Optnet: Differentiable optimization a layer in neural networks. In ICML, 2017. 3
Brandon Amos, Lei Xu, and J. Zico Kolter. Input convex neural networks. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 146-155 PMLR, 06-11 Aug 2017.
URL
ngs ml wtt.r..press/////e
7b. html
2
Shane T. Barratt. On the differentiability of the solution to convex optimization problems. arXiv: Optimization and Control, 2018.
Mathieu Blondel, Quentin Berthet, Marco Cuturi, Roy Frostig, Stephan Hoyer, Felipe Llinares-
Lopez, Fabian Pedregosa, and Jean-Philippe Vert. Efficient and modular implicit differentiation. ArXiv, abss/215..1118.. 2021.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman- and Qiao Zhang. JAX: composable transformations Python+NumPy programs, 2018. URL http:
//github. cttp crtoom//ooggee///// jax
5
Michael M. Bronstein, Joan Bruna, Yann LeCun, Arthur D. Szlam, and Pierre Vandergheynst. Ge- ometric deep learning: Going beyond euclidean data. IEEE Signal Processing Magazine, 34: 18 18-42, 2017.
Tian Qi Chen, Yulia Ruhanova, Jesse Bettencourt, and David Kristjanson Duvenaud. Neural ordi- nary differential equations. ArXiv, abs/1806.07366, 2018.
Taco Cohen and Max Welling. Group equivariant convolutional networks. In ICML, 2016. 2
Henry Darcy. Les fontaines publiques de la ville de Dijon exposition et application des principes à suivre des formules à employer dans les questions de distribution d'eau. Victor Dalmont, 1856.
Filipe de Avila Belbute-Peres, Kevin Smith, Kelsey Allen, Josh Tenenbaum, and Zico Kolter. End- to-end differentiable physics for learning and control. Advances in neural information processing systems, 31:7178-7188 2018. g
Jonas Degrave, Michiel Hermans, Joni Dambre, and Francis wyffels. A differentiable physics engine for deep learning in robotics. Frontiers in Neurorobotics, 13, 2019. ISSN 1662-5218. doi: 10. 20899/fnbot.2199.000
frontiersinn http://.......a.ttteee......... htLhttss:: htttpp/:/ww..oomierr......
Enbot 2019. 99, 00006
Priya L Donti, David Rolnick, and J Zico Kolter. De3: A learning method for optimization with
hard constraints. arXiv preprint 122:5 2021.
A Driscoll, N. Hale, and L. N. Trefethen. Chebfun Guide. Pafnuty Publications, 2014. URL http:/ www. chebfun. org/ /wcs/ /uide/
Chris Edwards. Neural networks learn to speed up simulations. Communications of the ACM, 65
(5):27-29, 2022. 2
Laurent Ghaoui, Fangda Gu, Bertrand Travacca, and Armin Askari. Implicit deep learning, SIAM J. Math. Data Sci., 3:930-958, 2021. 3
Johannes N. Hendriks, Carl Jidling, Adrian G. Wills, and Thomas Bo Schön. Linearly constrained neural networks. ArXiv, abs/2002.01600, 2020.
Magnus R. Hestenes and Eduard Stiefel. Methods of conjugate gradients for solving linear systems. Journal of research of the National Bureau of Standards, 49:409-433 1952. 5
11

---

Published as a conference paper at ICLR 2023
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Computation, 9:1735- 1780, 1997. 2
Carl Jidling, Niklas Wahlstrom, Adrian G. Wills, and Thomas Bo Schön. Linearly constrained gaussian processes. In NIPS, 2017.
Steven George Krantz and Harold R Parks. The implicit function theorem: history, theory, and applications. Springer Science & Business Media, 2002.
Aditi Krishnapriyan, Amir Gholami, Shandian Zhe, Robert Kirby, and Michael W Mahoney. Char- acterizing possible failure modes in physics-nformed neural networks. Advances in Neural Information Processing Systems, 34, 2021. 2
Aditi S Krishnapriyan, Alejandro F Queiruga, N Benjamin Erichson, and Michael W Mahoney. Learning continuous models for continuous physics. arXiv preprint ArX1v:2202.08444, 2022.
Markus Heggemanggrmann.... Algorithmic linearly constrained gaussian processes. NeurIPS, 2018.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient- learning applied to document recognition. Proc. IEEE, 86:2278-2324, 1998. 2
Kenneth Levenberg. A method for the solution of certain non linear problems least squares.
Quarterly of Applied Mathematics, 2: 2:164-168, 1944. 5
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 arXiv:2010.08895 2020.
1.3
Zongyi Li, Hongkai Zheng, Nikola Kovachki, David Jin, Haoxuan Chen, Burigede Liu, Kamyar Azizzadenesheli, and Anima Anandkumar. Physics-informed neural operator for learning partial differential equations. arXiv preprint arX.v:2111.03794. 2021. 17
Lu Lu, Pengzhan Jin, Guofei Pang, Zhongqiang Zhang, and George Em Karniadakis. Learning nonlinear operators via deeponet based on the universal approximation theorem of operators, Nat. Mach. Intell., 3:218-229, 2021a. 3
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, 4 2021b.
Julien Mairal, Francis Bach, Jean Ponce, and Guillermo Sapiro. Online dictionary learning for sparse coding. Proceedings of the 26th Annual International Conference on Machine Learning, ICML. '09, pp. 689-696, New York, NY, USA, 2009. Association for Computing Machinery. ISBN URL
1781605585161. doi: 10.1145/553774.155466 1553374.15530
1145/
Donald W. Marquardt, E. I. duPont, Bennett R, and George C Burrell, An algorithm for least-squares estimation of nonlinear parameters. 1963.
Lev Semenovich Pontryagin, EF Mishchenko, VG Boltyanskii, and RV Gamkrelidze. The mathe- matical theory of optimal processes. 1962.
M. Raissi, Perdikaris, and G.E. Karniadakis. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equa- tions. Journal of Computational Physics, 378:686 707, 2019. ISSN 0021-9991. doi: https://doi.
1gg/10.1016//.cpp.000000 10.045. URL httpa: www. www.:ssciieddiirettccomm
axticleep11/5002199111
2
Youcef Saad and Martin H. Schultz. Gmres: generalized minimal residual algorithm for solving nonsymmetric linear systems. Siam Journal on Scientific and Statistical Computing, 7:856-869, 1986. 16
12

---

Published as a conference paper at ICLR 2023
Samuel S. Schoenholz and Ekin D. Cubuk. Jax m.d. framework for differentiable physics. Advances in Neural Information Processing Systems, volume 33. Cur- Associates, Inc., 2020. URL https::/papprss wttps/ccp.ppee/////////////////. 83d3d4b6:957751 5e1 11aaaaaccca0333Paapppe pdf
Joseph Sill. Monotonic networks. In M. Jordan, M. Kearns, and S. Solla (eds.),
Advances in Neural Information Processing Systems. volume 10. MIT Press,
1997.
URL httpo::/prrooedlngg..ee..e. hww.pps.......pie.co.........m.o httpa///sssccc/pp/////////1/ 2r pdff
Kiwon Um, Robert Brand, Yun Raymond Fei, Philipp Holl, and Nils Thuerey, Solver-in-the-loop: Learning from differentiable physics to interact with iterative pde-solvers. Advances in Neural Information Processing Systems, 33:6111-6122, 2020.
Sifan Wang, Hanwen Wang, and Paris Perdikaris. Learning the solution operator of parametric partial differential equations with physics-informed deeponets. Science Advances, 7(40), 2021.
14
Jared Willard, Xiaowei Jia, Shaoming Xu, Michael Steinbach, and Vipin Kumar. Integrating
physics-based modeling with machine learning: A survey, arXiv preprint arXiv:2003.04919,
1(1):1-34, 2020. 2
T. Zhang, Z. Yao, A. Gholami, K. Keutzer, J. Gonzalez, G. Biros, and M. W. Mahoney. ANODEV2: A coupled Neural ODE evolution framework. Technical Report Preprint: arXiv:1906.0459, 2019. 3
13

---

Published as a conference paper at ICLR 2023
015
3.0
06
900
9.55
02
-013
02
035
(a) Target
(b) Hard-constrained difference
(c) Soft-constranned difference
Figure A.I: Heatmaps of ID convection example test set predictions. We compare hard- constrained model and the baseline soft-constrained model on a test set of new wavespeed parame- ters B. Both architectures are the same, except for additional PDE-CL in the hard-constrained model. (a) Target solutions of a subset of PDEs in the test set. (b) Difference between the predic-
tions of hard-constrained model and the target solution. (c) Difference between the predictions of the baseline soft-constrained model and the target solution. Over the test dataset, our model achieves 1.32% 0.02% relative error and 9.84 2.15 PDE residual test loss. In contrast, the soft-constrained model only reaches 2. 2.59% + 0.15% relative error and 774 1.2 PDE residual test loss. Our model achieves 49% less relative error than the soft-constrained model. The errors in both models are concentrated around the "sharp" features in the solution, but these errors have higher magnitude in the coft-constrained model.
A ID CONVECTION
We study a ID convection problem, describing transport phenomena. The problem can be formu- lated follows:
Ou(r,t)
Ou(x,t)
+ B(x))
(0, (), (0, 1),
Ōt
Ox
h(x) sin(TTx), TT g(t) sin
E (0,
(12)
t (0,1).
2
Here, h(x) is the initial condition represents the variable coefficients (denoted by Φ in Section 33. Given set of variable coefficients, 8(x). and spatiotemporal points (z,, tx), the goal is to predict the correct solution u(x, The 3(x) values are generated the same manner as in Wang et al. 2021) via (xx) v(x) min, v(x) + 1, where u(x:) is generated from Gaussian random field with length scale of 0.2. We physics- informed DeepONet baseline model (Wang et al. 2021), , trained with the PDE residual loss. Our hard-constrained model is composed of stacked dense layers and our PDE-CL, which allows for fair comparison. We provide more details on our setup and experiments in Appendix
0), g()) is the boundary condition (at x 0), and ((x)
Results. We plot example heatmaps from the test set in Figure A.1 We compare how close our hard-constrained model is to the target solution (Figure Ab and similarly for the soft-constrained baseline model (Figure A. The solution found by our card-constrained model much closer to the target solution than the solution found by the baseline model, and our model captures "sharp" features in the solution visibly better than the baseline model.
During the training procedure for both hard- and coft-constrained models, we track error an unseen test set of PDE solutions with different PDE parameters from the training set. We show these error plots in Figure In Figure the PDE residual loss for the hard-constrained model starts close to six orders magnitude lower than for the soft-constrained model, and it continues to remain low. In Figure A.2b we track the relative eтror with respect to the target solution obtained via a Lax- Wendroff scheme. Similarly, we see that our model starts at much smaller relative error immediately and also continues to decrease. Our model achieves 1.32% + 0.02% relative error and 9.84 + 2.15 PDE residual test loss, versus 2.59% + 0.15% and 774 1.2 for the soft-constrained baseline model. On the relative error metric, our model achieves 49% improvement the soft-
14

---

Published as a conference paper at ICLR 2023
107
constraint (aurs) Soft constraint
Hard constraint (ours)
Soft constraint
Lowest relative error for constraint
Interier test loss
10
Relative test eroor
10
10
10"
10a
10-1
10a
10P
10
10a
10a
10
10a
10a
10'
102
10"
10P
Batch
Batch
(a) PDE residual loss on test set
(b) Relative error on test set
Figure A.2: 1D convection: Error on test set during training. We train NN with the PDE residual loss function ("soft constraint" baseline) and the same NN architecture with our PDE-CL ("hard constraint"). During training, we track error on the test set, which we plot on log-log scale. (a) PDE residual loss on the test set, during training. We observe that the NN starts by fitting the initial and boundary condition regression loss during training, which explains why the PDE residual loss seems to go up initially. (b) Relative error on the test set, during training. Both measures show that our hard-constrained model starts at a much lower error on the test set at the very start of training. The grey, dashed line shows that the hard-constrained model achieves the same relative error as the soft-constrained model in over 100x fewer iterations, and ultimately achieves lower relative error. Wall-time comparison figures are given in Appendix B
10
Hard constraint (ours) Soft constra int
Hand constraint lours)
Soft constraint
10a
relative crror constraint
Vnterfor test 1600
10
Relative tast enter
307
101
101
10:3
10.
10a
10:1
10a
101
10"
10'
Runtime
Runtima
(a) PDE residual loss on test set
(b) Relative error test set
Figure B.I: Walltime plots for 1D convection. During the training procedure, we track error on an unseen test set. Our hard- constrained model reaches the optimal accuracy of the soft-constrained model in 10x less time.
constrained model. The standard deviation of our model for the relative error metric over the test dataset is also small. Our model trains faster than the soft-constrained method, due to much higher gains in accuracy per batch, even though each batch is slower (see Figure B.I.
B DETAILS ON THE ID CONVECTION PROBLEM
Experiment setup and implementation details. In this setting, the inputs to the models are two sets of ((x1,tt),... ti)...., (In, t,)), and ((<), "........ (2, (Iar,,,)) (,))) sampled points within the do- main X (interior points) and the corresponding 3(x) values. We use the former for fitting the PDE-CL, and the latter for computing the residual loss function. We also require a set
(( (n+1): -.., (Intn's ()) of sampled points the initial condition (t = 0) and boundary
15

---

Published as a conference paper at ICLR 2023
0.045
0.040
0.035
0.030
8
0.025
0.020
0.015
0.010
0.005
Soft constraint
Hard constraint (ours)
Figure B.2: 1D convection: Box plots showing
test set. We show the distribution of errors over the test set, at the end training. Our hard-constrained model has both a lower error and lower standard deviation compared the soft-constrained model.
condition (a 0). The training optimization problem is formulated as follows:
min Σ Lu( + [|Fa(uo: ww)) s.t. V3, F1 ..,)),,,))). = =0, =0 (13)
with,
Cp(ue)
(ual (3, Intis bnts) - u(3, ants, (n+:)) Σ i=1
Ful (uul(B,01.t1) (0 + 3(x1) aug(3, 21, at
@us(3, In. tn)
Ous (B, xn, tn)
B(xx)
Ox
Øl
Or
where 0 corresponds to parameters of the NN, u(x, is the solution the initial and boundary con- ditions, and F(ua)
PDE constraint that must be satisfied. The loss term L(us) is a regression loss over the initial and boundary conditions. The forward pass of the PDE-CL solves the following equality constrained problem,
min ((fow) s.t. ww -
(14)
where fo refers to the outputs of the base NN, on which we stack the PDE-CL. Since the ini- tial/boundary condition regression loss uses a quadratic penalty, this equality constrained problem in fact a convex equality constrained quadratic problem (EqQP), which is equivalent to linear sys- tem. We solve this linear system using GMRES (Saad Schultz 1986). compute the Jacobian via implicit differentiation with respect to Equation [4
In practice, we sample 750 points for the PDE-CL, and sample a separate 250 points for computing the residual in the loss function. To ensure fairness, sample 1000 points for the soft-constrained method, which are all used to compute the residual in the loss function. We use N=600 for the number of basis functions in the PDE-CL.
We show plots against wall time in Figure B.1 and the distributions of errors over the test set in Figure B.2
C DETAILS ON THE DARCY FLOW PROBLEM
Experiment setup and implementation details.
Our goal is to find parameters 0, which solve,
min Σ (FF) f(x))11 6
(15)
s.t. Vv, Fuo) f(x)-
16

---

Published as a conference paper at ICLR 2023
10
Hard coratroint (ours) Soft constraint
10
Hard constraint lours)
Soft constraint
Lowest relative
Interir test lozs
Reletive test eror
104
10-
10
10
10
10
10a
Runtime
Rumbrme
(a) PDE residual loss on test set
(b) Relative error on test set
Figure C.1: Walltime plots for Darey Flow. During the training procedure, we track error on an unseen test set. Our hard- constrained model achieves higher accuracy much more quickly than the soft-constrained model.
0.12
0.08
0.06
0.04
0.02
Soft constraint
Hard constraint (ours)
Box plots showing error over test set. We show the distribution of errors over the test set, at the end of training. Our hard-constrained model has both a lower error, as well as significantly lower standard deviation as compared to the soft-constrained model.
Figure C.2: 2D Darcy Flow:
By design, the objective function for feasible parameters is zero While numerical issues may prevent exact feasibility, solving the equality constrained problem by additionally minimizing the PDE residual helps the training procedure. The soft-constrainet training method is trained only by minimizing the PDE residual. We train the FNO model using the same hyperparameters as Li et al. (02TT. We denote the FNO model part of the architecture as fa. The PDE-CL constrains the output of the FNO model by solving the linear system,
Vu, Fo(fo) (xx)
(16)
To train our model, we compute the Jacobian of this layer via implicit differentiation, with respect this linear system equation.
We sample 3721 points for the PDE-CL. We also sample 3721 points for the soft-constrained method, which are all used compute the residual the loss function. We use N=4000 for the number of basis functions in the PDE-CL.
We show plots against wall time Figure and the distributions of errors over the test set in Figure C.2
D HARD CONSTRAINTS BOUND
To provide an evaluation of how hard the hard constraints are, we conduct an additional study where look at the error in prediction for points sampled and used to fit the PDE-CL. against points that
17

---

Published as a conference paper at ICLR 2023
100
Error for points used in PDE-CL Error for points NOT used in PDE-CL
80
60
40
-0.10
-0.05
0.00
0.05
0.10
0.15
Figure D.I: Histogram of errors: Error for points sampled by the PDE-CL, versus error for points not sampled. We consider a trained model, and perform inference on random PDE instance. In this plot, consider the ID convection setting. The histogram shows that the enror for points used in the PDE-CL (1000 points) is about the same as error for points not used for the PDE-CL (9000 points). This demonstrates that do not need to fit the PDE-CI EEECLLod on all points of the grid.
not sampled. With trained hard-constrainee model for ID convection, we sample batch of points and create density plot showing the error as a function of points used for fitting the PDE-CL and points not used for fitting the PDE-CL. The histogram in Figure D. shows that the emrors qualitatively the same between points used for fitting the PDE-CL and those not used. There are 1000 points used for fitting the PDE-CL, and 9000 points not used. Our results show that model achieves low error, even outside of the points used for fitting the PDE-CL.
E ABLATION: EVALUATING THE QUALITY OF THE LEARNED BASIS FUNCTIONS
We implement experiment to evaluate the quality (and the advantage) of our learned basis func- tions, compared to cubic interpolation. This experiment aims to understand whether our learned functions are useful outside of the points used for the constrained problem.
Problem setup. We start with model trained on the ID convection problem. The model was trained by sampling 750 points for fitting the PDE-CL, and 250 different points for the residual in the objective functions. The points were sampled from 100x 100 grid (10,000 points total). We sample 750 points from the grid, and solve the corresponding PDE-CL, which gives us candidate PDE solution. For baseline, We interpolate the solution find on these 750 points to the 10,000 points using a scipy. opt imi meiz''s 's cubic interpolation. We also interpolate to a 1000x1000 grid to see how our model's performance scales with resolution.
We compare the above results with the output of trained hard-constrained model. Once the linear combination weights are fit (same as the baseline), we now use our learned basis functions perform inference test dataset in Figure F. la] The figure shows that using our model reduces the error, as compared to using a standard interpolation on the hard-constrained points.
Results.
all 10,000 (or I million) points. We plot our results the
18

---

Published as a conference paper at ICLR 2023
0.07
0.07
Relalve
0.03
0.0
0.00
Intorporiedd
0.0
Laamned
(a) 100 grid
(b) 1000x 1000 grid
Figure E.1: Quality of learned basis functions. We compare the interpolation from the hard- constrained points against our model's learned prediction. Our learned basis functions have lower error, compared to the baseline interpolation. The enror gap increases with higher resolution the grid of interest (i.e., a finer discretization). Our learned basis functions are 36% more accu- rate than the baseline interpolation for the 100x100 grid, and 37% more accurate than the baseline interpolation for the 1000x1000 grid.
F COMPARISON TO NUMERICAL SOLVERS
We compare the complexity of our PDE-CL framework against numerical methods.
Problem setup. We define x n4 grid over 1D) domain with with samples. We have time horizon of [0, T), with n, samples. this case, we assume that the PDE linear. Let us suppose that we set the number of basis functions to N and the number of sampled points for solving the PDE-CL to n Tinterior + NJc + BBc. The PDE-CL will then solve (n + NJC + NBC) x N linear system. Our method currently results in dense linear system. Solving this linear system has complexity O(max(Panterior + dlE + nBc, NY x min(n dic + nBc, N)). This is added to forward pass using the NN on the whole grid, once the optimal linear combination has been computed. Fortunately, this forward pass is embarrassingly parallel.
RBC,N)),
On the other hand, finite difference method such Crank-Nicolson or Len-rof requires solv- ing a tri-diagonal system of size n. at each f step. This yields an overall complexity of O(nz x n.). Comparing the complexity of both methods, the PDE-CL asymptotically faster than numerical methods when,
max N)² min(n, N) < ny n4y
(17)
In our current framework, inference marginally slower on par with numerical solvers. How- ever, as we increase the resolution of our grid (finer discretization), our method compares favorably the numerical solver-our computational cost increases more slowly than the numerical solver. Additionally, the operations in our PDE-CL are poorly optimized for current hardware, i.c., GPU uti- lization is low. Our method will greatly benefit from future improvements in hardware acceleration, which is still nascent field in the context of linear solvers on GPUs.
19


Reviews:
[{"review_id": "yLOyVBA0qeQ", "reviewer": "Reviewer_aAWN", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "The reviewed work proposes a novel approach for representing solution operators of linear partial differentials as neural networks without the need for training data provided by solution pairs. Given an input function (right-hand side, initial conditions, boundary conditions, etc.) the proposed work uses a Fourier neural operator (FNO) to learn basis functions. Then, a classical linear solver is used to compute a linear combination of these basis functions that best solves the PDE. At training time, this linear system solve is treated as an \"implicit layer,\" and gradients are propagated through it. Thus, the FNO essentially solves a dictionary learning problem where it learns to automatically generate the best basis functions in which to express the solution. \n\nThis approach is compared to a physics-informed deep-O-net, showing minor improvements in accuracy vs. training time. ", "summary_of_the_review": "At present, I believe that the limitations of the reviewed work compared to standard physics-informed FNO/DeepONet, notably the possibly increased inference time and the limitation to linear problems, are excessive given the relatively modest performance gains. Therefore, I recommend the rejection of the paper in its present form. \nHowever, given how difficult it is to train physics-informed neural operators without training data, an improved version of the manuscript that addresses these issues could merit acceptance.", "strength_and_weaknesses": "### Strengths\nUsing neural networks for dictionary learning of optimal basis functions is an interesting idea that deserves further investigation. To the best of my knowledge, this idea is novel and given how well-understood the solution of linear systems is, it seems like a good idea to leverage these tools more when using neural networks to solve PDEs.\n\n### Weaknesses\nIn my opinion, the paper has three main weaknesses. \n\nThe first is the restriction of the proposed approach to linear PDEs. This loses the generality of other neural-network-based approaches, which is arguably one of their greatest strength. \n\nThe second is the requirement to perform a linear solve at test time, which potentially disposes with the cheap online cost of existing physics-informed methods. (It would be helpful if the authors could provide a comparison of the inference cost).\n\nThe third weakness is the overall minor improvement over the baseline, especially when accounting for the cost of the implicit layer. If I understand the results in figure D.1 correctly, they show also show a minimal improvement of the learned basis functions over interpolation.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "w8n_44Df4lJ", "reviewer": "Reviewer_xqsb", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper proposes a neural network approach to solve PDEs. The method relies on two key points. The first is outputting coefficients for an overcomplete basis of functions which parameterise the solution. The second is a PDE-constrained layer (e.g. solving a linear system) for incorporating the PDE constraints. The authors explain how to differentiate through the PDE-constrained layer using IFT. Their approach is demonstrated on some simple 1D and 2D problems. Notably, while their method can make use of solution data of the PDE, it doesn’t explicitly require it. Additionally, their approach provides guarantees about satisfying the constraints, rather than just penalising deviations from the constraints.", "summary_of_the_review": "The work is clear, high-quality, and novel. Based on reading the appendix, sufficient details are provided to reproduce the experiment.\n\nThe paper clearly explains the related work and how their approach differs. The paper clearly defines their model setup, and the loss that they train on. The authors explain why using Implicit Differentiation through the PDE-CL layer is crucial, and explain how it is implemented. Tradeoffs of the method are discussed. Special care is taken to explain how the loss function is designed with respect to which residual to penalise to encourage generalisation. Details of the experiments are provided to make it reproducible. Care is taken to show how their method generalises, and how it improves upon the baseline comparison.\n\nI don’t have any major concerns about the work, but this is partially due to my being not very familiar with the PDE + ML literature. This is why I put a lower confidence on my score.\n", "strength_and_weaknesses": "The paper clearly explains the related work and how their approach differs. The paper clearly defines their model setup, and the loss that they train on. The authors explain why using Implicit Differentiation through the PDE-CL layer is crucial, and explain how it is implemented. Tradeoffs of the method are discussed. Special care is taken to explain how the loss function is designed with respect to which residual to penalise to encourage generalisation. Details of the experiments are provided to make it reproducible.\n\nCare is taken to show how their method generalises, and how it improves upon the baseline comparison.\n\nI don’t have any major concerns about the work, but this is partially due to my being not very familiar with the PDE + ML literature. \n\nQuestion: Are there ways to bound the relative error of the solution using the PDE residual loss in your method? For example, is it possible for your method to have very low PDE residual loss, but for the relative solution error to not decrease?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "pRTYPAtPuS5", "reviewer": "Reviewer_3XGD", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "The authors consider a setting where the goal is to solve the parameterized partial PDE $\\mathcal{F}_\\phi(u) = \\mathbf{0}$, where $\\mathcal{F}_\\phi$ is an affine differential operator, the PDE parameter $\\phi$ is a function over some domain $\\mathcal{X}$, and the solution $u$ is also a function over some domain $\\mathcal{X}$. The aim is to train a neural network-based model $u_\\theta$ that approximates the mapping from parameters $\\phi$ to solutions $u(\\phi)$. In this paper, the authors introduce a PDE-constrained layer (implicit layer) within the neural network training process to enforce the PDE constraints over a given set of points in $\\mathcal{X}$, and thereby promote constraint satisfaction. Specifically, they let $u_\\theta = \\sum \\omega_i f_\\theta^i$, where $f_\\theta$ is a neural network and $\\omega$ is a weight vector provided/solved for by the implicit layer, and where $f_\\theta$ is trained end-to-end with the implicit layer. They show that this approach yields outputs that are close to the true solutions, for the 1D convection and 2D Darcy Flow settings.", "summary_of_the_review": "This paper provides a technically interesting concept, but the need for this concept is unclear without explicit wall clock time comparisons to traditional non-NN baselines (that is, does the new method beat those standard baselines?). A more thorough assessment of training choices (e.g. training/testing data-points and the choice of $N$), as well as (potentially) larger-scale experiments, are also necessary to assess the efficacy of the approach.", "strength_and_weaknesses": "The overall concept of the paper is interesting: Use a neural network to represent a basis, then explicitly solve for basis weights that attempt to satisfy the PDE constraints over a set of (train or test) points - and during training, train this system end-to-end. Overall, the paper is well-written, the idea is clear, and the experimental validation is illustrative.\n\nThe main weakness is that the experimental demonstrations are somewhat incomplete. Notably:\n* I was surprised not to see a wall clock time comparison to traditional solvers, as the lack of this comparison makes it difficult to assess whether this particular solution is actually needed / whether it improves upon the non-NN standard for PDEs characterized by affine or linear differential operators. \n* The choice of training points and the choice of testing points, as well as the choice of $N$, presumably have a large effect on performance - however, this is not really discussed or shown.\n* The experimental settings seem rather small in scale, and I would have liked to see larger-scale experiments.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "QZNK_87ZigS", "reviewer": "Reviewer_dRtC", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper presents a method for learning solutions to PDE that are guaranteed to satisfy the constraints. This is done in eq. (5) by learning a basis for the solutions and finding the best linear combination of them. This results in the overall learning objective in eq. (8) that learns the basis that best-minimizes the PDE residual on the training set. The method is evaluated on 1d convection problems (Section 4.1) and 2D Darcy Flow problems (Section 4.2). ", "summary_of_the_review": "I recommend to accept the paper as it's a reasonable idea and evaluation that will be influential to help the PDE solving community better-enforce constraints in the predictions.", "strength_and_weaknesses": "Strengths\n+ The standard PDE learning methods have no guarantees on how well the predicted solutions match the PDE constraints and having a model class that always satisfies them seems useful.\n+ The experimental results on the convection and Darcy Flow problems clearly show the residual losses and relative errors improve on the problems they care about.\n\nWeaknesses\n+ The biggest weakness is that the problem sizes are relatively lower that other ones considered in the literature, for example in [PDEBench](https://arxiv.org/pdf/2210.07182.pdf), and the comparisons to the soft-constrained PDE baseline do not appear to be established results. The paper would be stronger if it evaluated on the exact experimental setting from prior work on learning PDE solutions \n+ The inference procedure in eq. (5) seems computationally expensive to solve for every instance", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster