Title: Stochastic Taylor Derivative Estimator: Efficient amortization for arbitrary differential operators
Venue: NeurIPS 2024
Accepted: yes
Keywords: AI for Science

Automatic Differentiation

Deep Learning

Randomization
Abstract: Optimizing neural networks with loss that contain high-dimensional and high-order differential operators
  is expensive to evaluate with back-propagation due to $\mathcal{O}(d^{k})$ scaling of the derivative tensor size and the $\mathcal{O}(2^{k-1}L)$ scaling in the computation graph, where $d$ is the dimension of the domain, $L$ is the number of ops in the forward computation graph, and $k$ is the derivative order. In previous works, the polynomial scaling in $d$ was addressed by amortizing the computation over the optimization process via randomization. Separately, the exponential scaling in $k$ for univariate functions ($d=1$) was addressed with high-order auto-differentiation (AD). In this work, we show how to efficiently perform arbitrary contraction of the derivative tensor of arbitrary order for multivariate functions, by properly constructing the input tangents to univariate high-order AD, which can be used to efficiently randomize any differential operator.
  When applied to Physics-Informed Neural Networks (PINNs), our method provides >1000$\times$ speed-up and >30$\times$ memory reduction over randomization with first-order AD, and we can now solve 1-million-dimensional PDEs in 8 minutes on a single NVIDIA A100 GPU. This work opens the possibility of using high-order differential operators in large-scale problems.

Stochastic Taylor Derivative Estimator: Efficient amortization for arbitrary differential operators

Abstract

Optimizing neural networks with loss that contain high-dimensional and high-order differential operators is expensive to evaluate with back-propagation due to O ( d k ) scaling of the derivative tensor size and the O (2 k -1 L ) scaling in the computation graph, where d is the dimension of the domain, L is the number of ops in the forward computation graph, and k is the derivative order. In previous works, the polynomial scaling in d was addressed by amortizing the computation over the optimization process via randomization. Separately, the exponential scaling in k for univariate functions ( d = 1 ) was addressed with high-order auto-differentiation (AD). In this work, we show how to efficiently perform arbitrary contraction of the derivative tensor of arbitrary order for multivariate functions, by properly constructing the input tangents to univariate high-order AD, which can be used to efficiently randomize any differential operator. When applied to Physics-Informed Neural Networks (PINNs), our method provides >1000 × speed-up and >30 × memory reduction over randomization with first-order AD, and we can now solve 1-million-dimensional PDEs in 8 minutes on a single NVIDIA A100 GPU 1 . This work opens the possibility of using high-order differential operators in large-scale problems.

1 Introduction

In many problems, especially in Physics-informed machine learning [19, 33], one needs to solve optimization problems where the loss contains differential operators:

In this above, D α = ∂ | α | ∂x α 1 1 ,...,∂x α d d , α = ( α 1 , α 2 , . . . , α d ) is a multi-index, u θ is some neural network parameterized by θ , and f is some cost function. When either the differentiation order k or the dimensionality d is high, the objective function above is expensive to evaluate with back-propagation (backward mode AD) in both memory and computation: the size of the derivative tensor has scaling O ( d k ) , and the size of the computation graph has scaling O ( 2 k -1 L ) where L is the number of ops in the forward computation graph.

There have been several efforts to tackle this curse of dimensionality. One line of work uses randomization to amortize the cost of computing differential operators with AD over the optimization process so that the d in the above scaling becomes a constant for the case of k = 2 . Stochastic Dimension Gradient Descent (SDGD) [13] randomizes over the input dimensions where in each iteration, the partial derivatives are only calculated for a minibatch of sampled dimensions with back-propagation. In [12, 21, 15], the classical technique of Hutchinson Trace Estimator (HTE) [16] is used to estimate the trace of Hessian or Jacobian to inputs. Others choose to bypass AD completely to reduce the complexity of computation. In [31], the finite difference method is used for estimating the Hessian trace. Randomized smoothing [11, 14] uses the expectation over Gaussian random variable as ansatz, so that its derivatives can be expressed as another expectation Gaussian random variable via Stein's identity [39]. However, compared to AD, the accuracy of these methods is highly dependent on the choice of discretization.

In this work, we address the scaling issue in both d and k for the optimization problem in Eq. 1 at the same time, by proposing an amortization scheme that can be efficiently evaluated via high-order AD, which we call Stochastic Taylor Derivative Estimator (STDE) . Our main contributions are:

We demonstrate how Taylor mode AD [6], a high-order AD method, can be used to amortize the optimization problem in Eq. 1. Specifically, we show that, with properly constructed input tangents, the univariate Taylor mode can be used to contract multivariate functions' derivative tensor of arbitrary order;

We provide a comprehensive procedure for randomizing arbitrary differential operators with STDE, while previous works mainly focus on the Laplacian operator, and we provide abundant examples of STDE constructed for operators in common PDEs;

STDE encompass and generalizes previous methods like SDGD [13] and HTE [16, 12]. We also prove that HTE-type estimator cannot be generalized beyond fourth order differential operator;

We determine the efficacy of STDE experimentally. When applied to PINN, our method provides a significant speed-up compared to the baseline method SDGD [13] and the backward-free method like random smoothing [11]. Due to STDE's low memory requirements and reduced computation complexity, PINNs with STDE can solve 1-million-dimensional PDEs on a single NVIDIA A100 40GB GPU within 8 minutes , which shows that PINNs have the potential to solve complex real-world problems that can be modeled as high-dimensional PDEs. We also provide a detailed ablation study on the source of performance gain of our method.

2 Related works

High-order and forward mode AD The idea of generalizing forward mode AD to high-order derivatives has existed in the AD community for a long time [5, 18, 40, 22]. However, accessible implementation for machine learning was not available until the recent implementation in JAX [6, 7], which implemented the Taylor mode AD for accelerating ODE solver. There are also efforts in creating the forward rule for a specific operator like the Laplacian [24, 23]. Randomization over the linearized part of the AD computation graph was considered in [30]. Forward mode AD can also be used to compute neural network parameter gradient as shown in [2].

Randomized Gradient Estimation Randomization [28, 29, 8] is a common technique for tackling the curse of dimensionality for numerical linear algebra computation, which can be applied naturally in amortized optimization [1]. Hutchinson trace estimator [16] is a well-known technique, which has been applied to diffusion model [37] and PINNs [12]. Another case that requires gradient estimation is when the analytical form of the target function is not available (black box), which means AD cannot be applied. The method of zeroth-order optimization [25] can be used in this case, as it only requires evaluating the function at arbitrary input. It is also useful when the function is very complicated like in the case of a large language model [27].

3 Preliminaries and discussions

3.1 First-order auto-differentiation (AD)

ADis a technique for evaluating the gradient of composition of known analytical functions commonly called primitives. In an AD framework, a neural network F θ : R d → R d ′ is constructed as the composition of primitives F i that are parameterized by some parameters θ i . In this section, we will

consider the neural networks with linear computation graphs like F = F L ◦ F L -1 ◦ · · · ◦ F 1 , but the results generalize to arbitrary directed acyclic graphs (DAGs). We will assume that all hidden dimensions are h . See Appendix B for more details on first-order AD.

Forward mode AD Each primitives F i is linearized as the Fréchet (directional) derivative ∂F i : R h → L( R h , R h ) , which computes the Jacobian-vector-product (JVP): ∂F i ( a )( v ) = ∂F ∂ x ∣ ∣ a v , where a is referred to as the primal and v the tangent. ∂F i form a linearized computation graph (third row in Fig. 3), that computes the JVP of the composition ∂F ∂ x v :

By setting the tangent to v one of the standard basis of R d , JVP computes one column of the Jacobian D F , so the full Jacobian can be computed with d JVPs. Each JVP call requires O (max( d, h )) memory as only the current activation y i and tangent v i are needed to carry out the computation, and the computation complexity is usually in the same order as the forward computation graph. In the case of MLP, both the forward and the linearized graph have a complexity of O ( dh +( L -1) h 2 ) .

Backward mode AD Each primitives F i is linearized as the adjoint of the Fréchet derivative ∂ ⊤ F i instead, which computes the vector-Jacobian-product (VJP): ∂ ⊤ F i ( a )( v ⊤ ) = v ⊤ ∂F ∂ x ∣ ∣ a where v ⊤ is the cotangent. The linearized computation graph now runs in the reverse order:

which is also clear from Fig. 3. Furthermore, due to this reversion, we first need to do a forward pass to obtain the evaluation trace { y i } L i =1 before we can invoke the VJPs ∂ ⊤ F i , which apparent as shown in Eq. 3. Hence the number of sequential computations is twice as much compared to forward mode. The memory requirement becomes O ( d +( L -1) h ) as we need to store the entire evaluation trace. Similar to JVP, VJP computes one row of J F at a time, so the full Jacobian ∂F ∂ x can be computed using d ′ VJPs. When optimizing scalar cost functions ℓ ( θ ) : R n → R of the network parameters θ , backward mode efficiently trades off memory with computation complexity as d ′ = 1 and only 1 VJP is needed to get the full Jacobian. Furthermore, all parameter θ i can use the same cotangent v ⊤ , whereas with forward mode, separate tangent for each parameter θ i is needed.

3.2 Inefficiency of the first-order AD for high-order derivative on inputs

second-order derivative by applying VJP twice

VJP

Forward

x

∈

R

d

y

1

∈

R

h

y

2

y

3

θ

1

θ

2

θ

3

v

⊤

1

∈

R

h

v

⊤

2

v

⊤

3

v

⊤

∈

R

d

′

v

⊤

∂F

∂

x

y

=

F

(

x

)

θ

4

˜

v

⊤

∈

R

d

˜

v

⊤

1

˜

v

⊤

2

˜

v

⊤

3

˜

v

⊤

4

˜

v

⊤

5

˜

v

⊤

6

∂

2

F

d

∂x

i

∂x

j

v

d

˜

v

j

High-order input derivatives ∂ k u θ ∂ x k for scalar u θ can be implemented as repeated applications of first-order AD, but this approach will exhibit fundamental inefficiency that cannot be remedied by randomization.

Repeating backward mode AD With each repeated application of backward mode AD, the new evaluation trace will include the cotangents from the previous application of backward AD, so the length of sequential computation doubles . Furthermore, the size of the cotangent also grows by d times. Therefore applying backward mode AD has additional memory cost of O ( d +( L -1) h ) and

