Collapsing Taylor Mode Automatic Differentiation
This addresses a bottleneck in scientific machine learning by making PDE operator computations more efficient, though it appears incremental as an optimization of existing Taylor mode approaches.
The paper tackles the computational expense of computing partial differential equation (PDE) operators via nested backpropagation in scientific machine learning by introducing a collapsing technique for Taylor mode automatic differentiation that rewrites the computational graph. The method accelerates Taylor mode and outperforms nested backpropagation on popular PDE operators.
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.