Year: 2021

Abstract:
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 efficiency of finding subsequent solutions. This offers a unique opportunity to combine machine learning (in particular, meta-learning) and scientific computing. To date, a variety of such domainspecific 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 unified 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.

heading: 1 INTRODUCTION
text: 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. (2019)). For another example, in computational fluid dynamics, many methods involve solving a Poisson equation to compute the pressure field in every time step of the simulation (Ajuria Illarramendi et al. (2020)). In these situations, we can utilize the data of the previously solved problems to solve the next similar but unseen problems more efficiently, 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. (2017); Özbay et al. (2021); Tompson et al. (2017); Hsieh et al. (2018); Huang et al. (2020)). For example, in Tang et al. (2017), Ajuria Illarramendi et al. (2020), and Özbay et al. (2021), convolutional neural networks are used to predict the solution of Poisson equations, and Ajuria Illarramendi et al. (2020) and Özbay 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-Fabà et al. (2020); Chen et al. (2020); Psaros et al. (2021); Guo et al. (2021)). Meta-learning, or learning to learn, leverages previous learning experiences to improve future learning performance (Hospedales et al. (2021)), which fits 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 offline discovery of physics-informed neural network loss functions.
Although many methods have been proposed in this direction, they are often problem-specific. In other words, there lacks both a unified 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. (2017)) - 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.

heading: 2 GRADIENT-BASED META-SOLVING
text: 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.