additional computation cost of O ( 2 dh +2( L -1) h 2 ) , which is clear from Fig. 1. In general, with k repeated applications of backward mode AD will incur O ( 2 k -1 ( d +( L -1) h ) ) memory cost and O ( 2 k ( dh +( L -1) h 2 ) ) computation cost. And O ( d k -1 ) calls are needed to evaluate the entire derivative tensor. So both memory and compute scale exponentially in derivative order k

Repeating forward mode AD Consider u θ : R d → R . The input tangent dimension is d on the first application of forward mode AD, but on the second application, it will become d × d since we are now computing the forward mode AD for ∇ u θ : R d → R d . So the size of the input tangent with k repeated application is O ( d k ) , so it grows exponentially . This is also inefficient.

Mixed mode AD schemes are also likely inefficient See more detail in Appendix C.

3.3 Stochastic Dimension Gradient Descent

SDGD [13] amortizes high-dimensional differential operators by computing only a minibatch of derivatives in each iteration. It replaces a differential operator D with a randomly sampled subset of additive terms, where each term only depends on a few input dimensions

where ˜ D J denotes the SDGD operator that approximates the true operator D , J is the sampled index set, and | J | is the batch size. For example, in d -dimensional Poisson equation, N D = d , D = ∑ d j =1 ∂ 2 ∂x 2 j , and the additive terms are D j = ∂ 2 ∂x 2 j .

˜ D J are cheaper to compute than D due to reduced dimensionality: for each sampled index, by treating all other input as constant we get a function with scalar input and output. For a given index set J , the memory requirements are reduced from O ( 2 k -1 ( d +( L -1) h ) ) to O ( | J | (2 k -1 (1 + ( L -1) h )) ) , and the computation complexity reduces to O ( | J | 2 k ( h +( L -1) h 2 ) ) . This reduction is significant when d ≫ h as in the experimental setting of SDGD [13], but the exponential scaling in k persists.

3.4 Univariate Taylor mode AD

One way to define high-order AD is by determining how the high-order Taylor expansion of a univariate function changes when mapped by primitives. Firstly, the Fréchet derivative ∂F can be rewritten to operate on a space curve g : R → R d that passes through the primal a , i.e. g ( t ) = a , and has tangent g ′ ( t ) = v :

This shows that the ∂ (JVP) is the same as the univariate chain rule. The tuple J g ( t ) := ( g ( t ) , g ′ ( t )) can be thought of as the first-order expansion of g which lives in the tangent bundle of F . Treating F as the smooth map between manifolds, we can define the pushforward d F which pushes the first order expansion of g (i.e. J g ( t ) ) forward to the first order expansion of F ◦ g (i.e. J F ◦ g ( t ) ):

Naturally, to extend this to higher orders, one can consider the k th order expansion of the input curve g , which is equivalent to the tuple J k g ( t ) := ( g ( t ) , g ′ ( t ) , g ′′ ( t ) , . . . , g ( k ) ( t )) = ( a , v (1) , v (2) , . . . , v k ) known as the k -jet of g where v j is called the j th order tangent of g . J k g lives in the k th order tangent bundle of F , and we can define the k th-order pushforward d k F :

which pushes the k th order expansion of g (i.e. J k g ) forward to the k th order expansion of F ◦ g (i.e. J k F ◦ g ). ∂ k F = ∂ k ∂t k [ F ◦ g ]( t ) is the k -th order Fréchet derivative, whose analytical formula is given by the high-order univariate chain rule known as the Faa di Bruno's formula (Eq. 43).

Since J k g contains all information needed to evaluate ∂ j ∂t j [ F ◦ g ]( t ) for any j ≤ k , the map d k F is well-defined. d k defines a high-order AD: we can compute d k F of arbitrary composition F from the k th-order pushforward of the primitives d k F i , since d k is an homomorphism of the group ( { F i } , ◦ ) :

This approach of composing d k of primitives is also known as the Taylor mode AD. For more details on Taylor mode AD, see Appendix D.

4 Method

From the previous discussion, it is clear that the exponential scaling in k for the problem described in Eq. 1 cannot be mitigated by amortization alone. Although high-order AD methods like Taylor mode AD [6] can address this scaling issue, it is only defined for univariate functions. In this section, we describe a method that addresses the scaling issue in k and d simultaneously when amortizing Eq. 1 by seeing univariate Taylor mode AD as contractions of multivariate derivative tensor.

4.1 Univariate Taylor mode AD as contractions of multivariate derivative tensor

d F projects the Jacobian of F to R d ′ with a 1-jet J g ( t ) . Similarly, d k F contracts a set of derivative tensors to R d ′ with a k -jet J k g . We can expand ∂ k ∂t k F ◦ g with Eq. 43 to see the form of the contractions. For example, ∂F is JVP, and ∂ 2 F contains a quadratic form of the Hessian D 2 F :

From Eq. 43, one can always find a J l g with large enough l ≥ k such that there exists k ≤ l ′ ≤ l with ∂ l ′ F ( J l ′ g ) = D k F ( a ) · ⊗ k i =1 v ( v i ) where v i ∈ [1 , k ] , by setting some tangents v ( v i ) to the zero vector. That is, arbitrary derivative tensor contraction is contained within a Fréchet derivative of high-order, which can be efficiently evaluated through Taylor mode AD.

How large l should be depends on how off-diagonal the operator is. If the operator is diagonal (i.e. contains no mixed partial derivatives), l = k is enough. If the operator is maximally non-diagonal, i.e. it is a partial derivative where all dimensions to be differentiated are distinct, then the minimum l needed is (1 + k ) k/ 2 . For more details, please refer to Appendix F where a general procedure for determining the jet structure is discussed.

x

y

1

y

2

y

3

y

=

F

(

x

)

d

2

F

1

d

2

F

2

d

2

F

3

d

2

F

4

v

(1)

v

(1)

1

v

(1)

2

v

(1)

3

∂F

=

∂F

∂

x

v

(1)

v

(2)

v

(2)

1

v

(2)

2

v

(2)

3

∂

2

F

=

∂F

∂

x

v

(2)

+

∂

2

F

∂x

i

∂x

j

v

(1)

i

v

(1)

j

L

= 4

k

= 2

4.2 Estimating arbitrary differential operator by pushing forward random jets

Next, we show how to use the above facts to construct a stochastic estimator derivative operator. Differential operators can be evaluated through derivative tensor contraction. The action of the derivative D α = ∂ | α | ∂x α 1 1 ,...,∂x α d d on function u can be identified with the derivative tensor slice D | α | u ( a ) α . Differential operator L can be written as a linear combination of derivatives: L = ∑ α ∈I ( L ) C α D α , where I ( L ) is the set of tensor indices representing terms included in the operator L . For simplicity we only consider k th order differential operator, i.e. | α | = k ∈ N for all α . For scalar u : R d → R ,

we can identify a k th order differential operator L with the following tensor dot product

where d i ∈ [1 , d ] , i ∈ [1 , k ] is the tensor index on the i th axis, , and C ( L ) is a tensor of the same shape as D k u ( a ) that equals C α when d 1 , . . . , d k matches the multi-index α ∈ I ( L ) and 0 otherwise. We call C ( L ) the coefficient tensor of L . For example, the coefficient tensor of the Laplacian ∇ 2 is the d -dimensional identity matrix I . More complicated operators can be built as f ( x , u, D k 1 u, . . . , D k n u ) where f is arbitrary function.

Any derivative tensor contractions D k u ( a ) · C ( L ) can be estimated through random contraction, which can be implemented efficiently as pushing forward random jets from an appropriate distribution . With random ( v (1) , . . . , v ( k ) ) , we have

where ⊗ denotes Kronecker product, v ( v i ) d i ∈ [1 , k ] is the d i dimension of the v i th order tangent in the input k -jet. Eq. 11 is an unbiased estimator of the k th order operator L u = D k u ( a ) · C ( L ) when

For example, the condition for unbiasedness for the Laplacian ∇ 2 is E [ v ( a ) v ( b ) ⊤ ] = I . As discussed, one can always find a J l g with large enough l ≥ k such that ∂ l F ( J l g ) = D k F ( a ) · ⊗ k i =1 v ( v i ) , so with a distribution p over the input l -jet J l g that satisfies the unbiasedness condition (Eq. 12), we have

which means L u ( a ) can be approximated by the sample mean of the pushforwards of random l -jet drawn from p , which can be computed efficiently via Taylor mode AD. We call this method Stochastic Taylor Derivative Estimator (STDE) . The advantages of STDE are:

General: STDE can be applied to differential operators of arbitrary order and dimensionality.

Scalable: The scaling issue in the dimensionality d and the derivative order k are addressed at the same time. From the example computation graph (Fig. 2) we see that, for one call to d k F , the memory requirement has scaling of O ( kd ) and the computation complexity has scaling O ( k 2 dL ) . Like first-order forward mode AD, the derivative tensor D k u is never fully computed and stored. Combined with randomization, the polynomial scaling in d will be removed.

Parallelizable: The number of sequential computations does not grow with the order as can be seen in Fig. 2, and the computation can be trivially vectorized and parallelized since the pushforward of sample jets can be computed independently, and it uses the same computation graph ( d k u );

4.3 Constructing STDE for high-order differential operators with sparse random jets

Note that all coefficient tensor has the following additive decomposition:

where e i is the i th standard basis. For example, if the input dimension d is 3 , then e 2 = [0 , 1 , 0] ⊤ . As discussed before, there exists a J k g whose pushforward under ∂ l u is equivalent to contracting D k u with ⊗ k i =1 e d i . We call k -jet consisting of only standard basis and the zero vector 0 sparse . Therefore the discrete distribution p over the sparse k -jets in Eq. 14 satisfies the unbiasedness condition 12

where Z is the normalization factor and we identify ⊗ k i =1 e d i with the corresponding k -jet J k u .

4.3.1 Differential operator with easy to remove mixed partial derivatives

Next, we show some concrete examples for constructing STDE with sparse random jets.

Laplacian From Eq. 9 we know that the quadratic form of Hessian can be computed through ∂ 2 by setting v (2) = 0 and v (1) = e j . Therefore, the STDE of the Laplacian operator is given by

where J is the sampled index set, and the subscript [2] means taking the second-order tangent from the output jet. See example implementation in JAX in Appendix A.4.

High-order diagonal differential operators We call a differential operator diagonal if it is a linear combination of diagonal elements from the derivative tensor: L = ∑ d j =1 ∂ k ∂x k j . From Eq. 43 we see that setting the first-order tangent v (1) to e j and all other tangents v ( i ) to the zero vector gives the desired high-order diagonal element:

