Stochastic Taylor Derivative Estimator: Efficient amortization for arbitrary differential operators

TL;DR

This work tackles the computational bottleneck of optimizing losses with high-order differential operators in high dimensions by introducing the Stochastic Taylor Derivative Estimator (STDE). STDE recasts arbitrary differential operators as contractions of multivariate derivatives and leverages univariate Taylor mode AD through random jets to efficiently estimate these contractions, removing the exponential scaling with dimension $d$ and derivative order $k$ in a unified framework. The authors provide concrete sparse-jet constructions for common operators (e.g., Laplacian, diagonal high-order terms, mixed partials) and connect to Hutchinson’s trace estimator via dense jets, enabling broad applicability. Empirically, STDE yields substantial speedups and memory reductions in Physics-Informed Neural Networks (PINNs), solving 1-million-dimensional PDEs on GPUs within minutes, and demonstrates strong ablations across varying architectures, parallelization, and operator types. Overall, STDE broadens the practical use of high-order differential operators in large-scale scientific modeling by offering a general, scalable, and parallelizable amortization strategy that outperforms prior randomization-based approaches.

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 \emph{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.

Figures (5)

• Figure 1: The computation graph of computing second order gradient by repeated application of backward mode AD, for a function $F(\cdot)$ with $4$ primitives ($L=4$), which computes the Hessian-vector-product. Red nodes represent the cotangent nodes in the second backward pass. With each repeated application of VJP the length of sequential computation doubles.
• Figure 2: The computation graph of $\dd^{2}F$ for $F$ with $4$ primitives. Parameters $\theta _{i}$ are omitted. The first column from the left represents the input 2-jet $J_{g}^{2}(t)=(\vb{x} , \vb{v}^{(1)} , \vb{v}^{(2)})$, and $\dd^{2}F_1$ pushes it forward to the 2-jet $J_{F_1\circ g}^{2}(t)=(\vb{y}_{1} , \vb{v}_{1}^{(1)} , \vb{v}_{1}^{(2)})$ which is the subsequent column. Each row can be computed in parallel, and no evaluate trace needs to be cached.
• Figure 3: The computation graph of forward mode AD (left) and backward mode AD (right) of a function $F(\cdot)$ with $4$ primitives $F_{i}$ each parameterized by $\theta_{i}$. Nodes represent (intermediate) values, and arrows represent computation. Input nodes are colored blue; output nodes are colored green, and intermediate nodes are colored yellow.
• Figure 4: Convolutional weight sharing in the first layer, with input dimension $9$ and filter size $3$.
• Figure 5: Ablation on randomization batch size with Inseparable and effectively high-dimensional PDEs, dim=100k, 5 runs with different random seeds. Model converges when the difference of L2 error is below 1e-7.