Stochastic Taylor Derivative Estimator: Efficient amortization for arbitrary differential operators

Authors: zekun  shi
zheyuan  hu
min  lin
kenji  kawaguchi
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 08:29:46

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.

Main paper: 
 introduction: 
 In many problems, especially in Physics-informed machine learning [19, 32] , one needs to solve optimization problems where the loss contains differential operators: arg min θ f (x, u θ (x), D α (1) u θ (x), . . . ,EQUATIONIn this above,D α = ∂ |α| ∂x α 1 1 ,...,∂x α d d, α = (α 1 , α 2 , . . . , α d ) is a multi-index, u θ is some neural network parameterized by θ, and f is some cost function. When either the differentiation order k or the dimensionality d is high, the objective function above is expensive to evaluate with back-propagation (backward mode AD) in both memory and computation: the size of the derivative tensor has scaling O d k , and the size of the computation graph has scaling O 2 k-1 L where L is the number of ops in the forward computation graph.There have been several efforts to tackle this curse of dimensionality. One line of work uses randomization to amortize the cost of computing differential operators with AD over the optimization process so that the d in the above scaling becomes a constant for the case of k = 2. Stochastic Dimension Gradient Descent (SDGD) [13] randomizes over the input dimensions where in each iteration, the partial derivatives are only calculated for a minibatch of sampled dimensions with back-propagation. In [12, 21, 15] , the classical technique of Hutchinson Trace Estimator (HTE) [16] is used to estimate the trace of Hessian or Jacobian to inputs. Others choose to bypass AD completely to reduce the complexity of computation. In [30] , the finite difference method is used for estimating the Hessian trace. Randomized smoothing [11, 14] uses the expectation over Gaussian random variable as ansatz, so that its derivatives can be expressed as another expectation Gaussian random variable via Stein's identity [38] . However, compared to AD, the accuracy of these methods is highly dependent on the choice of discretization.In this work, we address the scaling issue in both d and k for the optimization problem in Eq. 1 at the same time, by proposing an amortization scheme that can be efficiently evaluated via high-order AD, which we call Stochastic Taylor Derivative Estimator (STDE). Our main contributions are:• We demonstrate how Taylor mode AD [6] , a high-order AD method, can be used to amortize the optimization problem in Eq. 1. Specifically, we show that, with properly constructed input tangents, the univariate Taylor mode can be used to contract multivariate functions' derivative tensor of arbitrary order;• We provide a comprehensive procedure for randomizing arbitrary differential operators with STDE, while previous works mainly focus on the Laplacian operator, and we provide abundant examples of STDE constructed for operators in common PDEs;• STDE encompass and generalizes previous methods like SDGD [13] and HTE [16, 12] . We also prove that HTE-type estimator cannot be generalized beyond fourth order differential operator;• We determine the efficacy of STDE experimentally. When applied to PINN, our method provides a significant speed-up compared to the baseline method SDGD [13] and the backward-free method like random smoothing [11] . Due to STDE's low memory requirements and reduced computation complexity, PINNs with STDE can solve 1-million-dimensional PDEs on a single NVIDIA A100 40GB GPU within 8 minutes, which shows that PINNs have the potential to solve complex real-world problems that can be modeled as high-dimensional PDEs. We also provide a detailed ablation study on the source of performance gain of our method. 
related works: 
 High-order and forward mode AD The idea of generalizing forward mode AD to high-order derivatives has existed in the AD community for a long time [5, 18, 39, 22] . However, accessible implementation for machine learning was not available until the recent implementation in JAX [6, 7] , which implemented the Taylor mode AD for accelerating ODE solver. There are also efforts in creating the forward rule for a specific operator like the Laplacian [23] . Randomization over the linearized part of the AD computation graph was considered in [29] . Forward mode AD can also be used to compute neural network parameter gradient as shown in [2] .Randomized Gradient Estimation Randomization [27, 28, 8] is a common technique for tackling the curse of dimensionality for numerical linear algebra computation, which can be applied naturally in amortized optimization [1] . Hutchinson trace estimator [16] is a well-known technique, which has been applied to diffusion model [36] and PINNs [12] . Another case that requires gradient estimation is when the analytical form of the target function is not available (black box), which means AD cannot be applied. The method of zeroth-order optimization [24] can be used in this case, as it only requires evaluating the function at arbitrary input. It is also useful when the function is very complicated like in the case of a large language model [26] .3 Preliminaries and discussions 3.1 First-order auto-differentiation (AD)AD is a technique for evaluating the gradient of composition of known analytical functions commonly called primitives. In an AD framework, a neural network F θ : R d → R d ′ is constructed as the composition of primitives F i that are parameterized by some parameters θ i . In this section, we will consider the neural networks with linear computation graphs likeF = F L • F L-1 • • • • • F 1 , butthe results generalize to arbitrary directed acyclic graphs (DAGs). We will assume that all hidden dimensions are h. See Appendix B for more details on first-order AD.Forward mode AD Each primitives F i is linearized as the Fréchet (directional) derivative ∂F i : R h → L(R h , R h ), which computes the Jacobian-vector-product (JVP): ∂F i (a)(v) = ∂F ∂x a v, where a is referred to as the primal and v the tangent. ∂F i form a linearized computation graph (third row in Fig. 3 ), that computes the JVP of the composition ∂F ∂x v:EQUATIONBy setting the tangent to v one of the standard basis of R d , JVP computes one column of the Jacobian D F , so the full Jacobian can be computed with d JVPs. Each JVP call requires O(max(d, h)) memory as only the current activation y i and tangent v i are needed to carry out the computation, and the computation complexity is usually in the same order as the forward computation graph. In the case of MLP, both the forward and the linearized graph have a complexity of O dh + (L -1)h 2 .Backward mode AD Each primitives F i is linearized as the adjoint of the Fréchet derivative ∂ ⊤ F i instead, which computes the vector-Jacobian-product (VJP):∂ ⊤ F i (a)(v ⊤ ) = v ⊤ ∂F∂x a where v ⊤ is the cotangent. The linearized computation graph now runs in the reverse order:EQUATIONwhich is also clear from Fig. 3 . Furthermore, due to this reversion, we first need to do a forward pass to obtain the evaluation trace {y i } L i=1 before we can invoke the VJPs ∂ ⊤ F i , which apparent as shown in Eq. 3. Hence the number of sequential computations is twice as much compared to forward mode. The memory requirement becomes O(d + (L -1)h) as we need to store the entire evaluation trace. Similar to JVP, VJP computes one row of J F at a time, so the full Jacobian ∂F ∂x can be computed using d ′ VJPs. When optimizing scalar cost functions ℓ(θ) : R n → R of the network parameters θ, backward mode efficiently trades off memory with computation complexity as d ′ = 1 and only 1 VJP is needed to get the full Jacobian. Furthermore, all parameter θ i can use the same cotangent v ⊤ , whereas with forward mode, separate tangent for each parameter θ i is needed.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes]Justification: We discussed the limitation of our work in section 6.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.• The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.• The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.• The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.• The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.• If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.• While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