Second-order parabolic PDEs Second-order parabolic PDEs are a large class of PDEs. It includes the Fokker-Planck equation in statistical mechanics to describe the evolution of the state variables in stochastic differential equations (SDEs), which can be used for generative modeling [38]. It also includes the Black-Scholes equation in mathematical finance for option pricing, the Hamilton-JacobiBellman equation in optimal control, and the Schrödinger equation in quantum physics and chemistry. Its form is given by

We have a second order derivative term 1 2 tr ( σ ( x , t ) σ ( x , t ) ⊤ ∂ 2 ∂ x 2 u ( x , t ) ) with off-diagonal term. The off-diagonals can be easily removed via a change of variable:

See derivation in Appendix E. Its STDE samples over the d terms in the expression above.

4.3.2 Differential operators with arbitrary mixed partial derivative

It is not always possible to remove the mixed partial derivatives but discussed in section 4.2, for an arbitrary k th order derivative tensor element D k u ( a ) n 1 ,...,n k , we can find an appropriate l -jet J l g ( t ) with g ( t ) = a such that ∂ l u ( J l g ) = D k u ( a ) n 1 ,...,n k . Here we show a concrete example.

2D Korteweg-de Vries (KdV) equation Consider the following 2D KdV equation

All the derivative terms can be found in the pushforward of the following jet:

where the subscript [ i ] means selecting the i th order tangent from the jet, and the prefactors are determined through Faa di Bruno's formula (Eq. 43). In this case, no randomization is needed since all the terms can be computed with just one pushforward. Alternatively, these terms can be computed with pushforwards of different jets of lower order (Appendix I.4). When input dimension d is high, randomization via STDE will provide significant speed up. We tested a few more high-order PDEs with irremovable mixed partial derivatives (see Appendix I.4), and the experimental results will be provided later.

4.4 Dense random jet and connection to HTE

In section 4.3 we show how to construct STDE with the pushforward of sparse random jets. It is also possible to construct STDE with dense random jets, i.e. jets with tangents that are not the standard basis. For example, the classical method of Hutchinson trace estimator (HTE) [16] can be implemented in the STDE framework as the pushforward of isotropic dense random jets, i.e. ( a , v , 0 ) ∼ δ a × p × δ with E p [ vv ⊤ ] = I .

We generalize the dense construction to arbitrary second-order differential operators using a multivariate Gaussian distribution with the eigenvalues of the corresponding coefficient tensor as its covariance. Suppose D is a second-order differential operator with coefficient tensor C . With the eigendecomposition C ′′ = 1 2 ( C + C ⊤ ) + λ I = UΣU ⊤ where -λ is smaller than the smallest eigenvalue of C , we can construct a STDE for D :

However, it is not always possible to construct dense STDE beyond the second order, even if we consider p with non-diagonal covariance. We prove this by providing a counterexample: one cannot construct an STDE for the fourth order operator ∑ d i =1 ∂ 4 ∂x 4 with dense jets. For more details on dense jets, see Appendix K. For specific high-order operators like the Biharmonic operator, it is still possible to construct STDE with dense jets which we show in Appendix J.

The main differences between the sparse and the dense version of STDE are:

sparse STDE is universally application whereas the dense STDE can only be applied to certain operators;

the source of variance is different (see Appendix K.3).

It is also worth noting that both the sparse and the dense versions of STDE would have similar computation costs if the batch size of random jets were the same. In general, we would suggest to use sparse STDE unless it is known a priori that the sparse version would suffer from excess variance and the dense STDE is applicable.

5 Experiments

We applied STDE to amortize the training of PINNs on a set of real-world PDEs. For the case of k = 2 and large d , we tested two types of PDEs: inseparable and effectively high-dimensional PDEs (Appendix I.1) and semilinear parabolic PDEs (Appendix I.2). We also tested high-order PDEs (Appendix I.4) that cover the case of k = 3 , 4 , which includes PDEs describing 1D and 2D nonlinear dynamics, and high-dimensional PDE with gradient regularization [42]. Furthermore, we tested a weight-sharing technique (Appendix G), which further reduces memory requirements (Appendix I.3). In all our experiments, STDE drastically reduces computation and memory costs in training PINNs, compared to the baseline method of SDGD with stacked backward-mode AD. Due to the page limit, the most important results are reported here, and the full details including the experiment setup and hyperparameters (Appendix H) can be found in the Appendix.

5.1 Physics-informed neural networks

PINN [33] is a class of neural PDE solver where the ansatz u θ ( x ) is parameterized by a neural network with parameter θ . It is a prototypical case of the optimization problem in Eq. 1. We consider PDEs defined on a domain Ω ⊂ R d and boundary/initial ∂ Ω as follows

where L and B are known operators, f ( x ) and g ( x ) are known functions for the residual and boundary/initial conditions, and u : R d → R is a scalar-valued function, which is the unknown solution to the PDE. The approximated solution u θ ( x ) ≈ u ( x ) is obtained by minimizing the mean squared error (MSE) of the PDE residual R ( x ; θ ) = L u θ ( x ) -f ( x ) :

where the residual points { x ( i ) } N r i =1 are sampled from the domain Ω . We use the technique from [26] that reparameterizes u θ such that the boundary/initial condition B u ( x ) = g ( x ) are satisfied exactly for all x ∈ ∂ Ω , so boundary loss is not needed.

Amortized PINNs PINN training can be amortized by replacing the differential part of the operator L with a stochastic estimator like SDGD and STDE. For example, for the Allen-Cahn equation, L u = ∇ 2 u + u -u 3 , the differential part of L is the Laplacian ∇ 2 . With amortization, we minimize the following loss

which is a modification of Eq. 24. Its gradient ∂ ˜ ℓ residual ∂θ is then an unbiased estimator to the gradient of the original PINN residual loss, i.e. E [ ∂ ˜ ℓ residual ∂θ ] = ∂ℓ residual ∂θ .

5.2 Ablation study on the performance gain

To ascertain the source performance gain of our method, we conduct a detailed ablation study on the inseparable Allen-Cahn equation with a two-body exact solution described in Appendix I.1. The results are in Table 1 and 2, where the best results for each dimensionality are marked in bold. All methods were implemented using JAX unless stated. OOM indicates that the memory requirement exceeds 40GBs. Since the only change is how the derivatives are computed, the relative L2 error is expected to be of the same order among different randomization methods, as seen in Table 3 in the Appendix. We have included Forward Laplacian which is an exact method. It is expected to perform better in terms of L2 error. However, as we can see in Table 3, the L2 error is of the same order, at least in the case where the dimension is more than 1000 .

JAX vs PyTorch The original SDGD with stacked backward mode AD was implemented in PyTorch. We reimplement it in JAX (see Appendix A.1). From Table 1 and 2, JAX provides ∼ 15 × speed-up and up to ∼ 4 × memory reduction.

Parallelization The original SDGD implementation uses a for-loop to iterate through the sampled dimension (Appendix A.1). This can be parallelized (denoted as 'Parallelized SDGD via HVP', details in Appendix A.2). Parallelization provides ∼ 15 × speed up and reduction in peak memory for the JIT compilation phase. We also tested mixed mode AD (dubbed as 'Forward-over-Backward SDGD'), which gives roughly the same performance as parallelized stacked backward mode, which is expected as explained in Appendix C.

Forward Laplacian Forward Laplacian [24] provides a constant-level optimization for the calculation of Laplacian operator by removing the redundancy in the AD pipeline, and we can see from

Table 1 and 2 that it is the best method in both speed and memory when the dimension is 100. But since it is not a randomized method, the scaling is much worse. Its computation complexity is O ( d ) , whereas a randomized estimator like STDE has a computation complexity of O ( | J | ) . Naturally, with a high enough input dimension d , the difference in the constant prefactor is trumped by scaling. When the dimension is larger than 1000, it becomes worse than even parallelized stacked backward mode SDGD.

STDE Compared to the best realization of baseline method SDGD, the parallelized stacked backward mode AD, STDE provides up to 10 × speed up and memory reduction of at least 4 × .

6 Conclusion

We introduce STDE, a general method for constructing stochastic estimators for arbitrary differential operators that can be evaluated efficiently via Taylor mode AD. We evaluated STDE on PINNs, an instance of the optimization problem where the loss contains differential operators. Amortization with STDE outperforms the baseline methods, and STDE also applies to a wider class of problems as it can be applied to arbitrary differential operators.

Applicability Besides PINNs, STDE can be applied to arbitrarily high-order and high-dimensional AD-based PDE solvers. This makes STDE more general than a branch of related methods. STDE is also more applicable than deep ritz method [41], weak adversarial network (WAN) [43], backward SDE-based solvers [3, 34, 10], deep Galerkin method [35], and the recently proposed forward Laplacian [24], which are all restricted to specific forms of second-order PDEs. STDE applies naturally to differential operators in PDEs, but it can also be applied to other problems that require input gradients. For example, adversarial attacks, feature attribution, and meta-learning, to name a few.

Limitations Being a general method, STDE forgoes the optimization possibilities that apply to specific operators. Furthermore, we did not consider variance reduction techniques that could be applied, which can be explored in future works. Also, we observed that lowering the randomization batch size improves both speed and memory profile, but the trade-off between cheaper computation and larger variance needs further analysis. Furthermore, the method is not suited for computing the high order derivative of neural network parameter as explained in Section 3.

Future works The key insight of the STDE construction is that the univariate Taylor mode AD contains arbitrary contraction of the derivative tensor and that derivative operators are derivative tensor contractions. This shows the connection between the fields of AD and randomized numerical linear algebra and indicates that further works in the intersection of these two fields might bring significant progress in large-scale scientific modeling with neural networks. One example would be the many-body Schrödinger equations, where one needs to compute a high-dimensional Laplacian. Another example is the high-dimensional Black-Scholes equation, which has numerous uses in mathematical finance.

References

Brandon Amos. Tutorial on amortized optimization, April 2023. arXiv:2202.00665 [cs, math].

Atılım Güne¸ s Baydin, Barak A. Pearlmutter, Don Syme, Frank Wood, and Philip Torr. Gradients without Backpropagation, February 2022. arXiv:2202.08587 [cs, stat].

Christian Beck, Sebastian Becker, Patrick Cheridito, Arnulf Jentzen, and Ariel Neufeld. Deep splitting method for parabolic PDEs. SIAM Journal on Scientific Computing , 43(5):A3135A3154, January 2021. arXiv:1907.03452 [cs, math, stat].

Sebastian Becker, Ramon Braunwarth, Martin Hutzenthaler, Arnulf Jentzen, and Philippe von Wurstemberger. Numerical simulations for full history recursive multilevel Picard approximations for systems of high-dimensional partial differential equations. Communications in Computational Physics , 28(5):2109-2138, June 2020. arXiv:2005.10206 [cs, math].

