Mixed Precision Training of Neural ODEs

AI summary

Overview

Research area: Numerical analysis and deep learning — specifically mixed precision arithmetic for continuous-time neural network architectures (Neural ODEs), sitting at the intersection of floating-point error analysis and efficient deep learning training.

Technical level: Advanced. The paper combines IEEE-754 floating-point format details, roundoff error analysis with formal assumptions and bounds, and custom adjoint (backward) recursion schemes.

Scope: The paper proposes, analyzes, and implements a mixed precision training framework for Neural ODEs with explicit time integrators, together with a dynamic adjoint scaling heuristic, a roundoff error analysis, and an open-source PyTorch package called rampde.

What This Paper Is About

Mixed precision training (using 16-bit formats for expensive computations while keeping weights in 32-bit) is standard practice in deep learning, but the paper argues and demonstrates that applying it naïvely to Neural ODEs is unreliable: roundoff errors accumulate during time integration, producing growing errors in the solution and gradients as the number of time steps increases. The authors design a mixed precision forward and backward scheme specifically for Neural ODEs, prove that its roundoff error stays on the order of the low-precision unit roundoff and does not grow uncontrollably with step count, and release a drop-in PyTorch package implementing it.

Key Contributions

1. Mixed precision ODE solvers and backpropagation. The authors design and analyze mixed precision explicit solvers for Neural ODEs with a custom backpropagation scheme that stores weights, states, and adjoints in high precision while evaluating the neural network in low precision. They report approximately 50% memory reduction and up to 2× speedup in their largest example.

2. Dynamic adjoint scaling. An adaptive scaling heuristic that maximally exploits the limited range of float16 during backpropagation, preventing overflow and underflow without requiring additional hyperparameter tuning.

3. Theoretical analysis. A proof that roundoff errors remain on the order of the unit roundoff of the low precision and do not grow uncontrollably with the number of time steps.

4. Implementation. The open-source PyTorch package rampde (robust automatic mixed precision for differential equations), hosted at https://github.com/EmoryMLIP/rampde, designed as a drop-in replacement for existing neural ODE implementations.

Main Findings

• Naïve mixed precision fails for Neural ODEs. The paper demonstrates theoretically and experimentally that straightforward application of mixed precision training can lead to accumulation of roundoff errors during time integration, growing errors in the solution and gradients as the number of time steps increases.

