GOTFlow: Learning Directed Population Transitions from Cross-Sectional Biomedical Data with Optimal Transport

GOTFlow is a novel framework that leverages graph-constrained optimal transport in a learned latent space to infer directed, interpretable population transitions and molecular drivers from cross-sectional biomedical data, overcoming limitations of existing methods in modeling non-linear, heterogeneous, and unbalanced biological dynamics.

The Gist

The Big Problem: The "Snapshot" Dilemma

Imagine you are trying to understand how a caterpillar turns into a butterfly. Ideally, you would film the whole process. But in biology and medicine, we rarely get to film the whole movie. Instead, we usually only have snapshots.

• We have a photo of a healthy tissue.
• We have a photo of a sick tissue.
• We have a photo of a tissue in the middle.

But these photos are taken from different people at different times. We don't know which specific healthy person turned into which specific sick person. It's like trying to figure out the plot of a movie by looking at 1,000 random still photos from the audience, where everyone is wearing different clothes and sitting in different seats.

Scientists have tried to guess the "movie plot" (the trajectory) before, but their methods often get confused, assume the story is too simple (like a straight line), or fail when the "cast" changes size (some people get sick and leave the study, new ones join).

The Solution: GOTFlow (The "Smart GPS")

The authors created a new tool called GOTFlow. Think of it as a Smart GPS for populations.

Instead of trying to track one specific person from start to finish, GOTFlow looks at the entire crowd at State A (e.g., Healthy) and the entire crowd at State B (e.g., Sick) and asks: "If we had to move the Healthy crowd to look like the Sick crowd, what is the most efficient way to do it?"

It uses a mathematical concept called Optimal Transport.

• The Analogy: Imagine you have a pile of sand (the healthy population) and you need to reshape it into a new pile (the sick population).
• The Old Way: You might just guess which grain of sand goes where, or assume the sand moves in a straight line.
• The GOTFlow Way: It calculates the absolute most efficient path for every single grain of sand to get to its new destination, minimizing the "energy" or effort required.

How It Works (The Three Magic Steps)

1. The "User-Defined Map" (The Graph)

In real life, biology isn't always a straight line. Sometimes a disease branches out (like a tree), or different treatments lead to different outcomes.

• GOTFlow's Trick: You tell the computer the rules of the road. You draw a map (a graph) saying, "Healthy can go to Early Disease, and Early Disease can go to Late Disease."
• Why it helps: It stops the computer from making impossible guesses (like assuming a late-stage cancer patient suddenly went back to being healthy). It respects the biological story you already know.

2. The "Shape-Shifting Lens" (Latent Space)

Biological data is messy. It's like trying to compare two clouds; they look different, but maybe they are made of the same water vapor.

• GOTFlow's Trick: It uses a "lens" (a neural network) to squish and stretch the data into a new, simpler shape where the differences between states become obvious.
• The Result: It finds the hidden geometry. It realizes that even though the data looks chaotic, the "Healthy" group and the "Sick" group are actually sitting on two distinct islands in a hidden landscape, and it figures out the bridge between them.

3. The "Flexible Crowd" (Unbalanced Transport)

In real biology, populations change size. Cells die, new ones are born, or people drop out of a study.

• The Old Problem: Old math said, "You must move 100% of the healthy people to the sick group, one-for-one." This is unrealistic.
• The GOTFlow Fix: It uses Unbalanced Optimal Transport. It allows the crowd to grow or shrink. It says, "Okay, 80% of the healthy people moved to the sick group, 10% died off, and 10% new cells appeared." This makes the model much more realistic.

What Did They Find? (The Three Stories)

The team tested GOTFlow on three real-world biological mysteries:

1. The Womb's Monthly Remodeling (Endometrium)

• The Story: The lining of the uterus changes every month to prepare for a baby. If it doesn't change correctly, pregnancy fails.
• The GOTFlow Discovery: They found that women who suffered miscarriages had a "sluggish" transition. Their uterine lining tried to change, but the "drift" (the movement from one state to the next) was weak and slow. GOTFlow identified exactly which genes were failing to "turn on" or "turn off" to cause this delay.