Claus Bendtsen and Ole Stauning. Tadiff , a flexible c + + package for automatic differentiation using taylor series expansion. 1997.

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

James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018.

Benyamin Ghojogh, Ali Ghodsi, Fakhri Karray, and Mark Crowley. Johnson-Lindenstrauss Lemma, Linear and Nonlinear Random Projections, Random Fourier Features, and Random Kitchen Sinks: Tutorial and Survey, August 2021. arXiv:2108.04172 [cs, math, stat].

Andreas Griewank and Andrea Walther. Evaluating Derivatives . Society for Industrial and Applied Mathematics, second edition, 2008.

Jiequn Han, Arnulf Jentzen, and Weinan E. Solving high-dimensional partial differential equations using deep learning. Proceedings of the National Academy of Sciences , 115(34):85058510, Aug 2018.

Di He, Shanda Li, Wenlei Shi, Xiaotian Gao, Jia Zhang, Jiang Bian, Liwei Wang, and Tie-Yan Liu. Learning Physics-Informed Neural Networks without Stacked Back-propagation, February 2023. arXiv:2202.09340 [cs].

Zheyuan Hu, Zekun Shi, George Em Karniadakis, and Kenji Kawaguchi. Hutchinson Trace Estimation for High-Dimensional and High-Order Physics-Informed Neural Networks. Computer Methods in Applied Mechanics and Engineering , 424:116883, May 2024. arXiv:2312.14499 [cs, math, stat].

Zheyuan Hu, Khemraj Shukla, George Em Karniadakis, and Kenji Kawaguchi. Tackling the curse of dimensionality with physics-informed neural networks. Neural Networks , 176:106369, 2024.

Zheyuan Hu, Zhouhao Yang, Yezhen Wang, George Em Karniadakis, and Kenji Kawaguchi. Bias-Variance Trade-off in Physics-Informed Neural Networks with Randomized Smoothing for High-Dimensional PDEs, November 2023. arXiv:2311.15283 [cs, math, stat].

Zheyuan Hu, Zhongqiang Zhang, George Em Karniadakis, and Kenji Kawaguchi. Score-Based Physics-Informed Neural Networks for High-Dimensional Fokker-Planck Equations, February 2024. arXiv:2402.07465 [cs, math, stat].

M.F. Hutchinson. A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Communications in Statistics - Simulation and Computation , 18(3):10591076, January 1989.

Martin Hutzenthaler, Arnulf Jentzen, Thomas Kruse, Tuan Anh Nguyen, and Philippe von Wurstemberger. Overcoming the curse of dimensionality in the numerical approximation of semilinear parabolic partial differential equations, July 2018.

Jerzy Karczmarczuk. Functional differentiation of computer programs. In Proceedings of the Third ACM SIGPLAN International Conference on Functional Programming , ICFP '98, pages 195-203, New York, NY, USA, 1998. Association for Computing Machinery.

George Em Karniadakis, Ioannis G. Kevrekidis, Lu Lu, Paris Perdikaris, Sifan Wang, and Liu Yang. Physics-informed machine learning. Nature Reviews Physics , 3(6):422-440, Jun 2021.

Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun, editors, 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings , 2015.

Chieh-Hsin Lai, Yuhta Takida, Naoki Murata, Toshimitsu Uesaka, Yuki Mitsufuji, and Stefano Ermon. Regularizing score-based models with score fokker-planck equations. In NeurIPS 2022 Workshop on Score-Based Methods , 2022.

Jacob Laurel, Rem Yang, Shubham Ugare, Robert Nagel, Gagandeep Singh, and Sasa Misailovic. Ageneral construction for abstract interpretation of higher-order automatic differentiation. Proc. ACM Program. Lang. , 6(OOPSLA2), oct 2022.

Ruichen Li, Chuwei Wang, Haotian Ye, Di He, and Liwei Wang. DOF: Accelerating high-order differential operators with forward propagation. In ICLR 2024 Workshop on AI4DifferentialEquations In Science , 2024.

Ruichen Li, Haotian Ye, Du Jiang, Xuelan Wen, Chuwei Wang, Zhe Li, Xiang Li, Di He, Ji Chen, Weiluo Ren, and Liwei Wang. Forward Laplacian: A New Computational Framework for Neural Network-based Variational Monte Carlo, July 2023. arXiv:2307.08214 [physics].

Sijia Liu, Pin-Yu Chen, Bhavya Kailkhura, Gaoyuan Zhang, Alfred Hero, and Pramod K. Varshney. A Primer on Zeroth-Order Optimization in Signal Processing and Machine Learning, June 2020. arXiv:2006.06224 [cs, eess, stat].

Lu Lu, Raphaël Pestourie, Wenjie Yao, Zhicheng Wang, Francesc Verdugo, and Steven G. Johnson. Physics-informed neural networks with hard constraints for inverse design. SIAM Journal on Scientific Computing , 43(6):B1105-B1132, 2021.

Sadhika Malladi, Tianyu Gao, Eshaan Nichani, Alex Damian, Jason D. Lee, Danqi Chen, and Sanjeev Arora. Fine-Tuning Language Models with Just Forward Passes, January 2024. arXiv:2305.17333 [cs].

Per-Gunnar Martinsson and Joel Tropp. Randomized Numerical Linear Algebra: Foundations &Algorithms, March 2021. arXiv:2002.01387 [cs, math].

Riley Murray, James Demmel, Michael W. Mahoney, N. Benjamin Erichson, Maksim Melnichenko, Osman Asif Malik, Laura Grigori, Piotr Luszczek, Michał Derezi´ nski, Miles E. Lopes, Tianyu Liang, Hengrui Luo, and Jack Dongarra. Randomized Numerical Linear Algebra : A Perspective on the Field With an Eye to Software, April 2023. arXiv:2302.11474 [cs, math].

Deniz Oktay, Nick McGreivy, Joshua Aduol, Alex Beatson, and Ryan P. Adams. Randomized Automatic Differentiation, March 2021. arXiv:2007.10412 [cs, stat].

Tianyu Pang, Kun Xu, Chongxuan Li, Yang Song, Stefano Ermon, and Jun Zhu. Efficient Learning of Generative Models via Finite-Difference Score Matching, November 2020. arXiv:2007.03317 [cs, stat].

Juncai Pu and Yong Chen. Lax pairs informed neural networks solving integrable systems, January 2024. arXiv:2401.04982 [nlin].

M. Raissi, P. Perdikaris, and G.E. Karniadakis. Physics-informed neural networks: a deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational Physics , 378:686-707, February 2019.

Maziar Raissi. Forward-Backward Stochastic Neural Networks: Deep Learning of Highdimensional Partial Differential Equations, April 2018. arXiv:1804.07010 [cs, math, stat].

Justin Sirignano and Konstantinos Spiliopoulos. Dgm: a deep learning algorithm for solving partial differential equations. Journal of computational physics , 375:1339-1364, 2018.

Maciej Skorski. Modern analysis of hutchinson's trace estimator. In 2021 55th Annual Conference on Information Sciences and Systems (CISS) . IEEE, March 2021.

Yang Song, Sahaj Garg, Jiaxin Shi, and Stefano Ermon. Sliced Score Matching: A Scalable Approach to Density and Score Estimation, June 2019. arXiv:1905.07088 [cs, stat].

Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-Based Generative Modeling through Stochastic Differential Equations, February 2021. arXiv:2011.13456 [cs, stat].

Charles M. Stein. Estimation of the Mean of a Multivariate Normal Distribution. The Annals of Statistics , 9(6):1135 - 1151, 1981.

MuWang. High Order Reverse Mode of Automatic Differentiation . PhD thesis, 2017. Copyright - Database copyright ProQuest LLC; ProQuest does not claim copyright in the individual underlying works; Last updated - 2023-03-04.

E Weinan and Ting Yu. The deep ritz method: a deep learning-based numerical algorithm for solving variational problems. Communications in Mathematics and Statistics , 6:1 - 12, 2017.

Jeremy Yu, Lu Lu, Xuhui Meng, and George Em Karniadakis. Gradient-enhanced physicsinformed neural networks for forward and inverse PDE problems. Computer Methods in Applied Mechanics and Engineering , 393:114823, April 2022. arXiv:2111.02801 [physics].

Yaohua Zang, Gang Bao, Xiaojing Ye, and Haomin Zhou. Weak Adversarial Networks for Highdimensional Partial Differential Equations. Journal of Computational Physics , 411:109409, June 2020. arXiv:1907.08272 [cs, math].

A Example implementations

A.1 PyTorch implementation of SDGD-PINN using backward mode AD

The original implementation of SDGD-PINN [13] computes the SDGD estimation of derivatives using a for-loop that iterates over the sampled PDE term/dimension. For example, given a function f representing the MLP PINN, the computation of SDGD for the Laplacian operator can be implemented in PyTorch as follows:

After computing the PDE differential operator, it is plugged into the residual loss, and then backwardmode AD is employed to produce the gradient for optimization concerning θ .

A.2 JAX implementation of SDGD Parallelization via HVP

A.3 JAX implementation of Forward-over-backward AD

The forward-over-backward AD In JAX mentioned in Appendix C can be implemented as follows:

A.4 JAX implementation of STDE for the Laplacian operator

The jet.jet function from JAX implements the high-order pushforward d n of jets in Eq. 7. It decomposes the input function into primitives, which have analytical derivatives derived up to arbitrary order, and uses the generalized chain rule (see section D.2) to compose the primitives into the pushforward of jets. Note that in the API of jet.jet , all the high-order tangents of the input jet are specified via the series argument.

B Further details on first-order auto-differentiation

B.1 Computation graph of first-order AD

x

y

1

y

2

y

3

y

=

F

(

x

)

F

1

F

2

F

3

F

4

v

v

1

v

2

v

3

∂F

∂

x

v

∂F

1

∂F

1

∂F

2

∂F

2

∂F

3

∂F

3

∂F

4

∂F

4

θ

1

θ

2

θ

3

θ

4

F

1

F

2

F

3

F

4

x

y

1

y

2

y

3

y

=

F

(

x

)

F

1

F

2

F

3

F

4

v

⊤

∂F

∂

x

v

⊤

3

v

⊤

2

v

⊤

1

v

⊤

∂

⊤

F

4

∂

⊤

F

4

∂

⊤

F

3

∂

⊤

F

3

∂

⊤

F

2

∂

⊤

F

2

∂

⊤

F

1

∂

⊤

F

