Identifying Drift, Diffusion, and Causal Structure from Temporal Snapshots

TL;DR

This paper introduces a novel method to identify the drift, diffusion, and causal structure of stochastic differential equations from temporal data, with applications to gene networks and other dynamic systems.

Contribution

It provides the first comprehensive approach for jointly estimating SDE components from marginal distributions, including an algorithm that handles anisotropic diffusion.

Findings

Theoretical conditions for identifiability of SDE parameters from marginals.

An iterative algorithm (APPEX) effectively estimates SDE parameters from data.

Successful demonstration on simulated linear additive noise SDEs.

Abstract

Stochastic differential equations (SDEs) are a fundamental tool for modelling dynamic processes, including gene regulatory networks (GRNs), contaminant transport, financial markets, and image generation. However, learning the underlying SDE from data is a challenging task, especially if individual trajectories are not observable. Motivated by burgeoning research in single-cell datasets, we present the first comprehensive approach for jointly identifying the drift and diffusion of an SDE from its temporal marginals. Assuming linear drift and additive diffusion, we show that non-identifiability can only arise if the initial distribution possesses generalized rotational symmetries. We further prove that even if this condition holds, the drift and diffusion can almost always be recovered from the marginals. Additionally, we show that the causal graph of any SDE with additive diffusion can be recovered from the identified SDE parameters. To complement this theory, we adapt entropy-regularized optimal transport to handle anisotropic diffusion, and introduce APPEX (Alternating Projection Parameter Estimation from $X_0$), an iterative algorithm designed to estimate the drift, diffusion, and causal graph of an additive noise SDE, solely from temporal marginals. We show that APPEX iteratively decreases Kullback-Leibler divergence to the true solution, and demonstrate its effectiveness on simulated data from linear additive noise SDEs.