L EARNING DIFFERENTIABLE SOLVERS FOR SYSTEMS WITH HARD CONSTRAINTS
A BSTRACT
We introduce a practical method to enforce partial differential equation (PDE) constraints for functions deﬁned by neural networks (NNs), up to a desired tolerance.
By combining methods in differentiable optimization and applications of the implicit function theorem to NN models, we develop a differentiable PDEconstrained layer that can be incorporated into a NN.
Inspired by dictionary learning, our model learns a family of functions, each of which deﬁnes a mapping from PDE parameters to PDE solutions.
At inference time, the model ﬁnds 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.
I NTRODUCTION
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., 2021).
Consider a parameterized partial differential equation (PDE), F φ ( u ) = 0 .
F φ is a differential operator, and the
PDE parameters φ and solution u are functions over a domain X .
Let Φ be a distribution of PDEparameter functions φ .
The goal is to solve the following feasibility problem by training a NN with parameters θ ∈ R p , i.e., ﬁnd θ such that, for all functions φ sampled from Φ , the NN solves the feasibility problem,
[Equation]
Training such a model requires solving highly nonlinear feasibility problems in the NN parameter space, even when F φ describes a linear PDE.
Current NN methods use two main training approaches to solve Equation 1.
The ﬁrst 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; it does not appear explicitly in the training algorithm.
The second approach (Raissi et al., 2019) aims to solve the feasibility problem 1 by considering the relaxation,
[Equation]
This second approach does not require access to any PDE solution data.
These two approaches have also been combined by having both a data ﬁtting loss and the PDE residual loss (Li et al., 2021).
However, both of these approaches come with major challenges.
The ﬁrst approach requires potentially large amounts of PDE solution data, which may need to be generated through expensive numerical simulations or experimental procedures.
It can also be challenging to generalize outside the training data, as there is no guarantee that the NN model has learned the relevant physics.
For the second approach, recent work has highlighted that in the context of scientiﬁc modeling, the relaxed feasibility problem in Equation 2 is a difﬁcult optimization problem (Krishnapriyan et al.,
; Wang et al., 2021; Edwards, 2022).
There are several reasons for this, including gradient imbalances in the loss terms (Wang et al., 2021) and ill-conditioning (Krishnapriyan et al., 2021), as well as only approximate enforcement of physical laws.
In numerous scientiﬁc domains including ﬂuid mechanics, physics, and materials science, systems are described by well-known physical laws, and breaking them can often lead to nonphysical solutions.
Indeed, if a physical law is only approximately constrained (in this case, “soft-constrained,” as with popular penalty-based optimization methods), then the system solution may behave qualitatively differently or even fail to reach an answer.
In this work, we propose a method to overcome these challenges by solving the PDE-constrained problem in Equation 1 directly.
We only consider the data-starved regime, i.e., we do not assume that any solution data is available on the interior of the domain (however, note that when solution data is available, we can easily add a data ﬁtting loss to improve training).
To solve Equation 1, we design a 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 a method to enforce hard PDE constraints by creating a differentiable layer, which we call PDE-Constrained-Layer or PDE-CL.
We make the PDE-CL differentiable using implicit differentiation, thereby allowing us to train our model with gradient-based optimization methods.
This layer allows us to ﬁnd the optimal linear combination of functions in a learned basis, given the PDE constraint.
• At inference time, our model only requires ﬁnding the optimal linear combination of the ﬁxed basis functions.
After using a small number of sampled points to ﬁt this linear combination, we can evaluate the model on a 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 a stationary (steady-state) spatial domain, the 1D Burger’s problem is a non-linear PDE on a spatiotemporal domain, and the 1D convection problem is a 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.
B ACKGROUND AND R ELATED 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 a 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 scientiﬁc modeling problems has gained interest in recent years (Willard et al., 2020).
NN architectures can also be designed such that they are tailored to a speciﬁc problem structure, e.g. local correlations in features (LeCun et al., 1998;
Bronstein et al., 2017; Hochreiter & Schmidhuber, 1997), symmetries in 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 scientiﬁc 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 et al., 2019), using NNs to learn differential operators in
PDEs such that many PDEs can be solved at inference time (Li et al., 2020; Lu et al., 2021a), and incorporating numerical solvers within the framework of NNs (Um et al., 2020).
It is sometimes possible to directly parameterize Gaussian processes (Lange-Hegermann, 2018; Jidling et al., 2017) or NNs (Hendriks et al., 2020) to satisfy PDEs, and ﬁt some desired loss function.
However, in the
PDEs we study, we cannot have a closed-form parameterization for solutions of the PDE.
Previous work in PDE-solving has tried to enforce hard constraints by enforcing boundary conditions (Lu et al., 2021b).
We instead enforce the PDE constraint on the interior domain.
Implicit layers.
A deep learning layer is a differentiable, parametric function deﬁned as f θ : x (cid:55)→ y .
For most deep learning layers, the two Jacobians ∂f∂x and ∂f∂θ are computed using the chain rule.
For these explicit layers, f θ is usually deﬁned 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.
Speciﬁcally, if the layer has input x ∈ R d in , output y ∈ R d out and parameters θ ∈ R p , we suppose that y solves the following nonlinear equation g ( x, y, θ )
=
for some g .
Under mild assumptions, this deﬁnes an implicit function f θ : x
(cid:55)→ y .
In our method, the forward function solves a constrained optimization problem.
When computing the Jacobian of the layer, it is highly memory inefﬁcient 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 (2017);
Barratt (2018); Blondel et al.
; Agrawal et al.
El Ghaoui et al.
(2021) and the Deep Implicit Layer NeurIPS 2021 tutorial 1 for more details).
Implicit layers have been leveraged in many applications, including solving ordinary differential equations (ODEs) (Chen et al., 2018), optimal power ﬂow Donti et al.
(2021), and rigid many-body physics (de Avila Belbute-Peres et al., 2018).
Differentiable physics.
In a different setting, recent work has aimed to make physics simulators differentiable.
The adjoint method (Pontryagin et al., 1962) is classically used in PDE-constrained optimization, and it has been incorporated into NN training (Chen et al., 2018; Zhang et al., 2019).
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; Schoenholz & 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.
M ETHODS
We describe the details of our method for enforcing PDE constraints within a NN model.
3. P ROBLEM SETUP
Our goal is to learn a mapping between a PDE parameter function φ : X → R and the corresponding
PDE solution u ( φ ) : X → R , where the domain X is an open subset of R d 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 et al., 2021).
Let F φ be a functional operator such that for all PDE parameter functions φ sampled from Φ , the solution u ( φ ) satisﬁes F φ ( u ( φ ))
= 0 .
The inputs to our NN vary depending on the domain of interest and the PDE parameters.
In the simplest case, the input is a pair ( x, φ ( x )) , where x ∈ X and φ ( x ) is the value of the PDE parameter at x .
The output of the NN is the value of the corresponding approximated solution u θ ( φ ) , for a given x .
We want to learn the mapping,
[Equation]
We show an example of such a mapping in Figure 1.
Importantly, we consider only the unsupervised learning setting, where solution data in the interior domain of the PDE is not available for training the model.
In this setting, the training is done by only enforcing the PDE, and the boundary conditions.
3. A DIFFERENTIABLE CONSTRAINED LAYER FOR ENFORCING PDE S
There are two main components to our model.
The ﬁrst component is a NN parameterized by θ , denoted by f θ .
The NN f θ takes the inputs described in Section 3.1 and outputs a vector in R N .
The output dimension N is the number of functions in our basis, and is a hyperparameter.
The second component of our model, the PDE-constrained layer or PDE-CL, is our main design contribution.
We implement a layer that performs a linear combination of the N outputs from the ﬁrst component, such that the linear combination satisﬁes the PDE on all points x i in the discretized domain.
Speciﬁcally, let ω be the weights in the linear combination given by the PDE-CL.
The output of our system is u θ =
(cid:80) Ni =1 ω
i f iθ , where f iθ is the i -th coordinate output of f θ .
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-ﬁnder for
PDEs; here we focus on both linear and nonlinear PDEs.
As an example, we describe a system based on an afﬁne PDE, where F φ is an afﬁne operator, which depends on φ .
In our experiments, we study an inhomogeneous linear system in Section 4.1, a homogeneous non-linear system in Section 4.2, and a homogeneous linear system in Section 4.3.
The operator G φ is linear when, for any two functions u, v from X to R , and any scalar λ ∈ R , we have that,
[Equation]
The operator F φ is afﬁne when there exists a function b such that the shifted operator F φ − b is linear.
Let G φ be the linear part of the operator:
G φ =
F φ − b .
We deﬁne the PDE-CL to ﬁnd the optimal linear weighting ω of the N 1D functions encoded by the ﬁrst NN component, over the set of sampled inputs x 1 , . . .
, x n .
The vector ω ∈ R N solves the linear equation system,
[Equation]
This linear system is a discretization of the PDE F φ ( u θ ) = 0 ; we aim to enforce the PDE at the sampled points x 1 , . . .
, x n .
The linear system has n constraints and N variables.
These are both hyperparameters, that can be chosen to maximize performance.
Note that once N is ﬁxed, it cannot be changed.
On the other hand, n can be changed at any time.
When the PDE is nonlinear, the linear system is replaced by a non-linear least-squares system for which efﬁcient solvers are available (Levenberg, 1944; Marquardt et al., 1963).
Backward pass of the differentiable constrained layer.
To incorporate the PDE-CL into an endto-end differentiable system that can be trained by ﬁrst-order optimization methods, we need to compute the gradients of the full model using an autodiff system (Bradbury et al., 2018).
To do this, we must compute the Jacobian of the layer.
The PDE-CL solves a linear system of the form g ( ω, A, b ) =
Aω − b = 0 in the forward pass, where
A ∈ R n × N , ω ∈ R N , b ∈ R n .
Differentiating g with respect to A and b using the chain rule gives the following linear systems, which must be satisﬁed by the Jacobians ∂ω∂A and ∂ω∂b ,
[Equation]
where 1 i = j 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 (Saad & Schultz, 1986)) for both the forward system
Aω = b 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 ablation study in E.1 to evaluate the quality of functions in our basis.
We also include an analysis and additional information on the enforcement of the hard constraints in D.
Loss function.
Our goal is to obtain a NN parameterized function which veriﬁes 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 x 1 , . . .
, x n is zero, up to numerical error of the linear solver used in 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 x (cid:48) 1 , . . .
, x (cid:48) n (cid:48) and build the corresponding linear system A (cid:48) , b (cid:48) .
Our loss value is (cid:107) A (cid:48)
ω − b (cid:48) (cid:107) 22 , where ω comes from the PDE-CL and depends on A and b , not A (cid:48) , b (cid:48) .
We compute gradients of this loss function using the Jacobian described above.
Another possibility is to use n > N .
In this case, the residual (cid:107)
Aω − b (cid:107) 22 will be non-zero, and while the “hard constraints” will not be satisﬁed during training, we can minimize this residual loss directly.
Let U ( X ) denote the uniform distribution over our (bounded) domain X .
Formally, our goal is to solve the bilevel optimization problem,
[Equation]
[Equation] [Equation]
n
[Equation] [Equation]
, x (cid:48) 1 , . . .
, x (cid:48) n (cid:48)
[Equation]
s.t. ω =
arg min
w
[Equation]
We approximate this problem by replacing the expectations by sums over ﬁnite samples.
The matrices A , A (cid:48) , and vectors b , b (cid:48) are built by applying the differential operator F φ to each function in our basis f iθ , 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 a new PDE parameter test point φ , the weights θ are ﬁxed as our function basis is trained.
In this paragraph, we discuss guarantees in the linear PDE case.
Suppose that we want the values of u θ over the (new) points x test 1 , . . .
, x test n test .
If n test < N , we can ﬁt ω in the PDE-CL using all of the test points.
This guarantees that our model satisﬁes the PDE on all test points: the linear system in the PDE-CL is underdetermined.
In practice, n test is often larger than N .
In this case, we can sample J ⊂ { 1 , . . . , n test } , | J | < N , and ﬁt the PDE-CL over these points.
Over the subset { x test j , j ∈ J } , the PDE will be satisﬁed.
Over the other points, the residual may be non-zero.
Another option is to ﬁt the PDE-CL using all the points x test 1 , . . .
, x test n test , in which case the residual may be non-zero for all points, but is minimized on average over the
[Figure]
sampled points by our 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 is enforced.
A smaller N allows for faster linear system solves.
Once ω is ﬁt, we can query u θ for any point in the domain.
In practice, we can choose | J | to be much smaller than n test ; here, the linear system we need to solve is much smaller than the linear system required by a numerical solver.
E XPERIMENTAL RESULTS AND IMPLEMENTATION
We test the performance of our model on three different scientiﬁc problems: 2D Darcy Flow (Section 4.1), 1D Burgers’ equation (Section 4.2), and 1D convection (Section 4.3).
In each case, the model is trained without access to any solution data in the interior solution domain.
The training set contains 1000 PDE parameters φ .
The model is then evaluated on a 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 L 2 error 1 M (cid:80) Mi
=1 (cid:107) u
θ ( φ i ) − u
( φ i ) (cid:107) 2 (cid:107) u
( φ i ) (cid:107); and the PDE residual loss 1 M (cid:80) M i =1 (cid:107)F φ i ( u θ ) (cid:107) 2 , 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. 2D D ARCY F LOW
We look at the steady-state 2D Darcy Flow problem, which describes, for example, ﬂuid ﬂow through porous media.
In this section, the PDE parameter (denoted by φ in the previous sections as a general variable) is ν ∈ L ∞ ((0 , 1) 2 ; R + ) , a diffusion coefﬁcient.
The problem is formulated as follows:
[Equation]
Here, f is the forcing function ( f ∈ L 2 ((0 , 1) 2 ; R ), and u ( x ) = 0 is the boundary condition.
The differential operator is then F ν ( u ) = −∇ · ( ν ( x ) ∇ u ( x )) .
Given a set of variable coefﬁcients ν , the goal is to predict the correct solution u .
The ν ( x ) values are generated from a Gaussian and
[Figure]
then mapped to two values, corresponding to the two different materials in the system (such as the ﬂuid and the porous medium).
We follow the data generation procedure from Li et al.
We use the Fourier Neural Operator (FNO) (Li et al., 2020) architecture, trained using a PDE residual loss as the baseline model (“soft-constrained”).
Our model uses the FNO architecture and adds our
PDE-CL (“hard-constrained”).
The domain (0 , 1) 2 is discretized over n x × n y points.
For each point on this grid, the model takes as input the coordinates, x ∈ (0 , 1) 2 , and the corresponding values of the diffusion coefﬁcients, ν ( x ) .
The boundary condition is satisﬁed by using a molliﬁer function (Li et al., 2020), and so the only term in the loss function is the PDE residual.
We use a constant forcing function f equal to 1.
We provide more details on our setup and implementation in Appendix C.
Results.
We plot example heatmaps from the test set in Figure 2.
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 2c).
Our hard-constrained model is much closer to the target solution, as indicated by the difference plots mostly being white (corresponding to zero difference).
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 to decrease as training proceeds, remaining signiﬁcantly lower than the baseline.
Similarly, in Figure 3b, we show the curves corresponding to the relative error and the PDE residual loss metric on the test dataset.
Our model starts at a much smaller relative error immediately and continues to decrease, achieving a ﬁnal lower test relative error.
On the test set, our model achieves 1 .
82 % ± 0 .
04 % relative error, versus 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 . 0433 .
Our model has almost two orders of magnitude lower PDE residual test loss, and it has signiﬁcantly lower standard deviation.
On the relative error metric, our model achieves a 71 % improvement over the soft-constrained 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.
[Figure]
4. 1D B URGERS ’ EQUATION
We study a non-linear 1D PDE, Burgers’ equation, which describes transport phenomena.
The problem can be written as,
[Equation]
Here, u 0 is the initial condition.
We consider periodic boundary conditions.
We aim to map the initial condition u 0 to the solution u .
We consider problems with ﬁxed viscosity parameter ν = 0 .
For this problem, we have a periodic boundary condition.
We follow the data generation procedure from Li et al.
(2020), which can be found here.
We use a physics-informed DeepONet baseline model (Wang et al., 2021) with regular multi-layer perceptrons as the base NN architecture, trained using the PDE residual loss.
Our hard-constrained model is composed of stacked dense layers and our PDE-CL, which allows for a fair comparison.
Because this PDE is non-linear, our PDE-CL solves a 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 set in Figure 4.
We compare how close our 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.
Note that our soft-constrained results are different from Wang et al.
(2021); this is because Wang et al.
(2021) did not use only dense layers.
Wang et al.
(2021) report a relative L 2 error of 1 . 5% after 200 , 000 steps of training, using a modiﬁed multi-layer perceptron architecture.
Our hard-constrained model achieves this level of accuracy after less than 5 , 000 steps of training, using just multi-layer perceptrons.
During the training procedure for both hard- and soft-constrained models, we track error on a validation set of PDE solutions with different PDE parameters from the training set.
We show these error plots in Figure 5.
In Figure 5b, the PDE residual loss for the hard-constrained model starts close to two orders of magnitude lower than for the soft-constrained model, and it continues to remain low.
In Figure A.2b, we track the relative error with respect to the target solution, obtained using the Chebfun package (Driscoll et al., 2014), following Wang et al.
Similarly, we see that our model starts at much smaller relative error immediately, and also continues to decrease.
The
[Figure]
large ﬂuctuations are due to the log-log plotting, and small batches used by our method for memory reasons.
Our model achieves 1 .
5 % ± 0 .
02 % relative error on a hidden test set.
4. 1D CONVECTION
We study a 1D convection problem, describing transport phenomena.
The problem can be formulated as follows:
[Equation]
Here, h ( x ) is the initial condition (at t = 0 ), g ( t ) is the boundary condition (at x = 0 ), and β ( x ) represents the variable coefﬁcients (denoted by φ in Section 3).
Given a set of variable coefﬁcients, β ( x ) , and spatiotemporal points ( x i , t i ), the goal is to predict the correct solution u ( x, t ) .
We provide results and more details in Appendix A.
C ONCLUSIONS
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 a method to enforce hard PDE constraints, when training NNs, by designing a differentiable PDE-constrained layer (PDE-CL).
We can add our layer to any NN architecture to enforce PDE constraints accurately, and then train the whole system endto-end.
Our method provides a means to control the trade-off between speed and accuracy through two hyperparameters.
We evaluate our proposed method on three problems representing different physical settings: a 2D Darcy Flow problem, which describes ﬂuid ﬂow through a porous medium; a 1D Burger’s problem, which describes viscous ﬂuids and a dissipative system; and a 1D convection problem, which describes transport phenomena.
Compared to the baseline soft-constrained model, our model can be trained in fewer iterations, achieves lower or competitive PDE residual error (measuring how well the PDE is enforced on the interior domain), and achieves lower or competitive relative error with respect to target solutions generated by numerical solvers.
Acknowledgements.
The authors would like to thank Quentin Berthet, David Duvenaud, Romain
Lopez, Dmitriy Morozov, Parth Nobel, Daniel Rothchild, Hector Roux de B´ezieux, and Alice Schoenauer 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 supported by the Intelligence Advanced Research Projects Agency (IARPA) and Army Research Ofﬁce (ARO) under Contract No.
W911NF-20-C-0035.
ASK and MWM would like to acknowledge the U.S. Department of Energy, Ofﬁce of Science, Ofﬁce of Advanced Scientiﬁc Computing Research, Scientiﬁc Discovery through Advanced Computing (SciDAC) program under contract No.
DE-AC02-05CH11231.
A 1D CONVECTION
We study a 1D convection problem, describing transport phenomena.
The problem can be formulated as follows:
[Equation]
Here, h ( x ) is the initial condition (at t = 0 ), g ( t ) is the boundary condition (at x = 0 ), and β ( x ) represents the variable coefﬁcients (denoted by φ in Section 3).
Given a set of variable coefﬁcients, β ( x ) and spatiotemporal points ( x i , t i ), the goal is to predict the correct solution u ( x, t ) .
The β ( x ) values are generated in the same manner as in Wang et al.
(2021) via β ( x ) = v ( x ) −
min x v ( x ) + 1 , where v ( x ) is generated from a Gaussian random ﬁeld with a length scale of 0 . 2 .
We use a physicsinformed 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 a fair comparison.
We provide more details on our setup and experiments in Appendix B.
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 A.1b), and similarly for the soft-constrained baseline model (Figure A.1c).
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.
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 A.2.
In Figure A.2a, the PDE residual loss for the hard-constrained model starts close to six orders of magnitude lower than for the soft-constrained model, and it continues to remain low.
In Figure A.2b, we track the relative error 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 a 49 % improvement over the soft-
[Figure] [Figure]
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.1).
B D ETAILS ON THE 1 D C ONVECTION PROBLEM
Experiment setup and implementation details.
In this setting, the inputs to the models are two sets of (( x 1 , t 1 ) , . . .
, ( x n , t n )) , and (( x (cid:48) 1 , t (cid:48) 1 ) , . . .
, ( x (cid:48) n (cid:48) , t (cid:48) n (cid:48) ))
sampled points within the domain X (interior points) and the corresponding β ( x ) values.
We use the former for ﬁtting the PDE-CL, and the latter for computing the residual loss function.
We also require a set (( x n +1 , t n +1 ) , . . .
, ( x n + n (cid:48) , t n + n (cid:48) )) of sampled points on the initial condition ( t = 0 ) and boundary condition ( x = 0 ).
The training optimization problem is formulated as follows:
[Equation]
with,
[Equation]
where θ corresponds to parameters of the NN, u ( x, t ) is the solution at the initial and boundary conditions, and F ( u θ ) is the PDE constraint that must be satisﬁed.
The loss term L ( u θ ) is a regression loss over the initial and boundary conditions.
The forward pass of the PDE-CL solves the following equality constrained problem,
[Equation]
where f θ refers to the outputs of the base NN, on which we stack the PDE-CL.
Since the initial/boundary condition regression loss uses a quadratic penalty, this equality constrained problem is in fact a convex equality constrained quadratic problem (EqQP), which is equivalent to a linear system.
We solve this linear system using GMRES (Saad & Schultz, 1986).
We compute the Jacobian via implicit differentiation with respect to Equation 14.
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, we 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 D ETAILS ON THE D ARCY F LOW PROBLEM
Experiment setup and implementation details.
Our goal is to ﬁnd parameters θ , which solve,
[Equation]
[Figure]
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-constrained training method is trained only by minimizing the PDE residual.
We train the FNO model using the same hyperparameters as Li et al.
We denote the FNO model part of the architecture as f θ .
The PDE-CL constrains the output of the FNO model by solving the linear system,
[Equation]
To train our model, we compute the Jacobian of this layer via implicit differentiation with respect to 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 to compute the residual in the loss function.
We use N=4000 for the number of basis functions in the PDE-CL.
We show plots against wall time in Figure C.1, and the distributions of errors over the test set in
Figure C.2.
D H ARD CONSTRAINTS BOUND
To provide an evaluation of how hard the hard constraints are, we conduct an additional study where we look at the error in prediction for points sampled and used to ﬁt the PDE-CL against points that were not sampled.
With a trained hard-constrained model for 1D convection, we sample a batch of points and create a density plot showing the error as a function of points used for ﬁtting the PDE-CL and points not used for ﬁtting the PDE-CL.
The histogram in Figure D.1 shows that the errors are qualitatively the same between points used for ﬁtting the linear combination and those not used.
There are 1000 points used for ﬁtting the PDE-CL, and 9000 points not used.
Our results show that our model achieves low error, even outside of the points used for ﬁtting the PDE-CL.
[Figure]
E A BLATION STUDIES E. E VALUATING THE QUALITY OF THE LEARNED BASIS FUNCTIONS
We add an experiment to evaluate the quality (and the advantage) of our learned basis functions, compared to cubic interpolation.
We want to understand whether our learned functions are useful outside of the points used for the constrained problem.
We start with a model trained on the 1D convection problem.
The model was trained by sampling 750 points for ﬁtting the PDE-CL, and 250 different points for the residual in the objective functions.
The points were sampled from a 100x100 grid (10,000 points total).
We sample 750 points from the grid, and solve the corresponding PDE-CL, which gives us a candidate PDE solution.
For our baseline, We interpolate the solution we ﬁnd on these 750 points to the 10,000 points using scipy.optimize ’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 our trained hard-constrained model: once the linear combination weights are ﬁt (same as the baseline), we now use our learned basis functions to perform inference over all 10,000 (or 1 million) points.
We plot our results over the test dataset in Figure E.1a.
The ﬁgure shows using our model reduces the error vs using interpolation on the hard-constrained points.
F C OMPARISON TO NUMERICAL SOLVERS
We compare complexities of our PDE-CL vs numerical methods.
Suppose that we wish to perform inference on a n x × n t grid over a 1-d domain, with n x samples, over a time span [0 , T ] , with n t samples.
Here we also suppose that the PDE is 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 = n interior + n IC + n BC .
The PDE-CL will then solve a ( n + n IC + n BC ) × N linear system.
Our method currently results in a dense linear system.
Solving this linear system has complexity O (max( n interior + n IC + n BC , N ) 2 × min( n + n IC + n BC , N )) .
We need to add to this a forward pass using the network on the whole grid, once the optimal linear combination has been computed.
Fortunately, this forward pass is embarrassingly parallel.
On the other hand, a ﬁnite difference method such as Crank-Nicolson or Lax-Wendroff requires to solve at each t step a tri-diagonal system of size n x , therefore yielding an overall complexity of
O ( n x × n t ) .
Comparing the complexity of both methods, the PDE-CL is asymptotically faster than numerical methods when,
[Equation]
In our current framework, inference is marginally slower or on par with numerical solvers.
However, as we increase the resolution of our grid (ﬁner discretization), our method compares favorably to the numerical solver, i.e., our computational cost increases more slowly than the numerical solver.
Additionally, the operations in our PDE-CL are poorly optimized for current hardware, i.e., GPU utilization is low.
Our method will greatly beneﬁt from future improvements in hardware acceleration, which is still a nascent ﬁeld in the context of linear solvers on GPUs.