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

arXiv: https://arxiv.org/abs/2502.19049 · HTML: https://mlsanity.com/paper/2502.19049/

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Builds on 38 of its 65 references that are in the map.
- First posted on arXiv 26 Feb 2025.

## Builds on (most cited first, 30 of 38)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [On the Opportunities and Risks of Foundation Models](https://mlsanity.com/paper/2108.07258/) (2021)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators](https://mlsanity.com/paper/1910.03193/) (2019)
- [TabPFN: A Transformer That Solves Small Tabular Classification Problems in a Second](https://mlsanity.com/paper/2207.01848/) (2022)
- [Conditional Neural Processes](https://mlsanity.com/paper/1807.01613/) (2018)
- [Analyzing and Improving the Training Dynamics of Diffusion Models](https://mlsanity.com/paper/2312.02696/) (2023)
- [Efficient Attention: Attention with Linear Complexities](https://mlsanity.com/paper/1812.01243/) (2018)
- [Transformers Can Do Bayesian Inference](https://mlsanity.com/paper/2112.10510/) (2021)
- [Neural Processes](https://mlsanity.com/paper/1807.01622/) (2018)
- [Attentive Neural Processes](https://mlsanity.com/paper/1901.05761/) (2019)
- [Scalable Gradients for Stochastic Differential Equations](https://mlsanity.com/paper/2001.01328/) (2020)
- [Deep Learning for Symbolic Mathematics](https://mlsanity.com/paper/1912.01412/) (2019)
- [Neural SDEs as Infinite-Dimensional GANs](https://mlsanity.com/paper/2102.03657/) (2021)
- [Towards a Neural Statistician](https://mlsanity.com/paper/1606.02185/) (2016)
- [ForecastPFN: Synthetically-Trained Zero-Shot Forecasting](https://mlsanity.com/paper/2311.01933/) (2023)
- [Neural SDE: Stabilizing Neural ODE Networks with Stochastic Noise](https://mlsanity.com/paper/1906.02355/) (2019)
- [Theoretical guarantees for sampling and inference in generative models with latent diffusions](https://mlsanity.com/paper/1903.01608/) (2019)
- [Kernels for sequentially ordered data](https://mlsanity.com/paper/1601.08169/) (2016)
- [Sparse learning of stochastic dynamic equations](https://mlsanity.com/paper/1712.02432/) (2017)
- [Efficient and Accurate Gradients for Neural SDEs](https://mlsanity.com/paper/2105.13493/) (2021)
- [ODEFormer: Symbolic Regression of Dynamical Systems with Transformers](https://mlsanity.com/paper/2310.05573/) (2023)
- [Modeling Temporal Data as Continuous Functions with Stochastic Process Diffusion](https://mlsanity.com/paper/2211.02590/) (2022)
- [Learning interpretable continuous-time models of latent stochastic dynamical systems](https://mlsanity.com/paper/1902.04420/) (2019)
- [The Variational Homoencoder: Learning to learn high capacity generative models from few examples](https://mlsanity.com/paper/1807.08919/) (2018)
- [SDE Matching: Scalable and Simulation-Free Training of Latent Stochastic Differential Equations](https://mlsanity.com/paper/2502.02472/) (2025)
- [Inference Networks for Sequential Monte Carlo in Graphical Models](https://mlsanity.com/paper/1602.06701/) (2016)
- [Dual Parameterization of Sparse Variational Gaussian Processes](https://mlsanity.com/paper/2111.03412/) (2021)
- [Latent SDEs on Homogeneous Spaces](https://mlsanity.com/paper/2306.16248/) (2023)
- [Amortized In-Context Bayesian Posterior Estimation](https://mlsanity.com/paper/2502.06601/) (2025)

## Built on by (most cited first, 7 of 7)

- [In-Context Learning of Temporal Point Processes with Foundation Inference Models](https://mlsanity.com/paper/2509.24762/) (2025)
- [Domain-Adaptation through Synthetic Data: Fine-Tuning Large Language Models for German Law](https://mlsanity.com/paper/2601.14160/) (2026)
- [Towards Fast Coarse-graining and Equation Discovery with Foundation Inference Models](https://mlsanity.com/paper/2510.12618/) (2025)
- [Towards Foundation Inference Models that Learn ODEs In-Context](https://mlsanity.com/paper/2510.12650/) (2025)
- [EVIL: Evolving Interpretable Algorithms for Zero-Shot Inference on Event Sequences and Time Series with LLMs](https://mlsanity.com/paper/2604.15787/) (2026)
- [Prior-Fitted Functional Flow: In-Context Generative Models for Pharmacokinetics](https://mlsanity.com/paper/2604.17670/) (2026)
- [Learning Continuous Neural Representation of Stochastic Hybrid Systems](https://mlsanity.com/paper/2609.38893/) (2026)

## Abstract (arXiv)

Stochastic differential equations (SDEs) describe dynamical systems where deterministic flows, governed by a drift function, are superimposed with random fluctuations, dictated by a diffusion function. The accurate estimation (or discovery) of these functions from data is a central problem in machine learning, with wide application across the natural and social sciences. Yet current solutions either rely heavily on prior knowledge of the dynamics or involve intricate training procedures. We introduce FIM-SDE (Foundation Inference Model for SDEs), a pretrained recognition model that delivers accurate in-context (or zero-shot) estimation of the drift and diffusion functions of low-dimensional SDEs, from noisy time series data, and allows rapid finetuning to target datasets. Leveraging concepts from amortized inference and neural operators, we (pre)train FIM-SDE in a supervised fashion to map a large set of noisy, discretely observed SDE paths onto the space of drift and diffusion functions. We demonstrate that FIM-SDE achieves robust in-context function estimation across a wide range of synthetic and real-world processes -- from canonical SDE systems (e.g., double-well dynamics or weakly perturbed Lorenz attractors) to stock price recordings and oil-price and wind-speed fluctuations -- while matching the performance of symbolic, Gaussian process and Neural SDE baselines trained on the target datasets. When finetuned to the target processes, we show that FIM-SDE consistently outperforms all these baselines.