1

θ

1

θ

2

θ

3

θ

4

F

1

F

2

F

3

F

4

B.2 Derivative via composition

First-order AD is based on a simple observation: for a set of functions L , the set of tuples of functions f and its Jacobian J f is closed under composition :

where ◦ denotes both function composition and the composition of the tuple ( f, J f ) . If we have the analytical formula of the Jacobian J f for every f ∈ L , then we can calculate the Jacobian of any composition of functions from L using the above composition rule for the tuple ( f, J f ) . The set L of functions are usually called the primitives .

B.3 Fréchet derivative and linearization

Given normed vector spaces V, W , the Fréchet derivative ∂f of a function f : V → W is a map from V to the space of all bounded linear operators from V to W , denoted as L( V, W ) , that is

such that at a point a ∈ V it gives the best linear approximation ∂f ( a )( · ) of f , in the sense that

Therefore, it is also called the linearization of f at point a . Concretely, consider a function in Euclidean spaces f : R n → R m . At any point a ∈ R n , the Fréchet derivative ∂f can be seen as the directional derivative of f :

where J f ( a ) ∈ R m × n denote the Jacobian of f at point a called the primal , and v ∈ R n , also called the tangent is a vector representing the direction. Therefore the Fréchet derivative is also called Jacobian-vector-product (JVP) . And we can write the truncated Taylor expansion as

Many operators have efficient JVP implementation due to sparsity. For example, element-wise application of scalar function (e.g. activation in neural networks) has diagonal Jacobian, and its JVP can be efficiently implemented as a Hadamard product. Another prominent example is discrete convolution, whose JVP has efficient implementation via FFT.

B.4 Adjoint of the Fréchet derivative

Given two topological vector spaces X,Y , the linear map u : X → Y has an adjoint t u : Y ′ → X ′ where X ′ , Y ′ are the dual spaces. The adjoint satisfies the following

In the finite-dimensional case, the dual space is the space of row vectors, and any linear map can be written as u ( x ) = A x . One can easily verify that the adjoint is the transpose: t u ( y ⊤ ) = y ⊤ A . The adjoint (transpose) of the Fréchet derivative of f : R n → R m , denoted as ∂ ⊤ f , is thus defined as

where v ⊤ is the cotangent which lives in the dual space of the codomain. Note that the adjoint is taken to v only where a is kept fixed. ∂ ⊤ f is also called vector-Jacobian-product (VJP) .

C Why mixed mode AD schemes like the forward-over-backward might not be better than stacked backward mode AD in the case of PINN

In AD literature [9], the second order derivative is recommended to be computed via forward-overbackward AD, i.e., first do a backward mode AD to get the first order derivative, then apply forward mode AD to the first order derivative to obtain the second order derivative. Usually, we will expect that forward-over-backward AD gives better performance in memory usage over stacked backward AD since the outer differential operator has to differentiate a larger computation graph than the inner one, and forward AD has less overhead as explained in section B.2. Essentially, forward-overbackward reverses the arrows in the third row in Fig. 1, therefore reducing the number of sequential computations and also the size of the evaluation trace. However, in the case of PINN, yet another differentiation to the network parameters θ needs to be taken. So, computing the second-order differential operator here with forward-over-backward AD might not yield any advantage.

D Taylor mode AD

D.1 High-order Fréchet Derivatives

The k th order Fréchet derivative of a function f : R n → R m at a point a is the multi-linear map with k arguments around point a that best approximates f . For example, when k = 2 , we have

where H f ( a ) ∈ R m × n × n denote the Hessian of f at point a , and v , v ′ ∈ R n . We can now write the second-order truncated Taylor series with it

For the more general case, we have

High-order Fréchet derivative can be seen as the best k th order polynomial approximation of f by taking all input tangents to be the same v ∈ R n :

D.2 Composition rule for high-order Fréchet derivatives

Next, we derive the higher-order composition rule by repeatedly applying the usual chain rule.

For composition f ( g ( x )) of scalar functions, we can generalize the chain rule for high-order derivatives by iteratively applying the chain rule to lower-order chain rules:

where we give the example of up to the third order. For arbitrary k , the k th order derivative of the composition is given by the Faa di Bruno's formula (scalar version)

where the outermost summation is taken over all partitions of the derivative order k . Here a partition of k is defined as a tuple ( p 1 , . . . , p k ) ∈ N k that satisfies

For vector-valued functions g : R n → R m , f : R m → R l , let

we can derive the following composition rule similarly

where again we give the example of up to the third order, and repeated indexes are summed as in Einstein notation. The general formula is again given by the multivariate version of the Faa di Bruno's formula. Note that in the multivariate version of the Faa di Bruno's formula, it is possible to take a derivative to distinguishable variables, but here we just present the version with indistinguishable input variables. This gives the composition rule for k th order total derivative.

The composition of the high-order Fréchet derivative ∂ k is the case of n = 1 , as the contraction with the input tangents v ( i ) ∈ R d is the same as composing with a scalar input function g : R → R d with v ( i ) = D i g . All derivative tensors of f ( g ( x )) can be represented using a R l vector, and similarly all derivative tensor v ( i ) of g can be represented using a R m vector. Then, the above chain rule can be simplified to

The Faa di Bruno's formula again gives the general formula for arbitrary derivative order

which is written in the perspective of input primal a and tangents v ( i ) .

E Removing the mixed partial derivatives term from second order semilinear parabolic PDE

F Evaluating arbitrary mixed partial derivatives

F.1 A concrete example

Let's first consider a concrete case. Suppose the domain is D -dimensional we want to compute the mixed derivative ∂ ∂x 2 i ∂x j . The naive approach would be to compute the entire third order derivative tensor D 3 f , which is a tensor of shape D × D × D , then extract the element at index ( j, i, i ) . However note that from Eq. 43, for any k > 3 , the pushforward of k -jet under d k f contains contractions of D 3 f . Although in the case of k = 3 , the only contraction of D 3 f is in the ∂ 3 f :

which can only be used to compute the diagonal or the block diagonal elements, when k > 3 , we will have a contraction that computes off-diagonal terms, i.e. the mixed partial derivatives. For example, in d 4 f , if all input tangents are set to zero except for v (1) and v (2) , ∂ 4 f becomes:

which contains the contraction of D 3 f that we want:

However, there are extra terms. We can remove them by doing two extract pushforwards. We can compute the desired mixed partial derivative with the following pushforward of standard basis:

If we go to higher-order jets, we can use more flexible contractions, and we can compute the mixed derivative with fewer terms to correct, hence less pushforwards. For example, the pushforward of the fifth-order tangent is

if all input tangents are set to zero except for v (1) and v (3) . With this we only need to remove one term:

Similarly, by going to the seventh-order tangent, we can compute this mixed derivative with only one pushforward. d 7 f contains ∂ 7 , and when all input tangents are set to zero except for v (2) and v (3) , ∂ 7 equals

which is the exact contraction we want. With this we have

F.2 Procedure for finding the right pushforwards for arbitrary mixed partial derivatives

More generally, consider the case where we need to compute arbitrary mixed partial derivative

where T is the number of different input dimensions in the mixed partial derivative, and q i t is the order. To compute it with k -jet pushforward, one needs to find:

a derivative order k ∈ N ,

a sparsity pattern for the tangents v ( i ) of the input jet, which is defined as the tuple of T integers J = ( j 1 , . . . , j T ) where v ( j ) = 0 when j ̸∈ J and j t < k for all t ∈ [1 , T ] ,

such that when setting

( p 1 , p 2 , . . . , p k ) ∈ N k is a partition of k as defined in Eq. 39.

Let's use the concrete example ∂ ∂x 2 i ∂x j again. In this case T = 2 , q i 1 = 2 and q i 2 = 1 . We demonstrated that this can be computed with one 7 -jet pushforward, which is equivalent to setting J = (2 , 3) , k = 2 j 1 + j 2 = 7 , and the partition (0 , 2 , 1 , 0 , 0 , 0 , 0) . The Faa di Bruno's formula (Eq. 43) ensures that the pushforward of the k th order tangent contains a contraction that can be used to compute the desired mixed partial derivative.

Furthermore, if there are no other partitions with a sparsity pattern that is the subset of the sparsity pattern of the partition in consideration, there are no extra terms to remove. Intuitively, if a partition has a sparsity pattern that is not a subset, it will vanish when we set the input tangents to zero according to the sparsity pattern of the partition in consideration. To understand this point better, let's look at the concrete example with the 5 -jet pushforward demonstrated above. (2 , 0 , 1 , 0 , 0) and (5 , 0 , 0 , 0 , 0) are both valid partition of k = 5 , and the sparsity pattern of (5 , 0 , 0 , 0 , 0) is the subset of that of (2 , 0 , 1 , 0 , 0) : p 1 are non-zero in both partition. Therefore the pushforward contains extra terms that can be removed with another pushforward. In the example with 7 -jet pushforward, no other partition has the sparsity pattern that is the subset of that of the partition (0 , 2 , 1 , 0 , 0 , 0 , 0) . This is equivalent to say, 2 + 2 + 3 is the only way to sum up to 7 when you can only use 2 and 3 , which can be verified easily.

With this setup, it is clear why the diagonal terms can always be computed with pushforward of the lowest possible order: ( k, 0 , . . . , 0) ∈ N k is always a valid partition k , and no other partition has sparsity pattern that is a subset of it.

For mixed partial derivatives, the difficulty scales the total order of the operator ∑ T t =1 q i t , and T which can be interpreted as the degree of the 'off-diagonalness' of the operator. For example, consider the case where T = 3 and q i 1 = 3 , q i 1 = 2 , q i 1 = 1 . This corresponds to the operator ∂ ∂x 3 i ∂x 2 j ∂x k . To avoid overlapping with the diagonal sparsity pattern ( k, 0 , . . . , 0) and to keep the order of derivative low, one might try k = 16 and the partition (0 , 3 , 2 , 1 , 0 , . . . ) ∈ N 16 . However, with higher k , there is more chance that other partitions will have a subset sparsity pattern. In this case (0 , 8 , 0 , 0 , 0 , . . . ) ∈ N 16 is one such example. One will need to either find all the partitions with subset sparsity pattern and remove them with multiple pushforward, or further increase the derivative order to find a pattern with no extra term.

G Further memory reduction via weight sharing in the first layer

When dealing with high-dimensional data, the parameters of the model's first layer in a conventional fully connected network would grow proportionally with the input dimension, resulting in a significant increase in memory requirements and forming a memory bottleneck due to massive model parameters. To address this issue, convolutional networks are often employed in deep learning for images to reduce the number of model parameters. Here, we adopt a similar approach to mitigate the memory cost of model parameters in high-dimensional PDEs, called weight sharing in the first layer.

