G RADIENT - BASED META - SOLVING AND ITS APPLICATIONS TO ITERATIVE METHODS FOR SOLVING DIFFERENTIAL EQUATIONS
Anonymous authors
Paper under double-blind review
A BSTRACT
In science and engineering applications, it is often required to solve similar computational problems repeatedly.
In such cases, we can utilize the data from previously solved problem instances to improve efﬁciency of ﬁnding subsequent solutions.
This offers a unique opportunity to combine machine learning (in particular, meta-learning) and scientiﬁc computing.
To date, a variety of such domainspeciﬁc methods have been proposed in the literature, but a generic approach for designing these methods remains under-explored.
In this paper, we tackle this issue by formulating a general framework to describe these problems, and propose a gradient-based algorithm to solve them in a uniﬁed way.
As an illustration of this approach, we study the adaptive generation of initial guesses for iterative solvers to speed up the solution of differential equations.
We demonstrate the performance and versatility of our method through theoretical analysis and numerical experiments.
I NTRODUCTION
It is common and important in science and engineering to solve similar problems repeatedly.
For example, in material science, a tremendous amount of physical and numerical experiments are conducted to discover and characterize new materials (Schmidt et al.
For another example, in computational ﬂuid dynamics, many methods involve solving a Poisson equation to compute the pressure ﬁeld in every time step of the simulation (Ajuria Illarramendi et al.
In these situations, we can utilize the data of the previously solved problems to solve the next similar but unseen problems more efﬁciently, and machine learning is a natural and effective approach for this.
Thus, in recent years, many learning-based methods have been proposed for repeated solutions of computational problems such as partial differential equations (PDEs) (Tang et al.
¨ Ozbay et al.
; Tompson et al.
; Hsieh et al.
; Huang et al.
For example, in Tang et al.
(2017), Ajuria Illarramendi et al.
(2020), and ¨Ozbay et al.
(2021), convolutional neural networks are used to predict the solution of Poisson equations, and Ajuria Illarramendi et al.
(2020) and ¨Ozbay et al.
(2021) propose to use the predicted solution as an initial guess of traditional numerical methods.
Hsieh et al.
(2018) combines a neural network and an iterative solver to accelerate it with maintaining the convergence guarantee.
In addition to the regular supervised learning approaches, there are several works where meta-learning approaches are taken to solve computational problems (Feliu-Fab`a et al.
; Chen et al.
; Psaros et al.
; Guo et al.
Meta-learning, or learning to learn, leverages previous learning experiences to improve future learning performance (Hospedales et al.
(2021)), which ﬁts the motivation utilizing the data from previously solved equations for the next one.
For example, Chen et al.
(2020) use meta-learning to generate a smoother of the Multi-grid Network for parametrized PDEs.
Psaros et al.
(2021) propose a meta-learning technique for ofﬂine discovery of physics-informed neural network loss functions.
Although many methods have been proposed in this direction, they are often problem-speciﬁc.
In other words, there lacks both a uniﬁed framework to explain them and a general design pattern to adapt them to new problem settings.
On the other hand, in the machine learning literature, there exists a general methodology - gradient-based meta-learning (Finn et al.
- that covers a variety of meta-learning problem settings.
In this paper, we generalize this approach to yield a framework, which we call gradient-based meta-solving (GBMS), that encompasses both learning and computational problems.
This offers a general means to understand and develop learning-based algorithms to speed up computation.
As an illustration of our approach, we apply GBMS to accelerate the solution of differential equations with iterative methods through learning.
We show the advantage of the proposed algorithm over the baseline classical and learning-based approaches through theoretical analysis and numerical experiments.
Finally, we incorporate the algorithm into a practical application and demonstrate its versatility and performance.
G RADIENT - BASED META - SOLVING
In this section, we introduce our core idea of gradient-based meta-solving.
First, we formulate a class of problems, which we call meta-solving, that includes both ordinary meta-learning and learning-based computational problems.
We then propose a general gradient-based algorithm for solving them.
Under our formulation, many methods proposed in related works can be regarded as special cases of the GBMS algorithm.
2. G ENERAL FORMULATION OF META - SOLVING
Let us now introduce the general formulation of meta-solving.
We ﬁrst ﬁx the required notations.
A task τ is a tuple τ =
( D τ , U τ , L τ ) consisting of a dataset D τ , a solution space U τ , and a loss function L τ .
A solution space U τ is a set of parametric candidate solutions, which is usually a subset of R N for some N ∈ N .
A loss function L τ is a function from U τ to R ≥ 0 that measures the quality of solution candidates.
To solve a task τ means to ﬁnd an approximate solution ˆ u ∈ U τ which minimizes L τ ( · ) .
Meta-solving considers the solution of not one, but a distribution of tasks by a learnable solver.
Thus, we consider a task space ( T , P ) as a probability space that consists of a set of tasks T and a task distribution P , which is a probability measure deﬁned on a suitable σ -algebra on T .
A solver Φ is a function from T × Θ to U , where U =
(cid:83) τ
∈T U τ .
θ ∈ Θ is the parameter of Φ , and Θ is its parameter space.
Here, θ may or may not be trainable, depending on the problem.
Then, solving a task τ ∈ T by an algorithm Φ with a parameter θ is denoted by Φ( τ ; θ )
= ˆ u .
A meta-solver Ψ is a function from T × Ω to Θ , where ω ∈ Ω is a parameter of Ψ and Ω is its parameter space.
A meta-solver Ψ parametrized by ω ∈ Ω is expected to generate an appropriate parameter θ τ ∈ Θ for solving a task τ ∈ T with a solver Φ , which is denoted by Ψ( τ ; ω )
= θ τ .
Then, by using the notations above, our meta-solving problem is deﬁned as follows:
Deﬁnition 1 (Meta-solving problem) .
For a given task space ( T , P ) , solver Φ , and meta-solver Ψ , ﬁnd ω ∈ Ω which minimizes E τ ∼ P
[ L τ (Φ( τ ; Ψ( τ ; ω ))] .
We present some familiar examples, which can be regarded as special cases of the meta-solving problem.
First, we can see that conventional meta-learning problems fall in this formulation.
Example 1 (Few-shot learning) .
As an example of the meta-learning problem, we take the fewshot regression problem with MAML (Finn et al.
The components of the problem are the following.
The task τ is to learn a regression model from data.
The dataset is D τ = { ( x i , y i )
Ki =1 } , which satisﬁes y =
f τ ( x ) for an unknown function f τ .
D τ is divided into the training set D train τ and validation set D val τ .
The solution parameter space U τ is a weights space of a neural network ( NN ) that models f τ .
Note that U
= U τ because the architecture of NN is shared across all tasks.
The approximate solution ˆ u ∈ U is the trained weights of NN , which is obtained by training on
D train τ .
The loss function L τ : U → R ≥ 0 is the mean squared error (MSE) on D val τ .
The task distribution ( T , P ) is determined by the distribution of the target function f τ and distribution of samples ( x, y ) .
The solver Φ : T × Θ → U is the single step gradient descent to minimize L train τ ( u ) =|
D train τ | (cid:80) (
x,y ) ∈ D train τ (cid:107) NN( x ; u ) − y (cid:107) 2 .
Its parameter θ ∈ Θ is initial weights u (0) of NN , so Θ = U .
Thus, Φ( τ ; θ ) = u (0) − α ∇ u (0)
L train τ ( u (0) )
= ˆ u , where α is a learning rate.
The meta-solver Ψ : T × Ω → Θ is considered as a constant function that returns its parameter ω ∈ Ω for any task τ ∈ T .
The parameter ω ∈ Ω is expected to be an appropriate initial weights for all τ ∈ T to be ﬁne-tuned easily.
Thus, Ω = Θ = U , and Ψ( τ ; ω )
=
ω
= θ = u (0) .
Note that the output of the meta-solver, the initial weights u (0) , does not depend on τ .
Then, the few-shot learning problem is deﬁned as meta-solving problem, which is min
[Equation]
In addition to the conventional learning problems, we can regard other computational problems, such as solving a differential equation, as a task of the meta-solving problem.
In constrast with
MAML, the inner-loop learning is now replaced with a family of iterative solvers for differential equations.
This necessitates the distinction of the meta-solver parameter space Ω and the solution space U .
Moreover, the meta-solver has to produce a task-speciﬁc parameter for the inner solver.
Example 2 (Solving differential equations) .
Suppose that we need to repeatedly solve similar instances of a class of differential equations with a given numerical solver.
The solver has a number of hyper-parameters, which sensitively affect accuracy and efﬁciency depending on the problem instance.
Thus, ﬁnding a strategy to optimally select solver hyper-parameters given a problem instance can be viewed as a meta-solving problem.
The components of the problem are the following.
The task τ is to solve a differential equation.
In this example, suppose that the target equation is the
Poisson equation − ∆ u
= f τ with Dirichlet boundary conditions u = g τ .
The dataset D τ contains data of the differential equation, D τ = { f τ , g τ } .
The solution parameter space U τ is R N τ .
The loss function L τ : U τ → R ≥ 0 measures the accuracy of ˆ u ∈ U τ .
In this example, the (cid:96) 2 -norm of the residuals obtained by substituting the approximate solution ˆ u into the equation can be used.
The task distribution ( T , P ) is the joint distribution of f τ and g τ .
The solver Φ : T × Θ → U is a numerical solver with a parameter θ ∈ Θ for the differential equation.
In this example, suppose that Φ is the
Jacobi method (Saad (2003)) and θ is its initial guess.
The meta-solver Ψ : T × Ω → Θ is a strategy characterized by ω ∈ Ω to select a parameter, an initial guess, θ τ ∈ Θ for each task τ ∈ T .
Note that the output of the meta-solver depends on τ , which is different from the case of Example 1 .
Then, ﬁnding the strategy to select parameters of the numerical solver becomes meta-solving problem.
The above examples explain why we describe our problem as meta-solving instead of meta-learning.
In Example 1, the task τ is learning from data, and the solver Φ is gradient descent.
In Example 2, the task τ is solving a differential equation, and the solver Φ is an iterative differential equation solver.
Regardless of the type of task and algorithm, in both cases, the solver Φ solves the task τ , and we do not distinguish whether Φ is a learning algorithm or other numerical solver.
In other words, learning algorithms such as gradient descent are also a type of numerical solvers, and we regard learning as a special case of solving.
It is also true for the outer learning algorithm to learn how to solve the task τ with the solver Φ , so learning to solve is a special case of solving to solve.
In this sense, we call it meta-solving.
2. G RADIENT - BASED META - SOLVING
In the previous section, we deﬁned meta-solving problems as a generalization of meta-learning problems, but how can we solve them effectively?
For meta-learning problems such as Example 1, general methodologies in the form of gradient-based meta-learning (Hospedales et al.
(2021)), e.g.
MAML (Finn et al.
(2017)), have been proposed.
In Example 1, MAML solves the minimization problem (1) by updating ω using the outer gradient descent algorithm ω
← ω − β ∇ ω L τ (ˆ u ) .
MAML is an algorithm to ﬁnd good initial weights of the neural network for conventional meta-learning problems, but it can be generalized to the meta-solving problems.
In meta-solving problems, tasks may not be learning problems, Φ may not be gradient descent, and its parameter θ may not be initial weights of a neural network.
However, we can still employ the same update rule as MAML, as long as L τ , Φ , and Ψ are differentiable.
Thus, for differentiable solvers Φ and Ψ , we propose gradientbased meta-solving algorithm (Algorithm 1) as a generalization of MAML.
This is a form of datadriven algorithm design, and differs from previous works in this direction (e.g. Hutter et al. Mitzenmacher & Vassilvitskii (2021), and Balcan (2021)).
These focus on discontinuous problems such as combinatorial optimization, while our framework focuses on differentiable problems.
2. O RGANIZING RELATED WORKS
Owing to the general formulation presented above, we can organize several related works on learning-based methods for scientiﬁc computing and describe them in a uniﬁed way.
We highlight the advantages of our method using the examples in this section.
The detailed description of how the task, the solver and the meta-solver are deﬁned in each case are found in Appendix A.
First, let us review Feliu-Fab`a et al.
This work proposes a neural network architecture inspired by the nonstantard wavelet form with meta-learning approach to solve the equations containing partial differential or integral operators.
It can be regarded as a special case of GBMS, where
←
− α ∇ task τ is solving the equation, solver Φ is the forward computation of the trained neural network, θ ∈ Θ is its weights, and meta-solver Ψ is the constant function that returns the weights θ = ω ∈ Ω .
Note that θ does not depend on the task τ .
We also note that other works where neural networks replace a whole or part of numerical methods (Tang et al.
(2017); ¨Ozbay et al.
; Tompson et al.
Hsieh et al.
(2018)) can be organized in the same way.
Thus, the meta-solving formulation includes many learning-based methods where meta-learning techniques are not explicitly employed.
We take Psaros et al.
as another example.
In this work, meta-learning is used to learn a loss function of the physics-informed neural network, shortly PINN (Raissi et al.
for solving
PDEs.
This also can be considered as a special case of GBMS, where τ is training the PINN, the solver Φ is the gradient descent for training the PINN, θ ∈ Θ is the weights of another neural network used as the loss function in the training, and the meta-solver Ψ is the constant function that returns the weights θ = ω ∈ Ω .
As in the previous example, θ does not depend on the task τ .
We remark that this example and Example 1 are similar in the sense that Φ is gradient descent and Ψ is the constant function returning neural network’s weights in both examples, though the task τ is different.
The uniﬁed framework sheds light on a similarity in various methods for various tasks, which enables us to apply a technique developed for one problem to another easily.
Lastly, let us consider Chen et al.
In this work, meta-learning is used to generate a parameter of PDE-MgNet, a neural network representing the multigrid method, for solving parametrized PDEs.
This also can be regarded as a special case of GBMS, where τ is solving a PDE, the solver Φ is the
PDE-MgNet, whose iterative function is implemented by a neural network φ with weights θ , and the meta-solver Ψ is another neural network with weights ω to generate θ depending on task τ .
Note that the meta-solver Ψ is trained with single step of the solver Φ but tested with multiple steps of Φ .
In other words, this work does not consider the solver Φ itself but instead its iterative function φ in the training.
We will show the importance of this difference in section 3.1.
The above examples show the generality of our meta-solving formulation that organizes a variety of methods in the systematic way regardless of their type of algorithm.
This overall approach allows us to easily design learning-based methods for new problem settings, and we will demonstrate it in the next section.
GBMS FOR ITERATIVE METHODS
In this section, we demonstrate how GBMS can be used to yield effective meta-solvers for new problem settings.
In particular, we consider the problem of accelerating iterative methods by adaptively choosing initial conditions.
Iterative methods are powerful tools to solve computational problems.
For example, the Jacobi method and SOR method are used to solve PDEs (Saad (2003)).
In iterative methods, a function φ τ , which depends on the task τ , is iteratively applied to the current approximate solution to update it closer to the true solution until it reaches a criterion, such as a certain error tolerance or number of iterations.
These methods require an initial guess as the starting point of the iterations, and the performance of the method depends on this choice.
Thus, adaptively choosing initial guesses for solvers is an effective way to speed up the computation process.
Here, we apply GBMS to derive a meta-solver that produces effective initial guesses for iterative solvers.
The meta-solving problem here can be viewed as a generalized version of Examplebeyond Jacobi solvers.
The task τ is any computational problem which can be solved by iterative methods.
For example, τ is solving a PDE as described in Example 2.
The task distribution ( T , P ) is deﬁned according to each problem.
The solver Φ : T × Θ → U is an iterative method with an initial guess θ = u (0)
∈ Θ , iterative function φ τ , and the number of iterations k , so Φ( τ ; u (0) )
=
φ kτ ( u (0) )
=
u ( k ) and Θ = U .
The meta-solver Ψ : T × Ω → Θ is a function parametrized by ω ∈ Ω , which takes τ ∈ T as an input and generates an initial guess θ τ = u (0)
τ ∈ Θ for the solver Φ .
We implement Ψ by a neural network, so ω ∈ Ω is its weights.
Then, Ψ is trained to minimize the expectation of L τ ( u ( k ) ) by the gradient descent.
Since both Φ and Ψ are implemented in a deep learning framework, ∇ ω L τ ( u ( k ) ) can be computed by back-propagation.
The entire process of the algorithm is presented in Algorithm 2.
Hereafter, we consider solving a PDE as a task of the meta-solving problem, but the method is applicable to other tasks, such as root ﬁnding.
Note that although Huang et al.
(2020), Ajuria Illarramendi et al.
(2020), and ¨Ozbay et al.
(2021) propose to use initial guesses generated by neural networks, these initial guesses are independent of the solvers.
On the other hand, our initial guesses are optimized for each solver.
We will show its advantage in the following sections.
i to k do u i τ ← φ τ i − 1) τ ) ; end
Update
← ω − ∇ ω
L τ u k )
τ
; 10 end
3. T OY PROBLEM : 1D P OISSON EQUATIONS
To study the property of the proposed algorithm, we consider solving 1D Poisson equations as a toy problem.
First, we show a theorem that guarantees the improvement of the proposed meta-solving approach for the Jacobi method and linear neural networks.
Then, we demonstrate that the theorem numerically holds for another iterative method and practical nonlinear neural network.
3.1. T HEORETICAL ANALYSIS
Let us recall Example 2 and set the domain of interest D
= (0 , 1) .
Then, the target equation becomes the following 1D Poisson equation with Dirichlet boundary condition:
[Equation]
To solve this equation numerically, it is discretized with ﬁnite difference scheme, and we rewrite it as the matrix equation Au = f , where the domain [0 , 1] is discretized into N points, so A ∈ R N × N and u, f ∈ R N .
Suppose that we solve the equation Au = f with the Jacobi method under randomly sampled f τ .
Then, the meta-solving problem for solving 1D Poisson equations is deﬁned by the following.
The task τ is to solve Au = f τ .
The dataset is D τ = { f τ , u τ } , where u τ is the solution of Au = f τ .
The solution parameter space is U τ = R N .
The loss function is L τ (ˆ u ) = (cid:107) u τ − ˆ u (cid:107) 2 .
The task distribution ( T , P ) is determined by the distribution of f τ , denoted by P f .
It is assumed to be centered and normalized, i.e. the mean of f τ is 0 and the covariance is the identity matrix.
The solver Φ k : T × Θ → U is the Jacobi method with k iterations starting at an initial guess u (0)
= θ ∈ Θ .
Its iterative function is φ τ ( u ) =
Mu + 12 f τ , where M = I − 12 A .
To summarize, Φ k ( τ ; u (0) )
=
φ kτ ( u (0) )
=
ˆ u .
Note that Φ 0 is the identity map.
The meta-solver Ψ : T × Ω → Θ is a linear neural network with weights ω ∈ Ω .
It can be represented as matrix multiplication for some matrix W , so Ψ( τ ; ω ) =
Wf τ = u (0) .
Here, we assume rank( W ) <
N to avoid the trivial solution W =
A − 1 , i.e. the meta-solver does not have the capacity to produce an exact solution for each task.
Then, the meta-solving problem becomes
[Equation]
As for this meta-solving problem, the following theorem holds.
The proof is in Appendix B.
Theorem 1 (Guarantee of improvement by meta-solving) .
For any k ≥ 0 , (3) has the unique minimizer W k .
Furthermore, if k 1 < k 2 , then for all k ≥ k 2 ,
[Equation]
where the equality holds if and only if W k 1 = W k 2 .
Theorem 1 guarantees improvement of meta-solving for the considered setting.
For example, suppose k 1 =
(regular supervised learning, i.e. directly predicting the solution without considering the solver) and k 2 = 5 .
Then, the theorem implies that the meta-solver trained with 5 Jacobi iterations is expected to give an better initial guess than the meta-solver obtained by regular supervised learning for any number of Jacobi iterations larger than 5 .
More generally, Theorem 1 shows that meta-learning with a higher number of Jacobi iterations in the inner solver improves the performance, provided one carries out at least a greater number of Jacobi iterations during inference.
This shows the advantage of the meta-solving approach over regular supervised learning.
3.1. N UMERICAL EXAMPLES FOR MORE GENERAL CASES
The key insight from the previous theoretical analysis is that meta-solving leverages the properties of the solver and adapts the selection of initial conditions to it.
Here, we show using numerical examples that this is the case for more complex scenarios, involving different task distributions, different iterative solvers, and a practical nonlinear neural network.
The meta-solving problem in this section is deﬁned by the following.
The task τ is the same as the previous section 3.1.1.
Let the number of discretization points N be 512 .
The task distribution ( T , P ) is determined by the distribution of u .
We consider two distributions P s and P h , where the solution u consists of sine functions and hyperbolic tangent functions respectively.
Their details are listed in Appendix C.1.
For each distribution, we prepare 30,000 tasks for training, 10,000 for validation, and 10,000 for test.
To solve the tasks, we use the Jacobi method and the Red-Black ordering SOR method (Saad (2003)) with k iterations starting at an initial guess u (0)
= θ ∈ Θ , denoted by Φ Jac ,k and Φ SOR ,k respectively.
Note that they are implemented by convolutional layers.
We consider two meta-solvers.
One is a heuristic initial guess generator Ψ BL that takes τ as an input and gives the heuristic initial guess, which is the linear interpolation of the boundary condition.
Ψ BL does not have a parameter and is used as a baseline.
The other is a variant of 1D U-Net (Ronneberger et al.
Ψ NN with weights ω ∈ Ω , which takes f τ and the heuristic initial guess Ψ BL ( τ ) as inputs and generates an initial guess θ τ ∈ Θ for the solver Φ .
The meta-solver Ψ NN is trained with solvers Φ Jac ,k , Φ SOR ,k with k = 0 , 4 , 16 , 64 by using the Algorithm 2.
For each setting, Ψ NN is trained six times with different random seeds.
The details of the architecture and hyper-parameters of Ψ NN are found in Appendix D.1.
The trained meta-solver Ψ NN and baseline Ψ BL are tested with solvers Φ Jac ,k , Φ SOR ,k with k = 0 , 4 , 16 , 64 .
The performances of the meta-solvers are measured by the MSE on the test set and presented in Table 1.
Figure 1 shows the comparison of the initial guesses u (0) obtained by Ψ NN with Φ Jac , 0 and Φ Jac , 64 .
Figure 2 is the convergence plot of Ψ NN trained with Φ SOR ,k on P s .
The results show that the meta-solver Ψ NN is optimized for each corresponding solver Φ .
In Table 1, the best performance for a given test solver is achieved at the diagonal where the training and test solvers match.
We can explicitly investigate this adaptive phenomena by visualizing the meta-solver outputs for a representative problem instance.
Figure 1 shows that the meta-solver trained with more iterations tends to ignore high frequencies and focus more on low frequencies.
This is because the
Jacobi method converges fast in high frequencies but slow in low frequencies (Saad (2003)), and the trained meta-solver compensates the weakness of the solver.
These results show that the meta-solver actively adapts to the nature of the inner-loop solver.
Furthermore, we can observe that Theorem 1 holds numerically for the results.
Figure 2 illustrates that a meta-solver trained with more iterations is always better than those trained with fewer iterations, provided the number of iterations during testing is large.
In addition, at 10,000 iterations, the error of Ψ NN trained with Φ SOR , 64 (MSE of 0.0054 and relative error of 0.013) remains approximately 8% better than Ψ NN trained with Φ SOR , 0 , which shows that the GBMS keeps its advantage for a large number of iterations.
These results support that Theorem 1 holds for various practical settings with signiﬁcant improvement
[Table]
[Table]
[Table]
[Figure]
3. A PPLICATION : I NCOMPRESSIBLE FLOW SIMULATIONS
In this section, we introduce an application of GBMS for iterative solvers.
GBMS is developed for repeatedly solving task τ sampled from a given task distribution P .
One of the typical situations where GBMS is effective is that similar differential equations are repeatedly solved as a step of solving a time-dependent problem.
For example, in many methods of incompressible ﬂow simula- tions, Poisson equations are solved to compute the pressure ﬁeld at every time step, and this process occupies a large part of the computation time for ﬂuid simulations (Ajuria Illarramendi et al.
Problem deﬁnition.
Let us now describe the incompressible ﬂow simulation setup.
We consider 2D channel ﬂow with a ﬂuctuating inﬂow.
The velocity ﬁeld v and pressure ﬁeld p follow the
Navier-Stokes equations in the form:
[Equation]
where ρ is the density, and ν is the viscosity.
The equation (6) is called the pressure Poisson equation.
Let ρ = 1 , ν = 0 .
Let the domain of interest D be (0 , 1) × (0 , 1) and the time t ∈ [0 , 1] .
Let y =
and y
= 1 be the non-slip wall and ﬂuctuating inﬂow v in deﬁned on x = 0 .
By using the ﬁnite difference scheme, we discretize the equations into 128 by 128 spacial grids and 1,000 time steps.
Then, v and p are computed alternately by the discretized equations.
Details of the numerical scheme for solving Navier-Stokes equations follow Barba & Forsyth (2018).
GBMS for iterative method is used for solving the pressure Poisson equation (6).
Let us deﬁne the meta-solving problem for the above setting.
The task τ is to solve a pressure
Poisson equation.
The dataset D τ is { f τ , p τ , f τ − 1 , p τ − 1 , f τ − 2 , p τ − 2 } , where f τ is the right hand side and p τ is the solution of the pressure Poisson equation.
f τ − i , p τ − i are those at i th previous time step. Although p τ is determined by f τ theoretically, additional features of previous timesteps can provide useful information to determine a good guess for p τ .
The solution parameter space U τ is R 128 × 128 .
The loss function is the relative (cid:96) 2 -error.
The task distribution ( T , P ) is determined by the distribution of v in and the computation process of the simulation.
Let P in be the distribution of v in , whose detail is found in Appendix C.2.
We prepare 20 inﬂows for training, 10 for validation, and 10 for test, which are denoted by V trainin , V valin and V testin respectively.
Then, we generate datasets under each inﬂow by solving the Navier-Stokes equations with the SOR method without GBMS, where the relative error tolerance is 10 − 9 .
The solver Φ SOR ,k is the SOR method with k iterations starting at an initial guess θ ∈ Θ .
We consider two meta-solvers.
One is a baseline Ψ BL that takes τ as an input and gives the heuristic initial guess, which is the solution of the previous step Poisson equation p τ − 1 .
This choice is based on the domain knowledge of ﬂuid simulations and standard in the literature (Ferziger & Peri´c (2002)).
The meta-solver Ψ NN is a variant of 2D U-Net with weights ω ∈ Ω , which takes D τ \ { p τ } as an input and generates initial guess θ τ ∈ Θ for the solver Φ SOR ,k . Note that Ψ NN trained with Φ SOR , 0 (i.e. solver-independent initial guess) is similar to the method in Ajuria Illarramendi et al. and is considered as a data-driven baseline in this paper. We also note that the choice of Ψ NN is arbitrary, and any other problem-speciﬁc neural network architectures can be used as Ψ NN .
Training.
We train the meta-solver Ψ NN with the SOR solver Φ SOR ,k with k = 0 , 4 , 16 , 64 by using the Algorithm 2.
The details of the architecture and hyper-parameters of Ψ NN are found in Appendix
D.2.
In addition, we use data augmentation during training to prevent overﬁtting to the high-accuracy training data.
During the training, we augment data by proceeding 10 time steps with Φ SOR ,k and Ψ NN being trained.
Speciﬁcally, after computing the loss for p τ , we proceed one time step using Φ SOR ,k and Ψ NN , and compute the loss for next p τ +1 using the computed previous step information.
Then, it is repeated 10 times for each mini batch.
Without the data augmentation, the inputs of Ψ NN are from the prepared dataset and are always accurate during training.
However, during inference, the inputs of Ψ NN are the outputs of the simulation at previous time steps, which can be of a slightly different distribution than encountered during training (i.e. distribution shift) due to numerical errors and generalization gaps.
This then causes further accuracy problems for the simulation at the next step, and this issue compounds itself in time.
This eventually degrades the performance if the shift in distribution is not dealt with.
The data augmentation is used to prevent it.
We remark that this training process does not require the whole of the simulation to be differentiable.
This is practically important because we can employ GBMS with established solvers by only replacing the Poisson solver and can utilize the other parts without any modiﬁcation.
Evaluation.
We incorporate the trained meta-solvers into simulations and evaluate them by the relative (cid:96) 2 -error of the velocity ﬁeld on test inﬂows under a ﬁxed number of the SOR iterations.
More speciﬁcally, the performance metric is
[Figure]
, where v i is the ground truth of i th step velocity under the inﬂow v in , which is obtained by using Ψ BL and Φ SOR ,k with a large enough k to achieve the relative error tolerance of 10 − 9 , and ˆ v i is the i th step velocity obtained by using the trained Ψ NN and Φ SOR ,k with a ﬁxed k .
To stabilize the simulation, the ﬁrst 100 steps are excluded, and the simulation with the trained Ψ NN starts with 101 steps.
The results presented in Table 2 demonstrate the advantage of GBMS.
Although all Ψ NN diverge for Φ SOR , 0 because of error accumulation, for the other solvers, the best performance with signiﬁcant improvement is achieved at the diagonal where the training and test solvers match.
For example, for the test solver Φ SOR , 4 , the accuracy of Ψ NN trained with Φ SOR , 4 is approximately 50 times better than regular supervised learning ( Ψ NN trained with Φ SOR , 0 ) and 130 times better than the classical baseline Ψ BL .
Furthermore, GBMS keeps its advantage for a larger number of iterations than the number which it is trained with.
For example, for Φ SOR , 64 , the meta-solver trained with Φ SOR ,has the best performance, followed by one trained with Φ SOR , 16 , Φ SOR , 4 , and Φ SOR , 0 in that order.
As for computation time, Ψ NN trained with Φ SOR , 4 and tested with Φ SOR , 4 (GBMS approach) takes 11 sec to simulate the ﬂow for one second, while Ψ BL tested with Φ SOR , 128 takes 105 sec to achieve similar accuracy to the GBMS approach.
Also, Ψ NN trained with Φ SOR , 0 and tested with Φ SOR ,takes 73 sec to achieve the similar accuracy. Since it takes approximately 7 hours to train Ψ NN with a GeForce RTX 3090, our approach is effective if we simulate the ﬂow for more than 268 seconds, which can be easily satisﬁed in practical settings.
To summarize, GBMS performs best for a ﬁxed number of iterations and can achieve high accuracy within a smaller number of iterations and shorter computation time than the classical baseline and regular supervised learning.
[Table]
C ONCLUSION
In this paper, we proposed a formulation of meta-solving and a general gradient-based algorithm (GBMS) to solve this class of problems.
In the proposed framework, many related works that used neural networks for solving differential equations were organized in a uniﬁed way and regarded as variants of GBMS.
Thus, the GBMS approach offers a general design pattern to develop solution algorithms that blends machine learning and scientiﬁc computing.
As a concrete illustration, we applied GBMS to iterative methods for solving differential equations and showed its advantage over both classical numerical methods and regular supervised learning.
In particular, the proposed method was tested in the incompressible ﬂow simulation and the accuracy was improved by approximately 50 times compared to regular supervised learning and 130 times compared to a classical baseline for a ﬁxed number of iterative solver iterations.
We will study applications of GBMS to other types of problems such as nonlinear equations in future work.
R EPRODUCIBILITY S TATEMENT
To ensure the reproducibility, we provide the detailed proof of Theorem 1 in Appendix B.
As for the experiments, we provide the details of dataset generation in Appendix C and the details of network architecture and training hyper-parameters in Appendix D.
In addition, we conducted experiments with different random seeds to obtain robust results.
Finally, we will make our source code public at a later date.
Source code with limited documentation is available upon request.
A D ETAILS OF EXAMPLES
Example 3 (Feliu-Fab`a et al.
Feliu-Fab`a et al.
(2020) propose the neural network architecture with meta-learning approach that solves the equations in the form L η u ( x ) =
f
( x ) with appropriate boundary conditions, where L η is a partial differential or integral operator parametrized by a parameter function η ( x ) .
This work can be described as follows:
• Task τ : The task τ is to solve a L η u ( x ) =
f
( x ) for η = η τ : – Dataset D τ : The dataset D τ is D τ = { η τ , f τ , u τ } , where η τ , f τ , u τ ∈ R N are the [], right hand side, and solution respectively.
The reference solution u τ is obtained by [].
– Solution parameter space U τ : The solution parameter space U τ is a subset of R N for N ∈ N .
– Loss function L τ : The loss function L τ : U → R ≥ 0 is the mean squared error with the reference solution, i.e. L τ (ˆ u )
=
(cid:107) u τ − ˆ u (cid:107) 2 .
• Task space ( T , P ) :
The task distribution ( T , P ) is determined by the distribution of η τ and
f τ . • Solver Φ : The solver Φ : T × Θ → U τ is implemented by a neural network imitating the wavelet transform, which is composed by three modules with weights θ =
( θ 1 , θ 2 , θ 3 ) .
In detail, the three modules, φ 1 ( · ; θ 1 ) , φ 2 ( · ; θ 2 ) , and φ 3 ( · ; θ 3 ) , represent forward wavelet transform, mapping η to coefﬁcients matrix of the wavelet transform, and inverse wavelet transform respectively.
Then, Φ is represented by Φ( τ ; θ ) =
φ 3 (( φ 2 ( η τ ; θ 2 ) φ 1 ( f τ ; θ 1 )); θ 3 )
= ˆ u .
• Meta-solver Ψ :
The meta-solver Ψ : T × Ω → Θ is the constant function that returns its parameter ω , so Ψ( τ ; ω )
=
ω
= θ and Ω = Θ .
Note that θ does not depend on τ in this example.
Then, the weights ω
= θ is optimized by a gradient descent algorithm as described in Algorithm 1.
Example 4 (Chen et al.
The target equation in Chen et al.
(2020) is a linear systems of equations A η u = f obtained by discretizing parameterized steady-state PDEs, where u, f ∈ R N and A η ∈ R N × N is determined by η , a parameter of the original equation.
This work can be described as follows:
• Task τ : The task τ is to solve a linear system A η u = f for η = η τ : – Dataset D τ : The dataset D τ is { η τ , f τ } .
– Solution parameter space U τ : The solution parameter space U τ is R N .
– Loss function L τ : The loss function L τ : U → R ≥ 0 is an unsupervised loss based on the residual of the equation, L τ (ˆ u ) = (cid:107) f τ − A η τ ˆ u (cid:107) 2 / (cid:107) f τ (cid:107) 2 .
• Task space ( T , P ) :
The task distribution ( T , P ) is determined by the distribution of η τ and
f τ . • Solver Φ : The solver Φ : T × Θ → U is iterations of a function φ τ ( · ; θ ) :
U → U that represents an update step of the multigrid method.
φ τ is implemented using a convolutional neural network and its parameter θ is the weights corresponding to the smoother of the multigrid method.
Note that weights of φ τ other than θ are naturally determined by η τ and the discretization scheme.
In addition, φ τ takes f τ as part of its input at every step, but we write these dependencies as φ τ for simplicity.
To summarize, Φ( τ ; θ
= φ kτ
( u (0) ; θ )
=
ˆ u ,
where k is the number of iterations of the multigrid method and u (0) is initial guess, which is 0 in the paper.
• Meta-solver Ψ :
The meta-solver Ψ : T × Ω → Θ is implemented by a neural network with weights ω , which takes A η τ as its input and returns weights θ τ that is used for the smoother inspired by the subspace correction method.
Then, ω is optimized by a gradient decent algorithm with the number of multigrid iteration k = 1 as described in Algorithm 1.
Example 5 (Psaros et al.
In Psaros et al.
(2021), meta-learning is used for learning a loss function of the physics-informed neural network, shortly PINN (Raissi et al.
The target equations are the following:
[Equation]
where D ⊂ R M is a bounded domain, u : [0 , T ] × D → R N is the solution, F λ is a nonlinear operator containing differential operators, B λ is a operator representing the boundary condition, u 0 ,λ : D → R N represents the initial condition, and λ is a parameter of the equations.
• Task τ : The task τ is to solve a differential equation by PINN: – Dataset D τ : The dataset D τ is the set of points ( t, x ) ∈ [0 , T ] × D and the values of u at the points if applicable.
In detail, D τ = D f,τ ∪ D b,τ ∪ D u 0 ,τ ∪ D u,τ , where D f,τ , D b,τ , and D u 0 ,τ are sets of points corresponding to the equation (a), (b), and (c) respectively.
D u,τ is the set of points ( t, x ) and observed values u ( t, x ) at the points.
In addition,
each dataset
D
,τ is divided into training set D train
,τ and validation set D val
· ,τ .
– Solution parameter space U τ : The solution parameter space U τ is the weights space of PINN.
– Loss function L τ : The loss function L τ : U → R ≥ 0 is based on the evaluations at the points in D val τ .
In detail,
[Equation]
,τ (ˆ u ) + L val u 0 ,τ (ˆ u ) , where
L val f,τ = w f | D f,τ | (cid:88) ( t,x ) ∈
D
f,τ (cid:96) ( F λ [ˆ u ]( t, x ) , 0 )
L val b,τ = w b | D b,τ | (cid:88) ( t,x ) ∈ D b,τ (cid:96) ( B λ [ˆ u ]( t, x ) , 0 )
L val u 0 ,τ = w u| D u 0 ,τ | (cid:88) ( t,x ) ∈ D u 0 ,τ (cid:96) (ˆ u (0 , x ) , u 0 ,λ ( x )) ,
and (cid:96) :
R N × R N → R ≥ 0 is a function.
In the paper, the mean squared error is used as (cid:96) .
• Task space ( T , P ) :
The task distribution ( T , P ) is determined by the distribution of λ .
• Solver Φ : The solver Φ : T × Θ → U τ is the gradient decent for training the PINN.
The parameter θ ∈ Θ controls the objective of the gradient decent, L train τ (ˆ u ; θ ) =
L train f,τ (ˆ u ; θ ) + L train b,τ (ˆ u ; θ )
+ L train u 0 ,τ (ˆ u ; θ ) + L train u,τ (ˆ u ; θ ) , where the difference from L val τ is that parametrized loss (cid:96) θ is used in L train τ instead of the MSE in L val τ .
Note that the loss weights w f , w b , w u 0 , w u in L train τ are also considered as part of the parameter θ .
In the paper, two designs of (cid:96) θ are studied.
One is using a neural network, and the other is using a learned adaptive loss function.
In the former design, θ is the weights of the neural network, and in the latter design, θ is the parameter in the adaptive loss function.
• Meta-solver Ψ :
The meta-solver Ψ : T × Ω → Θ is the constant function that returns its parameter ω , so Ψ( τ ; ω )
=
ω
= θ and Ω = Θ .
Note that θ does not depend on τ in this example.
Then, the parameter ω
= θ is optimized by a gradient decent algorithm as described in Algorithm 1.
B P ROOF OF T HEOREM
Proof of Theorem 1.
Equation (3) can be represented as follows:
[Equation]
Since the mean of f τ is 0 and the covariance is the identity matrix, W is the minimizer of (8) if and only if W is the minimizer of the following:
[Equation]
Let λ i and v i ( i = 1 , 2 , . . .
, N ) be eigenvalues and eigenvectors of M .
Note that λ i = cos iπN +1 and 1 > λ 1 > λ 2 > · · · > λ N > − 1 .
Since A = 2 I − 2 M , eigenvalues and eigenvectors of A are 2 − 2 λ
i and v i .
By the eigenvalue decomposition, M and A can be written as
[Equation] [Equation]
11) where Λ = diag ( λ i ) and V = ( v 1 , v 2 , . .
. , v N ) .
By the decompositions,
[Equation]
Thus, the minimizer W k is
[Equation]
where r = rank ( W k ) <
N , and i ( j, k ) is the index i where | λ ki 2 − 2 λ i | takes the top j th value in i ∈ { 1 , 2 , . . .
, N } .
Then, the following lemma holds.
Lemma 2.
If k 1 < k 2 , then for all k ≥ k
[Equation]
where the equality holds if and only if W k 1 = W k 2 .
Since (8) and (9) are equivalent, if Lemma 2 holds, then Theorem 1 holds.
Proof of Lemma 2.
Let
[Equation]
I k is the set of N − r indices corresponding to the least
N − r values of | λ ki 2 − 2 λ i | .
We have
[Equation]
If W k 1 = W k 2 , then (19) equals 0.
Assume W k 1 (cid:54) =
W k 2 , so I k 1 (cid:54)
=
I k 2 .
Note that
I k
= { min I k , min I k + 1 , . .
, min I k
+ ( N − r ) − 1 = max
I k } , and for any i ∈ I k , we have |
λ i | ≤ | λ max I k | .
In addition, for k 1 < k 2 , we have max I k 1 ≥ max I k 2 .
Let
[Equation] [Equation]
[Equation]
Then, we have
[Equation]
[Equation] [Equation] The last inequality (25) is because for any i ∈ I k 2 , (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
26) and
[Equation]
This completes the proof of Lemma 2.
Lemma 2 holds.
Thus, Theorem 1 holds.
C D ETAILS OF TASK DISTRIBUTIONS C. D ISTRIBUTIONS OF 1D P OISSON EQUATION
In P s , u is represented by u ( x
[Equation]
=
[Equation]
c
[Equation]
, (28) where
[Equation] [Equation]
[Equation]
In P h , u is represented by
[Equation]
(32) where
[Equation] [Equation]
[Equation]
C. D ISTRIBUTION OF INFLOW
In P in , v in = ( v 1 , v 2 ) T is represented by
[Equation] [Equation]
) where
[Equation]
and
[Equation]
D D ETAILS OF NETWORK ARCHITECTURE AND HYPER - PARAMETERS D. D ETAILS OF SECTION 3.1.
In section 3.1.2, Ψ NN is a variant of 1D U-Net with a residual connection to leverage the heuristic initial guess u h
= Ψ BL ( τ ) , so Ψ NN ( τ ) =
Ψ NN ( f τ , u h ) =
u h + 1DUNet
( f τ , u h ) , where 1DUNet consists of four stages with halved resolutions.
Each stage has two convolutional layers with kernel size 11 and the activation function tanh .
In the ﬁrst stage, the number of channels is 8 , and it is doubled as the resolution is halved.
By utilizing the linearity of the Poisson equation, inputs are normalized by (cid:107) f τ (cid:107) before feeding them into Ψ NN , and the ﬁnal output ˆ u is denormalized by (cid:107) f τ (cid:107) .
The model is trained for 2000 epochs by Adam with the learning rate 0 . 0005 and the batch-size 512 .
The learning rate is decreased by 1 / 5 times when the validation loss does not improve forepochs, and the training is terminated when the validation loss does not improve for 250 epochs.
The model with the best validation loss is used for evaluation.
D. D ETAILS OF SECTION 3.
In section 3.2, Ψ NN is a variant of 2D U-Net.
The differences from Ψ NN in section 3.1.2 are its dimension, kernel size 5 , and the starting number of channels 16.
The model is trained for 50 epochs by Adam with the learning rate 0 . 0005 and the batch-size 64 , and the learning rate is decreased by 1 / 5 at 30 epoch.
Note that the meta-solver is trained for 10 time steps for each mini batch by the data augmentation.
The model with the best validation loss is used for evaluation.
E D ETAILS OF EXPERIMENT RESULTS
[Table]
[Table] [Table]
[Table] [Table]