Foundation Inference Models for Stochastic Differential Equations: A Transformer-based Approach for Zero-shot Function Estimation

2025-02-26

AI Summary

This paper addresses the challenging zero-shot estimation of drift and diffusion functions in stochastic differential equations (SDEs) from sparse, noisy multivariate time-series observations. We propose FIM-SDE, the first Transformer-based foundational inference model for SDEs, which integrates amortized variational inference with neural operator principles and is trained via supervised functional mapping. Unlike conventional symbolic regression or Bayesian approaches, FIM-SDE achieves cross-system, fine-tuning-free generalization—requiring no parameter updates on target systems. Evaluated on diverse real-world and synthetic benchmarks—including double-well dynamics, Hopf bifurcations, human motion capture, oil prices, and wind speed data—FIM-SDE delivers high-accuracy, robust zero-shot function estimation. It substantially overcomes the generalization limitations of existing methods, establishing a new paradigm for physics-informed, data-driven SDE inference under scarce and corrupted observations.

Technology Category

• Machine Learning: Other Foundations of Machine Learning
• Reasoning under Uncertainty: Stochastic Optimization
• Computer Vision: Diffusion Models for Vision

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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 natural and social sciences alike. Yet current solutions are brittle, and typically rely on symbolic regression or Bayesian non-parametrics. In this work, we introduce FIM-SDE (Foundation Inference Model for SDEs), a transformer-based recognition model capable of performing accurate zero-shot estimation of the drift and diffusion functions of SDEs, from noisy and sparse observations on empirical processes of different dimensionalities. Leveraging concepts from amortized inference and neural operators, we train FIM-SDE in a supervised fashion, to map a large set of noisy and discretely observed SDE paths to their corresponding drift and diffusion functions. We demonstrate that one and the same (pretrained) FIM-SDE achieves robust zero-shot function estimation (i.e. without any parameter fine-tuning) across a wide range of synthetic and real-world processes, from canonical SDE systems (e.g. double-well dynamics or weakly perturbed Hopf bifurcations) to human motion recordings and oil price and wind speed fluctuations.

Problem

• Estimate drift and diffusion functions
• Zero-shot function estimation
• Transformer-based recognition model

Innovation

• Transformer-based zero-shot estimation
• Amortized inference for SDE functions
• Robust across synthetic and real data