Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Identifying Drift, Diffusion, and Causal Structure from Temporal Snapshots

Vincent Guan, Joseph Janssen, Hossein Rahmani, Andrew Warren, Stephen Zhang, Elina Robeva, Geoffrey Schiebinger

TL;DR

This work addresses the challenge of learning drift and diffusion parameters of an SDE from temporal marginals when trajectories are unobserved. It establishes a sharp identifiability criterion: for linear additive-noise SDEs, drift $A$ and diffusion $H=GG^\top$ are identifiable from marginals if and only if the initial distribution $X_0$ is not auto-rotationally invariant, and it shows that the causal graph can be recovered from these parameters. The paper then introduces APPEX, an algorithm that alternates trajectory inference via anisotropic Schrödinger-bridge OT (AEOT) with maximum-likelihood parameter estimation, provably reducing KL divergence at each step and converging to the true SDE under suitable conditions. Empirical results on simulated data demonstrate that APPEX accurately recovers drift, diffusion, and causal structure and outperforms Waddington-OT, even in high dimensions and in the presence of latent confounders, highlighting its practical potential for systems biology and related domains.

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.

Identifying Drift, Diffusion, and Causal Structure from Temporal Snapshots

TL;DR

This work addresses the challenge of learning drift and diffusion parameters of an SDE from temporal marginals when trajectories are unobserved. It establishes a sharp identifiability criterion: for linear additive-noise SDEs, drift and diffusion are identifiable from marginals if and only if the initial distribution is not auto-rotationally invariant, and it shows that the causal graph can be recovered from these parameters. The paper then introduces APPEX, an algorithm that alternates trajectory inference via anisotropic Schrödinger-bridge OT (AEOT) with maximum-likelihood parameter estimation, provably reducing KL divergence at each step and converging to the true SDE under suitable conditions. Empirical results on simulated data demonstrate that APPEX accurately recovers drift, diffusion, and causal structure and outperforms Waddington-OT, even in high dimensions and in the presence of latent confounders, highlighting its practical potential for systems biology and related domains.

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 ), 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.

Paper Structure

This paper contains 31 sections, 16 theorems, 106 equations, 10 figures, 4 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Our contributions
  3. Background and related work
  4. Background and related work
  5. Stochastic differential equations (SDEs) and modeling real-world systems
  6. Inference from temporal marginals
  7. Schrödinger bridge problem
  8. Identifiability of SDE parameters
  9. System identification from temporal marginals
  10. SDE parameter identifiability
  11. Causal graph identification from SDE parameters
  12. Our parameter estimation method
  13. Trajectory inference via anisotropic entropy-regularized optimal transport
  14. Parameter estimation via MLE
  15. The APPEX algorithm
  16. ...and 16 more sections

Key Result

Theorem 3

Let $p_0$ be a probability distribution on $\mathbb R^d$ with finite moments, and let denote the set of drift-diffusion pairs under a linear additive noise model eq: linear_additive_noise_SDE. Then, given the initialization $X_0 \sim p_0$, any parameter set $(A,H) \in \Theta$ is identifiable from its marginals $\{p_t : t \ge 0\}_{A,H}$, in the sense that if and only if $p_0$ is not auto-rotation

Figures (10)

  • Figure 1: Outline of our theoretical (a) and algorithmic (b) contributions. In (b), our algorithm Alternating Projection Parameter Estimation from $X_0$ alternates between Anisotropic Entropy-regularized Optimal Transport (AEOT) and Maximum Likelihood Estimation (MLE).
  • Figure 2: Four examples of auto-rotationally invariant distributions are illustrated. By Theorem \ref{['thm: identifiability']}, these initial distributions lead to non-identifiability of linear additive noise SDEs from temporal marginals.
  • Figure 3: Visualization of our parameter estimation algorithm APPEX. Given observed temporal marginals from an underlying SDE (left), APPEX alternates between trajectory inference (middle) and MLE parameter estimation (right) in order to find the SDE parameters that best represent the temporal marginal observations.
  • Figure 4: The mean absolute error for estimates of $A$ and $H$ using APPEX is shown per iteration for all three pairs of SDEs from Section \ref{['sec:Noniden_examples']}
  • Figure 5: The correlation between the estimated and true SDE parameters is plotted per iteration across $10$ random linear additive noise SDEs for dimensions $3$ and $10$
  • ...and 5 more figures

Theorems & Definitions (28)

  • Definition 1
  • Definition 2
  • Theorem 3
  • Definition 4
  • Lemma 5
  • Remark 6: Application to empirical marginals
  • Proposition 7: MLE estimators for drift and diffusion of SDE \ref{['eq: linear_additive_noise_SDE']} from multiple trajectories
  • Remark 8: Application to causal discovery
  • Lemma 9
  • Example 1: Starting at stationary distribution lavenant2021towardsshen2024multi
  • ...and 18 more