DOF: Accelerating High-order Differential Operators with Forward Propagation

Authors: ruichen  li
chuwei  wang
haotian  ye
di  he
liwei  wang
Venue: ICLR_cc_2024_Workshop_AI4DiffEqtnsInSci
Published: 2024-02-05 07:21:24

Abstract: Solving partial differential equations (PDEs) efficiently is essential for analyzing complex physical systems. Recent advancements in leveraging deep learning for solving PDE have shown significant promise. However, machine learning methods, such as Physics-Informed Neural Networks (PINN), face challenges in handling high-order derivatives of neural network-parameterized functions. Inspired by Forward Laplacian, a recent method on accelerating Laplacian computation, we propose an efficient computational framework, Differential Operator with Forward-propagation (DOF), for calculating general second-order differential operators without losing any precision. We provide rigorous proof of the advantages of our method over existing methods, demonstrating two times improvement in efficiency and reduced memory consumption on any architectures. Empirical results illustrate that our method surpasses traditional automatic differentiation (AutoDiff) techniques, achieving 2x improvement on the MLP structure and nearly 20x improvement on the MLP with Jacobian sparsity.

Main paper: 
 introduction: 
 Partial differential equations (PDEs) play a pivotal role in understanding and predicting the behavior of physical systems. While classical numerical methods have proven effective in some cases, they can be prohibitively challenging when dealing with complicated problems, e.g., turbulence in fluid dynamics [1] and high dimensional equations [2] . Recently, the advent of deep learning [3] has spurred a wave of innovations in leveraging neural networks (NN) for numerical solutions of PDEs. These NN-based approaches have been applied to various problems, including fluid dynamics [4; 5] , high-dimensional optimal control problems [6; 7] , and quantum many-body problems [8; 9] . Notable works, such as the Physics-Informed Neural Network (PINN) [4; 10; 11] and Neural Operator [12; 13; 14; 15] , showcase the potential of neural networks in capturing the underlying physics of systems governed by PDEs. Among these works, the central idea is to parameterize the solution as a neural network, and optimize this expressive network with the guidance from PDE. The ever-growing AutoDiff packages [16; 17; 18] enables convenient calculation of associated quantities such as residual losses and derivatives, avoiding discretization errors in classical methods. However, unlike most computer vision or natural language processing tasks in which first-order derivatives are sufficient for optimizations, in PDErelevant problems one has to deal with high-order derivatives. This raises a significant challenge as common AutoDiff packages are computationally intensive under this circumstance [19; 20] . There have been attempts to address this issue [21; 22; 23; 24] . However, they resort to either randomized methods [21; 22] or numerical differentiation [23; 24] , which introduce unsatisfying statistical errors and are limited in problems where high precision is not demanded.Recently, Li et al. [25] proposed a computational framework, Forward Laplacian (FL). This framework is designed specifically for accelerating Laplacian operator computation and thus can significantly boost the computation of Laplacian-relevant PDE like Schrödinger equation in quantum chemistry [8] . Remarkably, Forward Laplacian is proven to be precision-preserved as it introduces no statistical errors at all. Consequently, it is natural to ask whether the idea of FL can be leveraged for other PDEs associated with high-order derivatives computation.To fill this gap, we develop a new computational framework named Differential Operator with Forward-propagation (DOF). DOF shares a similar computational procedure as FL, yet can be applied to compute general second-order differential operators. We demonstrate that DOF outperforms conventional AutoDiff methods in memory and computational cost by a large margin, both theoretically and practically. In practice, DOF can accelerate NN-based solvers across a wide range of PDEs such as the non-homogeneous heat equation and Klein-Gordon equation.The contribution of this paper is summarized as follows:• We generalize the Forward Laplacian method and propose DOF to precisely compute arbitrary second-order differential operators of neural networks (DOF). • We demonstrate that DOF improves computation efficiency and memory consumption simultaneously, regardless of the architecture of neural networks, both theoretically and empirically. The improvement can be significant in common architectures like MLP. 
