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

TL;DR

This paper introduces FIM-SDE, a pretrained model capable of accurately estimating the drift and diffusion functions of low-dimensional stochastic differential equations from noisy data, enabling rapid adaptation and superior performance over existing methods.

Contribution

The paper presents FIM-SDE, a novel foundation inference model that performs in-context estimation of SDE functions, reducing reliance on prior knowledge and complex training procedures.

Findings

FIM-SDE achieves robust in-context estimation across diverse synthetic and real-world SDEs.

FIM-SDE matches the performance of traditional baselines when used out-of-the-box.

Finetuning FIM-SDE improves its accuracy beyond existing methods.

Abstract

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.