One-sentence Summary

Researchers from the National University of Singapore and Sea AI Lab propose the Stochastic Taylor Derivative Estimator, which constructs input tangents to univariate high-order automatic differentiation to efficiently contract arbitrary-order derivative tensors of multivariate functions, yielding $>1000\times$>1000× speed-up over prior randomization methods when solving high-dimensional PDEs with Physics-Informed Neural Networks.

Key Contributions

• The paper introduces STDE, a general method for constructing stochastic estimators for arbitrary differential operators by using univariate Taylor-mode automatic differentiation to evaluate arbitrary contractions of high-order derivative tensors efficiently.
• STDE combines randomization with high-order automatic differentiation to reduce the cost of high-dimensional and high-order differential operators in neural-network losses, making it applicable to Physics-Informed Neural Networks and other input-gradient problems such as adversarial attacks, feature attribution, and meta-learning.
• Experiments on Physics-Informed Neural Networks show more than 1000x speed-up and more than 30x memory reduction over randomization with first-order automatic differentiation, and enable 1-million-dimensional PDEs to be solved in 8 minutes on a single NVIDIA A100 GPU.

Introduction

The authors work at the intersection of automatic differentiation and physics-informed machine learning, where training objectives often contain high-order derivatives of neural networks with respect to their inputs. Computing these derivatives with repeated backward-mode AD can scale exponentially in both derivative order and input dimension, while randomized approaches such as SDGD, Hutchinson trace estimators, finite differences, and randomized smoothing can reduce cost but often rely on approximations, work well mainly for low-order or Laplacian-like operators, and may be sensitive to discretization or estimator variance.

The authors leverage Taylor mode AD to propose Stochastic Taylor Derivative Estimator, a method that uses constructed input tangents to contract high-order derivative tensors and randomize arbitrary differential operators efficiently. Their approach generalizes previous randomized estimators, clarifies limits of Hutchinson-type methods beyond fourth order, and enables faster, memory-efficient training of PINNs, including million-dimensional PDE problems on a single GPU.

Method

The authors first introduce Stochastic Dimension Gradient Descent as an operator approximation step. A differential operator $\mathcal{D}$D is decomposed into additive terms, each depending on only a few input dimensions, and a random subset is sampled at each iteration:

$$\mathcal{D}:=\sum_{j=1}^{N_{\mathcal{D}}}\mathcal{D}_j \approx \frac{N_{\mathcal{D}}}{|J|}\sum_{j\in J}\mathcal{D}_j:=\tilde{\mathcal{D}}_J,$$D:=j=1∑ND​​Dj​≈∣J∣ND​​j∈J∑​Dj​:=D~J​,

where $J$J is the sampled index set. This reduces the cost of differentiation over large input dimensions $d$d, but it does not remove the exponential dependence on derivative order $k$k. The main method therefore builds on Taylor mode automatic differentiation and recasts high order derivative tensor contraction as an efficient jet propagation process.

The core mathematical idea is to view a univariate Taylor mode AD operator as a contraction of a multivariate derivative tensor. A jet of order $k$k at a point $\mathbf{a}$a can be written as

$$J_g^k(t)=\left(g(t),\mathbf{v}^{(1)},\ldots,\mathbf{v}^{(k)}\right),$$Jgk​(t)=(g(t),v(1),…,v(k)),

and pushing this jet through a function $F$F computes a high order Fréchet derivative. By setting some tangent vectors to $\mathbf{0}$0, one can select diagonal, mixed, or structured entries of the derivative tensor. Thus, a $k$kth order differential operator can be represented as a contraction between a derivative tensor and a coefficient tensor. For a scalar function $u:\mathbb{R}^d\rightarrow\mathbb{R}$u:Rd→R, a $k$kth order operator $\mathcal{L}$L can be written as

$$\mathcal{L}u(\mathbf{a}) = \sum_{d_1,\ldots,d_k} D_u^k(\mathbf{a})_{d_1,\ldots,d_k} C_{d_1,\ldots,d_k}(\mathcal{L}) = D_u^k(\mathbf{a})\cdot \mathbf{C}(\mathcal{L}),$$Lu(a)=d1​,…,dk​∑​Duk​(a)d1​,…,dk​​Cd1​,…,dk​​(L)=Duk​(a)⋅C(L),

