Foundation Inference Models for Ordinary Differential Equations

AI summary

Overview

Research area: Machine learning for scientific discovery, specifically amortised inference of ordinary differential equations (ODEs) from noisy trajectory data.

Technical level: Advanced. The paper assumes familiarity with dynamical systems, neural operators, attention mechanisms, and amortised/simulation-based inference.

Scope: The paper introduces FIM-ODE, a pretrained "Foundation Inference Model" that predicts the vector field of a low-dimensional ODE directly from noisy, subsampled trajectories in a single forward pass, and evaluates it against the pretrained symbolic baseline ODEFormer and classical per-dataset neural and Gaussian process methods.

What This Paper Is About

Recovering the governing equations of a dynamical system from noisy measurements is a core problem in scientific modelling, but existing methods either demand heavy system-specific prior knowledge or expensive, unstable per-dataset training. The authors ask whether a single neural network, pretrained once on a deliberately simple distribution of polynomial ODEs, can infer vector fields for new and even out-of-distribution systems without any retraining. They also ask whether representing the vector field locally (with a neural operator) rather than as a global symbolic expression gives better accuracy while remaining interpretable.

Key Contributions

1. A simple polynomial pretraining prior. The authors define a prior distribution over ODEs in dimensions one to three whose vector fields are sparse multivariate polynomials of total degree at most three with coefficients drawn independently from a standard normal distribution. They show a model pretrained on this prior estimates out-of-distribution vector fields, including systems modelling human-motion trajectories.

2. A local neural-operator representation of the vector field. Instead of predicting a symbolic expression, FIM-ODE represents the inferred field with neural operators and queries it at arbitrary state-space locations. The authors show this local representation matches and often outperforms global symbolic representations, and that it remains interpretable because it can be queried for equilibria, Jacobians, stability types, and phase-space geometry.

3. Pretraining as a finetuning initialisation. The pretrained model serves as a strong starting point, enabling fast and stable finetuning when target dynamics are far out of distribution, which the authors report outperforms modern neural and Gaussian process baselines without requiring machine learning expertise.

4. Open release. The pretrained model, code repository, and tutorials are made available online at a project page listed in the paper.

Main Findings

• Zero-shot trajectory reconstruction on ODEBench. FIM-ODE consistently outperformed ODEFormer across all noise and subsampling configurations. Reported fraction of systems with variance-weighted R² > 0.9: FIM-ODE 84.4% (ρ=0.0, σ=0.0), 80.3% (ρ=0.0, σ=0.03), 75.4% (ρ=0.0, σ=0.05), 82.8% (ρ=0.5, σ=0.0), 74.6% (ρ=0.5, σ=0.03), 72.1% (ρ=0.5, σ=0.05). ODEFormer scored 63.1%, 61.5%, 61.5%, 63.9%, 66.4%, and 61.5% on the corresponding settings.

• Comparable zero-shot generalisation. On trajectory generalisation (evaluated from new initial conditions), the two methods performed comparably. FIM-ODE scored 29.5%, 32.8%, 29.5%, 27.0%, 27.0%, 28.7% at the R² > 0.9 threshold versus ODEFormer's 27.9%, 27.9%, 25.4%, 31.1%, 32.8%, 27.0%. FIM-ODE showed a clearer advantage at the less stringent R² > 0.8 threshold reported in the appendix.

• Much smaller model, far fewer pretraining systems. FIM-ODE's vector field predictor has 8M parameters versus 86M for ODEFormer (8M versus 86M for the vector field predictor; total FIM-ODE model is 13M parameters including the 5M-parameter auxiliary head). It was pretrained on 0.6M synthetic ODE systems versus roughly 50M for ODEFormer, about eighty times fewer.

• A substantial portion of the benchmark is out of distribution. Roughly one third of ODEBench systems have non-polynomial vector fields, including trigonometric and rational components. The authors report that separating in-distribution from out-of-distribution systems shows the aggregate results are not driven only by in-distribution cases.

• Local versus global representation trade-off, shown on three systems. On the frictionless pendulum (ODE 28, containing a sine term and therefore out of distribution), ODEFormer preserved the conservative closed-orbit structure and recovered a centre near the origin with saddle-type equilibria, whereas FIM-ODE reproduced observed trajectories but biased the field near the origin toward weak unstable spirals. On the reduced CDIMA model (ODE 42, containing rational terms), the true system has an unstable spiral near (1.78, 4.17); FIM-ODE identified a nearby candidate equilibrium and correctly classified it as an unstable spiral, while ODEFormer found a nearby exact symbolic equilibrium but reversed the local stability, predicting a stable spiral. On the Lotka-Volterra competition model (ODE 26, in distribution), FIM-ODE split the true coexistence saddle at (1, 1) into two nearby saddle-like candidates and recovered a stable node near (0, 2), but did not recover the stable node at (3, 0) or the unstable equilibrium at (0, 0); ODEFormer also recovered part of the boundary structure but not the full coexistence geometry.

• Low-data OOD regimes are a weakness for pure zero-shot inference. In the Van der Pol forecasting and FitzHugh Nagumo imputation experiments, both pretrained models were typically suboptimal relative to classical per-dataset methods. Running the Van der Pol experiment over 100 independent noise realisations produced large standard deviations, indicating that a short, noisy context trajectory does not constrain the dynamics enough for reliable zero-shot forecasting. Augmenting the context with additional trajectories (adding 2, 8, and 49, with initial conditions sampled from a normal distribution with variance 0.1 around the original Van der Pol initial condition) substantially reduced variance.

