Mixed Precision Training of Neural ODEs

TL;DR

This work tackles the instability of naive mixed-precision training for Neural ODEs by introducing a framework that evaluates the neural ODE dynamics in low precision while accumulating states and adjoints in high precision, coupled with a dynamic adjoint scaling scheme. The authors prove that forward, adjoint, and gradient errors remain bounded at $O(u_{\mathrm{L}}) + O(u_{\mathrm{H}}/h)$, and they implement this in the rampde PyTorch package, providing a drop-in replacement with multiple backward variants. Empirically, the approach yields approximately 50% memory reduction and up to 2x speedups across challenging Neural ODE tasks, including CNFs, OT-Flow, and STL-10, while preserving accuracy comparable to single-precision baselines. The results demonstrate that mixed-precision Neural ODEs can scale to larger problems on tensor-core hardware, with practical guidance on when and how to apply scaling (bf16/float16) to balance stability, speed, and memory usage.

Abstract

Exploiting low-precision computations has become a standard strategy in deep learning to address the growing computational costs imposed by ever larger models and datasets. However, naively performing all computations in low precision can lead to roundoff errors and instabilities. Therefore, mixed precision training schemes usually store the weights in high precision and use low-precision computations only for whitelisted operations. Despite their success, these principles are currently not reliable for training continuous-time architectures such as neural ordinary differential equations (Neural ODEs). This paper presents a mixed precision training framework for neural ODEs, combining explicit ODE solvers with a custom backpropagation scheme, and demonstrates its effectiveness across a range of learning tasks. Our scheme uses low-precision computations for evaluating the velocity, parameterized by the neural network, and for storing intermediate states, while stability is provided by a custom dynamic adjoint scaling and by accumulating the solution and gradients in higher precision. These contributions address two key challenges in training neural ODE: the computational cost of repeated network evaluations and the growth of memory requirements with the number of time steps or layers. Along with the paper, we publish our extendable, open-source PyTorch package rampde, whose syntax resembles that of leading packages to provide a drop-in replacement in existing codes. We demonstrate the reliability and effectiveness of our scheme using challenging test cases and on neural ODE applications in image classification and generative models, achieving approximately 50% memory reduction and up to 2x speedup while maintaining accuracy comparable to single-precision training.

Figures (4)

• Figure 1: Log-scale plot of $|y(t)|$ (black solid line) and $|y'(t)|$ (blue solid line) used in \ref{['ex:test_adjoint_scaling']} and the numerical approximation of the solution using RK4 with 400 equidistant time steps in float16. The horizontal dashed lines indicate the range of the normal numbers in float16. In this set up, the ODE almost exhausts the range of float16, complicating the derivative computation.
• Figure 1: Comparison of continuous normalizing flow (CNF) sample quality in \ref{['sub:cnf_experiment']} across precision formats and ODE solvers on 2D datasets also used in GrathwohlEtAl2018. The columns contain the datasets: 2-spirals (left pair), 8-Gaussians (center pair), and checkerboard (right pair) datasets. For each dataset, there are two columns, comparing torchdiffeq (left) and rampde (right). Rows demonstrate different precision and scaling configurations. The bottom row shows ground truth targets with float16 dynamic scaling, which is only available in rampde. Across these tests, the rampde solver exhibits visually comparable sample quality even without loss scaling, which is important for float16 training using torchdiffeq for all the datasets in this example.
• Figure 2: Roundoff error experiment for the 2D Continuous Normalizing Flow training using the 2spirals dataset in \ref{['sub:cnf_experiment']}. We compare the accuracy of forward (dashed) and backward propagation (solid) in float 16 for torchdiffeq (left) and rampde (right) for forward Euler (top) and RK4 (bottom) as the number of time steps increases. This is done by the relative error of each quantity specified in the legend of the figure to the float64 computation. The accuracy of the solution and gradients for the initial value is consistent for torchdiffeq, while the gradients of the neural network weights become slightly less accurate as $N$ grows. Consistent with our analysis in \ref{['sec:roundoff']}, the relative errors for rampde remain stable.
• Figure 3: Training loss convergence for the STL-10 experiment in \ref{['sub:stl10_experiment']} comparing torchdiffeq (triangular markers, dashed lines) and rampde (square markers, solid lines) in different precisions. The plot shows training loss (log scale) versus wall-clock time in hours, demonstrating that mixed precision variants can achieve comparable final loss values more quickly. For clarity, we omit rampde in float16 with gradient scaling as its curve is similar to that without scaling. We also omit torchdiffeq at float16 as it failed with NaN after a few training steps even with gradient scaling activated.

Theorems & Definitions (9)

• Example 3.1
• Theorem 4.2: Main result
• Lemma 4.3: Forward error
• Proof 1
• Lemma 4.4: Adjoint error
• Proof 2
• Lemma 4.5: Gradient error
• Proof 3
• Proof 4: Proof of \ref{['theo:main_result']}