method: 
 In PINNs and many other NN-based PDE solvers, the solution of a PDE, ϕ(x), is parameterized as neural networks, ϕ(x) := ϕ(x; θ), where θ represent the NN parameters. These methods necessitate computing the high-order derivatives of a neural network. It has been shown that standard AutoDiff methods are not efficient for the high-order derivatives calculation [19; 20] .In this work, we focus on the calculation of the second-order differential operators, which have the formEQUATIONHere, a ij , b i , c : R N → R are coefficients in the second order operator, and N is the input dimension (the time variable is comprised in x for evolution equations). In practice, the first term dominates the computation cost. Thus, for brevity and clarity, we will only focus on the case when b i ≡ c ≡ 0 in the following discussions. We always denote the symmetric matrix a ij (x) i,j as A(x).In standard AutoDiff packages, the second-order operator calculation is based on the Hessian matrix H = (∂ ij ϕ(x)) i,j . We call those methods Hessian-based methods. They use multiple Jacobian calculations to derive the Hessian matrix, resulting in a huge computation cost. Recently, Li et al. [25] proposed a new computational framework, Forward Laplacian (FL), which primarily focuses on accelerating the calculation of Laplacian, i.e., A ≡ I N . Below we briefly review this method.FL computes the Laplacian with one efficient forward pass, avoiding redundant calculation in the Hessian-based approach. Following the notation in Li et al. [25] , we describe FL in a computation graph G. The node set V = {v i |i = 0, 1, ..., M } represents the operations or variables used in a neural network. We use the abbreviation i → j if there is a directed edge from v i to v j in G, and we denote operations as F , e.g., v j = F j ({v i : i → j}) for all j ≥ 0. Notice that the node indices are arranged according to the topological order, i.e., for all i → j, we have i < j. The output of ϕ is denoted by v M . In addition, ∇ and ∆ represent the gradient operator and Laplacian operator with respect to the input, respectively. Detailed notations can be found in the appendix A.Specifically, according to node dependency, FL sequentially computes the Laplacian tuple (v i , ∇v i , ∆v i ) associated with each node. In a simplified case where v i depends only on v i-1 , i.e., v i = F i (v i-1 ), the graph can be represented as a chain. We can compute the output tuple in a forward pass:EQUATIONThe propagation rule of Laplacian tuple is derived through the chain rule:EQUATIONFor the general computation graph, we can generalize eq. ( 3) to the following formula:EQUATIONEQUATIONThen we can sequentially compute the Laplacian tuple for each node in the topological order according to eqs. ( 4) to (6) . As shown in Li et al. [25] , this approach outperforms the Hessian-based approach in numerous types of computational graphs. It has been successfully applied to solving the Schrodinger equation, resulting in over a magnitude of acceleration.We now formally propose DOF to efficiently compute all kinds of second-order operators. For brevity, we denote A(x) as A and a ij (x) as a ij in the following discussion. To compute i,j a ij ∂ i ∂ j ϕ(x), we first decompose the coefficient matrix A into L ⊤ DL such that D is a diagonal matrix whose diagonal elements are all ±1 and 0. As A is a symmetric matrix, this decomposition can be done easily. For instance, we can eigen-decompose A = S ⊤ ΣS, where S is an orthogonal matrix and Σ is the diagonal eigenvalue matrix, and choose L = |Σ| 1/2 S and D = sgn(Σ).During the computation, for each node v k , we compute tuple (v j , g j , s j ) := (v j , L∇v j , Lv j ). We can derive the following formula by chain rule:EQUATIONEQUATIONEQUATIONWe can derive Lϕ by sequentially applying this propagation rule to each node in the topological order. We will prove that this method outperform the Hessian-based methods in both the memory usage and computation cost for any neural network architecture: Theorem 2.1. The computation cost (counted in FLOPs) of DOF is at most half that of Hessianbased methods for any neural network architecture. Theorem 2.2. The memory consumption of DOF (M 1 ) is smaller than that of Hessian-based methods (M 2 ) for any neural network architecture.Specifically, M 1 ≲ 2 L M 2 for an L-layer MLP.The proof for theorems 2.1 and 2.2 could be found in appendices B and D. We remark that the memory and computation consumption of DOF can be further reduced in some specific network architectures. See section 3.2 for details.To better understand the DOF method, we discuss the implementation of DOF on two special classes of second-order operators.Elliptic Operator. For elliptic operator, the coefficient matrix A is positive so we have D = I N . As a result, eqs. ( 7) to ( 9) are reduced to eqs. ( 4) to (6) . The only difference between DOF and Forward Laplacian here is the initial value of the tuple (i.e. the tuple at node v j for j = 0). Thus, we utilize some existing Forward Laplacian package [26; 27] to compute the elliptic operator.Low-rank Coefficient Matrix. It is well-known that computing the Hessian-vector product can be faster than computing the entire Hessian matrix in the standard AutoDiff package. Namely, computing Hessian-vector product takes O(1/N ) cost of computing the entire Hessian matrix. Thus, in the standard AutoDiff package, if the coefficient matrix is a low-rank matrix, the calculation of the second-order operator can be accelerated through multiple Hessian-vector products. Similarly, the DOF method also accomplishes this acceleration when dealing with a low-rank coefficient matrix.If the coefficient matrix A is a rank-r matrix, we can eliminate the columns and rows associated with the zero eigenvalues in L and D. Then we have L ′ ∈ R r×N and D ′ ∈ R r×r that still satisfy A = L ′⊤ D ′ L ′ . Thus, while the propagation rule of DOF remains the same, the dimension of g k is reduced from N to r. According to the analysis in the appendices B and D, this reduction in the dimension will lead to an O(r/N ) reduction in both the memory usage and the computation cost. 
