Simulation-free Structure Learning for Stochastic Dynamics

TL;DR

StructureFlow tackles the challenge of jointly uncovering the causal structure and stochastic dynamics of high-dimensional systems from partial, noisy observations. It achieves this with a simulation-free framework that blends a time-invariant neural graphical vector field for structure with a time-dependent score to capture stochasticity, all learned through a Schrödinger-bridge formulation via score-and-flow matching and entropic OT. The method supports interventional data through knockout masks, enabling simultaneous structure recovery and trajectory inference even under perturbations, and is demonstrated on synthetic linear and BoolODE systems, biologically plausible simulations, and real single-cell perturbation data. Empirically, StructureFlow shows strong performance on both structure discovery and dynamical inference tasks, scales favorably with dimension, and offers a practical tool for mechanistic understanding and intervention design in complex biological systems.

Abstract

Modeling dynamical systems and unraveling their underlying causal relationships is central to many domains in the natural sciences. Various physical systems, such as those arising in cell biology, are inherently high-dimensional and stochastic in nature, and admit only partial, noisy state measurements. This poses a significant challenge for addressing the problems of modeling the underlying dynamics and inferring the network structure of these systems. Existing methods are typically tailored either for structure learning or modeling dynamics at the population level, but are limited in their ability to address both problems together. In this work, we address both problems simultaneously: we present StructureFlow, a novel and principled simulation-free approach for jointly learning the structure and stochastic population dynamics of physical systems. We showcase the utility of StructureFlow for the tasks of structure learning from interventions and dynamical (trajectory) inference of conditional population dynamics. We empirically evaluate our approach on high-dimensional synthetic systems, a set of biologically plausible simulated systems, and an experimental single-cell dataset. We show that StructureFlow can learn the structure of underlying systems while simultaneously modeling their conditional population dynamics -- a key step toward the mechanistic understanding of systems behavior.

Figures (18)

• Figure 1: Overview of StructureFlow for joint structure learning and dynamical inference.
• Figure 2: StructureFlow yields improved structure learning performance when scaled to high-dimensional systems. We compare with NGM-NODE and RF on synthetic linear systems with varying dimensionality ($d$) and system (graph) sparsity levels (5%, 20%, 40%).
• Figure 3: StructureFlow is consistently a top performing structure learning method across simulated biological systems. Here, we use interventional (with knockouts) and observational (no knockouts) data, and report average precision (AP) and area under the ROC curve (AUROC) scores.
• Figure 4: Summary of inferred structure across multiple synthetic biological systems. Connectivity matrices (heatmaps) showing ground truth and inferred graphs for the TF (trifurcating), LL (long linear), and HSC (hematopoietic stem cell) systems. The x/y-axis labels correspond to system variables and shading indicates edge interaction strength.
• Figure 5: Visualization of model predicted dynamics. 2D PCA visualization comparing dynamical inference methods for trajectory inference using leave-one-timepoint-out evaluation on the TF (trifurcating) system. We show left-out timepoint prediction for the observational (wild-type) setting and for two seen interventions (knockouts) selected for their diverse trifurcating paths.