Learning biophysical models of gene regulation with probability flow matching

1. Problem Statement

Cellular differentiation is governed by high-dimensional, stochastic gene regulatory networks (GRNs). While single-cell RNA sequencing (scRNA-seq) provides quantitative snapshots of the transcriptome, inferring the underlying mechanistic dynamics from these cross-sectional data remains a significant challenge.

Current computational methods face a trade-off between scalability and biophysical interpretability:

• Descriptive/Statistical Models: Methods like TrajectoryNet or standard Optimal Transport (OT) project data into low-dimensional latent spaces. While scalable, they lack mechanistic interpretability, cannot easily perform in silico gene perturbations, and often fail to generalize to unseen conditions (out-of-distribution).
• Mechanistic Models: Approaches based on differential equations (ODEs) or stochastic differential equations (SDEs) offer interpretability but struggle with high-dimensional data due to the computational cost of solving ODEs and computing Wasserstein distances.
• The Gap: Existing methods often ignore intrinsic molecular stochasticity or assume decoupled gene dynamics, leading to models that interpolate data well but fail to capture the true biological mechanisms of lineage specification and fate decisions.

2. Methodology: Probability Flow Matching (PFM)

The authors introduce Probability Flow Matching (PFM), a scalable, simulation-free framework for learning biophysically consistent stochastic processes directly from multi-marginal single-cell data.

Core Framework

PFM builds upon Flow Matching (FM) and Score-based Generative Modeling. Instead of solving expensive ODEs or computing optimal transport couplings iteratively, PFM learns a velocity field $u_\theta(x, t)$ that transports a simple base distribution to complex target distributions observed at discrete time points.

Key Mathematical Formulation:
The dynamics are modeled via the Fokker-Planck (FP) equation:
$$\frac{\partial p_t(x)}{\partial t} = -\nabla \cdot (u_t(x) p_t(x)) + g_t(x) p_t(x)$$
Where:

• $u_t(x) = f(x) - \nabla \cdot D(x) - D(x)\nabla \log p_t(x)$ is the probability flow velocity.
• $f(x)$ encodes gene-regulatory interactions (drift).
• $D(x)$ is the diffusion matrix capturing intrinsic/extrinsic noise.
• $g_t(x)$ models cell proliferation and death (growth dynamics).

Key Technical Innovations

1. Direct Inference in Gene Space: Unlike methods that project to latent spaces, PFM operates directly in the high-dimensional gene space (specifically a reduced Transcription Factor space), preserving mechanistic structure.
2. Chebyshev Interpolation for Conditional Paths: To handle multiple time points (multi-marginal data), PFM constructs smooth conditional paths $p_t(x|z)$ between snapshots. The authors replace standard linear or cubic-spline interpolants with spectrally regularized Chebyshev interpolants. This ensures globally smooth time derivatives and superior convergence in high-dimensional, low-sample regimes.
3. Unbalanced Flow Matching: The framework extends to unbalanced settings to account for cell proliferation and death. It introduces an auxiliary mass density $m_t$ that evolves along the probability path, allowing simultaneous inference of regulatory dynamics ($f$) and growth rates ($g$).
4. Biophysical Parameterization: The framework supports various biophysical models, including the Chemical Langevin Equation (CLE), which explicitly models intrinsic noise as a function of reaction rates (production $h(x)$ and degradation $\ell x$).

3. Key Contributions

• Scalable Mechanistic Learning: PFM enables the inference of stochastic differential equations directly from high-dimensional scRNA-seq data without the computational bottleneck of repeated ODE simulations.
• Superior Interpolation via Chebyshev Polynomials: The paper demonstrates that Chebyshev interpolation significantly outperforms linear and cubic-spline methods in reconstructing drift fields and probability flows, particularly in high dimensions.
• Integration of Growth Dynamics: The framework uniquely allows for the joint inference of gene-regulatory dynamics and population-level birth/death rates from unbalanced data.
• Validation of Biophysical Consistency: The authors prove that while deterministic models (ODEs) may achieve low interpolation error, only biophysically consistent stochastic models (like CLE) accurately recover the underlying regulatory logic and generalize to unseen conditions.

4. Results

The authors validated PFM on both synthetic data and three real-world hematopoiesis datasets.

A. Synthetic Benchmarks

• Interpolation Accuracy: In high-dimensional (up to 50D) multistable systems, PFM with Chebyshev interpolation achieved significantly lower Root Mean Squared Error (RMSE) in reconstructing drift fields compared to linear or cubic-spline baselines.
• Efficiency: PFM reduced computational and memory requirements by an order of magnitude compared to simulation-based methods like Probability Flow Inference (PFI).

B. Hematopoietic Differentiation (Ex Vivo & In Vitro)

• Gene Regulatory Network (GRN) Inference: When applied to a 24-gene transcription factor (TF) space, PFM successfully inferred regulatory interactions.
  • Stochastic vs. Deterministic: Deterministic models (TrajectoryNet, ODE) interpolated the data well but failed to recover the correct signed regulatory interactions (activation vs. inhibition). Only the biophysically consistent Chemical Langevin Equation (CLE) model accurately recovered the causal structure consistent with ChIP-seq data.
• Generalization to Unseen Conditions:
  • Models trained on stem/progenitor cells were tested on unseen initial conditions (e.g., different differentiation trajectories).
  • Result: Stochastic models (CLE) maintained high correlation with experimental terminal states, whereas deterministic models failed to recover the correct attractors.
• In Silico Perturbations (Knockouts):
  • The authors simulated knockouts of key TFs (KLF1 for erythroid, FLI1 for megakaryocytic).
  • Outcome: CLE correctly predicted lineage-specific effects (e.g., KLF1 KO disrupts erythroid but spares megakaryocytes). Deterministic models incorrectly predicted the loss of both lineages, failing to capture the robustness of the stochastic landscape.

C. Growth Dynamics Inference

• Applied to a dataset with pronounced proliferation/death, PFM successfully inferred cell-state-specific growth rates ($g(x)$).
• The inferred fate probabilities from the stochastic CLE model were probabilistic and aligned better with clonal lineage tracing data than the deterministic assignments from TIGON++.

5. Significance and Conclusion

This work establishes Probability Flow Matching as a powerful bridge between first-principles biophysical modeling and data-driven single-cell omics.

• Mechanistic Insight: It moves beyond descriptive "pseudotime" ordering to provide a quantitative, mechanistic understanding of how noise and regulation drive cell fate decisions.
• Predictive Power: By capturing intrinsic stochasticity, PFM models generalize robustly to unseen experimental conditions and perturbations, a critical requirement for therapeutic applications and in silico drug screening.
• Scalability: The simulation-free nature of PFM makes it feasible to apply mechanistic models to genome-scale or large-scale multi-omics datasets, overcoming previous computational barriers.

In summary, the paper argues that biophysical consistency is not just a theoretical preference but a practical necessity for building models that can accurately predict cellular behavior under new conditions and perturbations. PFM provides the scalable infrastructure to achieve this.