🤖 machine learning

In-Context Learning of Stochastic Differential Equations with Foundation Inference Models

The paper introduces FIM-SDE, a pretrained foundation model that leverages in-context learning and supervised training on diverse SDE paths to achieve accurate, zero-shot estimation of drift and diffusion functions from noisy time series data, while consistently outperforming traditional baselines when finetuned on target datasets.

Original authors: Patrick Seifner, Kostadin Cvejoski, David Berghaus, Cesar Ojeda, Ramses J. Sanchez

Published 2026-06-09

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Patrick Seifner, Kostadin Cvejoski, David Berghaus, Cesar Ojeda, 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

Imagine you are trying to figure out the rules of a game just by watching someone play it, but there's a catch: the player is being pushed around by invisible, random gusts of wind. You can see where they go, but you can't see the wind itself, and your view of the game is sometimes blurry or interrupted.

This is the problem scientists face when studying Stochastic Differential Equations (SDEs). These are mathematical formulas that describe systems moving with a mix of predictable flow (like a river current) and random chaos (like the wind). The goal is to find the two hidden "functions" that control this:

The Drift: The predictable part (the river current).
The Diffusion: The random part (the wind).

Usually, figuring out these rules requires a lot of data, a lot of time, and a lot of guesswork about what the rules might look like. If the data is messy or sparse, old methods often fail or take forever to compute.

The New Solution: FIM-SDE

The authors of this paper introduce a new tool called FIM-SDE (Foundation Inference Model for SDEs). Think of FIM-SDE as a super-smart detective who has spent years studying millions of fake games in a simulation lab.

Here is how it works, using simple analogies:

1. The "Training Camp" (Pre-training)

Before the detective ever sees a real case, they go to a massive training camp.

The Simulation: The researchers generate 600,000 fake scenarios. They invent random "drift" and "diffusion" rules (mostly simple math formulas like polynomials), simulate the game, and then deliberately "corrupt" the data. They add noise, blur the view, and skip time steps.
The Lesson: The detective (the AI model) is trained to look at these messy, fake game logs and immediately guess the hidden rules that created them. It learns to recognize patterns in the chaos.

2. The "Zero-Shot" Superpower (In-Context Learning)

Once trained, the detective can look at a brand new game log (real-world data) and guess the rules immediately, without needing to study that specific game for hours.

The Analogy: Imagine you've seen thousands of different types of weather patterns. If someone shows you a new, slightly blurry photo of a storm, you can instantly say, "Ah, that's a low-pressure system with high wind," without needing to run a new computer simulation.
The Result: The paper shows that FIM-SDE can look at messy data from things like stock prices, oil prices, or wind speeds and guess the underlying rules almost as well as experts who have spent days analyzing that specific data.

3. The "Quick Study" (Fine-tuning)

Sometimes, the detective's first guess is good, but not perfect. Maybe the real game has a weird quirk the training camp didn't cover.

The Analogy: Instead of starting from scratch, the detective just needs a quick refresher course.
The Result: The paper claims that FIM-SDE can adapt to a specific new dataset in just 100 steps (iterations). Other methods might need 5,000 steps to get the same accuracy. It's like the detective already knows the language and just needs to learn a few new slang words, whereas others are trying to learn the whole language from zero.

What Did They Test?

The authors tested this detective on two types of challenges:

Classic Physics Problems: Systems like the "Double-Well" (a ball rolling between two hills) and the "Lorenz Attractor" (a famous chaotic weather model).
Real-World Data: Stock prices for Facebook and Tesla, oil price fluctuations, and wind speed recordings.

The Verdict

On Messy Data: FIM-SDE is much more robust than older methods. When the data is noisy or sparse, the old methods often give up or give wrong answers. FIM-SDE keeps working.
On Speed: It is incredibly fast to adapt. It can learn a new system 50 times faster than previous deep learning methods.
On Accuracy: When it gets a quick "fine-tuning" on the specific data, it beats all the other methods (including symbolic math and Gaussian processes) in accuracy.

The Limitations (The Fine Print)

The paper is honest about what the detective can't do yet:

Simple Rules Only: The detective was trained on relatively simple math rules (polynomials). If the real world uses incredibly complex, non-mathematical rules, the detective might struggle.
Small Systems: It works best on systems with 1 to 3 variables (like a ball moving in 3D space). Trying to apply it to systems with hundreds of variables is currently too difficult because the "training camp" simulation becomes too slow and prone to errors.

In summary: FIM-SDE is a pre-trained AI that learns the "grammar" of random motion from millions of simulations. It can instantly read a messy real-world story and guess the rules, and if you give it a tiny bit of extra help, it becomes the best at guessing those rules faster than anyone else.

Technical Summary: In-Context Learning of Stochastic Differential Equations with Foundation Inference Models

Problem Statement

Stochastic Differential Equations (SDEs) model dynamical systems where deterministic flows (drift) are superimposed with random fluctuations (diffusion). A central challenge in machine learning and scientific discovery is the accurate estimation or "discovery" of these drift and diffusion functions from noisy, time-series data.

Current approaches face significant limitations:

Symbolic Regression: Relies heavily on pre-specified dictionaries of basis functions and often requires dense, clean data.
Gaussian Process (GP) Regression: Uses non-parametric priors but is sensitive to prior knowledge and often struggles with convergence when data is sparse or noisy, frequently requiring variational approximations.
Neural SDEs: Represent drift and diffusion via neural networks but typically require costly training involving backpropagation through SDE solvers (often via adjoint sensitivity methods) and must be retrained from scratch for every new system.

