A Stable and Scalable Method for Solving Initial Value PDEs with Neural Networks

Unlike conventional grid and mesh based methods for solving PDEs, neural networks have the potential to break the curse of dimensionality, providing approximate solutions to high-dimensional PDEs. While global minimization of the PDE residual over the network parameters works well for boundary value problems, catastrophic forgetting limits its applicability to initial value problems. In an alternative local in time approach, the optimization problem can be converted into an ODE on the network parameters and the solution propagated forward in time; however, we demonstrate that current methods utilizing this idea suffer from two key issues. First, following the ODE produces an uncontrolled growth in the conditioning of the problem, ultimately leading to unacceptably large numerical errors. Second, as the ODE methods scale cubically with the number of model parameters, they are restricted to small neural networks, significantly limiting their ability to represent intricate PDE initial conditions and solutions. Building on these insights we develop Neural-IVP, an ODE based IVP solver which prevents the network from getting ill conditioned and runs in time linear in the number of parameters, enabling us to evolve the dynamics of challenging high-dimensional PDEs with neural networks.