L EARNING DIFFERENTIABLE SOLVERS FOR SYSTEMS WITH HARD CONSTRAINTS
Anonymous authors
Paper under double-blind review
A BSTRACT
We introduce a practical method to enforce linear 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
PDE-constrained 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 ✓ 2 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., 2021;
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.
When the PDE is linear, this amounts to solving a linear system.
We can use computationally efﬁcient, optimized linear solvers for this system at inference time.
After using a small number of sampled points to ﬁt this linear combination, we can evaluate the model on a much higher resolution grid.
• Our method controls the trade-off between speed and accuracy through two parameters: the number of functions in the learned basis, and the number of sampled points used to ﬁt the linear combination.
• We provide empirical validation of our method on two problems representing different classes of PDEs.
The 1D convection problem is a hyperbolic PDE on a spatiotemporal domain, and the 2D Darcy Flow problem is an elliptic PDE on a stationary (steady-state) spatial domain.
In both situations, we have no access to solution data on the interior spatial and/or temporal domain.
We show that our approach has lower error than the soft constraint approach when predicting solutions for new, unseen test cases.
Compared to the soft constraint approach, our approach takes fewer iterations to converge to the correct solution, and also requires less training time.
• Our method relies on solving batches of linear systems.
It will strongly beneﬁt from future improvements in linear solvers on hardware accelerators (e
g., GPUs, TPUs).
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,
, 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 to 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 7! 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 2 R d in , output y 2 R d out and parameters ✓ 2 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 7!
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.
(2019) 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.
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 2 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 initial and/or 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 ✓ =
[Equation] 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 linear and afﬁne PDEs, such that solving the PDE constraint can be done by efﬁciently solving linear systems.
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 a homogeneous system in Section 4.1 and an inhomogeneous system in Section 4.2.
The operator G  is linear when, for any two functions u, v from X to R , and any scalar  2 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 ! 2 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.
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 2 R n ⇥ N , ! 2 R N , b 2 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 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 D.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 C.
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 0 1 , . . .
, x 0 n 0 and build the corresponding linear system A 0 , b 0 .
Our loss value is k A 0 !
b 0 k 22 , where ! comes from the PDE-CL and depends on A and b , not A 0 , b 0 .
We compute gradients of this loss function using the Jacobian described above.
Another possibility is to use n > N .
In this case, the residual k A !
b k 2 2 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]
x 0 1 , . . .
, x 0 n
[Equation]
✓
[Equation]
(  , x 1 , . . . , x n ; ✓ )  b 0 (  , x 0 1 , . . .
, x 0 n 0 ; ✓ )
[Equation]
We approximate this problem by replacing the expectations by sums over ﬁnite samples.
The matrices A , A 0 , and vectors b , b 0 are built by applying the differential operator F  to each function in our basis f i ✓ , using the sampled gridpoints.
Inference procedure.
At inference, when given a new PDE parameter test point  , the weights ✓ are ﬁxed as our function basis is trained.
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
[Figure]
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 2 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 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 two different scientiﬁc problems: 1D convection (Section 4.1) and 2D Darcy Flow (Section 4.2).
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 Lerror 1 M
[Equation]
; and the PDE residual loss 1 M [Equation] , 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.
4. 1D CONVECTION
We study a 1D convection problem, describing transport phenomena.
The problem can be formulated as follows:
[Equation]
[Figure]
Wall ﬁgures are given in Appendix
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 A.
Results.
We plot example heatmaps from the test set in Figure 2.
We compare how close our hard-constrained model is to the target solution (Figure 2b), and similarly for the soft-constrained baseline model (Figure 2c).
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 3.
In Figure 3a, 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 3b, 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 softconstrained 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 A.1).
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
[Figure]
a general variable) is ⌫ 2 L 1 ((0 , 1) 2 ; R + ) , a diffusion coefﬁcient.
The problem is formulated as follows:
[Equation]
Here, f is the forcing function ( f 2 L 2 ((0 , 1) 2 ; R ) and u ( x ) = 0 is the boundary condition.
The differential operator is then F ⌫ ( u ) = r · ( ⌫ ( x ) r 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 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 2 (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 B.
Results.
We plot example heatmaps from the test set in Figure 4.
We compare visually how close our hard-constrained model is to the target solution (Figure 4b), and how close the soft-constrained baseline model is to the target solution (Figure 4c).
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 5.
In Figure 5a, 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 5b, 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.
[Figure]
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.
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 two problems representing different physical settings: a 1D convection problem, which describes transport phenomena; and a 2D Darcy Flow problem, which describes ﬂuid ﬂow.
Compared to 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.
Our approach for enforcing PDE hard constraints for NNs is tailored to linear PDEs.
However, future work could extend our method to nonlinear PDEs.
This would require developing an approach to implicitly differentiate through solvers for nonlinearly constrained optimization problems.
Our method also opens exciting avenues for combining numerical methods and NNs more generally.
In particular, our method would immediately beneﬁt from improvements from two areas: better neural architectures for representing functions; and better linear solvers on GPUs.
Localized neural architectures such as MoE (Shazeer et al., 2017; Wang et al., 2022) would allow making the linear system in our layer sparse.
Leveraging this structure when solving the linear system in PDE-CL would improve our method, in terms of both speed and accuracy.