inefficiency of the first-order ad for high-order derivative on inputs: 
 second-order derivative by applying VJP twice High-order input derivatives ∂ k u θ ∂x k for scalar u θ can be implemented as repeated applications of first-order AD, but this approach will exhibit fundamental inefficiency that cannot be remedied by randomization.VJP Forward x ∈ R d y 1 ∈ R h y 2 y 3 θ 1 θ 2 θ 3 v ⊤ 1 ∈ R h v ⊤ 2 v ⊤ 3 v ⊤ ∈ R d ′ v ⊤ ∂F ∂x y = F (x) θ 4 ṽ⊤ ∈ R d ṽ⊤ 1 ṽ⊤ 2 ṽ⊤ 3 ṽ⊤ 4 ṽ⊤ 5 ṽ⊤ 6 ∂ 2 Fd ∂xi∂xj v d ṽjRepeating backward mode AD With each repeated application of backward mode AD, the new evaluation trace will include the cotangents from the previous application of backward AD, so the length of sequential computation doubles. Furthermore, the size of the cotangent also grows by d times. Therefore applying backward mode AD has additional memory cost of O(d + (L -1)h) and additional computation cost of O 2dh + 2(L -1)h 2 , which is clear from Fig. 1 . In general, with k repeated applications of backward mode AD will incur O 2 k-1 (d + (L -1)h) memory cost and O 2 k (dh + (L -1)h 2 ) computation cost. And O d k-1 calls are needed to evaluate the entire derivative tensor. So both memory and compute scale exponentially in derivative order k Repeating forward mode AD Consider u θ : R d → R. The input tangent dimension is d on the first application of forward mode AD, but on the second application, it will become d × d since we are now computing the forward mode AD for ∇u θ : R d → R d . So the size of the input tangent with k repeated application is O d k , so it grows exponentially. This is also inefficient.Mixed mode AD schemes are also likely inefficient See more detail in Appendix C.SDGD [13] amortizes high-dimensional differential operators by computing only a minibatch of derivatives in each iteration. It replaces a differential operator D with a randomly sampled subset of additive terms, where each term only depends on a few input dimensionsEQUATIONwhere DJ denotes the SDGD operator that approximates the true operator D, J is the sampled index set, and |J| is the batch size. For example, in d-dimensional Poisson equation,N D = d, D = d j=1 ∂ 2 ∂x 2 j, and the additive terms areD j = ∂ 2 ∂x 2 j .DJ are cheaper to compute than D due to reduced dimensionality: for each sampled index, by treating all other input as constant we get a function with scalar input and output. For a given index set J, the memory requirements are reduced fromO 2 k-1 (d + (L -1)h) to O |J|(2 k-1 (1 + (L -1)h)) ,and the computation complexity reduces to O |J|2 k (h + (L -1)h 2 ) . This reduction is significant when d ≫ h as in the experimental setting of SDGD [13] , but the exponential scaling in k persists.One way to define high-order AD is by determining how the high-order Taylor expansion of a univariate function changes when mapped by primitives. Firstly, the Fréchet derivative ∂F can be rewritten to operate on a space curve g : R → R d that passes through the primal a, i.e. g(t) = a, and has tangent g ′ (t) = v:EQUATIONThis shows that the ∂ (JVP) is the same as the univariate chain rule. The tuple J g (t) := (g(t), g ′ (t)) can be thought of as the first-order expansion of g which lives in the tangent bundle of F . Treating F as the smooth map between manifolds, we can define the pushforward dF which pushes the first order expansion of g (i.e. J g (t)) forward to the first order expansion ofEQUATIONNaturally, to extend this to higher orders, one can consider the kth order expansion of the input curve g, which is equivalent to the tuple J k g (t) := (g(t), g ′ (t), g ′′ (t), . . . , g (k) (t)) = (a, v (1) , v (2) , . . . , v k ) known as the k-jet of g where v j is called the jth order tangent of g. J k g lives in the kth order tangent bundle of F , and we can define the kth-order pushforward d k F :EQUATIONwhich pushes the kth order expansion of g (i.e. J k g ) forward to the kth order expansion of F • g (i.e.J k F •g ). ∂ k F = ∂ k ∂t k [F • g](t)is the k-th order Fréchet derivative, whose analytical formula is given by the high-order univariate chain rule known as the Faa di Bruno's formula (Eq. 43).g contains all information needed to evaluate ∂ j ∂t j [F • g](t) for any j ≤ k, the map d k F is well-defined. d k defines a high-order AD: we can compute d k F of arbitrary composition F from the kth-order pushforward of the primitives d k F i , since d k is an homomorphism of the group ({F i }, •):EQUATIONThis approach of composing d k of primitives is also known as the Taylor mode AD. For more details on Taylor mode AD, see Appendix D.The number of sequential computations does not grow with the order as can be seen in Fig. 2 , and the computation can be trivially vectorized and parallelized since the pushforward of sample jets can be computed independently, and it uses the same computation graph (d k u);Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [Yes] Justification: Most of our theoretical results are asymptotic analyses on the computation complexity, and we have clearly stated the assumption we have made. Our claim on the non-generalizability of HTE construction is proved rigorously in the Appendix K.Guidelines:• The answer NA means that the paper does not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.• All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.• Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.• Theorems and Lemmas that the proof relies upon should be properly referenced. 