heading: 2.1 GENERAL FORMULATION OF META-SOLVING
text: Let us now introduce the general formulation of meta-solving. We first fix 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 RN 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 find 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 defined on a suitable σ-algebra on T . A solver Φ is a function from T × Θ to U , where U = ⋃ τ∈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 defined as follows: Definition 1 (Meta-solving problem). For a given task space (T , P ), solver Φ, and meta-solver Ψ, find ω ∈ Ω 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. (2017)). The components of the problem are the following. The task τ is to learn a regression model from data. The dataset is Dτ = {(xi, yi)Ki=1}, which satisfies y = fτ (x) for an unknown function fτ . Dτ is divided into the training set Dtrainτ and validation set Dvalτ . 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 Dtrainτ . The loss function Lτ : U → R≥0 is the mean squared error (MSE) on Dvalτ . 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 minimizeLtrainτ (u) =
1 |Dtrainτ | ∑ (x,y)∈Dtrainτ
‖NN(x;u)− y‖2. Its parameter θ ∈ Θ is initial weights u(0) of NN, so Θ = U . Thus, Φ(τ ; θ) = u(0) − α∇u(0)Ltrainτ (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 fine-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 defined as meta-solving problem, which is
min ω E τ∼P [Lτ (Φ(τ ; Ψ(τ ;ω))] = min ω E τ∼P ∑ (x,y)∈Dvalτ ∥∥NN(x;ω − α∇ωLtrainτ (ω))− y∥∥2 . (1)
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-specific 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 efficiency depending on the problem instance. Thus, finding 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 RNτ . The loss function Lτ : Uτ → R≥0 measures the accuracy of û ∈ Uτ . In this example, the `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, finding 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.

heading: 2.2 GRADIENT-BASED META-SOLVING
text: In the previous section, we defined 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 find 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. (2011), Mitzenmacher & Vassilvitskii (2021), and Balcan (2021)). These focus on discontinuous problems such as combinatorial optimization, while our framework focuses on differentiable problems.

heading: 2.3 ORGANIZING RELATED WORKS
text: Owing to the general formulation presented above, we can organize several related works on learning-based methods for scientific computing and describe them in a unified 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 defined in each case are found in Appendix A.
First, let us review Feliu-Fabà et al. (2020). 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
Algorithm 1: Gradient-based meta-solving algorithm Require: (P, T ): task space, Φ,Ψ: differentiable solver, S: stopping criterion, α: learning rate Result: ω
1 initialize ω; 2 while S is not satisfied do 3 Sample task τ ∼ P ; 4 Generate parameter of Φ using meta-solver: θτ (ω)← Ψ(τ ;ω); 5 Solve task τ using solver: û(ω)← Φ(τ ; θτ (ω)); 6 Update meta-solver parameter: ω ← ω − α∇ωLτ (û(ω)); 7 end
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); Özbay et al. (2021); Tompson et al. (2017); 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. (2021) 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. (2019)), 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 unified 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. (2020). 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.

heading: 3 GBMS FOR ITERATIVE METHODS
text: 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 Example 2 beyond 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 defined 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 finding. Note that although Huang et al. (2020), Ajuria Illarramendi et al. (2020), and Özbay 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.
Algorithm 2: Gradient-based meta-solving algorithm for iterative solvers Require: (P, T ): task space, Ψ: differentiable solver, S: stopping criterion, k: number of iterations of the iterative solver, α: learning rate Result: ω
1 initialize ω; 2 while S is not satisfied do 3 Sample task τ ∼ P ; 4 Form iterative function φτ depending on τ ; 5 Generate initial guess using meta-solver: u(0)τ (ω)← Ψ(τ ;ω); 6 for i← 1 to k do 7 Update approximate solution by iterative function: u(i)τ (ω)← φτ (u(i−1)τ (ω)); 8 end 9 Update meta-solver parameter: ω ← ω − α∇ωLτ (u(k)τ (ω));
10 end

heading: 3.1 TOY PROBLEM: 1D POISSON EQUATIONS
text: 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.

heading: 3.1.1 THEORETICAL ANALYSIS
text: Let us recall Example 2 and set the domain of interestD = (0, 1). Then, the target equation becomes the following 1D Poisson equation with Dirichlet boundary condition:
− d 2
dx2 u(x) = f(x), x ∈ (0, 1)
u(0) = a, u(1) = b.
(2)
To solve this equation numerically, it is discretized with finite difference scheme, and we rewrite it as the matrix equationAu = f , where the domain [0, 1] is discretized intoN points, soA ∈ RN×N and u, f ∈ RN . 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 defined 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τ = RN . The loss function is Lτ (û) = ‖uτ − û‖2. The task distribution (T , P ) is determined by the distribution of fτ , denoted by Pf . 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 + 12fτ , where M = I − 1 2A. 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
min ω∈Ω E τ∼P [Lτ (Φ(τ ; Ψ(τ ;ω))] = min W E fτ∼Pf ∥∥uτ − φkτ (Wfτ )∥∥2 . (3) 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 Wk. Furthermore, if k1 < k2, then for all k ≥ k2,
E fτ∼Pf ∥∥uτ − φkτ (Wk1fτ )∥∥2 ≥ E fτ∼Pf ∥∥uτ − φkτ (Wk2fτ )∥∥2 , (4) where the equality holds if and only if Wk1 = Wk2 .
Theorem 1 guarantees improvement of meta-solving for the considered setting. For example, suppose k1 = 0 (regular supervised learning, i.e. directly predicting the solution without considering the solver) and k2 = 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.

heading: 3.1.2 NUMERICAL EXAMPLES FOR MORE GENERAL CASES
text: 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 defined 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 Ps and Ph, 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. (2015)) Ψ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 Ps.
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 significant improvement.

heading: 3.2 APPLICATION: INCOMPRESSIBLE FLOW SIMULATIONS
text: 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 flow simula-
tions, Poisson equations are solved to compute the pressure field at every time step, and this process occupies a large part of the computation time for fluid simulations (Ajuria Illarramendi et al. (2020)).
Problem definition. Let us now describe the incompressible flow simulation setup. We consider 2D channel flow with a fluctuating inflow. The velocity field v and pressure field p follow the Navier-Stokes equations in the form:
∂v ∂t + (v · ∇)v = −1 ρ ∇p+ ν∆v (5)
∆p = −ρ∇ · ((v · ∇)v) (6)
where ρ is the density, and ν is the viscosity. The equation (6) is called the pressure Poisson equation. Let ρ = 1, ν = 0.01. Let the domain of interest D be (0, 1) × (0, 1) and the time t ∈ [0, 1]. Let y = 0 and y = 1 be the non-slip wall and fluctuating inflow vin defined on x = 0. By using the finite 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 define 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 ith 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 R128×128. The loss function is the relative `2-error. The task distribution (T , P ) is determined by the distribution of vin and the computation process of the simulation. Let Pin be the distribution of vin, whose detail is found in Appendix C.2. We prepare 20 inflows for training, 10 for validation, and 10 for test, which are denoted by V trainin ,Vvalin and V testin respectively. Then, we generate datasets under each inflow 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 fluid simulations and standard in the literature (Ferziger & Perić (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. (2020) 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-specific 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 overfitting to the high-accuracy training data. During the training, we augment data by proceeding 10 time steps with ΦSOR,k and ΨNN being trained. Specifically, 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 modification.
Evaluation. We incorporate the trained meta-solvers into simulations and evaluate them by the relative `2-error of the velocity field on test inflows under a fixed number of the SOR iterations. More specifically, the performance metric is 1|V testin | ∑ vin∈V testin 1 900 ∑1000 i=101 ‖vi−v̂i‖ ‖v̂i‖ , where vi is the ground truth of ith step velocity under the inflow vin, 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 ith step velocity obtained by using the trained ΨNN and ΦSOR,k with a fixed k. To stabilize the simulation, the first 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 significant 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,64 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 flow 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,64 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 flow for more than 268 seconds, which can be easily satisfied in practical settings. To summarize, GBMS performs best for a fixed 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.

heading: 4 CONCLUSION
text: 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 unified 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 scientific 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 flow 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 fixed number of iterative solver iterations. We will study applications of GBMS to other types of problems such as nonlinear equations in future work.

heading: 5 REPRODUCIBILITY STATEMENT
text: 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.

heading: A DETAILS OF EXAMPLES
text: Example 3 (Feliu-Fabà et al. (2020)). Feliu-Fabà 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τ ∈ RN 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 RN 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̂) = ‖uτ − û‖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 coefficients 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. (2020)). 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 ∈ RN and Aη ∈ RN×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 RN . – Loss function Lτ : The loss function Lτ : U → R≥0 is an unsupervised loss based on
the residual of the equation, Lτ (û) = ‖fτ −Aητ û‖ 2 /‖fτ‖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 takesAητ 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. (2021)). 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. (2019)). The target equations are the following:
Fλ[u](t, x) = 0, (t, x) ∈ [0, T ]×D (a) Bλ[u](t, x) = 0, (t, x) ∈ [0, T ]× ∂D (b)
u(0, x) = u0,λ(x), x ∈ D, (c)
where D ⊂ RM is a bounded domain, u : [0, T ] × D → RN is the solution, Fλ is a nonlinear operator containing differential operators, Bλ is a operator representing the boundary condition, u0,λ : D → RN 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τ = Df,τ ∪Db,τ ∪Du0,τ ∪Du,τ , where Df,τ , Db,τ , and Du0,τ are sets of points corresponding to the equation (a), (b), and (c) respectively. Du,τ 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 Dtrain·,τ 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 Dvalτ . In detail,
Lτ (û) = L val τ (û) = L val f,τ (û) + L val b,τ (û) + L val u0,τ (û),
where
Lvalf,τ = wf |Df,τ | ∑ (t,x)∈Df,τ ` (Fλ[û](t, x),0) Lvalb,τ = wb |Db,τ | ∑ (t,x)∈Db,τ ` (Bλ[û](t, x),0)
Lvalu0,τ = wu0 |Du0,τ | ∑ (t,x)∈Du0,τ ` (û(0, x), u0,λ(x)) ,
and ` : RN × RN → R≥0 is a function. In the paper, the mean squared error is used as `.
• 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, Ltrainτ (û; θ) = Ltrainf,τ (û; θ) + L train b,τ (û; θ) + L train u0,τ (û; θ) + L train u,τ (û; θ), where the difference from L val τ is
that parametrized loss `θ is used in Ltrainτ instead of the MSE in L val τ . Note that the loss weights wf , wb, wu0 , wu in L train τ are also considered as part of the parameter θ. In the paper, two designs of `θ 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.

heading: B PROOF OF THEOREM 1
text: Proof of Theorem 1. Equation (3) can be represented as follows:
min W E fτ∼Pf ∥∥uτ − φkτ (Wfτ )∥∥2 = min W E fτ∼Pf ∥∥Mk(uτ −Wfτ )∥∥2 (7) = min
W E fτ∼Pf ∥∥Mk(A−1fτ −Wfτ )∥∥2 . (8) 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:
min W
∥∥Mk(A−1 −W )∥∥2 F . (9)
Let λi and vi (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 = 2I − 2M , eigenvalues and eigenvectors of A are 2− 2λi and vi. By the eigenvalue decomposition, M and A can be written as
M = V ΛV T (10)
A = V (2I − 2Λ)V T , (11)
where Λ = diag(λi) and V = (v1, v2, . . . , vN ). By the decompositions,∥∥Mk(A−1 −W )∥∥2 F = ∥∥V ΛkV T (V (2I − 2Λ)−1V T −W )∥∥2 F (12)
= ∥∥V Λk(2I − 2Λ)−1V T − V ΛkV TW )∥∥2
F (13)
Thus, the minimizer Wk is
Wk = r∑ j=1 (2− 2λi(j,k))−1vi(j,k)vTi(j,k), (14)
where r = rank(Wk) < N , and i(j, k) is the index i where | λ k i
2−2λi | takes the top jth value in i ∈ {1, 2, . . . , N}. Then, the following lemma holds.
Lemma 2. If k1 < k2, then for all k ≥ k2∥∥Mk(A−1 −Wk1)∥∥2F ≥ ∥∥Mk(A−1 −Wk2)∥∥2F , (15) where the equality holds if and only if Wk1 = Wk2 .
Since (8) and (9) are equivalent, if Lemma 2 holds, then Theorem 1 holds.
Proof of Lemma 2. Let
Ik := {1, 2, . . . , N} \ {i(j, k) : j = 1, 2, . . . , r}. (16)
Ik is the set of N − r indices corresponding to the least N − r values of | λ k i
2−2λi |. We have∥∥Mk(A−1 −Wk1)∥∥2F − ∥∥Mk(A−1 −Wk2)∥∥2F (17) = ∑ i∈Ik1 ( λki 2− 2λi )2 − ∑ i∈Ik2 ( λki 2− 2λi )2 (18)
= ∑
i∈Ik1\Ik2
( λki
2− 2λi
)2 − ∑ i∈Ik2\Ik1 ( λki 2− 2λi )2 . (19)
If Wk1 = Wk2 , then (19) equals 0. Assume Wk1 6= Wk2 , so Ik1 6= Ik2 . Note that Ik = {min Ik,min Ik + 1, . . . ,min Ik + (N − r) − 1 = max Ik}, and for any i ∈ Ik, we have |λi| ≤ |λmax Ik |. In addition, for k1 < k2, we have max Ik1 ≥ max Ik2 . Let
ik1 := arg min i∈Ik1\Ik2
| λ k i
2− 2λi | (20)
ik2 := arg max i∈Ik2\Ik1
| λ k i
2− 2λi | (21)
C := |Ik1 \ Ik2 | = |Ik2 \ Ik1 |. (22)
Then, we have
(19) ≥C
( λkik1
2− 2λik1
)2 − C ( λkik2
2− 2λik2
)2 (23)
=C  λk2+(k−k2)ik1
2− 2λik1
2 − λk2+(k−k2)ik2
2− 2λik2
2  (24)
>0. (25)
The last inequality (25) is because for any i ∈ Ik2 ,∣∣∣∣∣ λ k2 ik1
2− 2λik1
∣∣∣∣∣ > ∣∣∣∣∣ λk2i2− 2λi ∣∣∣∣∣ (26) and |λik1 | > |λmax Ik2 | ≥ |λi|. (27) This completes the proof of Lemma 2.
Lemma 2 holds. Thus, Theorem 1 holds.

heading: C DETAILS OF TASK DISTRIBUTIONS
text: C.1 DISTRIBUTIONS OF 1D POISSON EQUATION
In Ps, u is represented by
u(x) = 20∑ i=1 ai sin(biπ(x− ci)), (28)
where
ai ∼ N (0, 1) (29) bi ∼ Unif(0, 128) (30) ci ∼ Unif(0, 1). (31)
In Ph, u is represented by
u(x) = 20∑ i=1 ai tanh(biπ(x− ci)), (32)
where
ai ∼ N (0, 10) (33) bi ∼ Unif(0, 30) (34) ci ∼ Unif(0, 1). (35)
C.2 DISTRIBUTION OF INFLOW
In Pin, vin = (v1, v2)T is represented by
v1(y, t) = w(y, t) cos z(t) (36) v2(y, t) = w(y, t) sin z(t), (37)
where
w(y, t) = 0.5 + 0.5 sin(ay − bt)π sin yπ sin tπ (38) z(t) = e sin(ct− d)π (39)
and
a ∼ Unif(0, 10) (40) b ∼ Unif(−5, 5) (41) c ∼ Unif(0, 5) (42) d ∼ Unif(0, 2) (43)
e ∼ Unif(0, π 2 ). (44)

heading: D DETAILS OF NETWORK ARCHITECTURE AND HYPER-PARAMETERS
text: D.1 DETAILS OF SECTION 3.1.2
In section 3.1.2, ΨNN is a variant of 1D U-Net with a residual connection to leverage the heuristic initial guess uh = ΨBL(τ), so ΨNN(τ) = ΨNN(fτ , uh) = uh + 1DUNet(fτ , uh), 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 first 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 ‖fτ‖ before feeding them into ΨNN, and the final output û is denormalized by ‖fτ‖. 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 for 200 epochs, 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.2 DETAILS OF SECTION 3.2
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 DETAILS OF EXPERIMENT RESULTS

Review summaries:
This paper proposes leveraging data from previous problem instances to improve efficiency of solving similar ones in the future. A general gradient-based method is proposed, which is applied to generating initial guesses to differential equation solutions. This problem is formulated as a meta-learning problem. 

The paper proposed a gradient-based algorithm GBMS to solve PDEs based on the solutions of other similar problems. In GBMS, a network is trained to produce good initial guess for the iterative solver of the PDE. Numerical experiments are performed to show the effectiveness of the method.