Kronecker-Factored Approximate Curvature for Physics-Informed Neural Networks
CoRR(2024)
Abstract
Physics-informed neural networks (PINNs) are infamous for being hard to
train. Recently, second-order methods based on natural gradient and
Gauss-Newton methods have shown promising performance, improving the accuracy
achieved by first-order methods by several orders of magnitude. While
promising, the proposed methods only scale to networks with a few thousand
parameters due to the high computational cost to evaluate, store, and invert
the curvature matrix. We propose Kronecker-factored approximate curvature
(KFAC) for PINN losses that greatly reduces the computational cost and allows
scaling to much larger networks. Our approach goes beyond the established KFAC
for traditional deep learning problems as it captures contributions from a
PDE's differential operator that are crucial for optimization. To establish
KFAC for such losses, we use Taylor-mode automatic differentiation to describe
the differential operator's computation graph as a forward network with shared
weights. This allows us to apply KFAC thanks to a recently-developed general
formulation for networks with weight sharing. Empirically, we find that our
KFAC-based optimizers are competitive with expensive second-order methods on
small problems, scale more favorably to higher-dimensional neural networks and
PDEs, and consistently outperform first-order methods and LBFGS.
MoreTranslated text
AI Read Science
Must-Reading Tree
Example
Generate MRT to find the research sequence of this paper
Chat Paper
Summary is being generated by the instructions you defined