method: 
 From the previous discussion, it is clear that the exponential scaling in k for the problem described in Eq. 1 cannot be mitigated by amortization alone. Although high-order AD methods like Taylor mode AD [6] can address this scaling issue, it is only defined for univariate functions. In this section, we describe a method that addresses the scaling issue in k and d simultaneously when amortizing Eq. 1 by seeing univariate Taylor mode AD as contractions of multivariate derivative tensor. EQUATIONd2 .From Eq. 43, one can always find a J l g with large enough l ≥ k such that there existsk ≤ l ′ ≤ l with ∂ l ′ F (J l ′ g ) = D k F (a) • ⊗ k i=1 v (vi) where v i ∈ [1, k],by setting some tangents v (vi) to the zero vector. That is, arbitrary derivative tensor contraction is contained within a Fréchet derivative of high-order, which can be efficiently evaluated through Taylor mode AD.How large l should be depends on how off-diagonal the operator is. If the operator is diagonal (i.e. contains no mixed partial derivatives), l = k is enough. If the operator is maximally non-diagonal, i.e. it is a partial derivative where all dimensions to be differentiated are distinct, then the minimum l needed is (1 + k)k/2. For more details, please refer to Appendix F where a general procedure for determining the jet structure is discussed.EQUATIONi v(1) j 1) , v (2) ), andL = 4 k = 2EQUATIONEQUATION1 , v1 ) which is the subsequent column. Each row can be computed in parallel, and no evaluate trace needs to be cached. we can identify a kth order differential operator L with the following tensor dot productEQUATIONd i ∈ [1, d], i ∈ [1, k]is the tensor index on the ith axis, , and C(L) is a tensor of the same shape as D k u (a) that equals C α when d 1 , . . . , d k matches the multi-index α ∈ I(L) and 0 otherwise. We call C(L) the coefficient tensor of L. For example, the coefficient tensor of the Laplacian ∇ 2 is the d-dimensional identity matrix I. More complicated operators can be built as f (x, u, D k1 u, . . . , D kn u) where f is arbitrary function.Any derivative tensor contractions D k u (a)•C(L) can be estimated through random contraction, which can be implemented efficiently as pushing forward random jets from an appropriate distribution. With random (v (1) , . . . , v (k) ), we haveE[D k u (a) d1,...,d k v (v1) d1 . . . v (v k ) d k ] = D k u (a) d1,...,d k E[v (v1) d1 . . . v (v k ) d k ] = D k u (a) • E ⊗ k i=1 v (vi) (11)where ⊗ denotes Kronecker product, v(vi) di ∈ [1, k]is the d i dimension of the v i th order tangent in the input k-jet. Eq. 11 is an unbiased estimator of the kth order operatorEQUATION)For example, the condition for unbiasedness for the Laplacian vi) , so with a distribution p over the input l-jet J l g that satisfies the unbiasedness condition (Eq. 12), we have∇ 2 is E[v (a) v (b)⊤ ] = I. As discussed, one can always find a J l g with large enough l ≥ k such that ∂ l F (J l g ) = D k F (a) • ⊗ k i=1 v (EQUATIONwhich 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 k are addressed at the same time. From the example computation graph (Fig. 2 ) we see that, for one call to d k F , the memory requirement has scaling of O(kd) and the computation complexity has scaling O k 2 dL . Like first-order forward mode AD, the derivative tensor D k u is never fully computed and stored. Combined with randomization, the polynomial scaling in d will be removed.Note that all coefficient tensor has the following additive decomposition:EQUATION)where e i is the ith standard basis. For example, if the input dimension d is 3, then e 2 = [0, 1, 0] ⊤ . As discussed before, there exists a J k g whose pushforward under ∂ l u is equivalent to contracting D k u with ⊗ k i=1 e di . We call k-jet consisting of only standard basis and the zero vector 0 sparse. Therefore the discrete distribution p over the sparse k-jets in Eq. 14 satisfies the unbiasedness condition 12EQUATION)where Z is the normalization factor and we identify ⊗ k i=1 e di with the corresponding k-jet J k u . 4.3.1 Differential operator with easy to remove mixed partial derivatives Next, we show some concrete examples for constructing STDE with sparse random jets.Laplacian From Eq. 9 we know that the quadratic form of Hessian can be computed through ∂ 2 by setting v (2) = 0 and v (1) = e j . Therefore, the STDE of the Laplacian operator is given byEQUATION)where J is the sampled index set, and the subscript [2] means taking the second-order tangent from the output jet. See example implementation in JAX in Appendix A.4.High-order diagonal differential operators We call a differential operator diagonal if it is a linear combination of diagonal elements from the derivative tensor:L = d j=1 ∂ k ∂x k j .From Eq. 43 we see that setting the first-order tangent v (1) to e j and all other tangents v (i) to the zero vector gives the desired high-order diagonal element:EQUATION)General nonlinear second-order PDEs Second-order parabolic PDEs are a large class of PDEs. It includes the Fokker-Planck equation in statistical mechanics to describe the evolution of the state variables in stochastic differential equations (SDEs), which can be used for generative modeling [37] . It also includes the Black-Scholes equation in mathematical finance for option pricing, the Hamilton-Jacobi-Bellman equation in optimal control, and the Schrödinger equation in quantum physics and chemistry. Its form is given byEQUATIONWe have a second order derivative term 1 2 tr σ(x, t)σ(x, t) ⊤ ∂ 2 ∂x 2 u(x, t) with off-diagonal term. The off-diagonals can be easily removed via a change of variable:EQUATION)See derivation in Appendix E. Its STDE samples over the d terms in the expression above.It is not always possible to remove the mixed partial derivatives but discussed in section 4.2, for an arbitrary kth order derivative tensor element D k u (a) n1,...,n k , we can find an appropriate l-jet J l g (t) with g(t) = a such that ∂ l u(J l g ) = D k u (a) n1,...,n k . Here we show a concrete example. 2D Korteweg-de Vries (KdV) equation Consider the following 2D KdV equationEQUATION)All the derivative terms can be found in the pushforward of the following jet: 1) , . . . , v (13) ),EQUATIONEQUATION)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. 43). In this case, no randomization is needed since all the terms can be computed with just one pushforward. Alternatively, these terms can be computed with pushforwards of different jets of lower order (Appendix I.4). When input dimension d is high, randomization via STDE will provide significant speed up. We tested a few more high-order PDEs with irremovable mixed partial derivatives (see Appendix I.4), and the experimental results will be provided later.In section 4.3 we show how to construct STDE with the pushforward of sparse random jets. It is also possible to construct STDE with dense random jets, i.e. jets with tangents that are not the standard basis. For example, the classical method of Hutchinson trace estimator (HTE) [16] can be implemented in the STDE framework as the pushforward of isotropic dense random jets, i.e.(a, v, 0) ∼ δ a × p × δ with E p [vv ⊤ ] = I.We generalize the dense construction to arbitrary second-order differential operators using a multivariate Gaussian distribution with the eigenvalues of the corresponding coefficient tensor as its covariance. Suppose D is a second-order differential operator with coefficient tensor C. With the eigendecomposition C ′′ = 1 2 (C + C ⊤ ) + λI = UΣU ⊤ where -λ is smaller than the smallest eigenvalue of C , we can construct a STDE for D:E v∼N (0,Σ) [∂ 2 u(a)(Uv, 0)] -λE v∼N (0,I) [∂ 2 u(a)(v, 0)] = D 2 u (a) • [C ′′ -λI] = D 2 u (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 a counterexample: one cannot construct an STDE for the fourth order operatord i=1 ∂ 4∂x 4 with dense jets. For more details on dense jets, see Appendix K. For specific high-order operators like the Biharmonic operator, it is still possible to construct STDE with dense jets which we show in Appendix J.The main differences between the sparse and the dense version of STDE are: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 K.3).It is also worth noting that both the sparse and the dense versions of STDE would have similar computation costs if the batch size of random jets were the same. In general, we would suggest to use sparse STDE unless it is known a priori that the sparse version would suffer from excess variance and the dense STDE is applicable.Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?Answer: [Yes]Justification: We have included sample implementations of key steps of our method in the first section in the Appendix A.Guidelines:• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.• If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.• Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.• While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example • The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.• The full details can be provided either with the code, in appendix, or as supplemental material. 
