Mixed-precision computation is increasingly vital for scaling deep learning. However, standard strategies that are successful in many settings often break down for continuous-time models like Neural ODEs, leading to instability and accuracy loss. We address this gap by designing and analyzing explicit mixed-precision ODE solvers and a custom backpropagation scheme, tailored for training Neural ODEs in scientific machine learning tasks.
Our approach leverages low-precision arithmetic for neural network evaluations and intermediate state storage, while maintaining solution stability and accuracy via dynamic adjoint scaling and high-precision accumulation. This hybrid strategy substantially reduces both computational cost and memory usage, enabling efficient training of deep continuous-time models even with limited hardware resources.
We demonstrate the potential of our approach for generative modeling tasks based on continuous normalizing flows and conditional transport. Our results show that these techniques enable the training of more complex models on resource-limited hardware.
Copyright. All Rights Reserved.