Denote the input dimension as d , which is potentially excessively high, and the hidden dimension of the MLP as h , and assume that d ≫ h . The first layer weight is an d × h dimensional matrix, whereas all subsequent layers have a weight matrix with a size of only h × h .

By introducing a weight-sharing scheme, one can reduce the redundancy in the parameters in the first layer. Specifically, we perform an additional 1D convolution to the input vectors x i before passing the input into the MLP PINN, as in Fig. 4. The 1D convolution has filter size B that divides D and stride size B , so the convolution output is non-overlapping, and the number of channels is set to 1 .

This weight-sharing scheme reduces the parameters by approximately 1 B . The number of parameters in the filters is B × 1 , and the subsequent fully connected layer will have a weight matrix of size d B × H . Therefore, the total number of the first layer is reduced from d × h to only d B × h + B , and we can see that with a larger block size B , we will have fewer parameters, and the reduction factor is approximately 1 B . More concretely, suppose d = 10 6 , h = 100 where one million ( 10 6 ) dimensional problems are also tested experimentally, so the number of parameters in the first layer is d × h = 100 × 10 6 . If we use a block size of B = 100 , we will reduce the number of parameters to d B × h + B = 10 6 + 100 . If the block size is B = 10 , the number of parameters will be d B × h + B = 10 × 10 6 +10 . In other words, with a larger block size of B , we significantly reduce the number of model parameters.

We will demonstrate the memory efficiency and acceleration thanks to weight-sharing in the experimental section.

x

1

x

2

x

3

x

4

x

5

x

6

x

7

x

8

x

9

θ

1

θ

2

θ

3

y

1

y

2

y

3

Convolution with stride of 3

output

H Experiment setup

Each experiment is run with five different random seeds, and the average and the standard deviations of these runs are reported.

To get an accurate reading of memory usage, we use a separate run where GPU memory pre-allocation for JAX is disabled through setting the environment variable XLA_PYTHON_CLIENT_ALLOCATOR=platform , and the test data set is stored on the CPU memory. The GPU memory usage was obtained via NVIDIA-smi and peak memory was reported.

All the experiments were done on a single NVIDIA A100 GPU with 40GB memory and CUDA 12.2. with driver 535.129.03 and JAX version 0.4.23.

Network architecture and training hyperparameters For the semilinear parabolic PDEs tested in Appendix I.2 we follow the network architecture of the original SDGD [13]:

The network is a 4-layer multi-layer perceptron (MLP) with 128 hidden units activated by Tanh.

The network is trained with Adam [20] for 10K steps, with an initial learning rate of 1e-3 that linearly decays to 0 in 10K steps, where at each step we calculate the model parameters gradient with 100 uniformly sampled random residual points.

The model is evaluated using 20K uniformly sampled random points fixed throughout the training.

The zero boundary condition is satisfied via the following parameterization

where u MLP θ is the MLP network, and u θ is the PDE ansatz, as described in [26].

For the semilinear parabolic PDEs tested in Appendix I.2, we made the following modifications:

Instead of using re-parameterization, the boundary/initial condition is satisfied by adding a regularization loss to the residual loss:

where g ( · ) is the initial data, N b is the batch size for boundary points, u θ is the PDE ansatz, C g is the coefficient for the first-order derivative boundary loss term, which we set to 0 . 05 . The total loss is

Instead of discretizing the time and sample residual points using the underlying stochastic process, we uniformly sample the time steps between the initial and the terminal time, i.e. t ∼ uniform [0 , T ] , and then sample x directly from the distribution of X t , i.e. x ∼ N (0 , ( T -t ) · I d × d ) . To match the original training setting of 100 SDE trajectories with 0 . 015 step size for time discretization, we use a batch size of 2000 for residual points and 100 for boundary/initial points.

We use a 4-layer multi-layer perceptron (MLP) with 1024 hidden units activated by Tanh. The network is trained with Adam [20] for 10K steps, with an initial learning rate of 1e-3 that exponentially decays with exponent 0 . 9995 .

To test the quality of the PINN solution, we measure the relative L1 error at the point ( x test , T ) against the reference value computed via multilevel Picard's method [3, 4, 17].

In all experiments, we use the biased version of Eq. 25:

as the bias in practice is very small and does not affect convergence.

I Experiments Results

I.1 Inseparable and effectively high-dimensional PDEs

The first class of PDEs is defined via a nonlinear, inseparable, and effectively high-dimensional exact solution u exact ( x ) defined within the d -dimensional unit ball B d :

where L is a linear/nonlinear operator and g ( x ) = L u exact ( x ) . The zero boundary condition ensures that no information about the exact solution is leaked through the boundary condition. We will consider the following operators:

Poisson equation: L u ( x ) = ∇ 2 u ( x ) .

Allen-Cahn equation: L u ( x ) = ∇ 2 u ( x ) + u ( x ) -u ( x ) 3 .

Sine-Gordon equation: L u ( x ) = ∇ 2 u ( x ) + sin( u ( x )) .

For the exact solution, we consider the following with all c i ∼ N (0 , 1) :

two-body interaction: u exact ( x ) = (1 -∥ x ∥ 2 2 ) ( ∑ d -1 i =1 c i sin( x i +cos( x i +1 ) + x i +1 cos( x i )) ) .

three-body interaction: u exact ( x ) = (1 -∥ x ∥ 2 2 ) ( ∑ d -2 i =1 c i exp( x i x i +1 x i +2 ) ) .

We tested the performance of STDE on these equations, and the results are presented in Table 3, 4, 5, 6. For the Allen-Cahn equation, we performed a detailed ablation study (Table 3), and we expect these results to generalize over these second-order PDEs.

I.1.1 Further details on ablation study

The gain by using JAX instead of PyTorch Since the original SDGD was implemented in PyTorch, we implemented the stacked backward mode without parallelization in SDGD dimensions in JAX for fair comparison (dubbed as 'Stacked Backward mode SDGD in JAX' in Table 3). The for-loop over SDGD dimension is implemented using jax.lax.scan . Table 3 shows that, even with the original stacked backward mode AD, the speed of JAX implementation can be more than 10 × faster when the dimension is high. The memory profile is similar. The difference could come from the fact that JAX uses XLA to perform Just-in-time (JIT) compilation of the Python code into optimized kernels. However, note that for the case of 100,000 dimensions, the JAX implementation of the stacked backward mode AD encountered an out-of-memory (OOM) error. This is because performing JIT compilation requires extra memory, and the peak memory requirement during JIT compilation is higher than that during training.

Randomization batch size We also tested the case where the STDE randomization batch size is reduced to 16 . As seen in Table 3, in the case of Allen-Cahn provides ∼ 2 × speed up, without hurting performance. However, theoretically lowering the randomization batch size leads to higher variance. The trade-off between computational efficiency and stability in convergence warrants further studies.

I.2 Semilinear Parabolic PDEs

The second class of PDEs is the semilinear parabolic PDEs, where the initial condition is specified:

where g ( x ) is a known, analytical, and time-independent function that specifies the initial condition, and T is the terminal time. We aim to approximate the solution's true value at one test point x test ∈ R d , at the terminal time t = T , i.e. at ( x test , T ) .

We will consider the following operators

Semilinear Heat Eq.

with initial condition g ( x ) = 5 / (10 + 2 ∥ x ∥ 2 ) ,

Allen-Cahn equation

with initial condition g ( x ) = arctan(max i x i ) ,

Sine-Gordon equation

with initial condition g ( x ) = 5 / (10 + 2 ∥ x ∥ 2 ) ,

All three equation uses the test point x test = 0 and terminal time T = 0 . 3 .

I.3 Weight sharing

We tested the weight-sharing technique mentioned in Section G.

In this section, we evaluate the performance of the weight-sharing scheme described in Appendix G. We tested the best-performing method from Table 3 (STDE with small randomization batch size of 16 ) with different weight-sharing block sizes, on the inseparable Allen-Cahn equation with the two-body exact solution.

From Table 10, we can see that weight sharing drastically reduces the number of network parameters and memory usage. With B = 50 , there is a 2 . 5 x reduction in memory and there is no performance loss in terms of L2 relative error.

However, from the experiments we can see that, in both the 1M and the 5M case, increasing the block size beyond 50 provides diminishing returns. For the 1M case, increasing B to 1000 affects the convergence quality, as the L2 relative error goes up by 100 x. For 5M, the maximum block size one can use before degrading performance is 500 , which is expected as the dimensionality of the problem is higher.

From Table 10 we can also see that in the 5 M-dimensional case, we will have an out-of-memory (OOM) error without weight sharing. With weight sharing enabled, we can effectively solve the 5 M-dimensional PDE with good relative L2 error, in around 30 minutes.

I.4 High-order PDEs

Here we demonstrate how to use STDE to calculate mixed partial derivatives in some actual PDE. We will consider the 2D Korteweg-de Vries (KdV) equation and the 2D Kadomtsev-Petviashvili equation from [32], and the regular 1D KdV equation with gPINN [42].

We will demonstrate that STDE increases the speed for computing the mixed partial derivatives, as it avoids computing the entire derivative tensor. Since these equations are low-dimensional we do not need to sample over the space dimension.

In this section, the equations are all time-dependent and the space is 2D, and we will omit the argument to the solution, i.e. we will write u ( x , t ) = u . To test the speed improvement, we run the STDE implementation against repeated backward mode AD on a Nvidia A100 GPU with 40GB memory. The results are reported in Table 11. From the Table we see that STDE provides around ∼ 2 × speed up compared to repeated application of backward mode AD across different network sizes.

I.4.1 High-order low-dimensional PDEs

Alternative way to compute the terms in 2D Korteweg-de Vries (KdV) equation The terms in the 2D KdV equation

can alternatively be computed with the pushforward of the following jets

All the derivative terms can be found in these output jets { J ( i ) } :

2D Kadomtsev-Petviashvili (KP) equation Consider the following equation

which can be expanded as

All the derivative terms can be computed with a 5-jet, 4-jet, and a 2-jet pushforward. Let

Then all required derivative terms can be evaluated as follows.

Gradient-enhanced 1D Korteweg-de Vries (g-KdV) equation Consider the following equation

Gradient-enhanced PINN (gPINN) [42] regularizes the learned PINN such that the gradient of the residual is close to the zero vector. This increases the accuracy of the solution. Specifically, the PINN loss (Eq. 24) is augmented with the term