experiments: 
 We applied STDE to amortize the training of PINNs on a set of real-world PDEs. For the case of k = 2 and large d, we tested two types of PDEs: inseparable and effectively high-dimensional PDEs (Appendix I.1) and semilinear parabolic PDEs (Appendix I.2). We also tested high-order PDEs (Appendix I.4) that cover the case of k = 3, 4, which includes PDEs describing 1D and 2D nonlinear dynamics, and high-dimensional PDE with gradient regularization [41] . Furthermore, we tested a weight-sharing technique (Appendix G), which further reduces memory requirements (Appendix I.3). In all our experiments, STDE drastically reduces computation and memory costs in training PINNs, compared to the baseline method of SDGD with stacked backward-mode AD. Due to the page limit, the most important results are reported here, and the full details including the experiment setup and hyperparameters (Appendix H) can be found in the Appendix.PINN [32] is a class of neural PDE solver where the ansatz u θ (x) is parameterized by a neural network with parameter θ. It is a prototypical case of the optimization problem in Eq. 1. We consider PDEs defined on a domain Ω ⊂ R d and boundary/initial ∂Ω as followsEQUATIONwhere L and B are known operators, f (x) and g(x) are known functions for the residual and boundary/initial conditions, and u : R d → R is a scalar-valued function, which is the unknown solution to the PDE. The approximated solution u θ (x) ≈ u(x) is obtained by minimizing the mean squared error (MSE) of the PDE residual R(x; θ) = Lu θ (x) -f (x):EQUATIONwhere the residual points {x (i) } Nr i=1 are sampled from the domain Ω. We use the technique from [25] that reparameterizes u θ such that the boundary/initial condition Bu(x) = g(x) are satisfied exactly for all x ∈ ∂Ω, so boundary loss is not needed. Amortized PINNs PINN training can be amortized by replacing the differential part of the operator L with a stochastic estimator like SDGD and STDE. For example, for the Allen-Cahn equation, Lu = ∇ 2 u + u -u 3 , the differential part of L is the Laplacian ∇ 2 . With amortization, we minimize the following lossEQUATIONwhich is a modification of Eq. 24. Its gradient ∂ lresidual ∂θ is then an unbiased estimator to the gradient of the original PINN residual loss, i.e. E[ ∂ lresidual ∂θ ] = ∂ℓresidual ∂θ .To ascertain the source performance gain of our method, we conduct a detailed ablation study on the inseparable Allen-Cahn equation with a two-body exact solution described in Appendix I.1. The results are in Table 1 and 2 , where the best results for each dimensionality are marked in bold. All methods were implemented using JAX unless stated. OOM indicates that the memory requirement exceeds 40GBs. Since the only change is how the derivatives are computed, the relative L2 error is expected to be of the same order among different randomization methods, as seen in Table 3 in the Appendix. We have included Forward Laplacian which is an exact method. It is expected to perform better in terms of L2 error. However, as we can see in Table 3 , the L2 error is of the same order, at least in the case where the dimension is more than 1000. Parallelization The original SDGD implementation uses a for-loop to iterate through the sampled dimension (Appendix A.1). This can be parallelized (denoted as "Parallelized SDGD via HVP", details in Appendix A.2). Parallelization provides ∼15× speed up and reduction in peak memory for the JIT compilation phase. We also tested mixed mode AD (dubbed as "Forward-over-Backward SDGD"), which gives roughly the same performance as parallelized stacked backward mode, which is expected as explained in Appendix C.Forward Laplacian Forward Laplacian [23] provides a constant-level optimization for the calculation of Laplacian operator by removing the redundancy in the AD pipeline, and we can see from Table 1 and 2 that it is the best method in both speed and memory when the dimension is 100. But since it is not a randomized method, the scaling is much worse. Its computation complexity is O(d), whereas a randomized estimator like STDE has a computation complexity of O(|J|). Naturally, with a high enough input dimension d, the difference in the constant prefactor is trumped by scaling. When the dimension is larger than 1000, it becomes worse than even parallelized stacked backward mode SDGD.STDE Compared to the best realization of baseline method SDGD, the parallelized stacked backward mode AD, STDE provides up to 10× speed up and memory reduction of at least 4×. 
conclusion: 
 We introduce STDE, a general method for constructing stochastic estimators for arbitrary differential operators that can be evaluated efficiently via Taylor mode AD. We evaluated STDE on PINNs, an instance of the optimization problem where the loss contains differential operators. Amortization with STDE outperforms the baseline methods, and STDE also applies to a wider class of problems as it can be applied to arbitrary differential operators.Applicability Besides PINNs, STDE can be applied to arbitrarily high-order and high-dimensional AD-based PDE solvers. This makes STDE more general than a branch of related methods. STDE is also more applicable than deep ritz method [40] , weak adversarial network (WAN) [42] , backward SDE-based solvers [3, 33, 10] , deep Galerkin method [34] , and the recently proposed forward Laplacian [23] , which are all restricted to specific forms of second-order PDEs. STDE applies naturally to differential operators in PDEs, but it can also be applied to other problems that require input gradients. For example, adversarial attacks, feature attribution, and meta-learning, to name a few.Limitations Being a general method, STDE forgoes the optimization possibilities that apply to specific operators. Furthermore, we did not consider variance reduction techniques that could be applied, which can be explored in future works. Also, we observed that lowering the randomization batch size improves both speed and memory profile, but the trade-off between cheaper computation and larger variance needs further analysis. Furthermore, the method is not suited for computing the high order derivative of neural network parameter as explained in Section 3. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes] Justification: We use the average of 5 random seeds for all our experiment results. We also reported the standard deviation for the relative error in PINN training in the Appendix I.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.• The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).• The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)• The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean.• It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.• For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).• If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes]Justification: We have included the hardware and software specifications we used to conduct our experiments in the Appendix H.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.• The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.• The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes]Justification: We have read the NeurIPS Code of Ethics and made sure that the paper conforms to it.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.• The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: Our work is not tied to particular applications, and there are no obvious paths that lead to potential harm.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.• The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.• If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: Our work is foundational and not tied to particular applications.• At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA] Justification: Our paper did not involve crowdsourcing and human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.• Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.• According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector. 
