Research map / Probabilistic and scientific modeling / Differential equations and PDEs
Stochastic Taylor Derivative Estimator: Efficient amortization for arbitrary differential operators
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Stochastic Taylor Derivative Estimator (STDE), which uses randomized input jets pushed forward through univariate Taylor-mode AD to estimate arbitrary differential operators, and applies it to PINNs for 1-million-dimensional PDEs.
Problem
Training neural networks with losses containing high-dimensional and high-order differential operators (e.g., PINNs) is expensive with backward-mode AD, scaling exponentially in derivative order k and polynomially in dimension d.
Approach
Construct sparse or dense random input jets so that univariate Taylor-mode AD pushforwards compute unbiased random contractions of the multivariate derivative tensor, enabling randomized estimation of arbitrary differential operators with efficient high-order AD and no full derivative tensor.
Key results
- STDE reaches 13.61 it/s at 1M dimensions where all SDGD baselines and Forward Laplacian run out of memory (OOM) on a 40GB A100.
- Compared with parallelized backward-mode SDGD, STDE gives up to 10x speed-up and at least 4x memory reduction, per the paper.
- STDE solves 1-million-dimensional PDEs within 8 minutes on a single NVIDIA A100 40GB GPU.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use STDE instead of parallelized backward mode SDGD (ablation) | PINN training on two-body Allen-Cahn · Two-body Allen-Cahn, 100K D | Speed (it/s): 29.24 -> 156.90 | Helped |
| Use STDE instead of parallelized backward mode SDGD (ablation) | PINN training on two-body Allen-Cahn · Two-body Allen-Cahn, 100K D | Memory (MB): 4931 -> 1073 | Helped |
| Use STDE instead of parallelized backward mode SDGD (ablation) | PINN training on two-body Allen-Cahn · Two-body Allen-Cahn, 10K D | Speed (it/s): 216.83 -> 454.16 | Helped |
| Use STDE instead of parallelized backward mode SDGD (ablation) | PINN training on two-body Allen-Cahn · Two-body Allen-Cahn, 10K D | Memory (MB): 1177 -> 795 | Helped |
| Use STDE instead of parallelized backward mode SDGD (ablation) | PINN training on two-body Allen-Cahn · Two-body Allen-Cahn, 1K D | Speed (it/s): 845.21 -> 1054.39 | Helped |
| Use STDE instead of parallelized backward mode SDGD (ablation) | PINN training on two-body Allen-Cahn · Two-body Allen-Cahn, 1K D | Memory (MB): 579 -> 537 | Helped |
| Use STDE instead of Forward Laplacian (ablation) | PINN training on two-body Allen-Cahn · Two-body Allen-Cahn, 100 D | Speed (it/s): 1974.50 -> 1035.09 | Hurt |
| Use JAX rather than PyTorch for backward mode SDGD (ablation) | Backward mode SDGD · Two-body Allen-Cahn | Speed: ~15× speed-up | Helped |
| Parallelize the sampled-dimension SDGD computation (ablation) | Backward mode SDGD · Two-body Allen-Cahn | Speed: ~15× speed up | Helped |
| Replace randomized SDGD with exact Forward Laplacian (ablation) | PINN training · Two-body Allen-Cahn, 100 D | Speed (it/s): 1035.09 -> 1974.50 | Helped |
| Use Forward Laplacian rather than STDE (ablation) | PINN training · Two-body Allen-Cahn, 1K D and above | Speed | Hurt |
| Lower the randomization batch size (ablation) | STDE · PINN training | Speed and memory | Helped |
Surprising findings
- Forward Laplacian, an exact method, is fastest at 100D but becomes worse than parallelized SDGD above 1000 dimensions due to its O(d) scaling.
- The paper proves that dense-jet HTE-type estimators cannot be generalized to fourth-order operators, whereas sparse jets can.
Limitations
- No variance reduction techniques were considered.
- Lowering randomization batch size improves speed and memory but the cost-variance trade-off needs further analysis.
- Not suited for high-order derivatives with respect to neural network parameters.
Experimental setup
Datasets: two-body inseparable Allen-Cahn equation (synthetic PDE residual points)
Benchmarks: inseparable Allen-Cahn equation (speed and memory ablation, 100D to 1M D)
Metrics: iterations per second peak memory (MB) relative L2 error
Baselines: SDGD (PyTorch backward mode) Backward mode SDGD Parallelized backward mode SDGD Forward-over-Backward SDGD Forward Laplacian random smoothing
Builds on: Taylor mode AD Stochastic Dimension Gradient Descent Hutchinson trace estimator Forward Laplacian PINN
Closest papers by meaning
- Collapsing Taylor Mode Automatic Differentiation (NeurIPS 2025 · similarity 0.87)
- Derivative-Informed Training of Neural Operators On-the-Fly via Sketched Tangent Consistency (NeurIPS 2026 · similarity 0.87)
- An Efficient High-dimensional Gradient Estimator for Stochastic Differential Equations (NeurIPS 2024 · similarity 0.86)
- Advancing the Lower Bounds: an Accelerated, Stochastic, Second-order Method with Optimal Adaptation to Inexactness (ICLR 2024 · similarity 0.85)
- Kronecker-Factored Approximate Curvature for Physics-Informed Neural Networks (NeurIPS 2024 · similarity 0.85)
- Derivative-enhanced Deep Operator Network (NeurIPS 2024 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.