# Collapsing Taylor Mode Automatic Differentiation

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 12 of its 32 references that are in the map.
- First posted on arXiv 19 May 2025.

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

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Ab-Initio Solution of the Many-Electron Schr"odinger Equation with Deep Neural Networks](https://mlsanity.com/paper/1909.02487/) (2019)
- [Deep neural network solution of the electronic Schr"odinger equation](https://mlsanity.com/paper/1909.08423/) (2019)
- [The Global Landscape of Neural Networks: An Overview](https://mlsanity.com/paper/2007.01429/) (2020)
- [Torch.fx: Practical Program Capture and Transformation for Deep Learning in Python](https://mlsanity.com/paper/2112.08429/) (2021)
- [Hutchinson Trace Estimation for High-Dimensional and High-Order Physics-Informed Neural Networks](https://mlsanity.com/paper/2312.14499/) (2023)
- [Kronecker-Factored Approximate Curvature for Physics-Informed Neural Networks](https://mlsanity.com/paper/2405.15603/) (2024)
- [Dynamical Measure Transport and Neural PDE Solvers for Sampling](https://mlsanity.com/paper/2407.07873/) (2024)
- [Stochastic Taylor Derivative Estimator: Efficient amortization for arbitrary differential operators](https://mlsanity.com/paper/2412.00088/) (2024)
- [A Physics Informed Neural Network Approach to Solution and Identification of Biharmonic Equations of Elasticity](https://mlsanity.com/paper/2108.07243/) (2021)
- [Rayleigh-Gauss-Newton optimization with enhanced sampling for variational Monte Carlo](https://mlsanity.com/paper/2106.10558/) (2021)
- [DOF: Accelerating High-order Differential Operators with Forward Propagation](https://mlsanity.com/paper/2402.09730/) (2024)

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

- [Unbiased and Second-Order-Free Training for High-Dimensional PDEs](https://mlsanity.com/paper/2605.14643/) (2026)

## Abstract (arXiv)

Computing partial differential equation (PDE) operators via nested backpropagation is expensive, yet popular, and severely restricts their utility for scientific machine learning. Recent advances, like the forward Laplacian and randomizing Taylor mode automatic differentiation (AD), propose forward schemes to address this. We introduce an optimization technique for Taylor mode that 'collapses' derivatives by rewriting the computational graph, and demonstrate how to apply it to general linear PDE operators, and randomized Taylor mode. The modifications simply require propagating a sum up the computational graph, which could -- or should -- be done by a machine learning compiler, without exposing complexity to users. We implement our collapsing procedure and evaluate it on popular PDE operators, confirming it accelerates Taylor mode and outperforms nested backpropagation.