institutional review board (irb) approvals or equivalent for research with human subjects: 
 Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?Answer: [NA] Justification: Our paper did not involve crowdsourcing and human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.• Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.• We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.• For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: The key insight of the STDE construction is that the univariate Taylor mode AD contains arbitrary contraction of the derivative tensor and that derivative operators are derivative tensor contractions. This shows the connection between the fields of AD and randomized numerical linear algebra and indicates that further works in the intersection of these two fields might bring significant progress in large-scale scientific modeling with neural networks. One example would be the many-body Schrödinger equations, where one needs to compute a high-dimensional Laplacian.Another example is the high-dimensional Black-Scholes equation, which has numerous uses in mathematical finance.A.1 PyTorch implementation of SDGD-PINN using backward mode ADThe original implementation of SDGD-PINN [13] computes the SDGD estimation of derivatives using a for-loop that iterates over the sampled PDE term/dimension. For example, given a function f representing the MLP PINN, the computation of SDGD for the Laplacian operator can be implemented in PyTorch as follows: After computing the PDE differential operator, it is plugged into the residual loss, and then backwardmode AD is employed to produce the gradient for optimization concerning θ.f_x =def hvp(f, x, v):"""stacked backward-mode Hessian-vector product""" return jax.grad(lambda x: jnp.vdot(jax.grad(f)(x), v))( The jet.jet function from JAX implements the high-order pushforward d n of jets in Eq. 7. It decomposes the input function into primitives, which have analytical derivatives derived up to arbitrary order, and uses the generalized chain rule (see section D.2) to compose the primitives into the pushforward of jets. Note that in the API of jet.jet, all the high-order tangents of the input jet are specified via the series argument. x y 1 y 2 y 3 y = F (x) F 1 F 2 F 3 F 4 v v 1 v 2 v 3 ∂F ∂x v ∂F 1 ∂F 1 ∂F 2 ∂F 2 ∂F 3 ∂F 3 ∂F 4 ∂F 4 θ 1 θ 2 θ 3 θ 4 F 1 F 2 F 3 F 4 x y 1 y 2 y 3 y = F (x) F 1 F 2 F 3 F 4 v ⊤ ∂F ∂x v ⊤ 3 v ⊤ 2 v ⊤ 1 v ⊤ ∂ ⊤ F 4 ∂ ⊤ F 4 ∂ ⊤ F 3 ∂ ⊤ F 3 ∂ ⊤ F 2 ∂ ⊤ F 2 ∂ ⊤ F 1 ∂ ⊤ F 1 θ 1 θ 2 θ 3 θ 4 F 1 F 2 F 3 F 4First-order AD is based on a simple observation: for a set of functions L, the set of tuples of functions f and its Jacobian J f is closed under composition:EQUATIONwhere • denotes both function composition and the composition of the tuple (f, J f ). If we have the analytical formula of the Jacobian J f for every f ∈ L, then we can calculate the Jacobian of any composition of functions from L using the above composition rule for the tuple (f, J f ). The set L of functions are usually called the primitives.Given normed vector spaces V, W , the Fréchet derivative ∂f of a function f : V → W is a map from V to the space of all bounded linear operators from V to W , denoted as L(V, W ), that isEQUATIONsuch that at a point a ∈ V it gives the best linear approximation ∂f (a)(•) of f , in the sense thatEQUATIONTherefore, it is also called the linearization of f at point a. Concretely, consider a function in Euclidean spaces f : R n → R m . At any point a ∈ R n , the Fréchet derivative ∂f can be seen as the directional derivative of f :EQUATIONwhere J f (a) ∈ R m×n denote the Jacobian of f at point a called the primal, and v ∈ R n , also called the tangent is a vector representing the direction. Therefore the Fréchet derivative is also called Jacobian-vector-product (JVP). And we can write the truncated Taylor expansion asEQUATION)Many operators have efficient JVP implementation due to sparsity. For example, element-wise application of scalar function (e.g. activation in neural networks) has diagonal Jacobian, and its JVP can be efficiently implemented as a Hadamard product. Another prominent example is discrete convolution, whose JVP has efficient implementation via FFT.Given two topological vector spaces X, Y , the linear map u : X → Y has an adjoint t u : Y ′ → X ′ where X ′ , Y ′ are the dual spaces. The adjoint satisfies the followingEQUATIONIn the finite-dimensional 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: t u(y ⊤ ) = y ⊤ A.The adjoint (transpose) of the Fréchet derivative of f : R n → R m , denoted as ∂ ⊤ f , is thus defined as∂ ⊤ f : R n → L(R m , R n ), ∂ ⊤ f (a)(v) = v ⊤ J f (a), v ∈ R m (32)where v ⊤ is the cotangent which lives in the dual space of the codomain. Note that the adjoint is taken to v only where a is kept fixed. ∂ ⊤ f is also called vector-Jacobian-product (VJP).C Why mixed mode AD schemes like the forward-over-backward might not be better than stacked backward mode AD in the case of PINNIn AD literature [9] , the second order derivative is recommended to be computed via forward-overbackward AD, i.e., first do a backward mode AD to get the first order derivative, then apply forward mode AD to the first order derivative to obtain the second order derivative. Usually, we will expect that forward-over-backward AD gives better performance in memory usage over stacked backward AD since the outer differential operator has to differentiate a larger computation graph than the inner one, and forward AD has less overhead as explained in section B.2. Essentially, forward-overbackward reverses the arrows in the third row in Fig. 1 , therefore reducing the number of sequential computations and also the size of the evaluation trace. However, in the case of PINN, yet another differentiation to the network parameters θ needs to be taken. So, computing the second-order differential operator here with forward-over-backward AD might not yield any advantage.The kth order Fréchet derivative of a function f : R n → R m at a point a is the multi-linear map with k arguments around point a that best approximates f . For example, when k = 2, we haveEQUATIONwhere H f (a) ∈ R m×n×n denote the Hessian of f at point a, and v, v ′ ∈ R n . We can now write the second-order truncated Taylor series with itEQUATION)For the more general case, we haveEQUATIONHigh-order Fréchet derivative can be seen as the best kth order polynomial approximation of f by taking all input tangents to be the same v ∈ R n :EQUATIOND.2 Composition rule for high-order Fréchet derivatives Next, we derive the higher-order composition rule by repeatedly applying the usual chain rule.For composition f (g(x)) of scalar functions, we can generalize the chain rule for high-order derivatives by iteratively applying the chain rule to lower-order chain rules:EQUATIONwhere we give the example of up to the third order. For arbitrary k, the kth order derivative of the composition is given by the Faa di Bruno's formula (scalar version)EQUATION)where the outermost summation is taken over all partitions of the derivative order k. Here a partition of k is defined as a tuple (p 1 , . . . , p k ) ∈ N k that satisfiesEQUATION)For vector-valued functions g :EQUATIONwe can derive the following composition rule similarlyEQUATIONEQUATIONEQUATIONEQUATIONwhere again we give the example of up to the third order, and repeated indexes are summed as in Einstein notation. The general formula is again given by the multivariate version of the Faa di Bruno's formula. Note that in the multivariate version of the Faa di Bruno's formula, it is possible to take a derivative to distinguishable variables, but here we just present the version with indistinguishable input variables. This gives the composition rule for kth order total derivative.The composition of the high-order Fréchet derivative ∂ k is the case of n = 1, as the contraction with the input tangents v (i) ∈ R d is the same as composing with a scalar input function g : R → R d with v (i) = D i g . All derivative tensors of f (g(x)) can be represented using a R l vector, and similarly all derivative tensor v (i) of g can be represented using a R m vector. Then, the above chain rule can be simplified toEQUATIONEQUATIONEQUATIONEQUATIONEQUATION)The Faa di Bruno's formula again gives the general formula for arbitrary derivative order∂ k ∂t k f (g(t)) = (p1,...,p k )∈N k , k i=1 i•pi=k k! k i p i !(i!) pi • D k i=1 pi f (a) l,m1,...,m k i=1 p i • k j=1 1 j! v (j) mj pj ∈ R l .(43) which is written in the perspective of input primal a and tangents v (i) .E Removing the mixed partial derivatives term from second order semilinear parabolic PDEEQUATIONF Evaluating arbitrary mixed partial derivatives F.1 A concrete example Let's first consider a concrete case. Suppose the domain is D-dimensional we want to compute the mixed derivative ∂ ∂x 2i ∂xj . The naive approach would be to compute the entire third order derivative tensor D 3 f , which is a tensor of shape D × D × D, then extract the element at index (j, i, i). However note that from Eq. 43, for any k > 3, the pushforward of k-jet under d k f contains contractions of D 3 f . Although in the case of k = 3, the only contraction of D 3 f is in the ∂ 3 f :EQUATIONEQUATIONwhich can only be used to compute the diagonal or the block diagonal elements, when k > 3, we will have a contraction that computes off-diagonal terms, i.e. the mixed partial derivatives. For example, in d 4 f , if all input tangents are set to zero except for v (1) and v (2) , ∂ 4 f becomes:3•D 2 f (a) l,m1,m2 v (2) m1 v (2) m2 +6•D 3 f (a) l,m1,m2,m3 v (2) m1 v (1) m2 v (1)m3 +D 4 f (a) l,m1,m2,m3,m4 v (1) m1 v (1) m2 v (1) m3 v (1) m4 .(46) which contains the contraction of D 3 f that we want:EQUATIONHowever, 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:EQUATIONIf we go to higher-order jets, we can use more flexible contractions, and we can compute the mixed derivative with fewer terms to correct, hence less pushforwards. For example, the pushforward of the fifth-order tangent isEQUATIONif all input tangents are set to zero except for v (1) and v (3) . With this we only need to remove one term:EQUATIONSimilarly, by going to the seventh-order tangent, we can compute this mixed derivative with only one pushforward. d 7 f contains ∂ 7 , and when all input tangents are set to zero except for v (2) and v (3) ,∂ 7 equals 105 • D 3 f (a) l,m1,m2,m3 v (3) m1 v (2) m2 v (2)(51) which is the exact contraction we want. With this we have ∂ ∂x 2 i ∂x j u θ (x) = ∂ 7 u θ (x)(0, e i , e j , 0, 0, 0, 0)/105.F.2 Procedure for finding the right pushforwards for arbitrary mixed partial derivatives More generally, consider the case where we need to compute arbitrary mixed partial derivativeEQUATION)where T is the number of different input dimensions in the mixed partial derivative, and q it is the order. To compute it with k-jet pushforward, one needs to find:1. a derivative order k ∈ N, 2. a sparsity pattern for the tangents v (i) of the input jet, which is defined as the tuple of T integers J = (j 1 , . . . , j T ) where v (j) = 0 when j ̸ ∈ J and j t < k for all t ∈ [1, T ],such that when settingEQUATION(p 1 , p 2 , . . . , p k ) ∈ N k is a partition of k as defined in Eq. 39.Let's use the concrete example ∂ ∂x 2 i ∂xj again. In this case T = 2, q i1 = 2 and q i2 = 1. We demonstrated that this can be computed with one 7-jet pushforward, which is equivalent to setting J = (2, 3), k = 2j 1 + j 2 = 7, and the partition (0, 2, 1, 0, 0, 0, 0). The Faa di Bruno's formula (Eq. 43) ensures that the pushforward of the 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 are no extra terms to remove. Intuitively, if a partition has a sparsity pattern that is not a subset, it will vanish when we set the input tangents to zero according to the sparsity pattern of the partition in consideration. To understand this point better, let's look at the concrete example with the 5-jet pushforward demonstrated above. (2, 0, 1, 0, 0) and (5, 0, 0, 0, 0) are both valid partition of k = 5, and the sparsity pattern of (5, 0, 0, 0, 0) is the subset of that of (2, 0, 1, 0, 0): p 1 are non-zero in both partition. Therefore the pushforward contains extra terms that can be removed with another pushforward. In the example with 7-jet pushforward, no other partition has the sparsity pattern that is the subset of that of the partition (0, 2, 1, 0, 0, 0, 0). This is equivalent to say, 2 + 2 + 3 is the only way to sum up to 7 when you can only use 2 and 3, which can be verified easily.With this setup, it is clear why the diagonal terms can always be computed with pushforward of the lowest possible order: (k, 0, . . . , 0) ∈ N k is always a valid partition k, and no other partition has sparsity pattern that is a subset of it.For mixed partial derivatives, the difficulty scales the total order of the operator T t=1 q it , and T which can be interpreted as the degree of the "off-diagonalness" of the operator. For example, consider the case where T = 3 and q i1 = 3, q i1 = 2, q i1 = 1. This corresponds to the operator ∂ ∂x 3 i ∂x 2 j ∂x k . To avoid overlapping with the diagonal sparsity pattern (k, 0, . . . , 0) and to keep the order of derivative low, one might try k = 16 and the partition (0, 3, 2, 1, 0, . . . ) ∈ N 16 . However, with higher k, there is more chance that other partitions will have a subset sparsity pattern. In this case (0, 8, 0, 0, 0, . . . ) ∈ N 16 is one such example. One will need to either find all the partitions with subset sparsity pattern and remove them with multiple pushforward, or further increase the derivative order to find a pattern with no extra term.When dealing with high-dimensional data, the parameters of the model's first layer in a conventional fully connected network would grow proportionally with the input dimension, resulting in a significant increase in memory requirements and forming a memory bottleneck due to massive model parameters. To address this issue, convolutional networks are often employed in deep learning for images to reduce the number of model parameters. Here, we adopt a similar approach to mitigate the memory cost of model parameters in high-dimensional PDEs, called weight sharing in the first layer.Denote the input dimension as d, which is potentially excessively high, and the hidden dimension of the MLP as h, and assume that d ≫ h. The first layer weight is an d × h dimensional matrix, whereas all subsequent layers have a weight matrix with a size of only h × h.By introducing a weight-sharing scheme, one can reduce the redundancy in the parameters in the first layer. Specifically, we perform an additional 1D convolution to the input vectors x i before passing the input into the MLP PINN, as in Fig. 4 . The 1D convolution has filter size B that divides D and stride size B, so the convolution output is non-overlapping, and the number of channels is set to 1.This weight-sharing scheme reduces the parameters by approximately 1 B . The number of parameters in the filters is B × 1, and the subsequent fully connected layer will have a weight matrix of size We will demonstrate the memory efficiency and acceleration thanks to weight-sharing in the experimental section.x 1 x 2 x 3 x 4 x 5 x 6 x 7 x 8 x 9 Each experiment is run with five different random seeds, and the average and the standard deviations of these runs are reported.To get an accurate reading of memory usage, we use a separate run where GPU memory pre-allocation for JAX is disabled through setting the environment variable XLA_PYTHON_CLIENT_ALLOCATOR=platform, and the test data set is stored on the CPU memory. The GPU memory usage was obtained via NVIDIA-smi and peak memory was reported.All the experiments were done on a single NVIDIA A100 GPU with 40GB memory and CUDA 12.2. with driver 535.129.03 and JAX version 0.4.23.For the semilinear parabolic PDEs tested in Appendix I.2 we follow the network architecture of the original SDGD [13] :• The network is a 4-layer multi-layer perceptron (MLP) with 128 hidden units activated by Tanh.• The network is trained with Adam [20] for 10K steps, with an initial learning rate of 1e-3 that linearly decays to 0 in 10K steps, where at each step we calculate the model parameters gradient with 100 uniformly sampled random residual points.• The model is evaluated using 20K uniformly sampled random points fixed throughout the training.We tested the performance of STDE on these equations, and the results are presented in Table 3 , 4, 5, 6. For the Allen-Cahn equation, we performed a detailed ablation study (Table 3 ), and we expect these results to generalize over these second-order PDEs.The gain by using JAX instead of PyTorch Since the original SDGD was implemented in PyTorch, we implemented the stacked backward mode without parallelization in SDGD dimensions in JAX for fair comparison (dubbed as "Stacked Backward mode SDGD in JAX" in Table 3 ). The for-loop over SDGD dimension is implemented using jax.lax.scan. Table 3 shows that, even with the original stacked backward mode AD, the speed of JAX implementation can be more than 10× faster when the dimension is high. The memory profile is similar. The difference could come from the fact that JAX uses XLA to perform Just-in-time (JIT) compilation of the Python code into optimized kernels. However, note that for the case of 100,000 dimensions, the JAX implementation of the stacked backward mode AD encountered an out-of-memory (OOM) error. This is because performing JIT compilation requires extra memory, and the peak memory requirement during JIT compilation is higher than that during training.We also tested the case where the STDE randomization batch size is reduced to 16. As seen in Table 3 , in the case of Allen-Cahn provides ∼2× speed up, without hurting performance. However, theoretically lowering the randomization batch size leads to higher variance. The trade-off between computational efficiency and stability in convergence warrants further studies.The second class of PDEs is the semilinear parabolic PDEs, where the initial condition is specified:EQUATIONwhere g(x) is a known, analytical, and time-independent function that specifies the initial condition, and T is the terminal time. We aim to approximate the solution's true value at one test point x test ∈ R d , at the terminal time t = T , i.e. at (x test , T ).We will consider the following operators• Semilinear Heat Eq.EQUATION)with initial condition g(x) = 5/(10 + 2∥x∥2 ),• Allen-Cahn equationEQUATION)with initial condition g(x) = arctan(max i x i ),• Sine-Gordon equationEQUATION)with initial condition g(x) = 5/(10 + 2∥x∥ 2 ),All three equation uses the test point x test = 0 and terminal time T = 0.3.We tested the weight-sharing technique mentioned in Section G.In this section, we evaluate the performance of the weight-sharing scheme described in Appendix G. We tested the best-performing method from Table 3 (STDE with small randomization batch size of 16) with different weight-sharing block sizes, on the inseparable Allen-Cahn equation with the two-body exact solution. From Table 10 , we can see that weight sharing drastically reduces the number of network parameters and memory usage. With B = 50, there is a 2.5x reduction in memory and there is no performance loss in terms of L2 relative error.However, from the experiments we can see that, in both the 1M and the 5M case, increasing the block size beyond 50 provides diminishing returns. For the 1M case, increasing B to 1000 affects the convergence quality, as the L2 relative error goes up by 100x. For 5M, the maximum block size one can use before degrading performance is 500, which is expected as the dimensionality of the problem is higher.From Table 10 we can also see that in the 5M-dimensional case, we will have an out-of-memory (OOM) error without weight sharing. With weight sharing enabled, we can effectively solve the 5M-dimensional PDE with good relative L2 error, in around 30 minutes.Here we demonstrate how to use STDE to calculate mixed partial derivatives in some actual PDE. We will consider the 2D Korteweg-de Vries (KdV) equation and the 2D Kadomtsev-Petviashvili equation from [31] , and the regular 1D KdV equation with gPINN [41] . We will demonstrate that STDE increases the speed for computing the mixed partial derivatives, as it avoids computing the entire derivative tensor. Since these equations are low-dimensional we do not need to sample over the space dimension.In this section, the equations are all time-dependent and the space is 2D, and we will omit the argument to the solution, i.e. we will write u(x, t) = u. To test the speed improvement, we run the STDE implementation against repeated backward mode AD on a Nvidia A100 GPU with 40GB memory. The results are reported in Table 11 . From the Table we see that STDE provides around ∼2× speed up compared to repeated application of backward mode AD across different network sizes.It is expensive to apply gradient enhancement for high-dimensional PDEs. For example, the gradient of the residual for the inseparable Allen-Cahn equation described in I.1 is given byEQUATIONWith STDE randomization, we randomized the second order term ∂ 2 ∂x 2 i with index i sampled from [1, d] . We can also sample the gPINN penalty terms. As mentioned in Appendix F.1, we have J = d 7 u(x, 0, e i , e j , 0, 0, 0, 0),EQUATIONWe further haveEQUATIONso the STDE of the Laplacian operator can be computed together with the above pushforward. With this pushforward, we can efficiently amortize the gPINN regularization loss by minimizing the following upperbound on the original gPINN loss with randomized Laplacian lgPINN ({x (i) } Nr i=1 , I, J)EQUATION)where J is an independently sampled index set for sampling the gPINN terms. The total loss is lresidual (θ; {x (i) } Nr i=1 , I) + lgPINN ({x (i) } Nr i=1 , I, J).We call this technique amortized gPINN. The above formula applies to all PDEs where the derivative operator is the Laplacian. For example, for the Sine-Gordon equation, we haveEQUATIONWe use c gPINN = 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 are reported in Table 12 . We implement the baseline method based on the best performing first-order AD scheme, the parallelized backward mode SDGD, which we denoted as JVP-HVP in the table. Specifically, to compute the residual gradient we apply one more JVP to the HVP-based implementation of Laplacian (Appendix A.2). From the table, we see that STDE-based amortized gPINN performs better than the JVP-HVP implementation, and both are more efficient than applying backward mode AD in a for-loop. Furthermore, through amortizing we can apply gPINN to high-dimensional PDE which was intractable.In this section we establish the connection between the classical technique of HTE [16] and STDE by demonstrating that HTE is a pushforward of dense isotropic random 2-jet.J.1 Review of HTE HTE provides a random estimation of the trace of a matrix A ∈ R d×d as follows:EQUATIONwhere p(v) is isotropic, i.e. E v∼p(v) [vv T ] = I. Therefore, the trace can be estimated by Monte Carlo:EQUATIONwhere each v i ∈ R d are i.i.d. samples from p(v).There are several viable choices for the distribution p(v) in HTE, such as the most common standard normal distribution. Among isotropic distributions, the Rademacher distribution minimizes the variance of HTE. The proof for the minimal variance is given in [35] . J.2 HTE as the pushforward of dense isotropic random 2-jets Note that both HTE and the STDE Hessian trace estimator (Eq. ) are computing the quadratic form of Hessian, a specific contraction that is included in the pushforward of 2-jet. In STDE, the random vectors are the unit vectors whose indexes are sampled from the index set without replacement. This can be seen as a discrete distribution p(v) such that v = √ de i for i = 1, 2, • • • , d with probability 1/d, which is isotropic. Hence HTE can also be defined as a push forward of random 2-jet that are isotropic.We can now write the computation of HTE as followsEQUATION)where ∇2 N is the STDE for Laplacian with random jet batch size N .It was shown in [12] that the Biharmonic operatorEQUATIONhas the following unbiased estimator:EQUATION)where p is the d-dimensional normal distribution. Therefore its STDE estimator isEQUATIONK STDE with dense jets K.1 STDE with second order dense jets as generalization of HTE Suppose D is a second-order differential operator with coefficient tensor C. If C is not symmetric, we can symmetrize it asC ′ = 1 2 (C + C ⊤ ), and D 2 u (a) • C = D 2 u (a) • C ′ since D 2 u (a) is symmetric.Furthermore, we can make C positive-definite by adding a constant diagonal λI where -λ is smaller than the smallest eigenvalue of C. The matrix C ′′ = 1 2 (C + C ⊤ ) + λI then has the eigen decomposition UΣU ⊤ where Σ is diagonal and all positive. Now we haveEQUATIONK.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 =d i=1 ∂ 4 u ∂x 4 i .The mask tensor of L is the rank-4 identity tensor I 4 ∈ R d×d×d×d , so the condition for unbiasedness isEQUATIONwhere δ ijkl = 1 when i = j = k = l, and is 0 otherwise.In the most general case where a ̸ = b ̸ = c ̸ = d, we can sample v ∈ R 4d and split it into four R d vectors. In this case we can define blocks of covariance as E v∼p [v (a) v (b) ] = Σ ab , and Σ = Σ ab ab . Denote the fourth-moment tensor of p as µ ijkl , then Eq. 90 states that the block µ abcd in the fourth moment tensor should match C. Fourth moments can always be decomposed into second moments:M ijkl = µ abcd ijkl = Σ ab ij Σ cd kl + Σ ac ik Σ bd jl + Σ ad il Σ bc jk(91) So finding the p that satisfies Eq. 90 is equivalent to finding a zero-mean distribution p with covariance that satisfies the above equation. In the case of L, the mask tensor is block-diagonal: M ijkl = σ ij δ ij,kl . So in the case where a ̸ = b, set a = 1, b = 2, we haveEQUATIONandΣ = Σ 11 Σ 12 Σ 21 Σ 22 ∈ R 2d×2d. Firstly, consider the diagonal entries of σ:EQUATIONThis can always be satisfied by setting the diagonal entries of both Σ aa and Σ aa block as follows:EQUATIONNext, consider the entire σ matrix. We haveEQUATIONIn the case of L, we have σ ij = δ ij , so for i ̸ = j we haveEQUATIONwhich is impossible to satisfy since entries in a covariance matrix must be real.The variance of the sparse STDE estimator comes from the variance of selected derivative tensor elements, whereas the variance of the dense estimator comes from the derivative tensor elements that are not selected. For example, in the case of Laplacian, as also discussed in [12] , the variance of the sparse STDE estimator comes from the diagonal element of the Hessian, whereas the variance of the dense STDE estimator comes from all the off-diagonal element of the Hessian. Justification: Our claims are backed by both theoretical and experimental evidence. The theoretical evidence is provided in the section 4, and the experimental evidence is provided in section 5.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper.• The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.• The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.• It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.our code will be available at https://github.com/sail-sg/stde 38th Conference on Neural Information Processing Systems (NeurIPS2024). 