results: 
 In this section, we compare DOF with the standard Hessian-based approach on different operators and network architectures. We study three kinds of second-order operators: elliptic operator, elliptic operator with low-rank coefficient matrix, and general operator. As for the architecture, we choose the standard MLP structure and the MLP with Jacobian sparsity. Details could be found in appendix E.The benchmark results for the MLP are shown in table 1 . For both the elliptic and general operators, DOF demonstrates significant efficiency improvements, halving the computation cost and reducing memory usage by a third compared to the Hessian-based method, aligning with our theoretical predictions. In the case of the low-rank operator, DOF achieves even greater acceleration, attributed to the dimensionality reduction of g.As discussed in Li et al. [25] , the cost of Forward Laplacian method is significantly reduced when the Jacobian of intermediate component is sparse. We notice that this sparse Jacobian property also exists in some advanced PINN network architecture [28] . As a forward AutoDiff method, DOF can also leverage this property to significantly accelerate the calculation. The memory and time comparison are listed in table 2 and the architecture details can be found in appendix E. Compared with the Hessian-based method, DOF can significantly reduce both the memory and computation consumption in the MLP with Jacobian sparsity, showing a great potential in applying our method to the advanced machine learning-based PDE solver.In this section, we briefly review computation framework in deep learning and auto differentiation, and introduce the notations we will use in the following sections.The computation graph serves as a descriptive language of deep learning models across various deep learning toolkits, including PyTorch [17] , TensorFlow [18] , and Jax [16] .In the computation graph G associated with a function ϕ(x), x = (x 1 , ...x N ) ∈ R N (we will always use N for input dimension), the edges represent function arguments, and nodes represent operations or variables. Note that G is always a directed acyclic graph. We use {v -1 , ...v -N } to represent the external nodes of G (i.e., the input x of the neural network-parameterized function ϕ), and use {v 0 , ...v M } to represent the internal nodes, sorted in topological orders. A node can be referred to as a neuron in the network. Specifically, we use v M to serve as the network output ϕ. We use the abbreviation i → j if there is a directed edge from v i to v j in G. Furthermore, we denote operations as F , e.g., v j = F j ({v i : i → j}) for all j ≥ 0.Example A.1. Take Multi-Layer Perceptron (MLP) function ϕ(x) as an example.ϕ has the form ϕ = F L • F L-1 • ... • F 0 , where F l = (F l,1 , ...F l,N l+1 ) : R N l → R N l+1are the mapping in each layer with N 0 = N and N L+1 = 1. Let u l = (u l 1 , u l 2 , ...u l N l ) ∈ R N l be the vector consist of the neurons in the l-th layer, we have u 0 = x andu l+1 = F l (u l ) = σ(W l u l + b l ) where W l ∈ R N l+1 ×N l , b l ∈ R N l+1 arenetwork parameters and σ is a nonlinear function operated element-wise.In this setting, the nodes {v i } i in the computation graph areEQUATIONand the operations generating the node representing u l i and u l-0.5 i are u l i = σ(u l-0.5 i) and u l-0.5 i= N l-1 j=1 W l-1 ij u l-1 j + b l-1 i , respectively.In most machine learning toolkits, auto differentiation(AutoDiff) implemented with back propagation algorithm is applied to compute the derivatives of neural network functions.For a function ϕ(x), this method first performs forward propagation to obtain the value of each variable v j , and construct the computation graph G. Next, a backward process is employed by creating a new computation graph Ĝ where node vi ∈ Ĝ represents the operation to calculate ∂ϕ ∂v i , i = M, ..., -N . The associated computations areEQUATIONB PROOF OF THEOREM 2.1Our proof follows similar idea to what is discussed in section 4.2 in Li et al. [25] , except that we give the proof for general second order operator here.Recall that we are going to compute eq. ( 1) and we only need to analyze the computation cost of its first term. In the following, we will use the notation described in appendix A for computation graphs.The previous auto differentiation method obtains Lϕ(x) through computing the Hessian matrix and its inner product with A. It first performs forward propagation to obtain the value of each variable v j . Next, a standard backward process is employed by creating a new computation graph Ĝ where node vi ∈ Ĝ represents the operation to calculate ∂ϕ ∂v i , i = M, ..., -N . The associated computations areEQUATIONThe Hessian matrix is then obtained through the following forward-mode Jacobian calculation along G and Ĝ, respectively:EQUATIONEQUATIONwhere we always use ∇ to denote ∇ x for simplicity.The bottleneck of Hessian computation comes from eq. ( 13) and eq. ( 14).For every connected nodes (j → i), it takes N float multiplication to yield ∂Fi ∂v j ∇v j . Therefore eq. ( 13) takes N |E| floating point operations (FLOPs) if we only count multiplications, where E denotes the set of edges in G. For eq. ( 14), the second term also takes N |E| FLOPs. To calculate the computational cost of the first term in eq. ( 14), we first introduce two notations, T and R, which are both sets of ordered tuples:EQUATION)The AutoDiff method sums over j first to obtain j:i→j,l→j ∂ 2 Fj ∂v l ∂v i ∂ϕ ∂v j for all (i, l) ∈ R, and then sums over l. For the first step, by leveraging the symmetry of Hessian matrix, it spends 1 FLOPS for a pair of (i, l, j) and (l, i, j) in T . Thus it spends 0.5|T | FLOPs in total. For the second step, for any (i, l) ∈ R, it takes N FLOPs to multiply a vector ∇v l with a scalar. Thus this steps takes N |R| FLOPs in total.Consequently, the total FLOPs for the previous method is about N (|R| + 2|E|) + 0.5|T |.Next we analyze the computation cost of DOF. For readers' convenience, we repeat the propagation scheme here:During the computation, for each node v k , we maintain a tuple(v k , g k , s k ) := (v k , L∇v k , Lv k ).The propagation rule of this tuple is:EQUATIONEQUATION)EQUATIONFor the proposed DOF method, we perform forward propagation along G to obtain ϕ(x), L∇ϕ(x) (recall that L is a matrix that comes from the decomposition of A) and Lϕ(x). The second term L∇ϕ(x) is calculated along G according to eq. ( 17). The third term, i.e., the target differential operator, is calculated along G according to eq. ( 18).The computational cost of the DOF method is dominated by eq. ( 17) and eq. ( 18). As previously discussed, eq. ( 17) takes N |E| FLOPs. For the first term in eq. ( 18), we decompose its calculation into two steps. First, we compute {g i⊤ Dg l } i≤l, (i,l)∈R . Since D is a diagonal matrix only with diagonal element in {0, ±1}, this computation takes 0.5rank(D)|R| ≤ 0.5N |R| FLOPs in total.Next, following the topological order of G, we sum over i ≤ l for each j, deriving the first term in eq. (18) . By leveraging the symmetry of the Hessian matrix and Gram matrix, we reduce this computation by a factor of 2, which is 0.5|T | FLOPs. The computational cost of the second term is negligible compared with the first term since s i is a scalar.Summing the computational cost of all the terms, we have that the DOF method uses at most 0.5N (|R| + 2|E|) + 0.5|T | FLOPs.In practice, a large percentage of operations are linear transformations, and for any linear operation F j , ∂ 2 Fj ∂v i ∂v l = 0 for any i → j, l → j. This means the value |T | is much smaller than N |R| and N |E|. Thus, our method is about two times faster than the previous Hessian-based methods for computing second order differential operators of general neural network functions.In this section, we follow the notation in appendix A for an MLP ϕ and show that there is further speedup comparing with the general 2x result stated in theorem 2.1.For MLP, We could explicitly computeEQUATIONTo compute the first term in eq. ( 18), note thatEQUATIONso we actually haveEQUATIONThis gives an alternative way to compute the first term in (18) , whose total computation cost is reduced from r(D)|R| to r(D) L l=0 N l+1 FLOPs, where r stands for rank. Recall that the first term of ( 18) is one of the dominant calculations, this manner certainly boost the efficiency.We focus on the peak memory usage caused by the forward-mode Jacobian calculation, i.e., the storage of ∇v i (or g i ), which usually dominates the memory cost. In a typical forward-mode Jacobian calculation, we write ∇v i into memory when the algorithm is applied to v i and remove it from memory when all the direct subnodes of v i have been computed.We denote τ (i) := max{j : i → j}.Then at the moment our DOF forward propagation comes to the node v j , the memory consumption is C(j) := N i:i≤j≤τ (i)1.(EQUATIONIf we further denote M 1 to be the peak memory usage when executing the forward-mode DOF, we haveEQUATIONFrom eq. ( 25) we clearly see that M 1 ≤ N |V |, here V is the set of nodes in the computation graph G. Furthermore, for any j < M, C(j) ≤ N (|V | -1). As a result, M 1 = N |V | if and only if every nodes is pointing to the end node v M , which corresponds to either one-layer linear model or a multivariate elementary function taking in input x and directly giving the output. These extreme neural network functions rarely occur in deep learning literature.By contrast, in the Hessian calculation in AutoDiff, the computation graph becomes a combination of G and Ĝ. As suggested by eq. ( 14), the node vi is a direct subnode of v i . Consequently, at the moment when ∇v M is written into memory, every ∇v i for v i ∈ G have been written into memory and could not be released since vi have not been computed yet. As a result, the peak memory usage(denoted as M 2 ) is strictly larger than N |V |, and therefore larger than the peak memory of Forward Laplacian.In the specific case when ϕ is MLP. For any node u k i , k ∈ 1 2 N, τ (#u k i ) = #u k+0.5 N k+1 , k = l #u k+0.5 i , k = l + 0.5, l ∈ N, (27) where # denotes the order index of the node representing certain neurons in G. Thus,EQUATIONThus we conclude that M 1 ≤ N maxl N l + N l-1 ≲ N 2 L l N l = N 2 L |V | ≤ 2 L M 2 . 
 
 Appendix:  


