Foundation Inference Models for Ordinary Differential Equations

TL;DR

This work addresses the challenge of inferring vector fields for autonomous ODEs from noisy trajectories by learning a reusable inference procedure. It introduces Foundation Inference Models for ODEs (FIM-ODE), which pair a simple, polynomial-degree-$3$ pretraining prior in dimensions $d\in\{1,2,3\}$ with a Transformer-based neural-operator that represents the vector field locally and predicts $\hat{\mathbf{f}}_{\theta}(\mathbf{x}|\tilde{\mathcal{D}})$. The model achieves strong zero-shot performance compared to ODEFormer, while being much smaller and pretrained on far fewer systems, and it supports rapid finetuning to adapt to out-of-distribution dynamics, including real human-motion trajectories. The work also analyzes the benefits and trade-offs of local vs global representations and outlines directions to extend to higher dimensions with alternative priors and architecture tweaks, aiming toward data-efficient, broadly applicable ODE inference. Overall, FIM-ODE offers a practical, scalable pathway for accurate, fast ODE inference across scientific domains.

Abstract

Ordinary differential equations (ODEs) are central to scientific modelling, but inferring their vector fields from noisy trajectories remains challenging. Current approaches such as symbolic regression, Gaussian process (GP) regression, and Neural ODEs often require complex training pipelines and substantial machine learning expertise, or they depend strongly on system-specific prior knowledge. We propose FIM-ODE, a pretrained Foundation Inference Model that amortises low-dimensional ODE inference by predicting the vector field directly from noisy trajectory data in a single forward pass. We pretrain FIM-ODE on a prior distribution over ODEs with low-degree polynomial vector fields and represent the target field with neural operators. FIM-ODE achieves strong zero-shot performance, matching and often improving upon ODEFormer, a recent pretrained symbolic baseline, across a range of regimes despite using a simpler pretraining prior distribution. Pretraining also provides a strong initialisation for finetuning, enabling fast and stable adaptation that outperforms modern neural and GP baselines without requiring machine learning expertise.

Figures (21)

• Figure 1: Synthetic data generation (left) and FIM-ODE architecture (right).
• Figure 2: Comparison of ODEformer and FIM-ODE on two ODEBench systems. Each model infers a vector field from a single corrupted context trajectory, obtained by subsampling the ground truth ($\rho=0.5$) and adding noise ($\sigma=0.03$). The clean context trajectory is shown in green in the left column. We then integrate the inferred vector fields from two initial conditions specified by ODEBench to assess reconstruction and generalisation.
• Figure 3: Magnitude statistics of vector field points as a function of relative distance to bounding box borders for 1D ODEs. The 1D case exhibits unique behavior due to outlier systems (see Figure \ref{['fig:1d_outlier']}).
• Figure 4: Magnitude statistics of vector field points as a function of relative distance to bounding box borders for 2D ODEs.
• Figure 5: Magnitude statistics of vector field points as a function of relative distance to bounding box borders for 3D ODEs.