🤖 machine learning

Foundation Inference Models for Ordinary Differential Equations

The paper introduces FIM-ODE, a pretrained foundation model that efficiently infers ODE vector fields from noisy trajectories in a single forward pass, achieving strong zero-shot performance and enabling rapid, expert-free finetuning that outperforms existing symbolic, neural, and Gaussian process baselines.

Original authors: Maximilian Mauel, Johannes R. Hübers, David Berghaus, Patrick Seifner, Ramses J. Sanchez

Published 2026-06-09

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Maximilian Mauel, Johannes R. Hübers, David Berghaus, Patrick Seifner, Ramses J. Sanchez

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Idea: Teaching a Detective to Solve Crimes Instantly

Imagine you are a detective trying to figure out the rules of a game just by watching someone play it. You see a ball rolling across a table, bouncing off a wall, and slowing down. Your job is to guess the invisible "force field" (the rules) that caused the ball to move that way.

In science, this is called solving an Ordinary Differential Equation (ODE). Scientists use these equations to describe everything from how planets orbit to how viruses spread. Usually, figuring out these rules from messy, noisy data is like trying to solve a puzzle while wearing thick gloves—it's slow, difficult, and requires a lot of expert training.

This paper introduces a new tool called FIM-ODE. Think of FIM-ODE as a super-detective who has already solved millions of practice puzzles before ever seeing your specific case. Because it has done so much "homework" (pretraining), it can look at your messy data and instantly guess the rules of the game in a single glance, without needing to re-learn everything from scratch.

How It Works: The "Training Camp" vs. The "Real Game"

1. The Training Camp (Pretraining)
Most AI models are trained on one specific problem at a time. If you want to predict the weather, you train one model. If you want to predict stock prices, you train another. This is like hiring a new employee for every single job and making them read the manual from page one.

The authors took a different approach. They built a massive "training camp" for their AI.

The Curriculum: They generated 600,000 fake, made-up physics problems. These were simple enough to be solvable (using low-degree polynomials, which are like basic algebraic curves) but complex enough to teach the AI how to spot patterns.
The Lesson: The AI learned a general "inference algorithm." Instead of memorizing specific answers, it learned how to figure out the rules from noisy observations.

2. The Real Game (Zero-Shot Inference)
Once the AI finished its training camp, they threw it into the real world.

The Test: They gave it real-world data it had never seen before, including data from complex chemical reactions and even human motion capture (people walking and running).
The Result: The AI didn't need to be retrained. It looked at the messy data and immediately produced a map of the "force field" driving the motion. In many cases, it performed better than the previous state-of-the-art models, even though it was much smaller and had seen far fewer training examples.

The Secret Sauce: Local vs. Global Maps

The paper makes a clever distinction between two ways of drawing a map:

