mlsanityData to 14 Oct 2026

Stochastic Taylor Derivative Estimator: Efficient amortization for arbitrary differential operators

arXiv 2412.00088 ↗First posted 27 Nov 2024cs.LGcited by 20

Not on the map. Why?

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.

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

21 of 42 references.

References · 21

  1. 2020Score-Based Generative Modeling through Stochastic Differential Equationscited by 9,122
  2. 2017DGM: A deep learning algorithm for solving partial differential equationscited by 678
  3. 2017The Deep Ritz method: A deep learning-based numerical algorithm for solving variational problemscited by 583
  4. 2017Solving high-dimensional partial differential equations using deep learningcited by 556
  5. 2019Sliced Score Matching: A Scalable Approach to Density and Score Estimationcited by 384
  6. 2023Fine-Tuning Language Models with Just Forward Passescited by 345
  7. 2021Physics-informed neural networks with hard constraints for inverse designcited by 280
  8. 2021Gradient-enhanced physics-informed neural networks for forward and inverse PDE problemscited by 222
  9. 2023Tackling the Curse of Dimensionality with Physics-Informed Neural Networkscited by 79
  10. 2018Forward-Backward Stochastic Neural Networks: Deep Learning of High-dimensional Partial Differential Equationscited by 55
  11. 2019Deep splitting method for parabolic PDEscited by 53
  12. 2020Efficient Learning of Generative Models via Finite-Difference Score Matchingcited by 49
9 more
  1. 2023Randomized Numerical Linear Algebra : A Perspective on the Field With an Eye to Softwarecited by 38
  2. 2022Learning Physics-Informed Neural Networks without Stacked Back-propagationcited by 23
  3. 2023Hutchinson Trace Estimation for High-Dimensional and High-Order Physics-Informed Neural Networkscited by 23
  4. 2024Score-Based Physics-Informed Neural Networks for High-Dimensional Fokker-Planck Equationscited by 22
  5. 2023Bias-Variance Trade-off in Physics-Informed Neural Networks with Randomized Smoothing for High-Dimensional PDEscited by 19
  6. 2020A Modern Analysis of Hutchinson's Trace Estimatorcited by 11
  7. 2021Johnson-Lindenstrauss Lemma, Linear and Nonlinear Random Projections, Random Fourier Features, and Random Kitchen Sinks: Tutorial and Surveycited by 8
  8. 2023Forward Laplacian: A New Computational Framework for Neural Network-based Variational Monte Carlocited by 6
  9. 2024DOF: Accelerating High-order Differential Operators with Forward Propagationcited by 5

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2025: 11; 2026: 9.

20 papers.

Sort:

By area:

  1. 2025CMT: Mid-Training for Efficient Learning of Consistency, Mean Flow, and Flow Map Modelscited by 18
  2. 2025OmniScientist: Toward a Co-evolving Ecosystem of Human and AI Scientistscited by 18
  3. 2025ProPINN: Demystifying Propagation Failures in Physics-Informed Neural Networkscited by 6
  4. 2025Enabling Automatic Differentiation with Mollified Graph Neural Operatorscited by 6
  5. 2025Dual Natural Gradient Descent for Scalable Training of Physics-Informed Neural Networkscited by 6
  6. 2025MeanFlow Transformers with Representation Autoencoderscited by 6
  7. 2025Adaptive Physics-informed Neural Networks: A Surveycited by 5
  8. 2025Scalable Back-Propagation-Free Training of Optical Physics-Informed Neural Networkscited by 2
  9. 2025Collapsing Taylor Mode Automatic Differentiationcited by 1
  10. 2025PMNO: A novel physics guided multi-step neural operator predictor for partial differential equationscited by 1
  11. 2026Deep Neural networks for solving high-dimensional parabolic partial differential equationscited by 1
  12. 2025Quasi-Random Physics-informed Neural Networkscited by 0
8 more
  1. 2026Unbiased Gradient Estimation for Event Binning via Functional Backpropagationcited by 0
  2. 2026Stochastic Dimension Zeroth-Order Estimator: Stable and Memory-Efficient Training of PINNscited by 0
  3. 2026Stochastic Dimension Implicit Functional Projections for Global Integral Conservation in High-Dimensional PINNscited by 0
  4. 2026Stochastic-Dimension Frozen Sampled Neural Network for High-Dimensional Gross-Pitaevskii Equations on Unbounded Domainscited by 0
  5. 2026Unbiased and Second-Order-Free Training for High-Dimensional PDEscited by 0
  6. 2026Parameterized Representations via Implicit Stochastic Modulation for High-Dimensional and High-Order Neural PDE Solverscited by 0
  7. 2026Tensor Train Diffusion: Leveraging Low-Rank Structures for High-Dimensional Score-Based Samplingcited by 0
  8. 2026Spectral-Informed Neural Networks Outperform Spectral Methods in High-dimensional PDEscited by 0