• Finetuning recovers strong performance. Finetuning FIM-ODE on the fixed context dataset by minimising mean absolute error between predictions and context data rapidly adapted the model and substantially improved both forecasting and imputation performance. The paper's Table 4 caption states that lower MSE is better and that FIM-ODE (Averaged) reports the mean over 100 independent noise realisations; the truncated content does not include the numeric entries of that table.

Methodology in Plain English

The work follows the amortisation idea behind simulation-based inference: instead of optimising a model separately for every dataset, train one model once on a large collection of synthetic problems so that it learns a reusable inference procedure.

The authors first design a synthetic world. They generate vector fields whose components are sparse multivariate polynomials of total degree at most three, with coefficients drawn independently from a standard normal distribution and random masking of degrees and individual monomials. They generate systems in one, two, and three dimensions. For each sampled field they draw initial conditions from a standard normal distribution, integrate forward over an observation window [0, 10] sampled at 200 equidistant points (Δt = 0.05), using Euler's method with 20 steps per observation interval (integration step size 0.0025). Systems whose trajectories exceed a magnitude of 10² are discarded as divergent. Trajectories are then corrupted with the same scheme used by ODEFormer: multiplicative Gaussian noise, y_i = (1 + ε)x_i with ε ~ N(0, σ²) and σ in [0, 0.06], plus an independent Bernoulli mask that removes observations with probability ρ in [0, 0.5]. All K trajectories from a system share a noise scale but use independent masks.

The model itself is an encoder-decoder neural operator built on attention. Rather than feeding raw trajectories, the encoder builds transition features from each consecutive pair: the current state y_i, the displacement Δy_i, the element-wise squared displacement Δy_i², and the inter-observation time Δτ_i. This is motivated by the fact that Δy_i / Δτ_i is a finite-difference estimate of the vector field. Each feature is projected independently to dimension n/4 and concatenated into an n-dimensional vector; two layers of linear self-attention then produce a permutation-invariant context representation. The decoder takes a query location x, embeds it, and passes it through M cross-attention blocks that attend to the context representation, followed by an MLP that outputs a d-dimensional vector. Because the decoder is queried at arbitrary points, it produces a local estimate of the vector field rather than only at observed states.

State dimensions are normalised to zero mean and unit variance and inter-observation times are re-centred, so the model works in a normalised space and predictions are mapped back using the chain rule. Training minimises a mean absolute error between predicted and ground-truth vector fields at sampled query locations, with half the queries drawn uniformly over the spatial extent of the observed data and half drawn from states along the simulated trajectories. Because flow magnitudes vary widely across state space, the authors add an auxiliary head that predicts a log-variance term U_θ and use uncertainty weighting, corresponding to a Laplace likelihood with heteroscedastic scale.

Pretraining used a single 13M-parameter model (8M for the vector field and 5M for the auxiliary head) on 600K polynomial ODE systems: 80K 1D, 210K 2D, and 310K 3D. The number of context trajectories per system was sampled uniformly between 1 and 9, with 100 to 200 noisy observations per trajectory. Evaluation compared FIM-ODE against ODEFormer and against per-dataset baselines including GP-DNF, Bayesian Neural ODE (BNeuralODE), NeuralODE, ODE2VAE, LatentSDE, GPODE, and npODE. Trajectories generated from inferred vector fields were computed with scipy.integrate.solve_ivp.

Why This Matters

Impact on research. The paper argues that an amortised, local vector-field estimator can rival a much larger symbolic model pretrained on a far more complex prior, which pushes back on the assumption that ODE foundation models need elaborate, heterogeneous pretraining distributions. It also reframes interpretability: the inferred neural vector field can be interrogated for equilibria, Jacobians, stability types, and phase-space geometry just as a symbolic expression can, but it makes only local commitments where data provide constraints rather than global commitments everywhere.

Real-world applications:

• Recovering governing equations in molecular reaction networks, where concentration dynamics are a classical ODE modelling target.
• Population dynamics in biology, including predator-prey and competition systems such as Lotka-Volterra.
• Characterising chaotic systems, with atmospheric convection as the canonical example.
• Modelling human-motion trajectories, which the paper reports as an out-of-distribution application of the pretrained model.

Industry relevance. The pitch is that practitioners without machine learning expertise can obtain a vector field estimate in a single forward pass, or finetune for a specific system, instead of building a symbolic regression, Gaussian process, or Neural ODE pipeline. The released pretrained model, code repository, and tutorials lower the barrier to applying amortised system identification in scientific and engineering workflows.

Future Directions

• Extending beyond three dimensions. The authors state that nothing in the architecture prevents training FIM-ODE on higher-dimensional systems; the restriction to one to three dimensions comes from the current data generation pipeline.

• Handling conservative and non-polynomial structure. The frictionless pendulum failure case, where FIM-ODE biased the field toward weak unstable spirals instead of preserving closed orbits, points to a need for priors or representations that respect conservative geometry.

• Improving low-data zero-shot inference. The large variance across noise realisations in the Van der Pol and FitzHugh Nagumo experiments raises the question of how to make pure zero-shot inference reliable when context windows are short and noisy.

• Addressing the nonstationary prior. The paper notes in its limitations discussion that the polynomial prior's pointwise variance grows with the norm of the state, scaling like ‖x‖^(2α) where α is the highest degree, which has practical consequences for data generation.

Target Audience

Researchers and graduate students in machine learning for science, dynamical systems, and system identification; practitioners in physics, biology, chemistry, and engineering who need to recover governing equations from noisy measurements; and anyone evaluating foundation-model approaches against classical per-dataset inference baselines. Readers need background in ODEs, neural networks, and attention mechanisms to follow the methodology in detail.

Authors’ 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.