STDE++: Polynomial-Time Amortization for Linear Differential Operators
Zekun Shi, Zheyuan Hu, Min Lin, Kenji Kawaguchi.
Year: 2026, Volume: 27, Issue: 114, Pages: 1−50
Abstract
Optimizing neural networks with losses that contain high-dimensional and high-order differential operators is expensive to evaluate with backpropagation 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 domain dimension, $L$ is the number of ops in the forward computation graph and $k$ is the derivative order. Previous works addressed the polynomial scaling in $d$ 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 contractions 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 randomize any differential operator efficiently. When applied to Physics-Informed Neural Networks (PINNs) and compared against the original PyTorch implementation of SDGD, our method yields about $1.34\times 10^{3}$ average speedup and $31.8\times$ average memory reduction across the three inseparable 100K-dimensional PDEs in our benchmark; the best case is $1.59\times 10^{3}$ speedup and $33.8\times$ memory reduction on Allen-Cahn. We can now solve 1-million-dimensional PDEs in 8 minutes on a single NVIDIA A100 GPU. Furthermore, we proposed new methods for computing mixed partial derivatives using Taylor mode AD, which scales polynomially with the derivative order. This work opens the possibility of using high-order differential operators in large-scale problems.