Decision: 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.

Reviews: review: paper_summary: paper_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.

In 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.
strength_weakness: strength_weakness: 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.

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.

Looking 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.

In 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.
questions: 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: limitations: Yes, they have.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.
questions: 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.
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.
limitations: limitations: Some limitations are discussed in section 6.
paper_summary: paper_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.
strength_weakness: strength_weakness: NA
questions: questions: NA
limitations: limitations: NA
paper_summary: paper_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.
strength_weakness: strength_weakness: 1.	Compared 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.  
Furthermore, 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.
2.	All 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.  
The 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. 
3.	Some of the highly related existing works are missing. E.g.,
-	Regarding L63, the forward rule for general different operators is done in [1].
-	Regarding 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.
4.	As 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.
5.	The 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.
6.	Some parts of the writing are confusing. See **Questions**.
Also, there are several misuses of terminologies. E.g.,
-	Fréchet derivative, not Frechet.
-	The 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.
-	In 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.
-	In 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’.



[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  
[2]Tan, Songchen. Higher-Order Automatic Differentiation and Its Applications. Diss. Massachusetts Institute of Technology, 2023.
questions: questions: 1.	Which 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.
2.	In 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.
3.	Although 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.
4.	Could 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.
5.	Does 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?
6.	Instead 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.  

On writing and notations:  
1. 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). 
2.	In Eqn.(19), it should be $\sum_{i=1}^d$.
3.	In L247, do you mean $v\sim p$,, instead of $(a,v,0)\sim p$?
4.	In the ‘100D’ column of table 2, 539 is not the smallest number. Why is it in bold?
5.	There needs to be more description in section F.1.  
It is good to remind the readers that Einstein notation is employed in Eq.(45) and that the indices $n,\ n’,\ n’’$ are omitted.
The 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.   

