Zero Coordinate Shift: Whetted Automatic Differentiation for Physics-informed Operator Learning
Automatic differentiation (AD) is a critical step in physics-informed machine learning, required for computing the high-order derivatives of network output w.r.t. coordinates of collocation points. In this paper, we present a novel and lightweight algorithm to conduct AD for physics-informed operator learning, which we call the trick of Zero Coordinate Shift (ZCS). Instead of making all sampled coordinates as leaf variables, ZCS introduces only one scalar-valued leaf variable for each spatial or temporal dimension, simplifying the wanted derivatives from "many-roots-many-leaves" to "one-root-many-leaves" whereby reverse-mode AD becomes directly utilisable. It has led to an outstanding performance leap by avoiding the duplication of the computational graph along the dimension of functions (physical parameters). ZCS is easy to implement with current deep learning libraries; our own implementation is achieved by extending the DeepXDE package. We carry out a comprehensive benchmark analysis and several case studies, training physics-informed DeepONets to solve partial differential equations (PDEs) without data. The results show that ZCS has persistently reduced GPU memory consumption and wall time for training by an order of magnitude, and such reduction factor scales with the number of functions. As a low-level optimisation technique, ZCS imposes no restrictions on data, physics (PDE) or network architecture and does not compromise training results from any aspect.
Code (1)
Tasks
GPUOperator learningPhysics-informed machine learningMethods 이 논문이 사용한 방법론
Similar Papers 제목 키워드 기반
Jet Functors and Weil Algebras in Automatic Differentiation: A Geometric Analysis
We present a differential-geometric formulation of automatic differentiation (AD) based on jet functors and Weil algebras. In this framework, forward- and reverse-mode differentiation arise naturally as pushforward and c…
Local learning for stable backpropagation-free neural network training towards physical learning
While backpropagation and automatic differentiation have driven deep learning's success, the physical limits of chip manufacturing and rising environmental costs of deep learning motivate alternative learning paradigms s…
Learning vertical coordinates via automatic differentiation of a dynamical core
Terrain-following coordinates in atmospheric models often imprint their grid structure onto the solution, particularly over steep topography, where distorted coordinate layers can generate spurious horizontal and vertica…
Automatic Differentiation: Theory and Practice
We present the classical coordinate-free formalism for forward and backward mode ad in the real and complex setting. We show how to formally derive the forward and backward formulae for a number of matrix functions start…
Neural Stochastic Differential Equations: Deep Latent Gaussian Models in the Diffusion Limit
In deep latent Gaussian models, the latent variable is generated by a time-inhomogeneous Markov chain, where at each time step we pass the current state through a parametric nonlinear map, such as a feedforward neural ne…
Variational Inference