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

TL;DR

This work introduces a few-shot Bayesian Optimization (BO) framework based on the system's prior information, and develops a two-stage BO framework to effectively incorporate search space constraints, enabling efficient optimization of both initial conditions and observation timings.

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.

Figures (17)

• Figure 1: Illustration of meta-learning-based few-shot Bayesian Optimization (BO) with time-delay constraints (detailed in Section. \ref{['Sec:td_opt']}) on the Lotka-Volterra (LV) system on different models: Gaussian Process (GP), Neural Process (NP) and System Aware Neural ODE Process (SANODEP). We start with one trajectory (marked as $\circ$). Each trajectory ensures that adjacent observations respect a minimum time delay $\Delta t$ constraint. The figure clearly shows that after evaluating just one trajectory (the first, left-most column of three rows), the meta-learned SANODEP model already resembles the original system more closely than the non-meta-learned GP, demonstrating its efficiency in guiding the search toward promising areas.
• Figure 2: Graphical Model of SANODEP, the model predicts any time point in the new trajectory by knowing both observations from new trajectories and from past trajectories. Depending on whether the $X_{new}^{{\mathbb{C}}}$ and $T_{new}^{{\mathbb{C}}}$ consists of more than the initial condition, the model focuses on forecasting or interpolating tasks. The solid and dashed lines represent the generative and inference processes, respectively.
• Figure 3: Model evaluation performance comparison on a different number of context trajectories. The first row corresponds to forecasting performance (prediction with only the initial condition known), and the second row represents an interpolating setting. It can be seen that SANODEP is either on-par or marginally better than NP for most problems, and the initial condition augmented encoder-based model variant provides more robust performance across problems.
• Figure 4: Comparison of few-shot BO performances in terms of mean $\pm$ standard deviation. Except for the Lotka-Voterra ($3d$) problem, SANODEP-based few-shot BO demonstrates competitive performance on all the rest of the problems.
• Figure 5: Model Comparison of SANODEP vs Physics Informed (PI) SANODEP, PI-SANODEP provides remarkable better prediction accuracy. In addition, PI-SANODEP provides a reasonable estimation (curve) on system parameters (dashed line).

Theorems & Definitions (7)

• Proposition 1
• Lemma C.1
• proof
• Corollary C.1.1
• Lemma C.2
• proof
• proof