Overall, 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: limitations: Yes.

Comments: title: I will keep my rating unless a major weakness is found by other reviewers
Keep Score
comment: Thanks.
Thanks. I have updated my score.
> Regarding the ablation study on randomization batchsize you added, I don’t think it is very appropriate since there is also randomness in the optimization.

Thank you for your comment on the additional experiment. Firstly we would like to emphasize that, one of the main claims of our paper is when solving optimization problems like equation (1), the cost of computing expensive derivative operators can be amortized over the stochastic optimization of the neural network ansatz. This is why we conduct all the experiments under a stochastic optimization. We regret that this point was not communicated more clearly, and we will try to make this point more prominent in the final version of the paper.

It is possible to conduct yet another experiment that directly compares the accuracy of differential operators under different randomization batch sizes, and we do expect to see that variance decreases proportional to the randomization batch size. However, this does not determine the effectiveness of STDE as explained above, and we think it would be more interesting to show the effect of randomization batch size on the convergence of various PDEs.

> Additionally, do you have any intuition why batchsize=16&64 get stuck whilst both a smaller and a larger batchsize do not?

This was puzzling for us as well. Upon conducting further experiments, we found that it is most likely since we had chosen a set of particularly bad random seeds ({1,2,3,4,5}). We are rerunning the experiments and will report the new result once it is ready.