where c gPINN is the g-PINN penalty weight. To perform gradient-enhancement we need to compute the gradient of the residual:

All the derivative terms can be computed with one 2-jet and two 7-jet pushforward. Let

Then all required derivative terms can be evaluated as follows.

I.4.2 Amortized gradient-enhanced PINN for high-dimensional PDEs

It is expensive to apply gradient enhancement for high-dimensional PDEs. For example, the gradient of the residual for the inseparable Allen-Cahn equation described in I.1 is given by

The total loss becomes

With STDE randomization, we randomized the second order term ∂ 2 ∂x 2 i with index i sampled from [1 , d ] . We can also sample the gPINN penalty terms. As mentioned in Appendix F.1, we have

We further have

so the STDE of the Laplacian operator can be computed together with the above pushforward. With this pushforward, we can efficiently amortize the gPINN regularization loss by minimizing the following upperbound on the original gPINN loss with randomized Laplacian

where J is an independently sampled index set for sampling the gPINN terms. The total loss is

We call this technique amortized gPINN . The above formula applies to all PDEs where the derivative operator is the Laplacian. For example, for the Sine-Gordon equation, we have

We use c gPINN = 0 . 1 , and to get better convergence, we train for 20 K steps instead of 10 K steps as in all other experiments in this paper. The results are reported in Table 12. We implement the baseline method based on the best performing first-order AD scheme, the parallelized backward mode SDGD, which we denoted as JVP-HVP in the table. Specifically, to compute the residual gradient we apply one more JVP to the HVP-based implementation of Laplacian (Appendix A.2). From the table, we see that STDE-based amortized gPINN performs better than the JVP-HVP implementation, and both are more efficient than applying backward mode AD in a for-loop. Furthermore, through amortizing we can apply gPINN to high-dimensional PDE which was intractable.

J Pushing forward dense random jets

In this section we establish the connection between the classical technique of HTE [16] and STDE by demonstrating that HTE is a pushforward of dense isotropic random 2-jet.

J.1 Review of HTE

HTE provides a random estimation of the trace of a matrix A ∈ R d × d as follows:

where p ( v ) is isotropic , i.e. E v ∼ p ( v ) [ vv T ] = I . Therefore, the trace can be estimated by Monte Carlo:

where each v i ∈ R d are i.i.d. samples from p ( v ) .

There are several viable choices for the distribution p ( v ) in HTE, such as the most common standard normal distribution. Among isotropic distributions, the Rademacher distribution minimizes the variance of HTE. The proof for the minimal variance is given in [36].

J.2 HTE as the pushforward of dense isotropic random 2-jets

Note that both HTE and the STDE Hessian trace estimator (Eq. ) are computing the quadratic form of Hessian, a specific contraction that is included in the pushforward of 2-jet. In STDE, the random vectors are the unit vectors whose indexes are sampled from the index set without replacement. This can be seen as a discrete distribution p ( v ) such that v = √ d e i for i = 1 , 2 , · · · , d with probability 1 /d , which is isotropic. Hence HTE can also be defined as a push forward of random 2-jet that are isotropic.

We can now write the computation of HTE as follows

where ˜ ∇ 2 N is the STDE for Laplacian with random jet batch size N .

J.3 Estimating the Biharmonic operator

It was shown in [12] that the Biharmonic operator

has the following unbiased estimator:

where p is the d -dimensional normal distribution. Therefore its STDE estimator is

K STDE with dense jets

K.1 STDE with second order dense jets as generalization of HTE

Suppose D is a second-order differential operator with coefficient tensor C . If C is not symmetric, we can symmetrize it as C ′ = 1 2 ( C + C ⊤ ) , and D 2 u ( a ) · C = D 2 u ( a ) · C ′ since D 2 u ( a ) is symmetric.

Furthermore, we can make C positive-definite by adding a constant diagonal λ I where -λ is smaller than the smallest eigenvalue of C . The matrix C ′′ = 1 2 ( C + C ⊤ ) + λ I then has the eigen decomposition UΣU ⊤ where Σ is diagonal and all positive. Now we have

K.2 Why STDE with dense jets is not generalizable

Specifically, we will prove that it is impossible to construct dense STDE for the fourth-order diagonal operator L u = ∑ d i =1 ∂ 4 u ∂x 4 i .

The mask tensor of L is the rank-4 identity tensor I 4 ∈ R d × d × d × d , so the condition for unbiasedness is

where δ ijkl = 1 when i = j = k = l , and is 0 otherwise.

̸

̸

̸

In the most general case where a = b = c = d , we can sample v ∈ R 4 d and split it into four R d vectors. In this case we can define blocks of covariance as E v ∼ p [ v ( a ) v ( b ) ] = Σ ab , and Σ = [ Σ ab ] ab . Denote the fourth-moment tensor of p as µ ijkl , then Eq. 90 states that the block µ abcd in the fourth moment tensor should match C . Fourth moments can always be decomposed into second moments:

̸

So finding the p that satisfies Eq. 90 is equivalent to finding a zero-mean distribution p with covariance that satisfies the above equation. In the case of L , the mask tensor is block-diagonal: M ijkl = σ ij δ ij,kl . So in the case where a = b , set a = 1 , b = 2 , we have

and Σ = [ Σ 11 Σ 12 Σ 21 Σ 22 ] ∈ R 2 d × 2 d . Firstly, consider the diagonal entries of σ :

This can always be satisfied by setting the diagonal entries of both Σ aa and Σ aa block as follows:

Next, consider the entire σ matrix. We have

̸

In the case of L , we have σ ij = δ ij , so for i = j we have

which is impossible to satisfy since entries in a covariance matrix must be real.

K.3 Sparse vs dense jets

The variance of the sparse STDE estimator comes from the variance of selected derivative tensor elements, whereas the variance of the dense estimator comes from the derivative tensor elements that are not selected. For example, in the case of Laplacian, as also discussed in [12], the variance of the sparse STDE estimator comes from the diagonal element of the Hessian, whereas the variance of the dense STDE estimator comes from all the off-diagonal element of the Hessian.

L Further ablation study

Allen-Cahn

Poisson

Sine-Gordon

100

4

16

L2 relative error

10-1

100

64

L2 relative error

10-1

L2 relative error

10-1

256

10-2

10-2

10-2

10-3

10-3

10-3

0

2000 4000 6000 8000

0

2000 4000 60008000

0

2000 4000 60008

8000

Epoch

Epoch

Epoch

1016

1016

1016.

1014

1014

1014

1012.

1012

1012

Residual Loss

Residual Loss

Residual Loss

1010

1010.

1010.

108.

108

108

106

106.

106

104-

104

104-

102

102

102.

0

2500

5000

750010000

0

2500

5000

7500

10000

0

2500

5000

750010000

Epoch

Epoch

Epoch

350

350

300

300

300

Iterations per second (its/s)

Iterations per second (its/s)

250

250

250

200

200

150

150

150

100

100

100

50

50

50

0

0

0

α

256

BatchSize

80

140

70

120

80

Convergence time (s)

Convergence time (s)

Convergence time (s)

100

60

50

80

40

30

60

40

20

40

20

10

20

0

0

Q

256

×

00

α

25°

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: Our claims are backed by both theoretical and experimental evidence. The theoretical evidence is provided in the section 4, and the experimental evidence is provided in section 5.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We discussed the limitation of our work in section 6.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best

judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: Most of our theoretical results are asymptotic analyses on the computation complexity, and we have clearly stated the assumption we have made. Our claim on the non-generalizability of HTE construction is proved rigorously in the Appendix K.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We have included sample implementations of key steps of our method in the first section in the Appendix A.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We have open-sourced our code.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We have included the setting for all the hyperparameters in the Appendix H.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We use the average of 5 random seeds for all our experiment results. We also reported the standard deviation for the relative error in PINN training in the Appendix I.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We have included the hardware and software specifications we used to conduct our experiments in the Appendix H.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We have read the NeurIPS Code of Ethics and made sure that the paper conforms to it.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: Our work is not tied to particular applications, and there are no obvious paths that lead to potential harm.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: Our work is foundational and not tied to particular applications.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [NA]