2. The Slow Climb of Cancer Risk (Breast Cancer)

• The Story: Breast cancer isn't just "cancer" or "no cancer." It's a sliding scale of risk.
• The GOTFlow Discovery: They mapped the molecular journey from low risk to high risk. They found specific genes that acted like "gas pedals" (getting worse as risk increased) and "brakes" (getting better as risk increased). This helps doctors understand how a tumor evolves, not just that it exists.

3. The Zombie Brain (Prion Disease)

• The Story: Prion diseases (like Mad Cow Disease) slowly destroy the brain.
• The GOTFlow Discovery: By looking at mouse brains at different stages, they saw the "drift" accelerate as the disease got worse. They identified specific genes that acted as alarms for inflammation. The tool showed that the brain's immune system was screaming "Help!" long before the mouse showed physical symptoms.

Why This Matters

Before GOTFlow, scientists had to guess the story of how diseases progress, often assuming it was a straight line or relying on data they didn't have (like tracking the same person for 10 years).

GOTFlow is like a detective who can look at a pile of evidence (snapshots) and reconstruct the crime scene (the progression) with high accuracy, even if the witnesses (the patients) are different people.

It gives us:

1. Direction: Which way is the disease going?
2. Speed: How fast is it changing?
3. The Culprits: Which specific genes or molecules are driving the change?

This helps researchers design better drugs to stop the "drift" before it's too late.

Technical Summary

1. Problem Statement

Biological and clinical systems are inherently dynamic (e.g., disease progression, tissue remodeling), yet most available datasets are cross-sectional. These datasets capture populations at discrete states rather than tracking individuals over time.

• The Challenge: Existing methods for trajectory inference (e.g., pseudotime, RNA velocity) often rely on fixed feature spaces, assume linear progression, or require dense longitudinal sampling (common in single-cell data). They struggle to model heterogeneous, unbalanced population transitions across diverse modalities (e.g., bulk RNA-seq, histology) without explicit temporal tracking.
• The Gap: There is a need for a flexible framework that can infer directed population-level changes from static, cross-sectional data while accommodating non-linear geometry, branching relationships, and changes in population mass (e.g., cell proliferation or death).

2. Methodology: The GOTFlow Framework

GOTFlow is a framework that combines representation learning with unbalanced optimal transport (UOT) to model transitions between biological states.

A. Problem Formulation

• Input: A dataset of $N$ observations $\{(x_i, s_i)\}$, where $x_i$ is a feature vector (e.g., gene expression) and $s_i$ is a discrete state label (e.g., disease stage, time point).
• Constraint: Transitions are governed by a user-defined directed weighted graph $G = \{(S_\ell, T_\ell)\}$. This graph encodes hypothesized admissible transitions (e.g., Early $\to$ Late, Control $\to$ Treatment) and allows for branching or merging paths.

B. Core Components

1. Latent Representation Learning:

   • GOTFlow learns a parametric encoder $\phi_\theta$ that maps high-dimensional inputs $x_i$ to a latent embedding $z_i$.
   • Whitening: To prevent degenerate solutions (e.g., encoder collapse) and handle scale/anisotropy, the latent embeddings are globally whitened before transport calculations.

2. Unbalanced Optimal Transport (UOT):

   • Unlike standard OT which enforces strict mass conservation, GOTFlow uses UOT to allow mass to be created or destroyed between states. This is crucial for biological systems where population sizes change due to proliferation or cell death.
   • The transport plan $P$ between source state $S$ and target state $T$ minimizes an Unbalanced Entropic OT functional:
     $$U(\tilde{\mu}_S, \tilde{\mu}_T; P) = \sum C_{ij}P_{ij} + \varepsilon \sum P_{ij}(\log P_{ij} - 1) + \lambda D_{KL}(P(x) \| a) + \lambda D_{KL}(P(y) \| b)$$
     • $C_{ij}$: Transport cost based on latent distance.
     • $\varepsilon$: Entropic regularization for numerical stability.
     • $\lambda$: KL divergence penalties allowing mass deviation.

