Identifying Drift, Diffusion, and Causal Structure from Temporal Snapshots

2024-10-30

arXiv.org

AI Summary

This work addresses the challenges of stochastic differential equation (SDE) modeling and causal inference from snapshot data—such as single-cell measurements—where only marginal temporal distributions are observed, without individual trajectories. We establish the first identifiability theory for linear SDEs: under non-rotationally symmetric initial distributions, the drift, diffusion, and underlying causal graph are uniquely identifiable, and the causal graph can be fully recovered from the SDE parameters. To operationalize this theory, we propose APPEX, an end-to-end algorithm integrating entropy-regularized optimal transport, alternating projection optimization, and KL-divergence minimization to jointly estimate drift, diffusion, and causal structure. Our theoretical analysis guarantees both parameter identifiability and causal recoverability. Empirically, APPEX monotonically reduces KL divergence and achieves high-precision recovery of drift, diffusion, and causal graphs on additive-noise SDEs across three benchmark settings.

Technology Category

• Machine Learning: Causal Learning
• Reasoning under Uncertainty: Causality
• Intelligent Robots: State Estimation

Application Category

• Graph Algorithms and Modeling for the Web: Representation, reconstruction, and subgraph or motif discovery in Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

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 prove that these parameters are identifiable from marginals if and only if the initial distribution lacks any generalized rotational symmetries. We further prove that the causal graph of any SDE with additive diffusion can be recovered from the 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.

Problem

• Identify SDE parameters from temporal marginals.
• Recover causal graph from SDE with additive diffusion.
• Develop APPEX algorithm for estimating SDE drift and diffusion.

Innovation

• Identifies SDE drift and diffusion from marginals
• Uses entropy-regularized optimal transport for anisotropic diffusion
• Introduces APPEX algorithm for parameter estimation