where $\mathbf{C}(\mathcal{L})$C(L) is zero except at the multi indices corresponding to terms in the operator. The Laplacian is a simple example because its coefficient tensor is the identity matrix $\mathbf{I}$I.

The proposed Stochastic Taylor Derivative Estimator replaces the exact tensor contraction with a random contraction. If the moments of the random jet match the coefficient tensor, that is,

$$\mathbb{E} \left[ v_{d_1}^{(v_1)}\cdots v_{d_k}^{(v_k)} \right] = C_{d_1,\ldots,d_k}(\mathcal{L}),$$E[vd1​(v1​)​⋯vdk​(vk​)​]=Cd1​,…,dk​​(L),

then the pushforward of a random higher order jet gives an unbiased estimator:

$$\mathbb{E}_{J_g^l\sim p} \left[ \partial^l u(J_g^l) \right] = \mathcal{L}u(\mathbf{a}).$$EJgl​∼p​[∂lu(Jgl​)]=Lu(a).

In practice, the authors estimate $\mathcal{L}u(\mathbf{a})$Lu(a) by averaging several pushforwards of independently sampled jets. The main computational modules are therefore: a random jet sampler, a Taylor mode AD pushforward, and a sample mean estimator. This design avoids explicitly constructing and storing the full $k$kth order derivative tensor $D_u^k$Duk​.

The computation graph of the pushforward is central to the scalability of the method. As shown in the figure below:

The figure illustrates a second order pushforward through a sequence of primitives. The left column contains an input $2$2 jet with primal value $\mathbf{x}$x and tangents $\mathbf{v}^{(1)}$v(1) and $\mathbf{v}^{(2)}$v(2). Each primitive maps one jet to the next jet, so the output of one stage becomes the input of the next stage. The first tangent branch corresponds to a Jacobian vector product, while the second tangent branch accumulates Hessian vector products and the quadratic Hessian form. Because all tangent orders at a given layer can be processed in parallel, the method does not require sequential expansion over the derivative order. For one call to $d^kF$dkF, the memory scaling is $\mathcal{O}(kd)$O(kd) and the computation scaling is $\mathcal{O}(k^2 d L)$O(k2dL), where $L$L is the number of primitives. Randomization further removes the polynomial dependence on input dimension $d$d when only a subset of derivative directions is needed.

For many operators, the authors construct sparse random jets. Any coefficient tensor can be decomposed into basis tensor terms:

$$\mathbf{C}(\mathcal{L}) = \sum_{d_1,\ldots,d_k} C_{d_1,\ldots,d_k} \mathbf{e}_{d_1}\otimes\cdots\otimes\mathbf{e}_{d_k},$$C(L)=d1​,…,dk​∑​Cd1​,…,dk​​ed1​​⊗⋯⊗edk​​,

where $\mathbf{e}_i$ei​ is the $i$ith standard basis vector. The method samples sparse jets built from these basis vectors with probabilities proportional to the corresponding coefficient entries:

$$p\left( \otimes_{i=1}^{k}\mathbf{e}_{d_i} \right) = \frac{C_{d_1,\ldots,d_k}}{Z},$$p(⊗i=1k​edi​​)=ZCd1​,…,dk​​​,

where $Z$Z is a normalization factor. The normalization cancels the additive form of the operator. For example, the Laplacian estimator samples indices $j$j and applies the second order pushforward with tangent $\mathbf{e}_j$ej​ and zero second tangent:

$$\tilde{\nabla}^2_J u_\theta(\mathbf{a}) = \frac{d}{|J|} \sum_{j\in J} \partial^2 u_\theta(\mathbf{a})(\mathbf{e}_j,\mathbf{0}).$$∇~J2​uθ​(a)=∣J∣d​j∈J∑​∂2uθ​(a)(ej​,0).

For a high order diagonal operator $\mathcal{L}=\sum_j \partial^k/\partial x_j^k$L=∑j​∂k/∂xjk​, the same pattern applies by setting the first order tangent to $\mathbf{e}_j$ej​ and all higher order tangents to zero:

$$\tilde{\mathcal{L}}_J u_\theta(\mathbf{a}) = \frac{d}{|J|} \sum_{j\in J} \partial^k u_\theta(\mathbf{a})(\mathbf{e}_j,\mathbf{0},\ldots).$$L~J​uθ​(a)=∣J∣d​j∈J∑​∂kuθ​(a)(ej​,0,…).