3. Joint Optimization Objective:

   • The encoder parameters $\theta$ and transport plans are learned jointly.
   • OT Energy Minimization: Encourages embeddings where valid transitions (edges in $G$) have low transport cost.
   • Contrastive Loss (InfoNCE): A negative sampling strategy ensures that valid transitions have lower energy than invalid ones (negative candidates). This enforces a discriminative structure in the latent space.
   • Total Loss: Aggregates OT energy and contrastive loss across all edges in the graph.

C. Interpretability

After training, GOTFlow derives interpretable summaries:

• Drift Vectors: For each sample, a barycentric projection onto the target state defines a drift vector ($\bar{z} - z$), quantifying the direction and magnitude of change.
• Feature-Level Attribution: By projecting the transport plan back to the original feature space, GOTFlow identifies specific molecular drivers (e.g., genes) that change most strongly along a transition.

3. Key Contributions

1. Graph-Constrained Representation Learning: Introduces a general OT framework that integrates a user-defined state graph with learned latent geometry, enabling cohort-level analysis of heterogeneous populations without requiring longitudinal data.
2. Unbalanced Transport for Biology: Explicitly models changes in population mass (proliferation/death) via UOT, avoiding unrealistic one-to-one matching assumptions.
3. Interpretable Dynamics: Provides quantitative summaries (drift vectors) and feature-level attribution to identify molecular drivers of progression.
4. Validation: Demonstrates robustness on synthetic data and applicability across three distinct biomedical domains.

4. Experimental Results & Applications

The authors validated GOTFlow on synthetic data and three real-world biological datasets:

• Synthetic Data:

  • GOTFlow successfully recovered known bifurcating and merging transition structures.
  • Inferred drift vectors showed high agreement with ground-truth displacements (Mean Cosine Similarity: $0.720 \pm 0.152$).

• Application 1: Endometrial Remodelling (Luteal Phase):

  • Context: Analyzed 664 biopsies to model the transition from early to late luteal phases.
  • Finding: GOTFlow identified coherent transcriptional trajectories. Crucially, it detected reduced drift magnitude in patients with a history of miscarriage compared to controls, indicating impaired or delayed decidualisation.
  • Drivers: Identified specific gene shifts (e.g., reduced PLA2G2A/DIO2 ratio, altered IGFBP1) associated with reproductive failure.

• Application 2: Breast Cancer Risk Progression (TCGA-BRCA):

  • Context: Analyzed 1,061 tumors to model transitions across 5 risk strata (low to high).
  • Finding: The model revealed a coherent directional flow from low to high risk.
  • Drivers: Identified genes with increasing expression in high-risk states (e.g., TCP1, GARS1 linked to aggression) and decreasing expression (e.g., SUSD3, NXNL2 linked to favorable prognosis), aligning with known prognostic signals.

• Application 3: Prion Disease (Mouse Brain):

  • Context: Modeled transcriptional dynamics across the incubation period of prion infection.
  • Finding: Drift magnitudes were small in asymptomatic phases but increased significantly as neuropathology developed.
  • Drivers: Identified infection-specific shifts in neuroinflammation markers (C1qa, Cd44, Gfap), distinguishing infected mice from controls and validating the model's ability to capture disease-specific pathology.

5. Significance and Conclusion

• Generalizability: GOTFlow moves beyond single-cell trajectory inference to handle cohort-level, cross-sectional data across various modalities (bulk RNA-seq, imaging).
• Hypothesis-Driven: By allowing users to define the transition graph, it supports hypothesis-driven modeling of specific biological progressions (e.g., specific disease stages or treatment responses).
• Interpretability: Unlike "black box" deep learning approaches, GOTFlow provides explicit, quantitative measures of population flow and feature-level attribution, making it highly suitable for biological discovery.
• Limitations: The method relies on user-defined state discretization and transition graphs (introducing assumptions). The inferred dynamics are observational and do not prove causality without experimental validation.

In summary, GOTFlow establishes a new paradigm for analyzing directed population dynamics from static biomedical data, offering a powerful tool to quantify how molecular profiles evolve across disease and physiological states.