System-Aware Neural ODE Processes for Few-Shot Bayesian Optimization

Qing, Jixiang and Langdon, Rebecca D. and Lee, Robert Matthew and Shafei, Behrang and Wilk, Mark van der and Tsay, Calvin and Misener, Ruth (2025) System-Aware Neural ODE Processes for Few-Shot Bayesian Optimization. Transactions on Machine Learning Research, 2025. ISSN 2835-8856

Text (3616_System_Aware_Neural_ODE_P)
3616_System_Aware_Neural_ODE_P.pdf - Accepted Version
Available under License Creative Commons Attribution.
Download (17MB)

Abstract

We consider the problem of optimizing initial conditions and termination time in dynamical systems governed by unknown ordinary differential equations (ODEs), where evaluating different initial conditions is costly and the state's value can not be measured in real-time but only with a delay while the measuring device processes the sample. To identify the optimal conditions in limited trials, we introduce a few-shot Bayesian Optimization (BO) framework based on the system's prior information. At the core of our approach is the System-Aware Neural ODE Processes (SANODEP), an extension of Neural ODE Processes (NODEP) designed to meta-learn ODE systems from multiple trajectories using a novel context embedding block. We further develop a two-stage BO framework to effectively incorporate search space constraints, enabling efficient optimization of both initial conditions and observation timings. We conduct extensive experiments showcasing SANODEP's potential for few-shot BO within dynamical systems. We also explore SANODEP's adaptability to varying levels of prior information, highlighting the trade-off between prior flexibility and model fitting accuracy.