For second order parabolic PDEs, the diffusion term may contain non diagonal Hessian entries. The authors remove this difficulty through a change of variable so that the operator becomes a sum over terms of the form

$$\partial^2 u(\mathbf{x},t)(\sigma(\mathbf{x},t)\mathbf{e}_i,\mathbf{0}).$$∂2u(x,t)(σ(x,t)ei​,0).

Then the stochastic estimator samples over the terms generated by the columns of $\sigma(\mathbf{x},t)$σ(x,t).

When mixed partial derivatives cannot be removed by such a reformulation, the authors increase the jet order and encode each desired tensor entry as the output of a higher order pushforward. The required order depends on how non diagonal the target tensor slice is. If all differentiated dimensions are distinct, a higher order jet is needed, and the corresponding derivatives are recovered by reading different tangent orders from the output jet and applying chain rule scaling factors. This allows arbitrary $k$kth order derivative tensor elements to be accessed from Taylor mode AD without constructing the full tensor.

The method also supports dense random jets. The classical Hutchinson trace estimator is recovered as a special case where the first order tangent is drawn from a distribution satisfying $\mathbb{E}[\mathbf{v}\mathbf{v}^\top]=\mathbf{I}$E[vv⊤]=I, and the second order pushforward is taken with a zero second tangent. For a general second order operator with symmetrically regularized coefficient tensor

$$\mathbf{C}'' = \frac{1}{2}(\mathbf{C}+\mathbf{C}^\top)+\lambda\mathbf{I} = \mathbf{U}\boldsymbol{\Sigma}\mathbf{U}^\top,$$C′′=21​(C+C⊤)+λI=UΣU⊤,

the authors construct an estimator by drawing Gaussian tangents with covariance $\boldsymbol{\Sigma}$Σ and subtracting a correction term associated with $\lambda\mathbf{I}$λI. Dense STDE can be more efficient for some operators, but it is not universally applicable beyond second order. The sparse version remains general, while the dense version can reduce variance in selected cases.

Experiment

The experiments evaluate STDE for amortized PINN training on inseparable, semilinear parabolic, high-order, and high-dimensional PDEs. They show that stochastic differential estimation preserves the PINN optimization objective while substantially reducing training computation and memory relative to standard backward-mode approaches. Ablation studies further attribute gains to implementation choices and favorable scaling, with exact derivative methods becoming less effective in high dimensions while STDE remains more efficient.

The speed ablation shows that implementation and parallelization choices strongly affect throughput for the two-body Allen-Cahn equation. STDE scales more robustly with dimension than PyTorch or sequential backward-mode SDGD, and it remains usable at very high dimension where most baselines run out of memory. Forward Laplacian performs best at low dimension but has poor scaling and fails at higher dimensions. Moving backward-mode SDGD from PyTorch to JAX gives a large speed improvement. Parallelizing backward-mode SDGD substantially increases speed over the sequential JAX version. Forward Laplacian is fastest in low dimension but degrades sharply as dimension grows. STDE shows the best high-dimensional scalability among the compared methods. STDE is reported to provide major speed and memory gains over the strongest SDGD baseline.

The memory comparison shows that STDE remains memory efficient as input dimension grows, while baseline stochastic gradient methods and exact forward Laplacian scale less favorably. Reimplementing the original method in JAX and parallelizing sampled-dimension work lowers peak memory relative to the original PyTorch loop, but still causes out of memory failures at very large dimensions. STDE uses much less memory than the strongest SDGD baseline in high dimensional settings and avoids the poor scaling of the exact Laplacian approach. STDE remains memory efficient at the largest tested dimension where baselines fail. JAX and parallelization reduce peak memory relative to the original PyTorch implementation. Forward Laplacian is memory competitive only at small dimension and scales worse at high dimension. Parallelized backward-mode and mixed mode variants have similar memory behavior. STDE achieves substantial memory savings over the strongest SDGD baseline.

Experiments evaluate the two-body Allen-Cahn case through speed ablations and memory comparisons, validating the impact of implementations, parallelization, and scaling behavior across methods. The speed results show that JAX and parallelization improve backward-mode SDGD, Forward Laplacian is fastest only at low dimension, and STDE scales most robustly and remains usable at very high dimensions. The memory results confirm that STDE avoids the poor scaling of exact Laplacian and achieves substantial memory savings over the strongest SDGD baseline, while JAX and parallelization reduce peak memory.