Discretize-Optimize vs. Optimize-Discretize for Time-Series Regression and Continuous Normalizing Flows

TL;DR

This work benchmarks Disc-Opt against Opt-Disc for neural ODEs in time-series regression and continuous normalizing flows, showing Disc-Opt can match or closely approach Opt-Disc performance while dramatically reducing training costs when guided by careful discretization and re-discretization. It demonstrates that gradient accuracy in Disc-Opt is robust to forward-solver tolerances, enabling substantial speedups (e.g., ~20x in time-series tasks and ~6x on CNF datasets) with controllable invertibility via step-size choices and post-training re-discretization. The study highlights the trade-offs between forward discretization fidelity, gradient computation, and invertibility, and it proposes practical strategies like multilevel training to further accelerate learning. Overall, Disc-Opt emerges as a practical and scalable alternative for training neural ODEs in applications where the continuous model is essential for inference and invertibility, with broad implications for scientific machine learning and probabilistic modeling.

Abstract

We compare the discretize-optimize (Disc-Opt) and optimize-discretize (Opt-Disc) approaches for time-series regression and continuous normalizing flows (CNFs) using neural ODEs. Neural ODEs are ordinary differential equations (ODEs) with neural network components. Training a neural ODE is an optimal control problem where the weights are the controls and the hidden features are the states. Every training iteration involves solving an ODE forward and another backward in time, which can require large amounts of computation, time, and memory. Comparing the Opt-Disc and Disc-Opt approaches in image classification tasks, Gholami et al. (2019) suggest that Disc-Opt is preferable due to the guaranteed accuracy of gradients. In this paper, we extend the comparison to neural ODEs for time-series regression and CNFs. Unlike in classification, meaningful models in these tasks must also satisfy additional requirements beyond accurate final-time output, e.g., the invertibility of the CNF. Through our numerical experiments, we demonstrate that with careful numerical treatment, Disc-Opt methods can achieve similar performance as Opt-Disc at inference with drastically reduced training costs. Disc-Opt reduced costs in six out of seven separate problems with training time reduction ranging from 39% to 97%, and in one case, Disc-Opt reduced training from nine days to less than one day.

Figures (8)

• Figure 1: Illustration of the Disc-Opt approach for the Gaussian mixture problem in Section \ref{['sec:sub_8gauss']}. As common in optimal control, we use different discretization points for the weights and the states. Following the terminology in optimal control, we refer to the time points as control and state layers, respectively. We discretize the state equation using a Runge-Kutta 4 scheme with constant step size $h=1$. Image is not to scale.
• Figure 2: Time-series regression training iterations 100 and 300 comparing the Disc-Opt and Opt-Disc approaches with the ground truth \ref{['eq:juliaODE']}. Comparative convergence video of the two methods is available at https://imgur.com/nWxwVoe. Discrepancies between convergence behavior of the approaches vary with initial parameterization (Section \ref{['sec:seeds']}).
• Figure 3: For time-series regression, the Disc-Opt method converges in fewer iterations, and each of its iterations requires less time. The mean iteration clocktimes are 2.0 ms for Disc-Opt and 80.5 ms for Opt-Disc. Around iteration 24, the ODE solvers struggle to improve the training (cf. Figure \ref{['fig:derivCheck']}), but the Opt-Disc approach appears to suffer more.
• Figure 4: The derivative check \ref{['eq:deriv_check']} for time-series regression iterations 14 and 24 shown on log-log plot. As expected, the gradients of the Disc-Opt approach (bottom row) are correct; note the faster decay of $E_1(h)$ (red line) compared to $E_0(h)$ (blue line) as $h$ decays. The gradient in the Opt-Disc approach (top row) is correct for iteration 14, but not for iteration 24. In this case (top right), the function $E_1$ (red line) is greater than $E_0$ (blue line).
• Figure 5: Extrapolation of time-series regression models. After training each model on the time interval $t \in [0,1.5]$, we visualize how the trained neural networks extrapolate up to $t=6$. Each neural ODE models behaves as a smooth ODE, by construction, though the learned models oscillate much more quickly than the ground truth ODE.