The paper identifies a gap for a method that can perform in-context (zero-shot) estimation of SDE functions from noisy data without extensive retraining, while remaining robust to sparse observations and noise.

Methodology: FIM-SDE

The authors introduce FIM-SDE (Foundation Inference Model for SDEs), a pretrained neural recognition model designed to map noisy, discretely observed SDE paths directly to their underlying drift and diffusion functions. The methodology consists of two primary components:

1. Data Generation Model

To train a model capable of generalizing to unseen systems, the authors construct a synthetic dataset based on a generative prior $p_{prior}(f, G)$ .

Function Space: Drift functions are sampled from polynomials of degree $\leq 3$ , and diffusion functions from positive polynomials of degree $\leq 2$ . This covers a broad range of physical models, including chaotic systems.
Simulation & Corruption: SDE paths are simulated using the Euler-Maruyama scheme. These paths are then "corrupted" to mimic real-world data:
- Noise: Additive Gaussian noise is applied.
- Irregular Sampling: Paths are subsampled to create irregular inter-observation times ( $\Delta \tau$ ), which can be significantly larger than the simulation step ( $\Delta t$ ).
Training Data: The dataset comprises tuples of corrupted time series and their corresponding ground-truth drift/diffusion pairs.

2. Neural Recognition Model Architecture

FIM-SDE is a Transformer-based architecture inspired by Neural Operators (specifically DeepONets) and amortized inference.

Input Representation: The input context consists of transition tuples $(y_i, \Delta y_i, \Delta y_i^2, \Delta \tau_i)$ derived from the noisy time series. The inclusion of $\Delta y^2$ allows the network to learn separate feature channels for drift and diffusion estimation.
Context Matrix: A self-attentive Transformer encoder ( $\Psi$ ) processes the embedded transition tuples to create a "Context Matrix." This acts as the equivalent of a DeepONet's Branch-net, encoding the entire observation history.
Functional Attention: To estimate the drift or diffusion at a specific query location $x$ , the model employs a "Functional Attention" mechanism (equivalent to a DeepONet's Trunk-net). This mechanism queries the Context Matrix using the location $x$ as the query, dynamically weighting transitions based on their spatial proximity to $x$ .
Training Objective: The model is trained to minimize a weighted Mean Squared Error (MSE) between the estimated functions and the ground truth. A learnable uncertainty estimator $U_\theta$ is used to down-weight regions where the model is uncertain (e.g., sparse or highly noisy areas), preventing numerical instability.

3. Finetuning Strategy

While FIM-SDE operates in a zero-shot mode, it supports rapid finetuning:

Dense Data: Finetuning maximizes the likelihood of short-time transitions.
Sparse Data: Finetuning simulates the inferred SDE over the observation gaps and minimizes the discrepancy between the simulated and observed states.
This allows the model to adapt to specific target datasets with significantly fewer iterations than training a Neural SDE from scratch.

Key Contributions

Data Generation Model: A novel model for sampling and corrupting SDE paths restricted to polynomial drift and diffusion. The authors empirically demonstrate that this model encodes a strong prior, enabling generalization to both in-distribution synthetic systems and real-world datasets.
First In-Context Neural Recognizer: FIM-SDE is the first neural model capable of estimating drift and diffusion functions of low-dimensional SDEs from noisy data in a zero-shot manner.
Rapid Adaptation: The model can be finetuned to target processes (both dense and sparse) converging up to 50 $\times$ faster than standard Neural SDE methods while ultimately outperforming them.

Experimental Results

The authors evaluated FIM-SDE on eight canonical SDE systems (e.g., Double-well, Lorenz attractor) and four real-world datasets (Facebook/Tesla stock prices, oil prices, wind speed).

Zero-Shot Performance: Across canonical systems with varying noise levels ( $\rho$ ) and observation gaps ( $\Delta \tau$ ), FIM-SDE achieved robust performance. In noisy or sparse settings, it consistently outperformed Symbolic Regression (BISDE) and Gaussian Process (SparseGP) baselines. On clean, dense data, it performed comparably to these baselines. Notably, FIM-SDE remained stable where symbolic and GP methods frequently produced invalid SDE estimates or failed to converge.
Real-World Data: On stock price data, FIM-SDE's zero-shot performance surpassed both BISDE and LatentSDE. On wind and oil data, while zero-shot performance lagged slightly behind BISDE, finetuning FIM-SDE yielded the best overall performance across all systems.
Efficiency: When finetuned on the Lorenz attractor, FIM-SDE surpassed the LatentSDE baseline (trained for 5000 iterations) in less than 100 iterations. In real-world datasets, FIM-SDE required only 100 finetuning iterations compared to 5000 for LatentSDE, representing a 50 $\times$ reduction in compute.

Significance and Claims

The paper claims that FIM-SDE functions as a true foundation model for SDE discovery. Its significance lies in:

Circumventing Training Bottlenecks: It eliminates the need for costly, system-specific training pipelines and backpropagation through SDE solvers for every new dataset.
Robustness: It handles noisy, sparse, and irregularly sampled data effectively without requiring explicit variational regularization between observations.
Scientific Automation: By providing a zero-shot solution that can be rapidly adapted, it offers a pathway toward automating scientific discovery in scenarios where prior knowledge of dynamics is scarce.

The authors acknowledge limitations, primarily that the current prior is restricted to low-degree polynomial drift/diffusion and low-dimensional systems (due to rejection rates in high-dimensional sampling). However, they posit that the strong inductive biases learned from this restricted class are sufficient to serve as a powerful initialization for real-world applications.