Decision: Both reviewers mark this paper as a strong accept. I recommend adding the references suggested by Reviewer nEYx and I vote for acceptance.

Reviews: review: title: title: This submission aims to propose the Differential Operator with Forward-propagation (DOF) for solving partial differential equations and more specifically for calculating general second-order differential operators without losing any precision. The authors’ empirical results illustrate that their method surpasses traditional automatic differentiation (AutoDiff) techniques, achieving 2x improvement on the MLP structure and nearly 20x improvement on the MLP with Jacobian sparsity.
main_review: main_review: A) One major limitation of this work is that it only refers to physics inspired neural networks and ignores much that has been done in the last 5 years such as MGNO, Fourier neural operators or the multiwavelet neural operators. Here are some recent examples that solve similar problems:
- “Multipole graph neural operator for parametric partial differential equations” In Advances in Neural Information Processing Systems, volume 33, pages 6755–6766, 2020.
- "Non-linear operator approximations for initial value problems." In International Conference on Learning Representations (ICLR). 2022.
- "Coupled Multiwavelet Operator Learning for Coupled Differential Equations." In The Eleventh International Conference on Learning Representations. 2022.
- "Multiwavelet-based operator learning for differential equations." Advances in neural information processing systems 34 (2021): 24048-24062.
- "Identifying arguments of space-time fractional diffusion: data-driven approach." Frontiers in Applied Mathematics and Statistics (2020): 14.
- "Fourier neural operator with learned deformations for pdes on general geometries." arXiv preprint arXiv:2207.05209 (2022).
B) The authors stated in the article that the use of DOF will not cause a loss in precision. Can this be specifically shown in the experimental results section?
C) The authors compared the proposed strategy with the AutoDiff method based on the Hessian matrix. Can more methods be added to prove the improvement of DOF?
D) The authors need to improve the rigor of the paper’ writing. Many of the citations that appear in the article are in the wrong order. Also as highlighted above they ignored major developments in this field of many papers that collected hundreds of citations, so I see no reason for which they ignore good prior work.
title: title: Technically solid idea with provable efficiency claims
main_review: main_review: **Summary:** The paper proposes a new way to compute second-order differential operators using automatic differentiation. Authors extend on the recent results of the Forward Laplacian method by Li et al. to add arbitrary second-order differential operators with provably faster and lower memory calculations. The paper showcases results on MLP with/without Jacobian sparsity.

**Strengths:**
1. The paper reads well and is easy to follow. I enjoyed reading the paper.
2. The authors address a significant concern on the efficient calculation of second-order differential operators.
3. Convicing analysis of performance in domains of both time and memory.

**Areas of Improvement:**
1. The paper does not mention efficient higher-order derivative calculations like Taylor mode AD [bettencourt2019, tan2023]. The works relate greatly, and I would like to see performance comparisons.
2. How well does this technique generalize to compute arbitrary higher-order derivatives?
3. Open-source software: Do the authors plan to present their algorithm as open-source software?

[bettencourt2019] Bettencourt, Jesse, Matthew J. Johnson, and David Duvenaud. "Taylor-mode automatic differentiation for higher-order derivatives in JAX." Program Transformations for ML Workshop at NeurIPS 2019. 2019.

[tan2023] Tan, Songchen. Higher-Order Automatic Differentiation and Its Applications. Diss. Massachusetts Institute of Technology, 2023.