Training Neural ODEs Using Fully Discretized Simultaneous Optimization

2025-02-21

AI Summary

Neural ordinary differential equations (Neural ODEs) suffer from high computational overhead due to repeated numerical integration during training. To address this, we propose a fully discretized co-optimization framework. Our method reformulates the continuous dynamics and parameter learning as a large-scale nonlinear programming problem via collocation-based full discretization—a novel simultaneous optimization paradigm. We further introduce an alternating direction method of multipliers (ADMM) decomposition strategy to enable distributed, inter-batch co-optimization of submodels. The discretized optimization problem is efficiently solved using the IPOPT solver. Evaluated on the Van der Pol oscillator benchmark, our approach achieves significantly faster convergence and higher training efficiency compared to standard Neural ODE training. Experimental results demonstrate superior computational efficiency, numerical stability, and scalability—validating the framework’s effectiveness for optimizing continuous-depth models without sacrificing accuracy or robustness.

Technology Category

• Search and Optimization: Mixed Discrete/Continuous Search
• Machine Learning: Optimization
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Neural Ordinary Differential Equations (Neural ODEs) represent continuous-time dynamics with neural networks, offering advancements for modeling and control tasks. However, training Neural ODEs requires solving differential equations at each epoch, leading to high computational costs. This work investigates simultaneous optimization methods as a faster training alternative. In particular, we employ a collocation-based, fully discretized formulation and use IPOPT--a solver for large-scale nonlinear optimization--to simultaneously optimize collocation coefficients and neural network parameters. Using the Van der Pol Oscillator as a case study, we demonstrate faster convergence compared to traditional training methods. Furthermore, we introduce a decomposition framework utilizing Alternating Direction Method of Multipliers (ADMM) to effectively coordinate sub-models among data batches. Our results show significant potential for (collocation-based) simultaneous Neural ODE training pipelines.

Problem

• Optimize Neural ODE training efficiency.
• Reduce computational costs in differential equations.
• Enhance model convergence with simultaneous methods.

Innovation

• Collocation-based discretized optimization
• IPOPT for nonlinear large-scale optimization
• ADMM for sub-model coordination