The Global Map (The Old Way): Previous models tried to write a single, perfect mathematical sentence (a symbolic formula) that described the entire universe of the problem. It's like trying to describe the entire ocean with one sentence. If the sentence is slightly wrong in one spot, the whole map is wrong.
The Local Map (FIM-ODE's Way): FIM-ODE doesn't try to write one perfect sentence for the whole world. Instead, it acts like a hologram. It looks at a specific spot where you have data and draws a tiny, accurate map just for that neighborhood. It knows that if you move a little bit, the rules might change slightly, so it updates its local map.

The Analogy: Imagine trying to describe the terrain of a mountain.

The Global approach tries to write one poem that describes the peak, the valley, and the slope all at once.
The Local approach (FIM-ODE) says, "I don't need a poem for the whole mountain. I just need to know exactly what the ground feels like right where you are standing." This turns out to be much more accurate when the data is messy.

Why This Matters

1. Speed and Simplicity
Usually, to get a good model, you need a team of experts, massive computers, and weeks of training. FIM-ODE is like a pre-loaded GPS. You don't need to build the map; you just turn it on, and it guides you. It allows scientists without deep machine learning expertise to infer complex dynamics quickly.

2. Handling Messy Data
Real-world data is never perfect. It has noise (static) and gaps (missing pieces). Because FIM-ODE was trained on millions of "corrupted" fake scenarios, it is incredibly good at ignoring the static and filling in the gaps.

3. The "Fine-Tuning" Option
Sometimes, the pre-trained detective is good, but not perfect for a very specific, weird case (like a specific person's unique walking style). The paper shows that you can take this pre-trained model and give it a tiny bit of extra training (fine-tuning) on just that one specific case. It adapts incredibly fast, becoming a specialist in minutes rather than days.

Limitations (The Catch)

The paper is honest about where this tool currently struggles:

The "Curse of Dimensionality": The training camp was mostly for 1, 2, or 3-dimensional problems (like a ball moving in a line, on a plane, or in a room). When things get very complex (high dimensions), the "fake" training data sometimes explodes or becomes unstable. The authors admit they are currently limited to 3 dimensions because generating stable, realistic data for higher dimensions is hard.
Not a Magic Bullet: If the data is too sparse or the system is completely unlike anything the AI saw in training, it might still get confused. It is a tool for hypothesis generation, not a replacement for human scientific judgment.

Summary

The authors built a "Foundation Inference Model" for physics equations. By training a neural network on a vast library of simple, synthetic physics problems, they created a model that can instantly guess the rules of motion from noisy data. It works like a local mapmaker rather than a global poet, allowing it to be more accurate and adaptable than previous methods, all while requiring less computing power and expertise to use.

Technical Summary: Foundation Inference Models for Ordinary Differential Equations (FIM-ODE)

1. Problem Statement

Ordinary differential equations (ODEs) are fundamental to scientific modeling, yet inferring their underlying vector fields from noisy, sparse, and irregularly sampled trajectory data remains a significant challenge. Traditional approaches fall into three categories:

Symbolic Regression: Methods like SINDy or genetic programming search for symbolic expressions but often require clean, densely sampled data, accurate derivative estimates, or strong prior knowledge about the functional form.
Gaussian Process (GP) Regression: These methods depend heavily on the choice of the GP prior and often involve complex variational approximations or gradient matching.
Neural ODEs: While flexible, these require expensive, unstable training pipelines involving backpropagation through numerical solvers or slow adjoint methods, demanding substantial machine learning expertise.

These classical approaches follow a "per-dataset" optimization paradigm, where a model is trained from scratch for each specific system. This paper proposes shifting to an amortized inference paradigm using Foundation Inference Models (FIMs). The goal is to pretrain a single model on a diverse synthetic dataset to learn a reusable inference procedure, enabling fast, zero-shot vector field estimation from noisy trajectories in a single forward pass.

2. Methodology: FIM-ODE

The proposed FIM-ODE framework consists of two main components: a pretraining prior distribution over ODEs and a neural inference model.

2.1 Pretraining Prior Distribution

The authors define a synthetic data generation pipeline based on a prior distribution over autonomous, first-order ODEs in dimensions $d \in \{1, 2, 3\}$ .

Vector Fields: The prior assumes vector fields are composed of sparse, multivariate polynomials of degree at most three. Coefficients are sampled from $\mathcal{N}(0, 1)$ , with random masking applied to degrees and monomials to induce sparsity and structural variety. This choice is motivated by the observation that many canonical ODEs (e.g., Lorenz system, biological oscillators) are low-degree polynomial systems, and polynomials guarantee the existence and uniqueness of solutions via the Picard-Lindelöf theorem.
Data Generation: For each sampled vector field, multiple trajectories are simulated from initial conditions drawn from $\mathcal{N}(0, 1)$ .
Corruption: To mimic real-world observations, trajectories are corrupted with multiplicative Gaussian noise and random subsampling (removing a fraction of observations). Systems producing divergent trajectories are discarded to focus training on bounded regimes.

2.2 Model Architecture: Transformer-based Neural Operator

FIM-ODE is a Transformer-based neural operator designed to map noisy trajectory contexts to local vector field estimates.

Input Representation: Instead of raw trajectories, the model processes "transition features" derived from consecutive state pairs $(y_i, y_{i+1})$ . These features include the current state, displacement, squared displacement, and inter-observation times. This design leverages the finite-difference approximation of the vector field.
Encoder: A context encoder uses self-attention to create a permutation-invariant representation of the $K$ noisy trajectories.
Decoder: A functional attention decoder queries this context representation at arbitrary spatial locations $x$ . It uses cross-attention to extract relevant information from the context to produce a local estimate of the vector field $\hat{f}_\theta(x)$ .
Scale Invariance: The model employs reversible instance normalization for spatial dimensions and delta log centering for temporal dimensions to handle varying scales across different ODE systems.
Loss Function: Training utilizes a mixed sampling strategy (half from observed trajectory points, half from uniform spatial sampling) and an uncertainty-weighted loss (Laplace likelihood with heteroscedastic scale) to prevent the optimization from being dominated by high-magnitude regions of the vector field.

3. Key Contributions

Simpler Pretraining Prior with Strong Generalization: The authors introduce a pretraining prior based on low-degree polynomial vector fields (degree $\le 3$ ) in dimensions 1–3. Despite this prior being significantly simpler than the complex prior used by ODEFormer (which includes trigonometric and rational functions), FIM-ODE achieves strong zero-shot performance, often matching or outperforming ODEFormer.
Local Neural Operator Representation: Unlike symbolic regression methods that seek a global functional form, FIM-ODE represents the vector field locally using a neural operator. The authors demonstrate that this local representation can match or exceed the accuracy of global symbolic representations, particularly in data-rich regimes. Crucially, they show this local representation is interpretable: it can be queried for equilibria, Jacobians, stability types, and phase-space geometry, similar to symbolic ODEs.
Effective Finetuning for Out-of-Distribution (OOD) Regimes: The paper demonstrates that the pretrained model serves as a strong initialization. When target dynamics are far out-of-distribution (e.g., real-world human motion or low-data regimes), modest finetuning on the specific context data allows for rapid and stable adaptation, outperforming modern neural and GP baselines without requiring extensive ML expertise.

4. Experimental Results

The authors evaluate FIM-ODE across three main experimental settings:

ODEBench (Zero-Shot Inference): Compared against ODEFormer (an 86M parameter model pretrained on 50M systems) on a benchmark of 61 autonomous ODEs.
- FIM-ODE (8M parameters, pretrained on 0.6M systems) consistently outperformed ODEFormer in trajectory reconstruction across all noise and subsampling configurations.
- In trajectory generalization (predicting from new initial conditions), FIM-ODE performed comparably to ODEFormer, with a clear advantage at less stringent thresholds.
- Interpretability Analysis: On specific OOD systems (e.g., frictionless pendulum, CDIMA reaction), FIM-ODE successfully identified local stability properties (e.g., unstable spirals) that global symbolic models sometimes misclassified, highlighting the benefit of local estimation over rigid global functional commitments.
Low-Data OOD Regime (Van der Pol & FitzHugh-Nagumo):
- In zero-shot settings with short, noisy context windows, FIM-ODE showed sensitivity to noise realizations, a limitation of purely zero-shot inference in data-scarce regimes.
- However, finetuning FIM-ODE on the specific context data rapidly adapted the model, yielding substantial performance gains and surpassing classical baselines (NeuralODE, GP-based methods).
Real-World Human Motion (CMU Motion Capture):
- Applied to 5D latent representations of walking and running sequences.
- Zero-shot performance varied by subject; it was excellent for some (Subject 09) but poor for others (Subjects 35, 39) whose dynamics were less aligned with the polynomial prior.
- Finetuning on subject-specific data dramatically recovered performance, demonstrating that the 3D ODE model class is sufficient to represent these dynamics once the specific vector field is adapted.

5. Significance and Claims

The paper claims that FIM-ODE lowers the barrier to data-driven ODE inference. By amortizing the inference cost through pretraining, it enables scientists to infer dynamical models without:

Training bespoke models for each dataset.
Relying on large, clean datasets.
Possessing substantial machine learning expertise or computational budgets.

The authors emphasize that while the model is not a "silver bullet" for all dynamical systems (noting limitations in the non-stationary scaling of the polynomial prior and fixed dimensionality), it provides a powerful tool for hypothesis generation and model-assisted analysis. The results suggest that a local, neural-operator approach, trained on a simple prior, can generalize effectively to complex, real-world systems, offering a compelling alternative to complex symbolic regression or unstable neural ODE training pipelines.

Limitations Acknowledged:

The polynomial prior induces non-stationary scaling, causing vector field magnitudes to grow with distance from the origin, which limits the generation of stable, high-dimensional trajectories (hence the restriction to $d \le 3$ ).
The architecture currently has a fixed maximum state dimension, requiring a priori selection of dimensionality.
Zero-shot inference can fail in low-data regimes or when the target dynamics are poorly represented by the prior, necessitating finetuning.