# Stochastic Taylor Derivative Estimator: Efficient amortization for arbitrary differential operators

arXiv: https://arxiv.org/abs/2412.00088 · HTML: https://mlsanity.com/paper/2412.00088/

## Key facts

- Cited by 20 papers in the map (data through 14 Oct 2026).
- Builds on 21 of its 42 references that are in the map.
- First posted on arXiv 27 Nov 2024.

## Builds on (most cited first, 21 of 21)

- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [DGM: A deep learning algorithm for solving partial differential equations](https://mlsanity.com/paper/1708.07469/) (2017)
- [The Deep Ritz method: A deep learning-based numerical algorithm for solving variational problems](https://mlsanity.com/paper/1710.00211/) (2017)
- [Solving high-dimensional partial differential equations using deep learning](https://mlsanity.com/paper/1707.02568/) (2017)
- [Sliced Score Matching: A Scalable Approach to Density and Score Estimation](https://mlsanity.com/paper/1905.07088/) (2019)
- [Fine-Tuning Language Models with Just Forward Passes](https://mlsanity.com/paper/2305.17333/) (2023)
- [Physics-informed neural networks with hard constraints for inverse design](https://mlsanity.com/paper/2102.04626/) (2021)
- [Gradient-enhanced physics-informed neural networks for forward and inverse PDE problems](https://mlsanity.com/paper/2111.02801/) (2021)
- [Tackling the Curse of Dimensionality with Physics-Informed Neural Networks](https://mlsanity.com/paper/2307.12306/) (2023)
- [Forward-Backward Stochastic Neural Networks: Deep Learning of High-dimensional Partial Differential Equations](https://mlsanity.com/paper/1804.07010/) (2018)
- [Deep splitting method for parabolic PDEs](https://mlsanity.com/paper/1907.03452/) (2019)
- [Efficient Learning of Generative Models via Finite-Difference Score Matching](https://mlsanity.com/paper/2007.03317/) (2020)
- [Randomized Numerical Linear Algebra : A Perspective on the Field With an Eye to Software](https://mlsanity.com/paper/2302.11474/) (2023)
- [Learning Physics-Informed Neural Networks without Stacked Back-propagation](https://mlsanity.com/paper/2202.09340/) (2022)
- [Hutchinson Trace Estimation for High-Dimensional and High-Order Physics-Informed Neural Networks](https://mlsanity.com/paper/2312.14499/) (2023)
- [Score-Based Physics-Informed Neural Networks for High-Dimensional Fokker-Planck Equations](https://mlsanity.com/paper/2402.07465/) (2024)
- [Bias-Variance Trade-off in Physics-Informed Neural Networks with Randomized Smoothing for High-Dimensional PDEs](https://mlsanity.com/paper/2311.15283/) (2023)
- [A Modern Analysis of Hutchinson's Trace Estimator](https://mlsanity.com/paper/2012.12895/) (2020)
- [Johnson-Lindenstrauss Lemma, Linear and Nonlinear Random Projections, Random Fourier Features, and Random Kitchen Sinks: Tutorial and Survey](https://mlsanity.com/paper/2108.04172/) (2021)
- [Forward Laplacian: A New Computational Framework for Neural Network-based Variational Monte Carlo](https://mlsanity.com/paper/2307.08214/) (2023)
- [DOF: Accelerating High-order Differential Operators with Forward Propagation](https://mlsanity.com/paper/2402.09730/) (2024)

## Built on by (most cited first, 20 of 20)

- [CMT: Mid-Training for Efficient Learning of Consistency, Mean Flow, and Flow Map Models](https://mlsanity.com/paper/2509.24526/) (2025)
- [OmniScientist: Toward a Co-evolving Ecosystem of Human and AI Scientists](https://mlsanity.com/paper/2511.16931/) (2025)
- [ProPINN: Demystifying Propagation Failures in Physics-Informed Neural Networks](https://mlsanity.com/paper/2502.00803/) (2025)
- [Enabling Automatic Differentiation with Mollified Graph Neural Operators](https://mlsanity.com/paper/2504.08277/) (2025)
- [Dual Natural Gradient Descent for Scalable Training of Physics-Informed Neural Networks](https://mlsanity.com/paper/2505.21404/) (2025)
- [MeanFlow Transformers with Representation Autoencoders](https://mlsanity.com/paper/2511.13019/) (2025)
- [Adaptive Physics-informed Neural Networks: A Survey](https://mlsanity.com/paper/2503.18181/) (2025)
- [Scalable Back-Propagation-Free Training of Optical Physics-Informed Neural Networks](https://mlsanity.com/paper/2502.12384/) (2025)
- [Collapsing Taylor Mode Automatic Differentiation](https://mlsanity.com/paper/2505.13644/) (2025)
- [PMNO: A novel physics guided multi-step neural operator predictor for partial differential equations](https://mlsanity.com/paper/2506.01598/) (2025)
- [Deep Neural networks for solving high-dimensional parabolic partial differential equations](https://mlsanity.com/paper/2601.13256/) (2026)
- [Quasi-Random Physics-informed Neural Networks](https://mlsanity.com/paper/2507.08121/) (2025)
- [Unbiased Gradient Estimation for Event Binning via Functional Backpropagation](https://mlsanity.com/paper/2602.12590/) (2026)
- [Stochastic Dimension Zeroth-Order Estimator: Stable and Memory-Efficient Training of PINNs](https://mlsanity.com/paper/2603.24002/) (2026)
- [Stochastic Dimension Implicit Functional Projections for Global Integral Conservation in High-Dimensional PINNs](https://mlsanity.com/paper/2603.29237/) (2026)
- [Stochastic-Dimension Frozen Sampled Neural Network for High-Dimensional Gross-Pitaevskii Equations on Unbounded Domains](https://mlsanity.com/paper/2604.09361/) (2026)
- [Unbiased and Second-Order-Free Training for High-Dimensional PDEs](https://mlsanity.com/paper/2605.14643/) (2026)
- [Parameterized Representations via Implicit Stochastic Modulation for High-Dimensional and High-Order Neural PDE Solvers](https://mlsanity.com/paper/2606.22150/) (2026)
- [Tensor Train Diffusion: Leveraging Low-Rank Structures for High-Dimensional Score-Based Sampling](https://mlsanity.com/paper/2607.06841/) (2026)
- [Spectral-Informed Neural Networks Outperform Spectral Methods in High-dimensional PDEs](https://mlsanity.com/paper/2607.13566/) (2026)

## Abstract (arXiv)

Optimizing neural networks with loss that contain high-dimensional and high-order differential operators is expensive to evaluate with back-propagation due to O(dk) scaling of the derivative tensor size and the O(2k-1L) 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× speed-up and >30× 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.