DOF: Accelerating High-order Differential Operators with Forward Propagation
Ruichen Li, Chuwei Wang, Haotian Ye, Di He, Liwei Wang
TL;DR
This work addresses the computational and memory bottlenecks of high-order derivatives in neural PDE solvers by introducing Differential Operator with Forward-propagation (DOF), a framework that generalizes Forward Laplacian to arbitrary second-order operators. By decomposing the operator coefficient matrix $A$ as $A=L^{\top} D L$ and propagating a specialized tuple through the computation graph, DOF computes $\mathcal{L}\phi$ with forward-only passes, achieving precision preservation and significant efficiency gains. Theoretical results show FLOP and memory savings relative to Hessian-based AutoDiff, with further improvements for elliptic (positive-definite) and low-rank coefficient cases, and empirical validation on MLPs confirms roughly 2x speedups and notable memory reductions, especially under Jacobian sparsity. The approach offers a practical path to scalable, accurate NN-based PDE solvers for forward and inverse problems in physics-informed learning.
Abstract
Solving partial differential equations (PDEs) efficiently is essential for analyzing complex physical systems. Recent advancements in leveraging deep learning for solving PDE have shown significant promise. However, machine learning methods, such as Physics-Informed Neural Networks (PINN), face challenges in handling high-order derivatives of neural network-parameterized functions. Inspired by Forward Laplacian, a recent method of accelerating Laplacian computation, we propose an efficient computational framework, Differential Operator with Forward-propagation (DOF), for calculating general second-order differential operators without losing any precision. We provide rigorous proof of the advantages of our method over existing methods, demonstrating two times improvement in efficiency and reduced memory consumption on any architectures. Empirical results illustrate that our method surpasses traditional automatic differentiation (AutoDiff) techniques, achieving 2x improvement on the MLP structure and nearly 20x improvement on the MLP with Jacobian sparsity.