• Dynamic adjoint scaling is essential in float16. On a constructed test problem (y'(t) = -λ(t)y(t) with λ(t) = θ₁t² + θ₂t + θ₃, solved with a 4th-order Runge-Kutta method using 400 equidistant time steps, T = 2.65, weights θ = (8, -11, 2⁻¹⁶), initial state x = 65504/180), float16 without scaling produced a relative error in ∂y₀ of 1.99e+3 and relative errors of 1.00e+0 for all three weight derivatives — indicating underflow. With dynamic scaling, the same float16 derivatives had relative errors of 5.89e-3, 6.05e-3, 5.96e-3, and 5.88e-3, comparable to the forward solve error of 3.67e-3.

• Scaling has little effect where the format range is already sufficient. In float32, relative errors were 7.01e-5 for y(T) with no scaling and 7.01e-5 with dynamic scaling; derivative errors ranged from 1.25e-4 to 1.40e-4 (none) and 1.28e-4 to 1.60e-4 (dynamic). In bfloat16, values were nearly identical with and without scaling (e.g., 3.65e-2 vs. 3.65e-2 for y(T), and 4.50e-2 vs. 4.49e-2 for ∂y₀), consistent with bfloat16 sharing the range of float32. The paper notes that bfloat16 errors were larger than float16 errors because more mantissa bits are available in float16.

• A single global loss scale is insufficient. The paper notes that simply scaling the loss function to avoid underflow would cause overflow at intermediate time steps in their test example, motivating the step-by-step adaptive scaling.

• Accuracy is preserved relative to single precision. The authors report that their scheme maintains accuracy comparable to single-precision training while delivering approximately 50% memory reduction and up to 2× speedup.

• Reliability and applicability. The scheme was demonstrated on challenging test cases and on neural ODE applications in image classification and generative models. (The specific datasets and benchmark names used in the experiments are not reported in the provided content.)

Methodology in Plain English

The authors restrict themselves to a discretize-then-optimize approach: the ODE is first discretized in time with an explicit single-step method (for example, forward Euler or the 4th-order Runge-Kutta method used in the test problem), and then gradients of the resulting discrete objective are computed by reversing the time integration. They avoid optimize-then-discretize (continuous adjoint) methods because adaptive solvers may fail in mixed precision when required step sizes round to zero or when accuracy is limited by low-precision evaluations of the network.

In the forward pass, the state is carried in high precision (float32), the expensive neural network increment is evaluated in low precision (float16 or bfloat16), the result is converted back up and accumulated in high precision, and the intermediate state is stored in low precision to save memory.

In the backward pass, the adjoint variable, gradient with respect to weights, and gradient with respect to time steps are initialized and accumulated in high precision, while the vector-Jacobian products through the neural network are computed in low precision. Because the adjoint tends to shrink over many steps, the authors introduce a dynamic adjoint scaling factor initialized as S_N = 2^⌊-log₂(u_L ‖a‖)⌋ so that the scaled adjoint has magnitude near 1/u_L, maximizing use of the format's range. The scale is rounded to a power of two so that scaling only changes the exponent and does not perturb the mantissa. If an overflow or NaN appears, the scale is halved and the computation retried on the same computational graph; if no rescaling was needed and the adjoint is sufficiently small, the scale is doubled for the next step. The unit roundoff is u_L := 2⁻¹¹ for float16 and u_L = 2⁻⁸ for bfloat16.

For the error analysis, the authors adopt Higham's convention of marking rounded quantities with a tilde, define quantization errors bounded by the low-precision unit roundoff, assume float32 addition introduces error bounded by its own unit roundoff, and state smoothness/boundedness assumptions on the increment function's first and second derivatives and on the cost function's Lipschitz constant. Under these assumptions they bound the relative errors of the forward solution, adjoint variables, and gradients.

Why This Matters

This work closes a gap between mainstream deep learning tooling and continuous-time models: automatic mixed precision contexts in PyTorch, JAX, and Keras whitelist operators assumed to be safe in low precision, but those assumptions break for Neural ODEs, where repeated network evaluations and adjoint recursions amplify small errors. By providing both a theoretical justification and a working drop-in package, the paper makes mixed precision usable for a class of models that is otherwise expensive in both compute and memory — memory cost grows with the number of time steps or layers, and compute cost grows with repeated network evaluations.

Real-world applications:

• Image classification with Neural ODE models, one of the two application areas the authors evaluate.
• Generative modeling — neural ODEs are used in continuous-time generative and score-based models, the second evaluation area in the paper.
• Scientific machine learning, including physics-informed neural networks and Deep Operator Networks (DeepONet) for solving partial differential equations, a setting the paper cites as already benefiting from mixed precision.
• Memory- and energy-constrained training generally, since the paper's stated benefits are memory reduction (~50%) and speedup (up to 2×), alongside reduced energy consumption as a motivation for mixed precision.

Industry relevance: The package is distributed as a drop-in replacement for existing PyTorch neural ODE code, so practitioners can adopt it without rewriting models. The speedups arise in part from hardware features the paper describes — tensor compute units and block fused multiply-adds that accumulate 16-bit products in single precision — which are standard in modern GPU accelerators, tying the work directly to commodity training hardware.

Future Directions

• Extending the principles to more advanced integrators. The paper states that its principles for simpler time integrators can be applied to derive effective mixed-precision backward propagation for implicit additive Runge-Kutta schemes (as in the cited work) and for explicit stabilized Runge-Kutta (RKC) methods, suggesting these as natural follow-ons.
• Combining the mixed precision error bound with discretization error bounds. The authors note that their result isolates the finite-precision error and that the total error also includes discretization error; integrating the two would describe overall convergence behavior until discretization error dominates.
• Whether optimize-then-discretize (continuous adjoint) approaches can be made reliable in mixed precision. The paper explicitly sets these aside because adaptive solvers may fail when step sizes round to zero or accuracy is limited by low-precision network evaluations.
• The paper's final section outlines potential future improvements to the framework; the specific items are not detailed in the provided content.

Target Audience

Researchers and practitioners working on Neural ODEs, continuous-time deep learning, and scientific machine learning who need to train such models efficiently on GPU hardware. It is also relevant to numerical analysts interested in floating-point roundoff error accumulation in time integration, and to deep learning engineers who want a drop-in mixed precision replacement for PyTorch neural ODE code. Readers should be comfortable with ODE solvers, adjoint methods, and floating-point formats; the error analysis sections are geared toward readers with numerical analysis backgrounds.

Authors’ 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 consisting of explicit ODE solvers and a custom backpropagation scheme and shows their effectiveness in 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 numerical reliability is provided by custom dynamic adjoint scaling and by accumulating the solution and gradients in higher precision. These contributions address two key challenges in training neural ODEs: 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 \texttt{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.