Gradient-flow SDEs have unique transient population dynamics

TL;DR

This paper proves that both drift and diffusion of gradient-flow SDEs are jointly identifiable from population data outside equilibrium, and introduces nn-APPEX, a method for their joint inference.

Contribution

It provides the first complete characterization of identifiability for gradient-flow SDEs and proposes a novel inference method that jointly learns drift and diffusion from marginals.

Findings

nn-APPEX accurately infers both drift and diffusion.

Previous methods produce biased drift estimates due to incorrect diffusion assumptions.

Identifiability holds only when observing outside equilibrium.

Abstract

Identifying the drift and diffusion of an SDE from its population dynamics is a notoriously challenging task. Researchers in machine learning and single-cell biology have only been able to prove a partial identifiability result: for potential-driven SDEs, the gradient-flow drift can be identified from temporal marginals if the Brownian diffusivity is already known. Existing methods therefore assume that the diffusivity is known a priori, despite it being unknown in practice. We dispel the need for this assumption by providing a complete characterization of identifiability: the gradient-flow drift and Brownian diffusivity are jointly identifiable from temporal marginals if and only if the process is observed outside of equilibrium. Given this fundamental result, we propose nn-APPEX, the first Schrodinger Bridge-based inference method that can simultaneously learn the drift and diffusion of a gradient-flow SDE solely from observed marginals. Extensive experiments show that nn-APPEX's ability to adjust its diffusion estimate enables accurate inference, while previous Schrodinger Bridge methods obtain biased drift estimates due to their assumed, and likely incorrect, diffusion.