PNODE: A memory-efficient neural ODE framework based on high-level adjoint differentiation
Neural ordinary differential equations (neural ODEs) have emerged as a novel
network architecture that bridges dynamical systems and deep learning. However,
the gradient obtained with the continuous adjoint method in the vanilla neural
ODE is not reverse-accurate. Other approaches suffer either from excessive
memory requirement due to deep computational graphs or from limited choices for
the time integration scheme, hampering their application to large-scale complex
dynamical systems. To achieve accurate gradients without compromising memory
efficiency and flexibility, we present a new neural ODE framework, PNODE, based
on high-level discrete adjoint algorithmic differentiation. By leveraging
discrete adjoint time integrators and advanced checkpointing strategies
tailored for these integrators, PNODE can provide a balance between memory and
computational costs, while computing the gradients consistently and accurately.
We provide an open-source implementation based on PyTorch and PETSc, one of the
most commonly used portable, scalable scientific computing libraries. We
demonstrate the performance through extensive numerical experiments on image
classification and continuous normalizing flow problems. We show that PNODE
achieves the highest memory efficiency when compared with other
reverse-accurate methods. On the image classification problems, PNODE is up to
two times faster than the vanilla neural ODE and up to 2.3 times faster than
the best existing reverse-accurate method. We also show that PNODE enables the
use of the implicit time integration methods that are needed for stiff
dynamical systems.