Scalable Learning of Drift and Diffusion in High-Dimensional Stochastic Differential Equations

Stochastic differential equations (SDEs) provide a fundamental framework for modeling high-dimensional stochastic dynamical systems. Their drift and diffusion functions govern the deterministic dynamics and stochastic fluctuations of the underlying system, respectively. Learning these functions from observed data is crucial for understanding and predicting system behavior. Despite the importance of high-dimensional SDEs ($d\gg 100$) in practical applications, existing methods mostly rely on prior knowledge of the functional forms or are difficult to scale to the high-dimensional case because of prohibitive time and space complexity. To address these limitations, we propose KoopSDE, an algorithm that uses the stochastic Koopman operator to learn drift and diffusion functions from a single discretely observed trajectory without prior knowledge of functional forms. To address the curse of dimensionality, we incorporate Hutchinson trace estimation (HTE), reducing the time and space complexity from $\mathcal{O}(d^2)$ to $\mathcal{O}(d)$. We also establish a set of theoretical guarantees for KoopSDE. Experiments demonstrate that KoopSDE can learn $100,000$-dimensional SDEs on a single NVIDIA A800 GPU. This work opens up the possibility for learning the drift and diffusion functions of high-dimensional SDEs.