Stochastic Taylor Derivative Estimator: Efficient amortization for arbitrary differential operators

2024-11-27

arXiv.org

AI Summary

High-dimensional, high-order differential operators in Physics-Informed Neural Networks (PINNs) incur prohibitive computational cost and poor scalability due to exponential complexity in both dimensionality and derivative order. Method: We propose a scalable Stochastic Taylor Derivative Estimator (STDE), the first method enabling efficient tensor contraction for arbitrary-order derivative tensors of multivariate functions. STDE unifies treatment of polynomial-in-dimension and exponential-in-order complexity by integrating higher-order automatic differentiation with customized tangent vectors and randomized estimation. Results: In PINNs, STDE accelerates computation over first-order AD-based randomization by >1000× and reduces memory consumption by 30×. Using a single NVIDIA A100 GPU, it solves million-dimensional PDEs in under eight minutes—dramatically enhancing feasibility and efficiency of high-order physics-constrained modeling.

Technology Category

• Reasoning under Uncertainty: Stochastic Optimization
• Machine Learning: Matrix & Tensor Methods
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines
• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

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$ imes$ speed-up and>30$ imes$ 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.

Problem

• High-dimensional Loss Functions
• Higher-order Differentiation
• Computational Efficiency in Neural Networks

Innovation

• Stochastic Taylor Derivative Estimator
• High-dimensional Differential Equations
• Physical Information Neural Networks