Research map / Probabilistic and scientific modeling / Differential equations and PDEs
In-Context Learning of Stochastic Differential Equations with Foundation Inference Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Pretrains a transformer/neural-operator recognition model (FIM-SDE) on 600K synthetic polynomial SDEs so it infers drift and diffusion of unseen low-dimensional SDEs in-context, and finetunes it rapidly on target data.
Problem
Estimating drift and diffusion functions of stochastic differential equations from noisy, sparse or irregularly observed time series, without per-system retraining or strong prior knowledge.
Approach
A pretrained recognition model with self-attentive context encoding and functional attention (DeepONet-style branch/trunk) maps a set of noisy SDE transition tuples to drift and diffusion functions, trained on a synthetic prior of polynomial SDEs with an adaptive uncertainty-weighted MSE loss; finetuning adapts it to dense or sparse target data.
Key results
- On canonical SDEs with noise and sparse observations, zero-shot FIM-SDE achieves the lowest MMD in most noisy settings (e.g., rho=0.05, Delta tau=0.002: 2.2 vs 9 for BISDE on 2D-Synt).
- On real-world datasets, finetuned FIM-SDE gives the best MMD on Oil (0.068) and Facebook (0.006), while zero-shot FIM-SDE beats BISDE and LatentSDE on the stock datasets.
- On Lorenz attractor, FIM-SDE finetuned for 100 iterations surpasses LatentSDE trained for 5000 iterations on all three initial-condition tests (e.g., 0.0038 vs 0.0517 on N(0,1)).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| finetune FIM-SDE on target data | FIM-SDE · Wind | MMD: 0.330 ± 0.114 -> 0.029 ± 0.034 | Helped |
| finetune FIM-SDE on target data | FIM-SDE · Oil | MMD: 0.199 ± 0.137 -> 0.068 ± 0.042 | Helped |
| finetune FIM-SDE on target data | FIM-SDE · Facebook | MMD: 0.012 ± 0.026 -> 0.006 ± 0.007 | Helped |
| finetune FIM-SDE on target data | FIM-SDE · Tesla | MMD: 0.000 ± 0.000 -> 0.004 ± 0.008 | Hurt |
| finetune FIM-SDE on context data | FIM-SDE · Lorenz attractor, initial condition N(0, 1) | MMD: 0.403933 -> 0.003838 | Helped |
| finetune FIM-SDE on context data | FIM-SDE · Lorenz attractor, initial condition (1, 1, 1) | MMD: 1.151080 -> 0.037159 | Helped |
| finetune FIM-SDE on context data | FIM-SDE · Lorenz attractor, initial condition N(0, 2) | MMD: 0.378515 -> 0.003335 | Helped |
| use FIM-SDE zero-shot instead of LatentSDE | Lorenz attractor · Lorenz attractor, initial condition N(0, 1) | MMD: 0.051708 -> 0.403933 | Hurt |
| use FIM-SDE zero-shot instead of LatentSDE | Lorenz attractor · Lorenz attractor, initial condition (1, 1, 1) | MMD: 0.135003 -> 1.151080 | Hurt |
| use FIM-SDE zero-shot instead of LatentSDE | Lorenz attractor · Lorenz attractor, initial condition N(0, 2) | MMD: 0.054335 -> 0.378515 | Hurt |
Surprising findings
- Symbolic (BISDE) and GP (SparseGP) baselines often produce invalid SDE estimates even on uncorrupted data.
- Zero-shot FIM-SDE performs poorly on the drift-dominated Lorenz attractor, biased toward non-zero diffusion, though finetuning corrects this.
Limitations
- Prior restricted to low-degree polynomial drift and diffusion and low-dimensional (up to 3D) systems.
- Rejection rates in data generation rise sharply with dimension (50% in 1D, 78% in 2D, 92% in 3D).
- Zero-shot performance is weak on drift-dominated systems such as the perturbed Lorenz attractor; diagonal diffusion only.
Experimental setup
Datasets: synthetic polynomial SDE prior (600K SDEs) double-well 2D synthetic (Wang) damped linear damped cubic Duffing glycolysis Hopf Lorenz attractor Facebook stock Tesla stock oil price
Benchmarks: double-well 2D-Synt (Wang) damped linear damped cubic Duffing glycolysis Hopf Lorenz attractor Facebook stock Tesla stock oil price wind speed
Metrics: MMD (signature kernel) MSE of drift/diffusion
Baselines: SparseGP BISDE LatentSDE
Builds on: DeepONet SINDy Neural SDE prior-fitted networks
Closest papers by meaning
- Neural Stochastic Flows: Solver-Free Modelling and Inference for SDE Solutions (NeurIPS 2025 · similarity 0.91)
- SING: SDE Inference via Natural Gradients (NeurIPS 2025 · similarity 0.88)
- Stable Neural Stochastic Differential Equations in Analyzing Irregular Time Series Data (ICLR 2024 · similarity 0.88)
- Efficient Training of Neural Stochastic Differential Equations by Matching Finite Dimensional Distributions (ICLR 2025 · similarity 0.88)
- Variational Inference for SDEs Driven by Fractional Noise (ICLR 2024 · similarity 0.88)
- Statistical Spatially Inhomogeneous Diffusion Inference (AAAI 2024 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.