Justification: Our work does not use existing assets.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: Our work does not release new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: Our paper did not involve crowdsourcing and human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: Our paper did not involve crowdsourcing and human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a new method, STDE, to compute high-order differential operators of (high-dimensional) neural network-represented functions. Two key ingredients include a generalization of Taylor-mode AD and randomness in the algorithm. This work shows impressive performance of STDE, in terms of both speed and memory.\n\nSTRENGTHS:\n-\tThis work addresses an important and challenging task in physics-informed machine learning.\n-\tThe illustration and analysis of previous AD algorithms and STDE are clear and comprehensive.\n-\tSTDE displays noticeable potential to improve the computation of arbitrary differential operators in both memory consumption and efficiency.\n-\tThe work provides new insight into introducing randomness other than SDGD or HTE whose basic ideas are employing Monte Carlo over dimensionality.\n\nWEAKNESSES:\n1.\tCompared with original back-propagation or forward-Laplacian(FL), STDE is a random algorithm that does not give **precise** derivatives. In precision-demanding tasks, e.g. solving many-electron Schrödinger equations, the inaccurate value of differential operators will possibly result in an unreliable solution. To make the result more convincing, a cost-accuracy curve (STDE with different number of random samplings, i.e. $|J|$ in Eqn.(17)) should be added.  \nFurthermore, in light of this, the comparison between STDE and FL in super high dimensions (in table 1,2) is unfair. A possible remedy is to compare STDE with modified FL, which employs randomness to be more scalable, for instance, Monte Carlo over dimensions. The author can refer to [1] for more technical and implementation details of FL.\n2.\tAll the test cases of high-dimensional high-order operators are Laplacians (including $\\Delta^2$). The high-order diagonal differential operator is mentioned in L216 but does not appear in the experiments. A comparison between STDE and direct back-propagation is expected.  \nThe applications are all coming from the small area of PINN. To demonstrate the broad impact of this work, the author should target more tasks of general interest, e.g. second order derivative of NN parameter, which plays an important role in accelerating optimizations. \n3.\tSome of the highly related existing works are missing. E.g.,\n-\tRegarding L63, the forward rule for general different operators is done in [1].\n-\tRegarding the geometric interpretation of Taylor-mode AD in sec 3.4, although not explicitly mentioned in Ref. [6] in the paper, is discussed in section 2 of [2]. The author should cite [6] or the first paper presenting similar ideas, and clarify their novelty if they still hope to claim generalizing univariate Taylor mode AD as their contribution.\n4.\tAs is discussed in Section 4.4, STDE can not be applied to *arbitrary* high-order operators. Obviously, this restricts the potential application scenarios of STDE.\n5.\tThe author didn’t discuss how large $l$ would be for large $k$ (notations from L165), which is necessary to assess the effectiveness of STDE for general high-order derivatives.\n6.\tSome parts of the writing are confusing. See **Questions**.\nAlso, there are several misuses of terminologies. E.g.,\n-\tFréchet derivative, not Frechet.\n-\tThe form of second order PDE in Eq.(18) is not ‘general’. For instance, you are assuming ellipticity in the second-order term. Please refer to Eq.(10) or Eq.(1) in [1] for *general* second-order operator.\n-\tIn Appendix I.1 and the main text, ‘inseparable and effectively high-dimensional’, I guess what the authors intend to say is ‘elliptic’. The linear part of all the test equations is merely a Laplacian, casting doubt on the effectiveness of STDE for general elliptic equations.\n-\tIn Appendix I.2 and the main text, these PDEs are not nonlinear Fokker-Planck Equations, I think the correct name would be ‘semilinear parabolic equation’.\n\n\n\n[1] Li R, Wang C, Ye H, He D, Wang L. DOF: Accelerating High-order Differential Operators with Forward Propagation. ICLR 2024 Workshop on AI4Differential Equations In Science  \n[2]Tan, Songchen. Higher-Order Automatic Differentiation and Its Applications. Diss. Massachusetts Institute of Technology, 2023.\n\nQUESTIONS:\n1.\tWhich function does ‘univariate’ refer to in Sec 3.4, $g$ or $F$? What is the first contribution (in the introduction), generalizing Taylor mode AD to the case where $g$ is multivariate, or $F$? I am skeptical of the authors’ claiming it as a contribution/novelty because it seems like this is merely about replacing scalar Faà di Bruno's formula with that of tensors.\n2.\tIn Appendix D, where $g$ defines an $n$-dimensional manifold, is $n$ always set to 1 in this work? If not, please point me to where the method with general $n$ is discussed.\n3.\tAlthough the authors claim to remove the exponential complexity w.r.t. the order of derivatives, for general operators the size of coefficient tensor **${C}$** in Eqn(10) grows exponentially and the sampling procedure in Eqn(15) will become hard.\n4.\tCould the authors think of any real application scenarios with high-order high-dim operators? As for the case of low-dim high-order operators, the authors need to check if the differential operator takes up a considerable proportion of the whole process, e.g. training PINN. If not, the benefit of accelerating differential operators might be very limited.\n5.\tDoes STDE have the potential to make use of sparsity in the neural network function, as shown in FL and [1]? For the specific case of second-order operators, how would STDE perform when $C$ is dense but low-rank?\n6.\tInstead of tackling a scalar function (e.g. Laplacian of a function), could STDE help with the computation of tensor-value functions, e.g. Hessian or curvature? Hessian appears in Monge-Ampère equation and Ricci curvature appears in Einstein equation.  \n\nOn writing and notations:  \n1. Can you clarify on ‘$[2]$’ in L214 with more details or concrete examples? For me, $d^2u_\\theta$ in Eqn.(16) is a (0,2)-tensor (Hessian), but why does it take three input $(a,e_j,0)$? Please also clarify on ‘$[3]$’ in Eqn(44). \n2.\tIn Eqn.(19), it should be $\\sum_{i=1}^d$.\n3.\tIn L247, do you mean $v\\sim p$,, instead of $(a,v,0)\\sim p$?\n4.\tIn the ‘100D’ column of table 2, 539 is not the smallest number. Why is it in bold?\n5.\tThere needs to be more description in section F.1.  \nIt is good to remind the readers that Einstein notation is employed in Eq.(45) and that the indices $n,\\ n’,\\ n’’$ are omitted.\nThe author should inform the readers that in the discussion $n=1$ for $n$ in Eq(40) and notations like $\\partial^4u(x)(e_i,e_j,0,0)$ refers to setting $v^{(1)}=e_i,\\ v^{(2)}=e_j$ instead of the taking $e_i,e_j,0,0$ as inputs for a (0-4) tensor.   \n\nOverall, this paper has the potential to be a good work. Hope the concerns above can be well addressed and then I will upgrade my assessment accordingly.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-22"}

comments: SUMMARY:
This paper proposes a stochastic optimization approach to find the minimizer of a cost function, which involves the complicated differential operators. The problem is computationally complex, hence, the authors propose to deal with a minibatch of derivatives in each iteration which reduces computational complexity. The method has application in various learning problems involving complicated differential operators such as physics-informed neural networks. Experiments are provided to evaluate the method.

STRENGTHS:
The paper seems to address a difficult optimization problem. The dimension reduction idea seems to be new and interesting. Application to PINN seems also novel.

WEAKNESSES:
NA

QUESTIONS:
NA

LIMITATIONS:
NA
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2024-06-30

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work proposes a scalable method for the optimization of loss functions including higher-order derivatives. The proposed method interprets arbitrary differential operators as derivative tensor contractions which are then estimated through random contractions. These random contractions can be computed efficiently using Taylor mode AD.\n\nSTRENGTHS:\n1. The proposed method is novel and interesting. I think the idea of random contractions of the derivative tensor is very intuitive and at least in principle should be scalable and a clear improvement over SDGD by ameliorating the exponential scaling in the order of the the differential operator. \n\n2. The idea is technically sound. There is sufficient empirical validation.\n\n2. The well-written presentation of the background material and JAX implementations of SDGD are also useful secondary contributions.\n\nWEAKNESSES:\n1. The presentation could be improved a tad bit. I think having 4 pages of background material is a bit unnecessary specially the general background on automatic differentiation could easily be moved to the appendix and space could be better utilized by explaining the method in more detail and discussing the experimental details. For example, a schematic diagram similar to Figure 2 illustrating STDE could be helpful in understanding the proposed method.\n\nQUESTIONS:\n1. In 291 it is stated \"Since the only change is how the derivatives are computed, the relative L2 error is expected to be of the same order among different methods\". Is this accurate? SDGD and STDE are stochastic approximations whereas Forward Laplace is not.\n2. I am a bit unclear about how to choose a distribution over the l-jets that satisfy the unbiasedness condition for any specific problem. It seems that the sparse random jets have the advantage of universally applicable but they also seem to involve a lot of redundant computations.\n\nLIMITATIONS:\nSome limitations are discussed in section 6.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper addresses the computational challenges of optimizing neural networks with loss functions that include high-dimensional and high-order differential operators. These challenges arise due to the scaling of the derivative tensor size with the dimension of the domain (d) and the computational graph's size with the number of operations (L) and the order of the derivative (k). Traditional methods either amortize the computational cost over the optimization process via randomization or use high-order auto-differentiation (AD) for univariate functions to tackle these issues.\n\nIn this work, the authors propose a method to efficiently perform arbitrary contraction of the derivative tensor for multivariate functions. This is achieved by constructing input tangents to univariate high-order AD, enabling efficient randomization of any differential operator. When applied to Physics-Informed Neural Networks (PINNs), this approach provides significant speed and memory efficiency improvements, achieving over 1000 times speed-up and 30 times memory reduction compared to randomization with first-order AD. The method allows solving 1-million-dimensional partial differential equations (PDEs) in just 8 minutes on a single NVIDIA A100 GPU, opening up the possibility of using high-order differential operators in large-scale problems.\n\nSTRENGTHS:\nThe paper introduces STDE, a general method for constructing stochastic estimators for arbitrary differential operators, which can be efficiently evaluated using Taylor mode auto-differentiation (AD). When evaluated on Physics-Informed Neural Networks (PINNs), a specific optimization problem where the loss function includes differential operators, STDE significantly outperforms baseline methods. Furthermore, STDE's applicability extends beyond PINNs to arbitrarily high-order and high-dimensional AD-based PDE solvers, making it more general than related methods.\n\nThe strengths of this paper are as follows:\n\nGenerality: STDE can be applied to a wide range of problems, including those involving arbitrarily high-order and high-dimensional differential operators. This broad applicability distinguishes STDE from other methods, which are often restricted to specific forms of second-order PDEs.\n\nEfficiency: The method enables efficient evaluation of stochastic estimators through Taylor mode AD, providing significant computational benefits.\n\nPerformance: In practical evaluations on PINNs, STDE outperforms baseline methods in terms of both speed and memory efficiency. It demonstrates over 1000 times speed-up and 30 times memory reduction compared to first-order AD randomization.\n\nScalability: STDE allows for the solution of extremely large-scale problems, such as 1-million-dimensional PDEs, in a matter of minutes on advanced hardware like the NVIDIA A100 GPU.\n\nVersatility: Beyond PINNs, STDE can be applied to various AD-based PDE solvers, making it a versatile tool for tackling a broad spectrum of differential operator-based optimization problems.\n\nOverall, STDE's generality, efficiency, performance, scalability, and versatility make it a powerful method for addressing high-dimensional and high-order differential operator challenges in neural network optimization.\n\nWEAKNESSES:\nWhile the paper demonstrates the significant strengths and broad applicability of STDE, it is important to acknowledge some limitations and areas for future improvement. As a general method, STDE may not leverage the specific optimization possibilities that are available for particular operators. Additionally, the study did not explore variance reduction techniques, which could potentially enhance the method's performance and could be a promising area for future research.\n\nAnother observation is that while reducing the randomization batch size improves both the speed and memory profile of STDE, this comes with a trade-off in the form of increased computational variance. Further analysis is required to understand and optimize this balance between computational efficiency and variance.\n\nLooking ahead, the paper identifies an intriguing connection between the fields of automatic differentiation (AD) and randomized numerical linear algebra, highlighting the potential for future work at this intersection. Such research could lead to significant advancements in large-scale scientific modeling with neural networks.\n\nIn summary, while there are areas for refinement, the contributions of this paper are substantial. The development of STDE as a general and efficient method for constructing stochastic estimators for arbitrary differential operators is a notable achievement, offering substantial benefits for high-dimensional and high-order differential operator problems in neural network optimization. The identified limitations and future research directions provide a clear path for further enhancing this already impressive work.\n\nQUESTIONS:\nGiven the importance of various other complex equations in scientific modeling, I am curious about the applicability of STDE to equations such as the Nonlinear Schrödinger Equation (NLS), the fourth-order NLS, and the Navier-Stokes equations. Could you elaborate on how STDE might perform or be adapted for these specific cases? Additionally, are there any preliminary results or theoretical considerations you could share regarding the application of STDE to these important equations?\n\nLIMITATIONS:\nYes, they have.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-24"}