Stochastic Taylor Derivative Estimator: Efficient amortization for arbitrary differential operators
Authors: Zekun Shi
Zheyuan Hu
Min Lin
Kenji Kawaguchi
Venue: NeurIPS 2024 oral (2024)
Keywords: AI for Science
Automatic Differentiation
Deep Learning
Randomization
Primary area: machine_learning_for_physical_sciences

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
Zekun Shi
Zheyuan Hu National University of Singapore
Min Lin Sea Al Lab
National University of Singapore
Sea Al Lab
e0792494 (11:44444uunus @u..nuseduu
linmin sea.com
shizk @ sea.com
Kenji Kawaguchi
National University Singapore
kenji @ nus.edu.sg
Abstract
Optimizing neural networks with loss that contain high-dimensional and high-order differential operators expensive to evaluate with back-propagation due (d*)) scaling of the derivative tensor size and the (22 scaling in the computation graph, where is the dimension of the domain, L is the number of ops in the forward computation graph, and & 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 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 efficiently randomize any differential operator. When applied to Physics-Infooe Informed Neural Networks (PINNs), our method provides > 1100x speed-up and >30x memory reduction over randomization with first-order AD. and we PDEs in minutes on a single NVIDIA A100 GPL This
solve
I-million dimensional L-miliion-dimensionad work opens the possibility of using high-order differential operators in large-scale problems.
1 Introduction
In many problems, especially in Physics-informee informed machine learning 119 33 optimization problems where the loss contains differential operators:
needs solve
arg min f(x, (x),,, Da (xx)......))) (a(")
un (x)),,
ug : Rd Rd
(1)
glal
a (a11, .., oa) is multi-index, up is some neural network parameterized by and is some 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 (dd)). and the size of the computation graph has scaling O (2k-11,, where L is the number of ops in the forward computation graph.
this above, D Dz 1
Our code is available mtlhtts::/ //gtt/ub/// www/saow.ccc.oooo..o...........a http///aallgg//ttde agde Received the Best Paper Award at NeurlPS 2024
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

There have been several efforts to tackle this curse of dimensionality. One line of work uses randomization 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 l = 2. Stochastic Dimension Gradient Descent (SDGD) [3] randomizes over the input dimensions where in each iteration, the partial derivatives are only calculated for a minibatch of sampled dimensions with pack-propagation. In [12] 21 [5]. the classical technique of Hutchinson Trace Estimator (HTE) [16 is used to estimate the trace of Hessian 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 [1144] 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 's identity [3]] 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 & for the optimization problem in Eq. a 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 66) 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 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 [6 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, method provides a significant speed-up compared to the baseline method SDGD [33 and the backward free method like random smoothing [11]. Due STDE's low memory requirements and reduced computation complexity, PINNS with STDE can solve --milion-dimensionn PDEs on a single NVIDIA A100 40GB GPU within 8 minutes, which shows that PINNs have the potential to solve complex world problems that can be modeled as high-dimensional PDEs. We also provide a detailed ablation study the source of performance gain of method.
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 [8
40 22]. However, accessible
implementation for machine learning was not available until the recent implementation in JAX (6] 77, 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 (4 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 22]
is a common technique for tackling the curse of dimensionality for numerical linear algebra computation, which can be applied naturally in amortized optimization I Hutchinson trace estimator [6] is a well-known technique, which has been applied to diffusion model 337 and PINNs [12. Another case that requires gradient estimation 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 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 22
Randomized Gradient Estimation Randomization 28
(25) be used in this case, as it only requires
3 Preliminaries and discussions
3.1 First-order auto-differentiation (AD)
AD a technique for evaluating the gradient of composition of known analytical functions commonly called primitives. In an AD framework, a neural network Fo Rd - Rdd is constructed as the composition of primitives F, that are parameterized by some parameters 0, In this section, we will

---

0x can be computed consider the neural networks with linear computation graphs like F = FL FL-1 o ... Fi. but the results generalize 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; is linearized as the Fréchet (directional) derivative OP: OF R + L(RR, RR), which computes the Jacobian- vector-produce (JVP): (a)(() = ax v, where is referred to the primal and v the tangent. 0F, form linearized computation graph (third row in Fig. that computes the JVP of the composition OF 8x
aF 0x = OF())v)) (OFL OF (0 (x))(x)(v).
(2)
By setting the tangent to v one of the standard basis of Rd, JVP computes one column of the Jacobian Dr, so the full Jacobian can be computed with d JVPs. Each JVP call requires O(max(d, h)) memory as only the current activation and tangent v, are needed to carry out the computation, and the computation complexity is usually in the same order 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²).
Backward mode AD Each primitives F; is linearized as the adjoint of the Fréchet derivative a F instead, which computes the pector-Jacobiat poroouuannppouuut (VJP): Ø (v) = of the cotangent. The linearized computation graph now runs in the reverse order: 8x la where v is
v OF 0x = a F(x)(v
= (i( ox..))...c.
=00
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:jez before can invoke the VJPs θ F;, 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 JF at a time, so the full Jacobian @F
using d' VJPs. When optimizing scalar cost functions {(0) R - R of the network parameters 0, backward mode efficiently trades off memory with computation complexity as d' 1 and only VJP is needed to get the full Jacobian. Furthermore, all parameter 6; can use the same cotangent v whereas with forward mode, separate tangent for each parameter 0; is needed.
3.2 Inefficiency of the first-order AD for high-order derivative on inputs
61
81
0₂
64
F(x)
Forward
x E Rd
Y1 E R
Y3
V
RA
V
8 Fi Br,dey Vadj
V
V
V E Rd
second-order derivative by applying VJP twice
Figure 1: The computation graph of computing second order gradient by repeated application of backward mode AD, for function F(-) with 4 primitives (L 4), which computes the Hessian- vector-product. Red nodes represent the cotangent nodes in the second backward pass. With each repeated application of VJP the length of sequential computation doubles.
@ us for scalar us can be implemented as repeated applications of High-order input derivatives ex" 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, 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 times. Therefore applying backward mode AD has additional memory cost of Old (L 1)h) and
3

---

1)A²), which is clear from Fig. In general, with ls repeated applications of backward mode AD will incur (
additional computation cost of (2dh + 2(L 1 2
1)k) memory cost and (2 (L 1)}²)) computation cost. And O 1d calls needed to evaluate the entire derivative tensor. So both memory and compute scale exponentially in derivative order l:
Repeating forward mode AD Consider ug E2 R. The input tangent dimension is d on the first application of forward mode AD, but on the second application, it will become d x d since we are now computing the forward mode AD for Vug R -> Rd. So the size of the input tangent with k repeated application is Oddh), so it grows exponentially. This is also inefficient.
Mixed mode AD schemes are also likely inefficient See more detail Appendix
C
3.3 Stochastic Dimension Gradient Descent
SDGD 13 amortizes high-dimensional differential operators by computing only 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
ND
ND J
D
Dj j=1
D
JEJ
where D 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, No = d,
o² and the additive terms are Dj D مj= 0z3.
8
Duy
DJ are cheaper to compute than D due to reduced dimensionality: for each sampled index, by treating
all other input as constant we get function with scalar input and output. For given index set J. the memory requirements are reduced from (2 (d + (L 1)h)) to (1J[(22-1() (1+ (L - 1))))), and the computation complexity reduces to [1 + (L 1)h²)). This reduction is signifi-
cant when d >>> h in the experimental setting of SDGD [3] but the exponential scaling in k persists.
3.4 Univariate Taylor mode AD
One way to define high-order AD by determining how the high-order Taylor expansion of univariate function changes when mapped by primitives. Firstly. the Fréchet derivative OF can be rewritten to operate on space curve g R - Iedd that passes through the primal a, i.e. g(l) = a, and has tangent g (t) V:
)F((())))) OF
d g'(t) - -t Irrog g(()).
-
(5)
x=g(()
This shows that the Ó (JVP) is the same as the univariate chain rule. The tuple J, = go
can be thought of as the forst-order expansion of which lives in the tangent bundle of F. Treating F as the smooth map between manifolds, we can define the pushforward dF which pushes the first order expansion of g (i.e. (!) forward to the first order expansion of F g (i.e JFog (t)):
d dt [F al(1)
(F) = (1) [F o (1),
(6)
Naturally, to extend this to higher orders, one can consider the kth order expansion of the input curve g. which is equivalent to the tuple (1) == (g((), g' (L),, (1),g"()), ( (1) = (a, v(), v()) ...,, v*) known as the k-jet of g where v is called the jth order tangent of 2h lives in the kth order tangent bundle of F. and we can define the kth-order pushforward d F:
à a F (0) = (t) = [FFg(()) 2 [F gl())
0²
a Du [F al(()
[ ol) 02
(7)
=(F(a), (F(a)(v())) 82
02 F(a)(v(L) (())
F(a(()( ".))...,'''(()((v
v(()))
which pushes the keth order expansion of g (i.c. 1hy forward to the kth order expansion of F o (i.c. Jb.o). F*F = ag [F g) (1) the k-th order Fréchet derivative, whose analytical formula given by the high-order univariate chain rule known the Faa di Bruno's formula (Eq. 43).

---

Since Jø contains all information needed to evaluate well-defined. d defines high-order AD; we can compute d F' of arbitrary composition F from the kth-order pushforward of the primitives d F; since dl is an homomorphism of the group ()
057 [ g](() for any S k, the map d F
1 (t) = J = dt F [d F2 o d FI)(
(8)
This approach of composing d² of primitives also known as the Taylor mode AD. For more details on Taylor mode AD, see Appendix D
Method
From the previous discussion, it is clear that the exponential scaling in k for the problem described in Eq. [ cannot be mitigated by amortization alone. Although high-order AD methods like Taylor mode AD 66] can address this scaling issue, it is only defined for univariate functions. In this section, we describe method that addresses the scaling issue in and simultaneously when amortizing Eq. by seeing univariate Taylor mode AD as contractions of multivariate derivative tensor.
4.1 Univariate Taylor mode AD as contractions of multivariate derivative tensor
dF projects the Jacobian of F' Rd with a 1-jet (t). Similarly, d contracts set of derivative tensors Rdd with k-jjt Jz. We can expand @* ar Fog with Eq. to see the form of the contractions.
For example, OF JVP, and F2F contains quadratic form of the Hessian Dp.
0² 0 FF(a)(v()) v(2) = Dt2
(1), F g = (a(a) + D2 de di d2
(9)
From Eq. 43 one can always find a J with large enough 1 2 1 such that there exists S < with 0" F(J)))
Df & ((()) where v, E kl. by setting some tangents v(t) 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 should be depends on how off- off-diagonal the operator is. If the operator is diagonal (i.e. contains no mixed partial derivatives). i.e. it is partial derivative where all dimensions to be differentiated are distinct, then the minimum / needed (1 + k/22 For more details, please refer to Appendix F where general procedure for determining the jet structure is discussed.
k is enough. the operator is maximally non- -indigannal,
d'P.
d'F2
d'Pss
d'F
V1
Y2
Ya
F(x)
vi
v2
via
oF (!) aF
(2)
(2
oF de
(2)
6 daiune
(1)
Figure The computation graph of d for F with primitives. Parameters 0 are omitted. The first column from the left represents the input 2-jet Jz (x, v(1), v(2),, and d F, pushes it forward which is the subsequent column. Each row can be computed in parallel, and no evaluate trace needs to be cached.
to the 2-jet J7 og (t) (y1, v") v(2))
4.2 Estimating arbitrary differential operator by pushing forward random jets
Next, show how to use the above facts to construct stochastic estimator derivative operator. Differential operators can be evaluated through derivative tensor contraction. The action of the derivative a
glal
on function u can be identified with the derivative tensor slice D (a)a.
0æ. ........ Bu
Differential operator L can be written as a linear combination of derivatives: L C == C. Da, where I(C) is the set of tensor indices representing terms included in the operator L. For simplicity only consider kth order differential operator, i.c. |a] k E N for all a. For scalar RRd R.
5

---

we can identify a kth order differential operator L with the following tensor dot product
Lu(a)= -EEEI(E) CaDou(a) Edh da (a) Ca (C) - Di C(C),
(10)
where d, E [1,d]ii E [1, kj is the tensor index the ith axis, and C(C) is tensor of the same shape as (a) that equals Ca when .j...., d2 matches the multi-inded a E I(C) and 0 otherwise. We call C(C) the coefficient tensor of L. For example, the coefficient tensor of the Laplacian V the d-dimensional identity matrix I. More complicated operators can be built as f(x, u, Da, u. Dr.. ")
where f is arbitrary function.
Any derivative tensor contractions (a) C(C) can be estimated through random contraction, which can be implemented efficiently as pushing forward random jets from an appropriate distribution. With random (v(1),.. ,..., ....,,((k) vill)), we have
(a)d ddi de
(v1) (r - (a((a)d Blogit) ...0oe - E (0
(1))
[
(11)
where @ denotes Kronecker product, v/z E [1, is the ds dimension of the e,th order tangent in the input k-jet. Eq. 11 is an unbiased estimator of the kth order operator Lu = (a) C(C) when
Elud,
(e1) Vdl - Ca
(vx)
(12)
For example, the condition for unbiasedness for the Laplacian V is (e) y [ As discussed, one can always find Jÿ with large enough / > k such that ø¹ FJJ)) = D (a) o v(et), with
distribution p over the input I-jet J'g that satisties the unbiasedness condition (Eq. 12 we have
(nmppl'u(J)) Es-ーplo' ut u(J)) = Elv (wx) ... dd [ = (a) C(C) Lu(a),
(13)
which means Lu(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:
1. General: STDE can be applied to differential operators of arbitrary order and dimensionality.
2. Scalable:
The scaling issue in the dimensionality d and the derivative order l are addressed at the same time. From the example computation graph (Fig. 2 we see that, for one call to d F, the memory requirement has scaling of O (kd) and the computation complexity has scaling OOk² dL).
Like first-order forward mode AD, the derivative tensor D is never fully computed and stored. Combined with randomization, the polynomial scaling in will be removed.
3. 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 (dfu);
4.3 Constructing STDE for high-order differential operators with sparse random jets
Note that all coefficient tensor has the following additive decomposition:
C(C)
Σ
Cdry...de de Cd,
Cd2
(14)
(EE(I)
where e; is the ith standard basis. For example, if the input dimension d is 3, then e₂ [0, 1, As discussed before, there exists J whose pushforward under o'u is equivalent to contracting DÊ; with
e 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. satisfies the unbiasedness condition 12
p Cdi.... de /Z. di,...,d4 d2 E I()),
(15)
where Z is the normalization factor and we identify ef_e, with the corresponding k-jet J.:
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 0² by ej. Therefore, the STDE of the Laplacian operator is given by
setting v and v
2 su)((a)) = d
a2 uo ua(
d [J
Σ doloaa)e 0) = [J|
Σ JEJ
u e, 0)
(16)
2x JEJ
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 Appendix A.4
High-order diagonal differential operators We call differential operator diagonal if is a linear combination of diagonal elements from the derivative tensor: L == - 8 From Eq. 43 we see
that setting the first-order tangent v!1) ej and all other tangents v(t) the zero vector gives the desired high-order diagonal element:
d
Σ (a) = |]|
uo(a)= = Eorosaale,,,... O....).
(17)
0x jEJ
|J| JEJ
Second-order parabolic PDEs are large class of PDEs. It includes the Bokker-Plancee equation 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 Mlck-SShhles equation mathematical finance for option pricing, the Hamilton-Jcchh Bellman equation in optimal control, and the Schrödinger equation in quantum physics and chemistry. Its form given by
Second-order parabolic PDEs
कर ux f) + 1 tr u(xx,) + Vu(x, f) 1(xx t) + f(t, x, u(x, f), T (xxx) t))
(18)
We have a second order derivative term tr (a (x, t)T The off-diagonals can be easily removed via a change of variable:
T o ux, with off-diagonal term
1
tr (xxx,))o(x,t))
0² Ox2 u(x,
1 E ()(xx)))(o(x 0).
(19)
2
See derivation in Appendix Its STDE samples over the d terms in the expression above.
4.3.2 Differential operators with arbitrary mixed partial derivative
is not always possible remove the mixed partial derivatives but discussed in section 4.2 for an arbitrary kth order derivative tensor element D (a) we can find an appropriate 7-jet J (t)
with g(t) such that ('u(J)) D. m........ Here show concrete example.
2D Korteweg-de Vries (KdV) equation Consider the following 2D KdV equation
Bty + Menwy + 3(uy - Uan + 22yy = 0.
(20)
All the derivative terms can be found in the pushforward of the following jet:
J = d ddd 11²uxx v v (13)), (3) ex, (4) = Eey, (7) = e,, V = Vi & (3,4,7),
Us -111/ Uy 3/2// Ugx 3jaj: Uay J/s /35, (21)
Myn /,0 Uky 3.0 /330, Uanayy J111/ /202200.
where the subscript [i] means selecting the ith order tangent from the jet, and the prefactors are determined through Faa di Bruno's formula (Eq. (3). 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[IA) I.4). When input dimension randomization via STDE will provide significant speed up. We tested a few more high-order PDEs with irremovable mixed partial derivatives (see Appendix [.4] and the experimental results will be provided later.
high,
7

---

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. is also possible to construct STDE with dense random jets, i.e. jets with tangents that not the standard basis. For example, the classical method of Hutchinson trace estimator (HTE) [[6] can be implemented in the STDE framework the pushforward of isotropic dense random jets, i.. (a, 0) ~ 6n x P x 8 with Ep[vv]] II
We generalize the dense construction to arbitrary second-order differential operators using 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 eigendeccomooition C" = 1 (C
AI UEU where -À smaller than the smallest
eigenvalue of C, we can construct STDE for D:
2) [0 (()) (Uv, 0)] VV
li (v, 0)] =
(a) C. (22)
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 counterexample one cannot construct an STDE for the fourth order operator Σ a² dz with dense jets. For more details 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
The main differences between the sparse and the dense version STDE are:
1. sparse STDE is universally application whereas the dense STDE can only be applied to certain operators;
2. the source of variance is different (see Appendix
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.
Experiments
We applied STDE amortize the training of PINNs 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 [.4] that cover the case of k = 3, 4, which includes PDEs describing ID and 2D nonlinear dynamics, and high-dimensional PDE with gradient regularization [42]. Furthermore, we tested weight-sharing technique (Appendix G which further reduces memory requirements (Appendix L.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 B3 is a class of neural PDE solver where the ansatz us (x) is parameterized by a neural network with parameter 0. It is a prototypical case of the optimization problem in Eq. [ We consider PDEs defined on a domain Ω c Rd and boundary/initial òΩ as follows
Lu(x) f(x),
Z Bu(x) = g(x), x
(23)
where L and B are known operators, f(x) and g(x) are known functions for the residual and boundary/initial conditions, and Rdd + IR is a scalar-valued function, which is the unknown solution to the PDE. The approximated solution up ux(x) ss u(x) is obtained by minimizing the mean squared error (MSE) of the PDE residual 0) Luo(x) -(x):
reesddual (8; 11
Σ (uo(x! f(x((())) N
1
(24)
where the residual points (x()).. are sampled from the domain Ω. We use the technique from 26 that reparameterizes Wi such that the boundary/initial condition Bu(x) = g(x) are satisfied exactly for all x E 02, so boundary loss is not needed.
8

---

Amortized PINNs PINN training can be amortized by replacing the differential part of the operator with stochastic estimator like SDGD and STDE. For example, for the Allen-Caihe equation. Lu = v + w², the differential part of L is the Laplacian V With amortization, we minimize the following loss
Nr N E
Presidiual (0; (x(1) 1= J,K) = - - -(x(()))] ( f(x(1)) (25)
which is modification of Eq. 24 Its gradient
oigsen is then an unbiased estimator to the gradient of 00
the original PINN residual loss, i.e. E
Dimeru 00
Denedan 08
5.2 Ablation study on the performance gain
To ascertain the source performance gain of our method, we conduct detailed ablation study on the inseparable Allen-Cahi equation with two-body exact solution described in Appendix 1.1 The results are in Table and where the best results for each dimensionality are marked bold. All methods 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, seen Table 3 in the Appendix. We have included Forward Laplacian which is an exact method. It is expected to perform better in terms of L.2 error. However, least in the case where the dimension is more than 1000.
can in Table 3 the I.2 error is of the same order,
Table 1: Speed ablation for the two-body Allen-Cahi equation.
Speed (it/s) f
100
IK
100K
IM
Backward mode SDGD (PyTorch) 13
55.56
3.70
1.85
0.23
OOM
Backward mode SDGD
40.63
37.04
29.85
OOM
OOM
Parallelized backward mode SDGD
1376.84
845.21
216.83
29.24
OOM
Forward-ove over- Backward SDGD
778.18 18
560.91
193.91
27.18
OOM
Forward Laplacian 24
1974.50
373.73
32.15
OOM
OOM
STDE
1035.09
1054.39
454.16
156.90
13.61
Table 2: Memory ablation for the two-body Allen- Cahn equation.
Memory (MB)
100 D
IK D
10K D
100K D
IM D
Backward mode SDGD (PyTorch) [3]
1328
1788
4527
32777
OOM
Backward mode SDGD
553
565
1217
OOM
OOM
Parallelized backward mode SDGD
539
579
1177
4931
OOM
Forward-over Backward Bar SDGD
537
579
1519
4929
OOM
Forward Laplacian 24
507
913
5505
OOM
OOM
STDE
537
795
1073
6235
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 and 2 speed-up and up 44x memory reduction.
JAX provides ~15x
Parallelization The original SDGD implementation uses for-loop iterate through the sampled dimension (Appendix A. This can be parallelized (denoted as "Parallelized SDGD via HVP". details in Appendix A.2). A..2). Parallelization provides ~15x speed up and reduction in peak memory for the JIT compilation phase. We also tested mixed mode AD (dubbed as "Forward-over. Bareer-awkwardd SDGD"), which gives roughly the same performance parallelized stacked backward mode, which is expected as explained in Appendix
Forward Laplacinn Forward Laplacian 24 provides a constant-level optimization for the caleu- lation of Laplacian operator by removing the redundancy in the AD pipeline, and we can see from

---

Table I and 2 that it is the best method in both speed and memory when the dimension is 100. But since not a randomized method, the scaling is much worse. Its computation complexity is O(d). whereas randomized estimator like STDE has a computation complexity of ((J]). Naturally, with a high enough input dimension d, the difference the constant prefactor trumped by scaling. When the dimension is larger than 1000, it becomes worse than even parallelized stacked backward mode SDGD.
STDE Compared the best realization of baseline method SDGD, the parallelized stacked back- ward mode AD, STDE provides up to 10x speed up and memory reduction of at least 4x.
6 Conclusion
We introduce STDE, general method for constructing stochastic estimators for arbitrary differential operators that can be evaluated efficiently via Taylor mode AD. We evaluated STDE PINNs, an instance of the optimization problem where the loss contains differential operators. Amortization with STDE outperforms the baseline methods, and STDE also applies 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- digh--imensional AD-based PDE solvers. This makes STDE more general than branch of related methods. STDE is also more applicable than deep ritz method [4]], weak adversarial network (WAN) [43] backward 3
SDE-based solvers
34 [0], deep Galerkin method [5]. and the recently proposed forward Laplacian [24), which are all restricted to specific forms of second-order PDEs. STDE applies
naturally differential operators in PDEs, but also be applied to other problems that require
input gradients. For example, adversarial attacks, feature attribution, and meta-learning, to name few.
Limitations Being 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, 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 B
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 large-scale scientific modeling with neural networks. One example would be the many-body Schrödinger equations, where one needs to compute a high-dimenssonal Laplacian. Another example the high-dimensional Mlack-Scholes equation, which has numerous uses in mathematical finance.
10

---

References
[1] Brandon Amos. Tutorial on amortized optimization, April 2023, arXiv:2202.0066) [cs, math].
[2] Atilum Güneş Baydin, Barak A. Pearlmutter, Don Syme, Frank Wood, and Philip Torr. Gradients without Backpropagation, February 2022. arXiv:2202.08587 [cs, stat].
[3] Christian Beck, Sebastian Becker, Patrick Cheridito, Arnulf Jentzen, and Ariel Neufeld. Deep splitting method for parabolic PDEs. SIAM Journal on Scientific Computing, 43(5):A3135-
A3154, January 2021. arXiv:1907.0444 [cs, math, stat].
[4] Sebastian Becker, Ramon Braunwarth, Martin Hutzenthaler, Arnulf Jentzen, and Philippe von Wurstemberger. Numerical simulations for full history recursive multilevel Picard approxi- mations for systems of high-dimensional partial differential equations. Communications in Computational Physics, 28(5):2109-213 June 2020. arXiv:2005.10200 [cs, math].
[5] Claus Bendtsen and Ole Stauning. Tadiff, flexible c + package for automatic differentiation using taylor series expansion. 1997.
[6] Jesse Bettencourt, Matthew J. Johnson, and David Duvenaud. Taylor-mode automatic differen- tiation for higher-order derivatives in JAX. In Program Transformations for MI. Workshop at NeurlPS 2019, 2019.
[7] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Minn and Qiao Zhang. JAX: composable transformations of Python+NumPP programs, 2018.
[8] Benyamin Ghojogh, Ali Ghodsi, Fakhri Karray, and Mark Crowley. Johnson-LLindenstrauss Lemma, Linear and Nonlinear Random Projections, Random Fourier Features, and Random Kitchen Sinks: Tutorial and Survey, August 2021. arXiv:2108.04177 [cs, math, stat].
[9] Andreas Griewank and Andrea Walther. Evaluating Derivatives. Society for Industrial and Applied Mathematics, second edition, 2008
[10] Jiequn Han, Aruulf Jentzen, and Weinan E. Solving high-dimensional partial differential equations using deep learning. Proceedings of the National Academy of Sciences, 15((4))8505- 8510, Aug 2018.
[11] Di He, Shanda Li, Wenlei Shi, Xiaotian Gao, Jia Zhang. Jiang Bian, Liwei Wang, and Tic- Yan Liu. Learning Physics-Infommn Neural Networks without Stacked Back-propagation February 2023. arXiv:2202.09344 09344 [cs].
[12] Zheyuan Hu, Zekun Shi, George Em Karniadakis, and Kenji Kawaguchi. Hutchinson Trace Es- timation for High-Dimensiona and High-Order Physics-Infomin Informed Neural Networks. Computer Methods in Applied Mechanics and Engineering, 424: 1 16883, May 2024. arXiv:2312.144499 [cs, math, stat].
[13] 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.
[14] Zheyuan Hu. Zhouhao Yang, Yezhen Wang, George Em Karniadakis, and Kenji Kawaguchi. VaaiaVa Trade-off in Physics-Informe Neural Networks with Randomized Smoothing for High-Dimensien PDEs, November 2023. arXiv:2311. 1. [cs, math, stat].
[15] Zheyuan Hu, Zhongqiang Zhang, George Em Karniadakis, and Kenji Kawaguchi. Score-Based Physics-Informet Neural Networks for figh-Diiensiooal Pokker-Plande Equations, February 2024. arXiv:2402.07466 [cs, math, stat].
[16] 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): 1059- 1076, January 1989.
[17] Martin Hutzenthaler, Arnulf Jentzen, Thomas Kruse, Tuan Anh Nguyen, and Philippe von Wurstemberger. Overcoming the curse of dimensionality the numerical approximation of semilinear parabolic partial differential equations, July 2018.
11

---

[18] Jerzy Karczmarezuk. 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.
[19] George Em Karniadakis, Ioannis G. Kevrekidis, Lu Lu, Paris Perdikaris, Sifan Wang. and Liu Yang. Physics-informa machine learning. Nature Reviews Physics, 3(6):422 440, Jun 2021.
[20] 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.
[21] Chieh-Hsin Lai, Yuhta Takida, Naoki Murata, Toshimitsu Uesaka, Yuki Mitsufuji, and Stefano Ermon. Regularizing score- score-based models with score fokker- Backer-planck equations. In NeurlPS 2022 Workshop on Score-Bace Based Methods, 2022.
[22] Jacob Laurel, Rem Yang. Shubham Ugare, Robert Nagel, Gagandeep Singh, and Sasa Misailovic. A general construction for abstract interpretation of higher-order automatic differentiation. Proc. ACM Program. Lang., 6(0OPSLA2), oct 2022.
[23] 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 AlADifferentilEEuattinn In Science, 2024.
[24] 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: New Computational Framework for Neural Network-bard based Variational Monte Carlo, July 2023. arXiv:2307.08214 [physics].
[25] 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.0663 066244 [cs, eess, stat].
[26] Lu Lu, Raphaël Pestourie, Wenjie Yao, Zhicheng Wang. Francesc Verdugo, and Steven G. Johnson. Physics-informeet neural networks with hard constraints for inverse design. SIAM 2021.
Journal on Scientific Computing. B
[27] Sadhika Malladi, Tianyu Gao, Eshaan Nichani, Alex Damian, Jason D. Lee, Danqi Chen, and Sanjeev Arora. Fine- ToninTTnig Language Models with Just Forward Passes, January 2024.
arXiv:2305.1 13.333335.17333111 [cs].
[28] Per-Gunnal Martinsson and Joel Tropp. Randomized Numerical Linear Algebra: Foundations & Algorithms, March 2021. arXiv:2002.01387 [es, math].
[29] Riley Murray, James Demmel, Michael W. Mahoney, N. Benjamin Erichson, Maksim Mel- nichenko, Osman Asif Malik, Laura Grigori, Piotr Luszezek, Michal Dereziński, 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. 1474 [cs, math].
[30] Deniz Oktay, Nick McGreivy, Joshua Aduol, Alex Beatson, and Ryan P. Adams. Randomized Automatic Differentiation, March 2021. arXiv:2007. 201.20007.10422 [cs, stat].
[31] Tianyu Pang, Kun Xu, Chongxuan Li, Yang Song, Stefano Ermon, and Jun Zhu. Effi- cient Learning of Generative Models via Difte-DDifeeeece Score Matching, November 2020. arXiv:2007.03311 [cs, stat].
[32] Juncai Pu and Yong Chen. Lax pairs informed neural networks solving integrable systems, January 2024. arXiv:2401 .04982 [nlin].
[33] M. Raissi, P. Perdikaris, and G.E. Karniadakis. Physies-informed neural networks: deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational Physics, 378:686-707, February 2019.
[34] Maziar Raissi. Forward-Backward Stochastic Neural Networks: Deep Learning of High- dimensional Partial Differential Equations, April 2018. 1rXiv:1804.07010 [cs, math, stat].
[35] Justin Sirignano and Konstantinos Spiliopoulos. Dgm: deep learning algorithm for solving partial differential equations. Journal of computational physics, 375:1339- 2018.
12

---

[36] Maciej Skorski. Modern analysis of hutchinson's 's trace estimator. In 2021 55th Annual Conference on Information Sciences and Systems (CISS). IEEE, March 2021.
[37] 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].
[38] 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:201 1. 13456 [cs, stat].
[39] Charles M. Stein. Estimation of the Mean of Multivariate Normal Distribution. The Annals of Statistics, 9()): 1151, 1981.
[40] Mu Wang. High Order Reverse Mode of Automatic Differentiation. PhD thesis, 2017. Copyright Datahase copyright ProQuest L.LC: ProQuest does not claim copyright in the individual underlying works; Last updated 2023-03-04.
[41] 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.
[42] Jeremy Yu, Lu Lu, Xuhui Meng, and George Em Karniadakis. Gradient-enhanced physics- informed neural networks for forward and inverse PDE problems. Computer Methods in Applied Mechanics and Engineering. 393: 14823, April 2022. arXiv:2111.0288 [physics].
[43] Yaohua Zang, Gang Bao, Xiaojing Ye, and Haomin Zhou. Weak Adversarial Networks for High- dimensional Partial Differential Equations. Journal of Computational Physics, 411:109409 June 2020. arXiv: 1907. 08272 [cs, math].
13

---

A Example implementations
A.1 PyTorch implementation of SDGD-PINN using backward mode AD
The original implementation of SDGD-PINN [3 computes the SDGD estimation of derivatives using a for-loop that iterates over the sampled PDE term/dimmnniion. For example, given function representing the MLP PINN, the computation of SDGD for the Laplacian operator can be implemented in PyTorch as follows:
f_x = torch. autograd ittograd.grad(f..... sum (), create_grape (reateggraph-Tre) [o]
idx_set np. np.random. choice (dim, adgd_batch_siae replace=Palse) hess_diag_va = 0.
for in idx_set:
dess_diage = torch. autograd grad(
f_x[:, i] sum(), x, greate-graph=Tiee [o] [:, i]
maas_diag_ve += hess gess_diag_ie i * dim adgd_batch, size
After computing the PDE differential operator, it plugged into the residual loss, and then backward- mode AD is employed produce the gradient for optimization concerning 0.
A.2 JAX implementation of SDGD Parallelization via HVP
def hvp(f, x, v):
"""stacked backward-mode Hessian-vector product""" return jax. grad (1ambda x: jmp. vdot (jax, grad(f) (x), v)) (x)
f. f_hees__dd __hess_daa_ffn lambda i: hvp(f-partial) x_i, jnp eye [i]) [i]
idx_set set jax random. choice(
key, dim, shape
replaceFaase
hess_diag_val jax. vmap whwss_eias__ff__ (idx_set)
A.3 JAX implementation of Forward-over-backward AD
The forward-over-backward AD In JAX mentioned in Appendix can be implemented fol- lows:
f_grad_fn fn jax. gr)
f_x, f f_hess_fr fn jaa. www.:innaizeef_granae x_i) jup over vÿp [i]) [i]
f chess_diag_ss fn lambda i: f hhess_fnjjnpeeeee....eeeeeeeeeee
hess hess_diag_val val = vexp Sift_hes__iiie fn) (idx_set)
A.4 JAX implementation of STDE for the Laplacian operator
idx_set set jax random. choice( key, dim, chape
replace=False
rand_jet jax. (1ambda i: jnp.eye (dim) [i]) (idx_set) pushfwd_2_nt 2__rdd___ff fn = lambda fun-fn, primals=(x,), series=((v, jup zeros zeros(dim)), (dim)),)
jet(
) pushforward of the 2-jet (a, v, 0), i.e. Idd-2 f(z, v, 0)
f_vals, (., vhv) jax. vmap (pushfwd_2_fn) (rand_jet) hess diss_diag_vll dim batch_size vhv
The jet function from JAX implements the ooghorrder pushforward d 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.2p to compose the primitives into the pushforward jets. Note that the API of specified via the series argument.
all the high-order tangents of the input jet
14

---

B
Further details on first-order auto-differentiation
B.1
Computation graph of first-order AD
8g
0,
0:
0
8g
0g
Fi
F2
F/
F1
F1
P2
Pg
F4
Fg
Es
F4
F(x)
F(x)
OF;
OF
0.F
BF4
0 F
F2
Fa
8
F4
0F 0.
8F,
F1
P2
F3
F4
Figure 3: The computation graph of forward mode AD (left) and backward mode AD (right) of function F()) with 4 primitives F; each parameterized by 0,, Nodes represent (intermediate) values, and arrows represent computation. Input nodes are colored blue; output nodes are colored green, and intermediate nodes are colored yellow.
B.2 Derivative via composition
First-order AD is based on simple observation: for set functions L, the set of tuples of functions f and its Jacobian J, is closed under composition:
Jf) Jg) = (f g, Jog).
(26)
where denotes both function composition and the composition of the tuple (J. If we have the analytical formula of the Jacobian J for every E L, then we can calculate the Jacobian any composition of functions from C using the above composition rule for the tuple (f) The set C 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 + is map from V to the space of all bounded linear operators from V to W, denoted L(V, W), that is
0f V + L(V,W) W),
(27)
such that at a point a V it gives the best linear approximation 0f(a)()) of the sense that
llf(a + h) f(a) - 0f (a)
lim |b|++0
(28)
|/blly
Therefore, it is also called the linearization of f at point Concretely, consider a function in Euclidean spaces _ R At any point IR, the Fréchet derivative Of can be seen as the directional derivative of
(29) where (a) R denote the Jacobian of f at point called the primal, and R" also called the tangent is a vector representing the direction. Therefore the Fréchet derivative is also called sacobian-vector-pronn (JVP). And we write the truncated Taylor expansion
: - + L(R",R") RT), 0f(a)(v) Jy(a)v
f(a + Ax) f(a) + 0f Ax) + OAx²).
(30)
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 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 X + Y has an adjoint 'u Y' + X" where X', Y the dual spaces. The adjoint satisfies the following
Vy E
('uly),=) (v, u(x))
(31)
15

---

In the mnite-dimensionaa case, the dual space is the space of row vectors, and any linear map can be written as u(x) Ax. One can easily verify that the adjoint is the transpose: 'u(y") y A. The adjoint (transpose) of the Fréchet derivative of IR" R denoted as 0J is thus defined as
f R2 +
o f(a)(v) v Ja()),
VER R
(32)
where T is the cotangent which lives in the dual space of the codomain. Note that the adjoint is taken to only where is kept fixed. 0" f is also called vector- Jacbraa--oban-rodceettt (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 [ the second order derivative is recommended be computed via forward-over- backward AD, i.c.. 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
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 explained in section B.2 Essentially, forward-over- backward reverses the arrows in the third row in Fig. [ 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 sccond-ordes differential operator here with forward-over www.aadddve----bawaateee AD might not yield any advantage.
D
Taylor mode AD
D.1 High-order Fréchet Derivatives
The kth order Fréchet derivative of a function f R" - IR" at a point a is the multi-linear map with k arguments around point that best approximates For example, when
have
8 f : + L(R x R"), 8 f())(v, (2)0)vvv)==- = (y(a)v's
Σ J.k
( (33)
where Hy(a) E RRxXX denote the Hessian of f at point a, and v, v' E R". We can now write the second-order truncated Taylor series with
f(a + Ax) f(a) + (Ax) + 8 (Ax) Ax) + +OAx'). For the more general case, we have
(34)
a* Ff:RD RD + L
O R R
0 f(a)(v()) ...., v(())=
Σ D)(a) Do 1
(1)
(35)
High-order Fréchet derivative can be seen as the best kth order polynomial approximation of by taking all input tangents be the same IR::
f(a + Ax) =f + (f(a)(v) +
1 2 of f(a(v.vv+++. +......))....
1 0h k! ot f(
+0(x)))))
(36)
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 (g((4) of scalar functions, tives by iteratively applying the chain rule to lower-order chain rules:
generalize the chain rule for high-order deriva-
dr [(((x)) = ()(x)) (x)
0²
dz2 /g((x) - (0) (z) + (g(x)) (x))²
03
3f (x) (x)) + (((()) ( (x) - (g(xx(
=
(37)
16

---

where we give the example of up to the third order. For arbitrary k, the kthh order derivative of the composition is given by the Faa di Bruno's formula (scalar version)
0h Oz (((())
Σ
k!
(f(E o (x)
II 1
(O))))) Ps
(38)
II. P
(P) PU)NN E &ppzk
where the outermost summation is taken over all partitions of the derivative order k. Here partition of l defined as a tuple (PI....,PP) E No that satisfies
i. pt = k.
(39)
For vector-valued functions g R
Rm
Rm IR,, let
v(I 0g(x)
a = g(x) E R"
E Rmxx
8x
(40)
(2) 8 g(x)
(3) 0 gg(())
R BXXXXX
E RR IXXXXX
Ox2
Ox3
we can derive the following composition rule similarly
f(((x)) (D (m) u m.n Rixn
Ox
dx2 f(((x)) =Dy (a)) I.m u mmn + D2(()) ( () mm mmmn nd E R
83
(41)
fx33(o(()) = ). (3)
m.n.n'.n"
+3 D7 aa) I. m.n mi m'.""
+D3 (a) m m mn ml ae m",n" E R plxxxnxa
(1) (1)
where again we give the example of up the third order, and repeated indexes are summed in Einstein notation. The general formula is again given by the multivariate version of the Faa di Bruno's s formula. Note that in the multivariate version of the Faa di Bruno's 's formula, it is possible take derivative to distinguishable variables, but here we just present the version with indistinguishable input variables. This gives the composition rule for kth order total derivative.
The composition of the high-order Fréchet derivative J is the case of n = 1, as the contraction with the input tangents (+) E Rd the same as composing with a scalar input function g IR + Rd with (+) D All derivative tensors of (g(()) can represented using a RC vector, and similarly all derivative tensor v() of g can be represented using a I vector. Then, the above chain rule can be simplified
0e ((((t) = =Dy (a) (1) E TIR
8²
0re [(((1) = =Dy (22 (7)a)). m. 1 ( TR!
03
0y 053 f(g(t)) (ay (3) + mm
o 2 (a) m. mn of (1) (1)
Unp un R².
(42)
The Faa di Bruno's formula again gives the general formula for arbitrary derivative order
dek {(((()) = Σ
ge
kd Di Pi (a),,
DF
k
(i)
E ER!
Tle p!!(ttp
ff
(PI...PP)ENN JEN S m
(43)
which is written in the perspective of input primal and tangents (1)
17

---

E Removing the mixed partial derivatives term from second order semilinear parabolic PDE
ir (o(x, ()((x))) (oo(x,))"(HHsss))) (Hess.u))xx ())
- tr (o(x, (x, ())
[x (Hess,u)) ()x 1)
(44)
Σ
c (xx (Hessxu) (x),
2
i=0
Ou((x, f), (es,, 1a]-
1=0
F Evaluating arbitrary mixed partial derivatives
E. A concrete example
Let's first consider a concrete case. Suppose the domain is D-dimensional we want to compute the mixed derivative āuraw; The naive approach would be to compute the entire third order derivative tensor D3. which is tensor of shape D D D, then extract the element index (j. i, However note that from Eq. 43 for any 3, the pushforward k-jet under d contains contractions of DJ. Although in the case of l 3, the only contraction of D3 is in the 0%
Im, m Um um Un2
(1) (1)
(45)
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-dia off-diagonal terms, i.c. the mixed partial derivatives. For example, in d f. all input tangents are set zero except for v()) and v(2), 8ff becomes:
3- D3 (m
(1) (2) (2) + BD) ()) I. my Un u mg Uma + (a)1, (ag.m..m.. ma.. my m2 my v (1) mi
(1) (1) (1)
ma." (46)
which contains the contraction of D that we want:
(1) my Ur Unng Piss
(47)
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:
On? (x = uo (x) e,, 0) وr (e; 0,0, 0) 338 (x)(e;, 0)]/6.
(48)
If we go 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
10 D3 (a) i. im my my
mg (3) ( (1)) + ())a). m.m.s V (1),,11 (1) (1)
(49)
mi ma mi "mng
if all input tangents set zero except for v(1) and v3) With this only need to remove one term:
Or,an, up (x) = (x) 0, ej, 0) Po 0,0,0 0)
(50)
Similarly, by going to the seventh-order tangent, we can compute this mixed derivative with only one pushforward. d contains 07, and when all input tangents are zero except for (22) and (3)
87 equals
105 (a) J.m. m...m.mmm Umg
(3) (2) (2)
(51)
which is the exact contraction we want. With this we have
i ug (x) a us (x) e, 0, e,, 0,000,000000
(52)
Ou?dx,
18

---

F.2 Procedure for finding the right pushforwards for arbitrary mixed partial
derivatives
More generally, consider the case where need to compute arbitrary mixed partial derivative
023 96
(53)
41 dr 9- Dz
where T' the number of different input dimensions in the mixed partial derivative, and qi, is the order. To compute it with k-jet pushforward, one needs to find:
1. a derivative order & E:
2. a sparsity pattern for the tangents v of the input jet, which is defined as the tuple of T integers J = (i....., where v = 0 when j $ J and je < k for all & E [1,T],
such that when setting
j J
Pj p=
(54)
Die: J.
(P1, (P1,P2... P.....,.,, Pk) ) No is partition of k defined in Eq. 39
Let's use the concrete example 02702z
a again. In this case T = 2, qi 2 and qi 1. We
demonstrated that this can be computed with one 7-jet pushforward, which is equivalent to setting
= (2, k 2j1 + j2 7, and the partition (0, 2, 1, O, 0, 0, 0). The Faa di Bruno's 's formula (Eq. 43 ensures that the pushforward of the kth 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 has a sparsity pattern that is not a subset, it will vanish when we set the input tangents to zero according the sparsity pattern of the partition in consideration. To understand this point better, let's ss look at the concrete example with the 5-jet pushforward demonstrated above. (2, 0, 1,0, 0,0 0) and (5, O, O, O, 0) are both valid partition of k 5. and the sparsity pattern of (5, O, O, O, is the subset of that of (2, 0, 1, 0, 0): 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). This equivalent say, 2 + 2 is the only way to sum up to when you can only use and which can be verified casily.
extra terms to remove. Intuitively, if a partition
With this setup, it is clear why the diagonal terms can always be computed with pushforward of the lowest possible order: (Ac, 0..... 0) E N* is always a valid partition k and no other partition has sparsity pattern that subset of it.
For mixed partial derivatives, the difficulty scales the total order of the operator Σt= qi., and T which can be interpreted as the degree of the "off-diagonalness" of the operator. For example, consider the case where T = and 9,, = 3, 3,4, = 2, 9i = 1. This corresponds to the operator 0=/0=fast To avoid overlapping with the diagonal sparsity pattern (k: ........., 0) and to keep the
order of derivative low, one might try k: = 16 and the partition 2,, 1,0..... E N16 However, with higher there is more chance that other partitions will have subset sparsity pattern. In this case (0, 8, O, 0, O....) E 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 nigh-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 similar approach to mitigate the memory cost of model parameters in nigh-dimensional PDEs, called weight sharing in the first layer.
19

---

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 dimensional matrix, whereas all subsequent layers have a weight matrix with size of only h x h.
reduce the redundancy in the parameters in the first layer. Specifically, we perform an additional ID convolution to the input vectors x; before passing the input into the MLP PINN, as in Fig. The ID 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.
By introducing a weight-sharing scheme,
This weight-sharing scheme reduces the parameters by approximately B. The number of parameters in the filters is B x 1, and the subsequent fully connected layer will have weight matrix of size d x H. Therefore, the total number of the first layer is reduced from d x h to only d and we can see that with a larger block size B, we will have fewer parameters, and the reduction factor is approximately B. More concretely, suppose d 10 h 100 where one million (10%) dimensional problems are also tested experimentally, so the number of parameters the first layer is d x h = 100 10 If we use block size of B 100, will reduce the number of parameters to B h d x h = 10 x 10 10. In other words, with larger block size of B, we significantly reduce B
B h + B,
B
10 + 100. If the block size is B = 10, the number of parameters will be
the number of model parameters.
We will demonstrate the memory efficiency and acceleration thanks to weight-sharing in the experi- mental section.
31
x22
T3
344
25
.T6
3/7
F8
Convolution with stride of
0
02
03
output
1/1
/2
ys
Figure Convolutional weight sharing in the first layer, with input dimension and filter size
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 memory usage, we use a separate run where GPU memory pre-allocation for JAX is disabled through setting the environment variable XLA_ PiL_PYTHO___ CLIIONNCLIITT_LL WWW.CCIITTAAALLCTTRRRRR.
and the test data set stored on the CPU memory. The GPU memory usage obtained via NVIDIA- smi and peak memory reported.
All the experiments were done on single NVIDIA A100 GPU with 40GB memory and CUDA 12.2.
with driver 535.129 129.03 and JAX version 0.4.23.
Network architecture and training hyperparameters For the semilinear parabolic PDEs tested in Appendix 1.2 we follow the network architecture of the original SDGD [3]
The network is a 4-layer multi-layer perceptron (MLP) with 128 hidden units activated by Tanh.
The network trained with Adam (00 for 10K steps, with an initial learning rate of le-3 that linearly decays 0 in 10K steps, where 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.
20

---

The zero boundary condition is satisfied via the following parameterization
(x) (1 |
(55)
where ugg MLP is the MLP network, and ug is the PDE ansatz, described in [26].
For the semilinear parabolic PDEs tested in Appendix 1.2 we made the following modifica- tions:
Instead of using oe-paaaetterization, the boundary/initial condition is satisfied by adding regularization loss to the residual loss:
N₂
No V (xg 0) - (g(xb, i)
fooundary (0; |
(x(.,. 0) g(x,,)))
No
No
(56)
where g(-) is the initial data, N₂ is the batch size for boundary points, is the PDE ansatz, Cg is the coefficient for the first-order derivative boundary loss term, which we set to 0.05. The total loss is
No Fresduual (0; {x., ) + 200ooundrry (0; {x6,
(57)
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. ~ uniform(0, I'. and then sample x directly from the distribution of X₁, i.e. To match the original training setting of 100 SDE trajectories with 0.015 step size for time discretization,
(N (T-()-Iaxa))
a batch size of 2000 for residual points and 100 for boundary/mntiall 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 LI error at the point (Nrest, T' against the reference value computed via multilevel Picard's method 3 4 17
experiments, we use the biased version of Eq. 25
Eesiduua px
Σ [x)o((())) - f(x(1)
-
(58)
N
as the bias in practice is very small and does not affect convergence.
I Experiments Results
L1 Inseparable and effectively high-dimensional PDEs
The first class of PDEs is defined via nonlinear, inseparable, and effectively high -igh-dimensional exact solution (x) defined within the d-dimensional unit ball Bd:
Cu(x) -x(),, x E Bd
(59)
u(x) =0, X dBd
where L is linear/nonlinear operator and g(x) ( The zero boundary condition ensures
that no information about the exact solution leaked through the boundary condition. We will consider the following operators:
Poisson equation: Lu(x) ux)(
Allen-Cahn equation: Lu(x) = V²u(x) (x) + u(x) - u(x)3.
Sine-Gordon equation: Lu(x) = v²u(x) + sin (xx)).
For the exact solution, consider the following with all e A/00 1):
two-body interaction: Weeaan x) = - |xx]2 D c sin(x; + cos(2:41) + ai41 cos(
three-body interaction: Mexacr(X (x) = (1 |x()) E exp((((4411))))
21

---

We tested the performance of STDE on these equations, and the results are presented in Table 3 S 6 For the Allen-Cahi equation, we performed detailed ablation study (Table 3). and we expect these results to generalize over these second-order PDEs.
11I1 Further details on ablation study
The gain by using JAX instead of PyTorch Since the original SDGD was implemented in Py Torch, we implemented the stacked backward mode without parallelization in SDGD dimensions JAX for fair comparison (dubbed as "Stacked Backward mode SDGD in JAX" in Table 3) The for-loop over SDGD dimension is implemented using j jax. lax.1a scam. Table 3 shows that, even with the original stacked backward mode AD, the speed of JAX implementation when the dimension is high. The memory profile 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 JIT compilation requires extra memory, and the peak memory requirement during JIT compilation is higher than that during training.
scan.
be more than 10x faster
out-of-memory (OOM) error. This is because performing
Randomization batch size We also tested the case where the STDE randomization batch size reduced to 16. As seen in Table in the case of Allen-Cahn provides ~2x 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.
1.2 Semilinear Parabolic PDEs
The second class of PDEs is the semilinear parabolic PDBs, where the initial condition is speci- fied:
a at u(x, () = Lu(xx() (x, 1) E Rd [0, (x, R x {0}
(60)
u(x,t) =x(x).
where g(x) is known, analytical, and time-independent function that specifies the initial condition, and T the terminal time. We aim to approximate the solution's 's true value at one test point Nest E Rd
at the terminal time t = T, i.e. at (Xest, T).
We will consider the following operators
Semilinear Heat Eq.
u(x,
Lu(x,t)= 1) = V2u(x,t)
(61)
1 u(x,t) t)²
with initial condition g(x) = 5/(00+1
Allen-Cahn equation
Cu(x,t) = V2u(x t) + uu(x f) u((x 2)3
(62)
with initial condition g(x) arctan m;;),
Sine-Gordon equation
Lu(x, () = V²u(x, () + sin(u(x, ()).
(63)
with initial condition g(x) = 5/ 1))
All three equation uses the test point X 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 Appendix
G We tested the best- best-performing method from Table 3I 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.
22

---

Table 3: Computational results for the Inseparable Allen-Cahn equation with the two-body exact solution, where the randomization batch size is set to 100 unless stated otherwise.
Method
Metric
100 D
IK D
1OK D
100K D
IM D
Speed
55.56it/s
3.70it/s
1.85i//s
0.23iu/s
OOM
Backward mode SDGD (Py (PyTorch)
Memory
1328MB
1788MB
4527MB
32777MB
OOM
Error
7.187F-03
5.615F-04
1.864F-03
2.178F-03
OOM
Speed
40.63iv's
37.04iv'
29.85it/s
OOM
OOM
Backward mode SDGD (JAX)
Memory
553MB
565MB
1217MB
OOM
0OM
3.51E-03
7.29F-04
3,46E-03
Error
18.47E-05
15.45E-06
12.01E-04
OOM
0OM
Speed
1376.84it/s
845.21it/s
216.83it/s
29.24it/s
OOM
Parallelized backward mode SDGD
Memory
539MB
579MB
1177MB
4931MB
00M
6,87E-03
3.12E-03
2.59E-03
1.60--03
Error
+6.97E-05
+7.04E-04
-2.20E-05
+ 1.11E-05
00M
Speed
178.18
560.91it/s
193.91it/s
27.18it/
OOM
Forward-over -Backward SDGD
Memory
537MB
579MB
1519MB
4929MBB
00M
4.07E-03
2.19E-03
5,47E-04
4.21E-03
Error
+7.42F--05
+2.03E-004
+7.48E-05
+2.53F-04
0OM
Speed
1974.50it/s
373.73it/s
32.15i/s
OOM
OOM
Forward
Laplacian
Memory
507MB
913MB
5505MB
OOM
00M
4.33E-03
5.50F-04
5,58E-03
Error
14.97E-05
14.60E-05
+2.73E-04
OOM
0OM
Speed
1035.09it/s
1054.39it/s
454.16it/s
156.90it/s
13.61it/s
STDE
Memory
543MB
537MB
795MB
1073MB
6235MB
1.03E-02
6,21E-04
3,45E-03
2.59E-03
1,38E-03
Error
+7.69E-05
+2.22E-04
+1.17E-05
+7.93--06
+3.34E-05
Speed
1833.78it/s
1559.36it/s
587.60it/s
283.33it/s
21.34it/s
STDE
(batch size=16)
Memory
457MB
481MB
741MB
1063MB
6295MB
1,89E-02
7,07E-04
8,33E-04
1,50E-03
3,99E-03
Error
+2.37E-04
+1.02E-05
#2.96E-04
+1.02E-05
+3.41E-05
Table 4: Computational results for the Inseparable Poisson equation with two-body exact solution.
Method
Metric
100D
IK D
10K D
100K D
IM D
Speed
55.56it/s
3.70it/s
1.85it/s
0.23it/s
OOM
Backward mode SDGD (PyTorch) 13
Memory
1328MB
1788MB
4527MB
32777MB
OOM
Error
7.189E-03
5.611E-04
1885E-03
2.175E-03
OOM
Speed
2020.05it/s
1649.20it/s
584.98it/s
281.78it/s
20.38it/s
STDE (batch size=16)
Memory
457MB
48 481MB
741MB
1063MB
6295MB
3,50E-03
4.91E-04
4,70E-03
3.49E-03
9.18E-04
11.44E-04
13.45E-05
10.10
12.14E-05
16.39E-06
23

---

Table 5: Computational results for the Inseparable Sine-Gordon equation with two-body exact
solution.
Method
Metric
100 D
1K D
10K D
100K D
IM D
Speed
55.56it/s
3.70iu/s
1.85i/s
0.23iu'
OOM
Backward mode SDGD (PyTorch)
Memory
1328MB
1788MB
4527MB
32777MB
OOM
7.192E-03
5.641E-04
1.854E-0
2.177E-03
OOM
Speed
1926.33it/s
1467.38its
566.26it/s
279.24ius
19.88it/s
STDE (batch size=16)
Memory
457MB
481MB
741MB
1063MB
6295MB
3.64E-03
5.40E-04
5.32E-03
9.56E-04
9.47E-04
Error
+1.46F-04
+7.21F-05
1..12
+8.03F--06
+8.30F-00
Table Computational results for the Inseparable Allen-Cahn, Poisson, and Sine-Gordon equation with the three-body exact solution, computed via STDE with randomization batch size | set to 16. *STDE with randomization batch size ([Jl) of 16 performs poorly on the M dimensional Inseparable Poisson equation with three-body exact solution: the 1.2 relative error is only 9.05B-02-+ 6.88E-04. To get better convergence, we increase the randomization batch size to 50 for the 1M case. This incurs no memory cost and is only slightly slower than the original setting (speed is 46.80it/s
when randomization batch size is 16).
Eq.
Metric
100 D
IK
10K D
100K D
1M D
Speed
1938.80it/s
1840.21it/s
1291.67it/s
356.76it/s
46.97it/s
Allen-Cahi
Memory
461MB
481MB
539MB
1055MB
6233MB
9.97E-03
1.43E-03
6.21E-04
1.56E-05
2.25E-06
Error
+38898-04
11.60E-04
18.15
13.28E-07
+1.8EE-07
Speed
1991.28iu/s
1872.31it/s
1276.21it/s
364.04it/s
31.73it/s
Poisson
Memory
473MB
481MB
539MB
1055MB
6233MB
1.00E-02
1,02E-03
1,01E-04
9.26E-02
4.82E-06
Error
+3.27B-04
+3.67E-05
-2.40B-07
+5.36E-04
2 16E-07
Speed
1938.80iu's
1840.21ius
1291.6ittss
356.76it/
46,88it/s
Sine-Gordon
Memory
475MB
479MB
539MB
1063MB
6233MB
9.97E-03
1,43E-03
6.21E-04
1.56E-05
2.31E-05
Error
3.89E-04
+ 1..60--04
-8.15E-00
13.28E-07
11.4EE006
Table 7: Computational results for the Time-dependent Semilinear Heat equation, where the number of SDGD sampled dimensions is set to 10.
Method
Metric
10
100 D
IK D
10K D
Speed
Backward mode SDGD (Py (orch)))
Memory
Error
1.052E-03
5.263E-04
6.910E-03
1.598F-03
Speed
211,63it/
207.66ius
188.31it/s
93.21iu/
BackwardBackward mode SDGD (JAX)
Memory
619MB
621MB
655MB
1371MB
8.55E-05
4.02E-04
3.81E-04
2.60E-03
Eror
16.75E-05
+2.07E-04
14.43E-04
11.38E-03
Speed
660.82it/s
635.16it/s
599.15it/s
361.11it/s
STDE
Memory
625MB
625MB
657MB
971MB
6.99E-05
3.69E-04
3.38E-04
6.08E-03
Error
+.5.788-05
12.194-04
+3.30EE-14
+7.47E-03

---

Table 8: Computational results for the Time-dependent Allen-Cahy equation, where the number of SDGD sampled dimensions is set 10.
Method
Metric
D0
100 D
IKD
10K D
Speed
BackwardBackward
mode SDGD (PyTorch) 13
Memory
Error
7.815E-04
3.142F-04
7.042F-04
2.477E-04
Speed
211 211.38ss/s
206.42it/s
188.02it/s
93.20it/s
Backward mode SDGD (JAX)
Memory
619MB
621MB
657MB
1371MB
6.31E-02
4.38E-03
1.35E-03
3.97E-04
Error
+3.79E-02
+2.48E-03
+1.33E-03
+3.03E-04
Speed
677.51it/s
650.98it/s
598.33it/s
361.31it/s
STDE
Memory 533MB
535MB
657MB
903MB
6.37E-02
4.38E-03
1.26E-03
3.79E-04
Error
+3.77E-02
+2,47E-03
+1.29E-03
+2.75E-04
Table 9: Computational results for the Time-dependent Sine-Gordon equation, where the number of SDGD sampled dimensions is set 10.
Method
Metric
D
100 D
IKD
10K D
Speed
BackwardBackward
mode SDGD (PyTorch) 1
Memory
Error
7.815E-04
3.142E-04
7.042E-04
2.477E-04
Speed
210.83it/s
207.44ius
187.98it/s
93.17iu'
BackwardBackward mode SDGD (JAX)
Memory
619MB
621MB
655MB
1371MB
5.39I-05
9.15E-05
4.19E-04
3.74E-02
Error
+0.-055
+6.06E-05
+2.18E-04
+4.15E-02
Speed
629.
629.04it/s
608.83iv/s
596.12it/s
365.09it's
STDE
Memory
525MB
539MB
655MB
971MB
4,15E-05
2.54E-04
4,05E-03
1.66E-02
Error
+3.21E-05
+1.76E-04
+1.44--02
5.955F-03
can see that weight sharing drastically reduces the number of network parameters and memory usage. With B 50, there 2.5x reduction in memory and there no performance loss in terms of L2 relative error.
From Table TO
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 100x. For 5M, the maximum block size one can use before degrading performance is 500, which is expected the dimensionality of the problem higher.
From Table 10 we can also see that in the 5M-dimensional case, we will have an out-of off-mmmooyyy (OOM) error without weight sharing. With weight sharing enabled, we can effectively solve the E.M-dmmeniional PDE with good relative L2 error, in around 30 minutes.
1.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 Badomtsev-Petviashhe equation from (32). and the regular ID KdV equation with gPINN [4]
25

---

Table 10: Effects of different weight sharing block sizes B for the Inseparable Allen- Cahn equation with two-body exact solution solved with STDE with randomization batch size of 16. B equals no weight sharing.
dim
10
50
100
500
1000
Speed
21.34it/s
16.67it/s
23.14it/
23.73it/s
25.47it/s
26.60it/s
Memory
6295MB
4819MB
2505MB
2461MB
2409MB
2403MB
IM
#Param.
128,033,281
12,833,292
2,593,332
1.313,382
289,782
162,282
3,99E-03
1.86E-02
4.76E-03
1.22E-03
2.57E-03
6.06E-01
Error
13.41E-05
13.13E-04
11.27E-04
6.00EE-0
11.15E-04
14.17E-04
Speed
0OM
3.16it/s
4.47it/s
4.74it/s
4.82it/s
4.76it/s
Memory
0OM
25023MB
10595MB
10359MB
10163MB
10143MB
5M
#Param.
640,033,281
64,033,292
12,833,332
6,433,382
1,313,782
674,282
5.11E-01
3.13E-03
3,94E-03
1.98E-03
6.27E-01
Error
OOM
14.01E-04
+2.34E-04
-2.22E-04
15,20E--05
+3.03E-04
We will demonstrate that STDE increases the speed for computing the mixed partial derivatives, avoids computing the entire derivative tensor. Since these equations are low-dimensional we do not need 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, () = 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 I1 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.41 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
Why Marray + 3()y Uxn 2uyy 0. can alternatively be computed with the pushforward of the following jets
(64)
3 ()
d (xx, O, ez ey. O,...), g(2) d u(x O, e.),, g(3) = uuu O, ey, 0).
(65)
All the derivative terms can be found in these output jets {3(0)]
U 3/) My ((() Uez = 3/4// /3, Uzy = (1) 110, Uyy = 3/2)
(66)
yog
(3) 3 Mizray (3 (1) 2804 /8400, Bty (3 (2)
Uyyy 1/3.
2D Kadomtsev-Pettiishhii (KP) equation Consider the following equation (u: + 6uu, + Upzz) a + 30 0.
(67)
which can be expanded as
Utx + 6u, Ur Guüzz + Urrzz 30 Uyp 0. All the derivative terms be computed with 5-jet, 4-jet, and 2-jet pushforward. Let
(68)
(1) ==d" e, ez, 0,
(2) =dd ez, 0,, 0,0,0)
(69)
(3) :=d2uuxx , 0).
Then all required derivative terms can be evaluated as follows.
Utr //// /10.
(2) Ur = -/11] Ugx = 3p2] Uneaz = 3(4]
(70)
Uyy 3
26

---

Gradient-enhanced ID Korteweg-de Vries (g-KdV) equation Consider the following equa- tion
Ut + W22 + zdugzz 0.
(71)
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 the solution. Specifically, the PINN loss (Eq. 24 is augmented with the term
PEPINN (fx(0)
Σ Nr j
Σ
a Dzy
R(x())
(72)
The total loss becomes
residuull CaPINN lyPINN
(73)
where CEPINN is the g-PINN penalty weight. To perform gradient-enhancement need compute the gradient of the residual:
R(x, t) == ut + udx + Ollaary
(74)
VR(a, t) = [wet + Ue
uite
Utx + Uz Mure + Otlyzzar
All the derivative terms can be computed with one 2-jet and two 7-jet pushforward. Let
( :=dd Cg., 0, 0,0,0,0,0000
( dx 0, e,, 0,0.0)
(75)
3 (3) ==dRu(x u(x, e,, 0)-
Then all required derivative terms can be evaluated as follows.
(1) = -1 Uza -12/ Uxzr
(1)
[3] Ugana 3 3[4]] Unzzzz
3
(1)
Utww (3 (11) 1/35, Uge (3 Uezzaw 1/5, Ut =3(3] Uzzzer
(76)
(3)
Uet
Table 11: Speed scaling for training low-dimensional high-order PDEs with different network sizes. The base network has depth L and width lower-order pushforwards.
128. STDE* the alternative scheme using
Speed (it/s)
network size
Base
L
16
256
512
1024
Backward
762.86
279.19
123.20
656.01
541.10
349.23
2D Kdv
STDE
1372.41
642.82
303.39
1209.30
743.75
418.13
STDE*
1357.64 64
606.43
272.01
1203.97
841.07
442.32
Backward
766.79
278.53
123.67
642.34
525.23
340.94
2D
STDE
1518.82
676.16
304.95
1498.61
1052.62
642.21
Backward 621.04
232.35
102.39
559.65
482.52
293.97
ID DDg-KA KdV
STDE
1307.27
593.21
253.48
1187.31
776.65
441.50
L.422 Amortized gradient-enhancee PINN for high-dimensional PDEs
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
a
82
R(x) 0r,
Σ
u(x)
u(x) =²(xx) f(x)
Oxz
Ox?
(77)
g3
0 0 <x() 3u² 33²((x) 2 u(x) Fxy
Σ
u(x) +
Ou, (xx)
tx,0n 0u?
27

---

With STDE randomization, we randomized the second order term 8² with index i sampled from
[1, d.. We can also sample the gPINN penalty terms. As mentioned Appendix F.1
we have
3
0,0,0,0,
dx,dr uxx) 1
(78)
We further have
8² dz?
u(x) = -34//3.
(79)
so the STDE of the Laplacian operator can be computed together with the above pushforward. With this pushforward, can efficiently amortize the gPINN regularization loss by minimizing the following upperbound on the original gPINN loss with randomized Laplacian
EYPINN (fx(()) ( IE.. J)
a
dz, u(x) 3x)(x) u(xx 0
ΣΣ Nr JEJ iEI
0z,0r u(x) + dz,
dx;" f(x)
(80)
03 dz 0z u(x) + θ u(x) 3x)(x) dz, dz,
u(x)
dz; f(x)
Nr JEJ NET
where J an independently sampled index set for sampling the gPINN terms. The total loss
Cresidual (0; [xxxxe (6) [ + EAPNN (1x ((x(1 (N N. I,J)
(81)
We call this technique amortized gPINN. The above formula applies all PDEs where the derivative operator is the Laplacian. For example, for the Sine-Gordon equation, we have
ÉGPINN (0) Nr ( (
8
2
(82)
ΣΣ JEJ iES
0x, u(x) + cos u(x) dx;
u(x) 0x, f(x)
N
We use CgPINN = 0.1, and to get better convergence, we train for 20K steps instead of 10K steps as in all other experiments in this paper. The results 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 JVP-HVP in the table. Specifically, compute the residual gradient apply one more JVP to the HVP-based implementation of Laplacian (Appendix (A.2). From the table, we that STDE-based amortized gPINN performs better than the JVP-HVP implementation, and both are more efficient than applying backward mode AD in for-loop. Furthermore, through amortizing we can apply gPINN high-dimensional PDE which was intractable.
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 E Iddd as follows:
tr(A) E E [V Av], ve Rd
(83)
where p(v) is isotropic, i.c. Ev-p(v) = I. Therefore, the trace can be estimated by Monte Carlo:
tr(A) 1 v Av;
(84)
where each V; E dd 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].
28

---

Table 12: Performance comparison of STDE-gPINN for digh-diimnsiooal inseparable PDEs. "None" the "gPINN method" column indicates that no gPINN loss was used.
Equation
gPINN method
Metric
100 D
IK D
10K D
100K D
Speed
256.75it/s
249.48it/s
108.80it's
61.04it/s
JVP-HVP
3.97B-02
1.02E-03
3.08E-04
1.39E-03
+3.98E-04
16.89E-05
#7.48E-06
+ +1.2E-05
Allen-
Speed
366.46it/s
324.60it/s
207.85it/s
155.40it/s
Cahn
STDE
4.34E-02
5,26E-04
1.25E-03
7.61E-04
Error
+3.72E-04
+2.26E-05
E4.07E-05
+ +1.03E-04
4.98F--02
6.32F-03
1.19E-04
5.43F-04
None
Error
+3.82F-04
+4.43E-05
+1.04F-05
+4.30F-06
Speed
1008.65ivs
188.10it/s
413.32iv/s
107.68iVs
JVP-HVP
1.85E-03
1.02E-03
1.79E-04
5.76E-04
Error
14.61E-05
16.89E-05
11006-05
11.37E-04
Sine- Gordon
Speed
1165.35it/s
948.99it/s
542.36it/s
210.75it/s
STDE
6.69E-03
1,12E-03
1.76E-04
1.55E-03
11.48E-04
+1.38E-05
-5.31B-06
+4.30E-05
4.74E-03
7.02E-04
1.31E-04
8.07E-04
None
Error
16.68E-05
11.691E05
+1,22E-05
14.01E-06
J.2 HTE 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, specific contraction that included 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 = vide for i 1, 2, d with probability 1/d, which isotropic. Hence HTE can also be defined as a push forward of random 2-jet that are isotropic.
We now write the computation of HTE as follows
d Σ (ux(()(v,,,)), 0), V, ~ p(v).
2
Ug
(85)
where 22, 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
J
ur OnzOn
uxx)
(86)
j=1
has the following unbiased estimator:
3 E Evpp(y) (v,00))
(87)
where is the d-dimensional normal distribution. Therefore its STDE estimator is
xxxx)
d
Σ 0,000)) Vy ~(0,
(88)
3N
K STDE with dense jets
K.1 STDE with second order dense jets as generalization of HTE
Suppose second-order differential operator with coefficient tensor C. If C not symmetric, can symmetrize it C' = CT), and (a) C (a) C' since D (a) is symmetric.
29

---

Furthermore, we can make C resitive-definie by adding a constant diagonal XI where -À is smaller than the smallest eigenvalue of C. The matrix C" 2 (C+ + C + AI then has the eigen decomposition UEUT where is diagonal and all positive. Now we have
[v--V(0 Uvv U*] = UEU = CC"
(89)
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 Lu Ei=1 8'u
The mask tensor of L. is the rank-4 identity tensor I E RRdxdxdxd, so the condition for unbiasedness
(1) (b) (e) Ex-ple (u,")) " 2 at MejKi deRRE, a,b,c,de (1,
(90)
where 6ight 1 when i =
and is otherwise.
In the most general case where / b / + d, we can sample E Ridd and split it into four Rd vectors. In this case we can define blocks of covariance as Ev ((())(()) = Lub, and Σ = [Xab] Denote the fourth-moment tensor of p as Hiikl, then Eq. 90 states that the block fa mbed in the fourth moment tensor should match C. Fourth moments can always be decomposed into second moments:
ab"
Mysai abodt E Mijkl = Lig 2ki + Eacyhd "Jl + Da 2jk
Dad
(91)
So finding the p that satisfies Eq. 90 is equivalent to finding a zero-mean distribution with covariance that satisfies the above equation. In the Mijui = Jiydig.tt So in the case where a & b, set a = 1, = 2, we have
of L, the mask tensor is black-diagonal:
Mejos
1212 = EE +
(92)
I1 221
1
and Σ
2dd2
Firstly, consider the diagonal entries of a:
2
agaa = 3( E [1,2}
(93)
This can always be satisfied by setting the diagonal entries of both Ma and ae block as fol- lows:
E = Von/3, a {1,2}
(94)
Next, consider the entire a matrix. We have
1212 = E 2リ + + Bi Mijo
3 Voonas + 2(
(95)
In the case of L, we have oij = óij. so for i
+ j we have
1 + 2(1))22
0 = 3
(96)
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 selected. For example, in 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 off-diagonal element of the Hessian.
of Laplacian, as also discussed in 112, the variance of the
L Further ablation study
30

---

Figure 5: Ablation on randomization batch size with Inseparable and effectively high-dimensional PDEs, 100= runs with different random seeds. Model converges when the difference of L.2 error is below le-7.
Allen-Cahn
Poisson
Sine-Gordon
10P
L2 relative error
10-1
100
L2 relative error
10-4
L2 relative etror
10"11
256
10
10-2
10
10
2000
4000
8000
2000
1000
8000
2000
1000
8000
Epoch
Epoch
Epoch
10a4
10
2014
10a4
10a4
1012
1012
10
1010
10
10
Resiuall
Residual
Residual
10a
10*
10a
109
10"
10*
10a
10a
2500
5000
7500
10000
2500
5000
7500
10000
2500
5000
7500
10000
Epoch
Epoch
Epoch
350
350
300
300
300
Nerrttoons per setand (itsis)
ler lterations pet secand (itss)
lterationss per second (itsi))
250
250
250
200
200
150
150
150
50
50
50
100
256
100
256
100
256
80
140
120
Convergence time (s)
Conversence Lime (s)
Canvergence time 6s)
60
1.00
50
80
90
60
30
40
20
10
20
100
256
100
256
255
31

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect paper's contributions and scope?
Question:
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 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 be expected to generalize 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 separate "Limitations" section their paper.
The paper should point 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 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, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails 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 fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best
32

---

judgment and recognize that individual actions in favor of transparency play an impor- tant 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 include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proof's 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 ex- perimental results of the paper to the extent that 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
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer this question will not 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 dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending 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 specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, provide access 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 replicate the results, access a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
33

---

(a) If the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., large language model), then there should either way to access this model for reproducing the results or 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 authors are welcome describe the particular way they provide for reproducibilityy 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 possible for other researchers have some path reproducing verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions 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 NeurlPS code and data submission guidelines (https://nips.cr
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply not including code, unless this is central to the contribution (e.g., for 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 httpppppwwwppppipeeeeeeeeeueeuue
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 subset of experiments reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification:
We have included the setting for all the hyperparameters in the Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
34

---

The experimental setting should be presented in the core of the paper to a level of detail that is necessary 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
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We the average of 5 random seeds for all our experiment results. We also reported the standard deviation for the relative error PINN training in the Appendix
Guidelines:
The answer NA means that the paper does include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, 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).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
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 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 range (e.g. negative етror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We have included the hardware and software specifications we used to conduct our experiments in the Appendix
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 't make into the paper).
35

---

9. Code Of Ethics
Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of httpp:///ppcrppss.cc////////laaa
Question:
Answer: [Yes]
Justification: We have read the NeurlPS Code of Ethics and made sure that the paper conforms to
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 deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations 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 tied to particular applications, and there are 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 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 direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that 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 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 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 monitor how 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 models that have 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.
36

---

Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have 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 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 challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
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?
Question:
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 dataset.
The authors should state which version of the asset is used and, if possible, include 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.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, wwpprswittcooeeeoooppppo thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and the documentation provided alongside the assets?
Answer: [NA]
Justification:
Our work does not release assets.
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details of the dataset/code/moodl 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.
37

---

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 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 research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According 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 research with human subjects.
Depending 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 the institution conducting the review.
38


Reviews:
[{"review_id": "EisrlY9oD0", "reviewer": "Reviewer_xY8N", "summary": "The 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.", "questions": "Given 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?", "limitations": "Yes, they have.", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The 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.", "weaknesses": "While 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Weaknesses**\n> While 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\nThanks for your suggestion on future improvement. You are right that there could exist a more efficient scheme tailored for specific operators like the Laplacian. However, to derive and implement a comprehensive set of optimized derivative estimators is a challenging task whose workload is akin to implementing an AD framework, therefore out of the scope of the current paper. Nevertheless, we hope that our paper can inspire future work in this direction. The same goes for variance reduction. We have done some preliminary work on variance reduction which will form a separate paper from this one.\n\n> Another 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\nThanks for pointing out this. We have conducted further ablation studies on the randomization batch size, please refer to the global rebuttal and the PDF attached.\n\n**Questions**\n> Given 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\nFor the NLS, the nonlinear term $|\\psi |^2\\psi$ does not contain derivatives, and the Laplacian term can be handled by STDE similar to the experiments we have done in this paper. As for the fourth-order NLS, the fourth-order term is a Biharmonic operator which can be handled by STDE as well as discussed in Appendix J.3, as well as in [12]. As for the Navier-Stokes equation, it is both low-dimensional and low-order, so STDE may not provide any significant acceleration. However, if we consider generalized Navier-Stokes equations that include higher-order spatial derivatives to account for dispersive effects, similar to generalized KdV equations discussed in Appendix I.4.1, then STDE might provide some acceleration.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MgSkBWw25n", "reviewer": "Reviewer_yjY8", "summary": "This 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.", "questions": "1. 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.", "limitations": "Some limitations are discussed in section 6.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. 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.", "weaknesses": "1. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Weaknesses**\n> 1. 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\nWe appreciate your input on writing. We agree that by moving the background section to the appendix we would have more space for experiment details, but doing this could also disturb the flow of the paper, as we discussed the motivation for our method when going through the background material.\n\n**Questions**\n> 1. 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.\n\nThanks for pointing out this issue. Yes, you are right, Forward Laplacian is an exact method, so it is expected to perform better in terms of L2 error. However, as we can see in Table 4, the L2 error is of the same order, at least in the case where the dimension is >1k. We will update the description of this section in the final version to reflect this point.\n\n> 2. 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\nRegarding computation cost, it is worth pointing out 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. The main differences between the sparse and the dense version of STDE are (1) sparse STDE is universally application whereas the dense STDE can only be applied to certain operators, as you've pointed out; (2) the source of variance is different (see Appendix K.3). In general, we would suggest to use sparse STDE unless we know a priori that the sparse version would suffer from excess variance and the dense STDE is applicable.\n\nFor constructing STDE for a specific problem, you could refer to the various examples provided in the paper, or follow the method outlined in Appendix F for sparse jet, or Appendix K.1 for dense jet.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Yzze56lKMv", "reviewer": "Reviewer_SkVJ", "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.", "questions": "NA", "limitations": "NA", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your positive review of our paper. We would appreciate it if you could provide any suggestions on potential improvements to our paper, and we welcome any future questions you might have on our paper.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ArQUq3L12m", "reviewer": "Reviewer_cUDy", "summary": "The 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.", "questions": "1.\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.", "limitations": "Yes.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "-\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.", "weaknesses": "1.\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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Weakness**\n> 1. Compared with original ...\n\nThanks for your insightful comment. We want to first clarify that, the main takeaway we wish to show out of this comparison between FL and STDE is to highlight the fact that randomization is important for scalability in dimensionality. While FL removes the redundancy in normal AD, being an exact method its complexity still grows with the dimensionality of the domain. We are not claiming that STDE beats FL, but rather to emphasize that randomization is key and there is potentially a lot of research to be done at the intersection of randomized linear algebra and automatic differentiation. In fact, from Table 2 one can see that in the low-dimensional regime, FL is clearly the best method both in terms of speed and memory.\n\nIt is certainly true that a randomized FL will perform much better in terms of scalability, but this entails non-trivial work. In general, for each specific differential operator, there could exist optimized forward rules like FL, as also pointed out by reviewer xY8N. We believe that developing a comprehensive set of such rules is a promising direction, but we think it is out of the scope of this paper.\n\nRegarding the accuracy reduction due to randomization, we would suggest the reviewer inspect the relative error data shown in Table 3. From the table, one can see that L2 error under randomization has the same order as exact methods like FL. However, to ascertain the stochastic error for solving a specific problem like the many-electron Schrödinger equations requires further investigation.\n\nFor the cost-accuracy curve you've mentioned, we have added additional experiments where the results are included in the rebuttal PDF. Please refer to the global rebuttal and the PDF attached.\n\n> 2. All the test ...\n\n- Besides the Laplacians and the Biharmonic operator, we also provide an example of amortized g-PINN in Appendix I.4.2 which is both high-dimensional and high-order. Furthermore, we have provided a general procedure for constructing sparse STDE for arbitrary derivative operators in Appendix F.\n- The high-order diagonal differential operator is a straightforward extension of the Laplacian operator.\n- We believe that the comparison between randomized AD and direct back-propagation is redundant as it is already done in previous work ([13]).\n- This work applies to any problem where one needs to compute a differentiable operator on the input of an NN, as explained in the introduction. STDE applies naturally to PINN 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. The method is not suited for computing the derivative of NN parameter as explained in Section 3.\n\n> 3. Some of the highly related existing works are missing...\n\nThank you for pointing out these relevant works that we are not aware of.\n- In [1], an exact AD method is proposed for computing arbitrary **second-order** differential operators. The idea is similar to our idea in equations (19) and (22) where matrix decomposition and change of basis are employed, but the overall algorithm is different. It is worth stressing again that our method applies to **arbitrary derivative order**. We will cite this paper in our final version since it is highly relevant.\n- Regarding section 3.4, this is just background material as it falls under the preliminaries section. We provided a detailed write-up here because it serves as an important context for our main contribution in section 4.1, and we want to set up the notation in a way that facilitates the explanation of our idea. It is true that [2] also discussed the geometric interpretation of Taylor-mode AD, but this is also discussed in earlier literature we cited in Section 2. Nevertheless, we will include [2] in related works in the final version for completeness' sake.\n\n> 4. As is discussed in Section 4.4...\n\nSection 4.4 section refers to the **dense** version of STDE which has limited applicability. The **sparse** STDE is universally applicable. Reviewer yjY8 also asked a question about the comparison between the two, please refer to my answer to that question for further clarification.\n\n> 5. The author didn’t ...\n\nThanks for bringing up this important question. How large $l$ should be depends on how off-diagonal the operator is. If the operator is diagonal as in the case of equation (17), $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 Section F where a general procedure for determining the jet structure is discussed.\n\n> 6. Some parts of the writing are confusing...\n\nThanks for your detailed checking on the correctness of the terminology.\n- When we say equation (18) is 'general' we mean that it includes a large class of commonly used second-order PDEs. We also mention that they are parabolic in the paragraph.\n- When we say 'inseparable and effectively high-dimensional', we are stressing the fact that the solution does not have strong symmetries such that the effective dimension can be drastically reduced via some change of variables. For example, if the solution has spherical symmetry, then the equation becomes a 1D under spherical coordinates.\n- For the equations in Appendix I.2, you are right that they are semilinear parabolic equations. However, they are also specifically driftless Fokker-Planck equations with semilinear extension, and Fokker-Planck equations are a subset of parabolic equations. Similarly, semilinear equations are nonlinear, but there are other types of nonlinearity like quasilinear. So I guess both terminologies are okay, and this is really just a personal preference.\n\n**Questions**\n\nThanks for these insightful questions. See our response to these questions in the global rebuttal.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
**Additional Experiments**

We have conducted further ablation studies on the randomization batch size (see the attached PDF). We ran all three equations from the Inseparable and effectively high-dimensional PDEs (Appendix I.1) with moderately high dimensions (100k), with different randomization batch sizes (1, 4, 16, 64, 100, 256). And we compared their final L2 error and convergence time. As expected, with a smaller batch size the iterations per second are higher (third row of the figure). Memory cost remains roughly the same since the computation graph was not changed. Rather surprisingly, the final L2 error and the time for convergence did not exhibit a linear relationship to randomization batch size, as can be seen in the first row and last row of the figure. Specifically, one can see that the batch size of 1 provides good L2 error and convergence time, regardless of the equation chosen. One explanation is that stochastic optimizers like Adam already have built-in mechanisms to control the variance and are robust to noise during training, so smaller batch sizes can perform well. This warrants further investigation that is out of the scope of this paper.

**Questions from reviewer cUDy**

Reviewer cUDy asked a lot of insightful questions. We answer them here as it would benefit all reviewers to better understand our papers.

> 1. Which function does 'univariate' refer to ...

Firstly we would like to clarify that, our contribution is the insight that the scalar version of Faà di Bruno's formula can be used to perform derivative tensor contractions for multivariate functions, and the vector-valued version of Faà di Bruno's formula is not needed, as discussed in section 4.1. Although the vector-valued Faà di Bruno's formula is introduced in Appendix D.2 (equation 41), it is **not** used in our method. We only use the univariate version (equation 42). In fact, one of our main insights was that for multivariate functions, the scalar version of Faà di Bruno's formula can be used to perform derivative tensor contractions,

As for your other question, in section 4.1, the multivariate function that we perform tensor contraction on is $F$. In section 3.4, the univariate function we are referring to is $F\circ g$, where the high-order chain rule is given by the scalar version of Faà di Bruno's formula.

> 2. In Appendix D ...

In Appendix D we discussed both the scalar and the vector version of Faà di Bruno's formula. The scalar version corresponds to the case of $n=1$ and the vector version corresponds to the case of $n>1$. In this work, we only consider $n=1$, which is enough for our method as discussed in section 4.1.

> 3. Although the authors claim...

Typically we would assume that the coefficient tensor has a certain structure. For example, in the case of equation 77, the operator $\frac{\partial^{3}}{\partial x_{j} \partial x_{i}^{2}}u$ is the diagonal of every rank-2 slice over the first axis of $D^3_u$. So one could sample the index $jii$ by sampling $j$ and $i$ separately from $[1,N]$. Even in the most general case where the index set to be sampled has no apparent structure, sampling the index set is still cheaper than computing the whole derivative tensor.

> 4. Could the authors ...

One example would be the many-body Schrödinger equations, where we need to compute a high-dimensional Laplacian. Another example is the high-dimensional Black-Scholes equation, which has numerous uses in mathematical finance.

> 5. Does STDE ..

In FL and [1], sparsity of intermediate derivative tensor are exploited to save memory and compute. The same idea can be applied to STDE, but it would require some modification to the JAX Taylor mode library. We have not yet done this, but it is an interesting future direction to work on.

> 6. Instead of tackling ...

In general, STDE provides a way to compute high-order derivative tensor elements with one forward pass. For high-order and low-dimensional operators, if the derivative tensor is sparse, then acceleration is possible, as shown by the experiments in section I.4.1.

In the case of Monge-Ampère equation, there is no sparsity since the operator $\det(\text{Hess} u)$ contains all entries from the second-order derivative tensor (Hessian matrix). One still might perform sampling among the additive terms in the determinant, which could provide acceleration when the dimension is high.

In the case of Ricci curvature, the operator is highly nonlinear since it involves computing the inverse of the metric tensor. In this case, STDE cannot be applied straightforwardly since obtaining an unbiased estimator usually requires linearity.

> On writing and notations:

> 1. Can you clarify ...

$d^2u_\theta$ does not denote the Hessian, but rather the second-order pushforward as described in section 3.4, equation 7. The input $(a,e_j,0)$ is a 2-jet, where $a$ is the primal, and $e_j$ and $0$ are the tangents.

The need for indexing notation like $[2]$ comes from the fact that the output is also a jet, which is a tuple that contains the primal and the tangents. We use the indexing notation to select a specific output tangent from the output jet.

> 2. In Eqn.(19), it should be $\sum_{i=1}^d$.

Yes, you are right, thanks for pointing out the typo.

> 3. In L247, do you mean $v\sim p$,, instead of $(a,v,0)\sim p$?

I think these two are synonymous, as $a$ here is not a R.V.

> 4. In the ‘100D’ column of table 2, 539 is not the smallest number. Why is it in bold?

Thanks for pointing out this typo, the smallest memory usage should be 507MB achieved by FL.

> 5. There needs to ...

In the Appendix, Einstein notation was first used and mentioned in equation 41.

As for equation 45, we are using the scalar version of Faà di Bruno's formula, so there are no $n,\ n’,\ n’’$ indices.

$\partial^4$ refers to the fourth-order Fréchet derivative, whose definition is given in equation 35. It is not a tensor. The inputs are the tangent vectors.

Meta-review: The authors describe an interesting combination of randomized approximation and auto-differentiation for high-order tensors arising from physics-informed neural network training.  The paper received positive reviews even before rebuttal, and the reviewers are clearly unanimous in their support at the end of the discussion period.  This seems like a clear accept.
Decision: Accept (oral)