> I would suggest adding the clarifications you replied in the final version.

Yes, we will be revising the final version to incorporate the clarifications we made in the above response. We would like to thank the reviewer again for his/her detailed checking on the terminologies, which greatly improved the quality of our paper.

> Fokker-Planck equation has a clear definition and derivation from evolution of probability densities. I could not agree with the usage of the terminology here and do think every word should be made precise (even though it is not relevant to the contribution of this work at all.)

Upon further consideration, we think that it would indeed be more precise to rename nonlinear Fokker-Planck in Appendix I.2 to semilinear parabolic equations, considering the interpretation of the Fokker-Planck equation. Initially, we used the name Fokker-Planck to be consistent with the baseline method SDGD (Section 5.2 in [1]). Thanks again for checking the terminologies used in our paper!

> Then you should use either $v\sim p$ or $\delta_{a}\times p\times \delta$ (or something like that).

Thank you for your suggestion. We think both of the suggested modifications are more precise than the current formula used in the paper. We will incorporate this change into the final version of our paper.
Thanks for your reply.

Regarding the ablation study on randomization batchsize you added, I don’t think it is very appropriate since there is also randomness in the optimization. I would instead suggest directly comparing the accuracy on differential operator itself, e.g. Laplacian of a complicated network function. Additionally, do you have any intuition why batchsize=16&64 get stuck whilst both a smaller and a larger batchsize do not?

I would suggest adding the clarifications you replied in the final version. Also, I want to comment on some notations and terminologies and do hope the author to make them precise given the diverse background of NeurIPS attendees.

>  However, they are also specifically driftless Fokker-Planck equations with semilinear extension.

Fokker-Planck equation has a clear definition and derivation from evolution of probability densities. I could not agree with the usage of the terminology here and do think every word should be made precise (even though it is not relevant to the contribution of this work at all.)

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

Then you should use either $v\sim p$ or $(a,v,0)\sim \delta_a\times p \times \delta$ (or something like that).
Thank you very much for your answer. I will keep my score.
Thanks for the clarification! Overall I like the paper and intend to keep the score.