Collapsing Taylor Mode Automatic Differentiation
accelerationautomatic differentiationcollapsing procedurecomputational efficiencycomputational graphforward laplacianlinear pde operatorsmachine learning compilernested backpropagationoptimization techniquepartial differential equationsperformance evaluationrandomized taylor modescientific machine learningtaylor mode
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