Research map / Probabilistic and scientific modeling / Differential equations and PDEs
Collapsing Taylor Mode Automatic Differentiation
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 collapsed Taylor mode, a graph-rewrite that propagates the sum of highest-order Taylor coefficients instead of propagating then summing, accelerating PDE operators (Laplacian, weighted Laplacian, biharmonic) relative to standard Taylor mode and nested first-order AD.
Problem
Computing PDE operators such as the Laplacian and biharmonic operator via nested backpropagation is expensive, and vanilla Taylor mode AD is not faster than nesting, limiting use in PINNs and variational Monte Carlo.
Approach
Because the highest-degree Taylor coefficient propagates linearly through Faà di Bruno's formula, the sum over multiple directional K-jets can be pushed inside the propagation, collapsing R jets into one; this recovers the forward Laplacian as a special case and extends to randomized Taylor mode and general linear operators via a mixed-partial interpolation family.
Key results
- Collapsed Taylor mode is roughly twice as fast as nested first-order AD for exact and stochastic (weighted) Laplacians while using only 40-50% of the memory.
- For the stochastic biharmonic operator, collapsed Taylor mode achieves up to 9x speed-up in time and 5x in differentiable memory over baselines.
- For the exact biharmonic, collapsing reduces the quadratic vector-count coefficient by 25% (from 6D^2 to 4.5D^2).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Collapse the highest Taylor coefficients instead of propagating each directional coefficient and summing afterward | Taylor mode for exact and stochastic Laplacians · Laplacian | runtime: roughly twice as fast | Helped |
| Collapse the highest Taylor coefficients instead of propagating each directional coefficient and summing afterward | Taylor mode for exact and stochastic Laplacians · Laplacian | peak memory (differentiable): 40-50% memory | Helped |
| Collapse the highest Taylor coefficients across sampled directions | Taylor mode for the stochastic biharmonic operator · Stochastic biharmonic | runtime: up to 9x in time | Helped |
| Collapse the highest Taylor coefficients across sampled directions | Taylor mode for the stochastic biharmonic operator · Stochastic biharmonic | peak memory (differentiable): 5x in memory | Helped |
Surprising findings
- Vanilla Taylor mode in JAX is about 50% slower than nested backpropagation for the Laplacian.
- JAX's Taylor mode was consistently slower than the authors' PyTorch Taylor mode implementation despite using jit.
Limitations
- Taylor arithmetic supports only a small number of primitives in the PyTorch implementation, implemented case by case.
- Implementation not fully optimized: no in-place operations and no efficient Taylor-mode schemes beyond Faà di Bruno.
- Applicability restricted to linear differential operators.
Experimental setup
Benchmarks: Laplacian runtime and memory benchmark weighted Laplacian benchmark biharmonic benchmark
Metrics: runtime (ms per datum) peak GPU memory (MiB)
Baselines: nested first-order AD standard Taylor mode JAX forward Laplacian JAX Taylor mode
Builds on: Taylor mode AD forward Laplacian randomized Taylor mode Faà di Bruno formula torch.fx
Closest papers by meaning
- Stochastic Taylor Derivative Estimator: Efficient amortization for arbitrary differential operators (NeurIPS 2024 · similarity 0.87)
- Efficient Learning of PDEs via Taylor Expansion and Sparse Decomposition into Value and Fourier Domains (AAAI 2024 · similarity 0.86)
- An Interpretable Approach to the Solutions of High-Dimensional Partial Differential Equations (AAAI 2024 · similarity 0.85)
- Toward Efficient Kernel-Based Solvers for Nonlinear PDEs (ICML 2025 · similarity 0.84)
- PRDP: Progressively Refined Differentiable Physics (ICLR 2025 · similarity 0.84)
- An operator preconditioning perspective on training in physics-informed machine learning (ICLR 2024 · similarity 0.84)
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.