🤖 machine learning

ContextFlow: Context-Aware Flow Matching For Trajectory Inference From Spatial Omics Data

ContextFlow is a novel context-aware flow matching framework that integrates local tissue organization and ligand-receptor communication patterns to guide the inference of biologically meaningful and statistically consistent trajectories from longitudinal spatially resolved omics data, outperforming existing state-of-the-art methods.

Original authors: Santanu Subhash Rathod, Francesco Ceccarelli, Sean B. Holden, Pietro Liò, Xiao Zhang, Jovan Tanevski

Published 2026-05-15

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Santanu Subhash Rathod, Francesco Ceccarelli, Sean B. Holden, Pietro Liò, Xiao Zhang, Jovan Tanevski

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: Predicting the Future of Cells

Imagine you are trying to understand how a city changes over time. You have a series of photos taken at 9:00 AM, 12:00 PM, and 3:00 PM. However, there's a catch: you can't watch the video. You only have the snapshots, and the people in the photos are unpaired (you don't know which person in the 9:00 AM photo became the person in the 12:00 PM photo).

In biology, scientists have this exact problem with spatial omics data. They have snapshots of tissues (like a brain or a liver) at different times, showing where cells are and what genes they are using. They want to reconstruct the "movie" of how these cells move, change, and interact to heal a wound or grow an organ.

The Problem: The "Stranger Danger" of Old Methods

Previous methods tried to guess which cell in the morning photo became which cell in the afternoon photo by looking only at their appearance (their genetic "face").

The Flaw: Imagine two people in a crowd. One is a baker, and the other is a firefighter. If they both happen to be wearing a red hat (a similar genetic trait), an old computer program might think, "Aha! The baker turned into the firefighter!"
The Reality: In biology, this is impossible. A baker cannot turn into a firefighter just because they wore the same hat. Similarly, an "excitatory" neuron cannot turn into an "inhibitory" neuron just because their genes look similar at a specific moment.
The Result: These old methods create "biologically implausible" movies where cells do things that nature forbids, leading to a confusing and incorrect story of how the tissue heals or grows.

The Solution: ContextFlow (The "Smart Detective")

The authors created ContextFlow, a new AI framework that acts like a smart detective. Instead of just looking at the cell's "face" (genes), it looks at the context around the cell to make a better guess.

They use two main clues (called "priors") to guide the AI:

1. The Neighborhood Clue (Spatial Smoothness)

The Analogy: Think of a neighborhood. If you see a family moving from one house to the house next door, it makes sense. If you see a family suddenly teleporting to a house on the other side of the city, it's suspicious.
How it works: Cells in tissues usually stay close to their neighbors. If a cell moves, it likely moves to a nearby spot, not across the whole organ. ContextFlow uses this "local neighborhood" rule to say, "It's highly unlikely this cell jumped across the brain; it probably stayed in the neighborhood."

2. The Conversation Clue (Ligand-Receptor Communication)

The Analogy: Imagine people in a room talking. A baker might talk to a delivery driver about bread. A firefighter might talk to a paramedic about safety. They don't usually swap roles mid-conversation.
How it works: Cells "talk" to each other using chemical signals (ligands and receptors). If Cell A is sending a "grow" signal and Cell B is listening, they are likely part of the same process. ContextFlow checks these "conversations." If the conversation patterns don't match a specific transition, the AI knows that transition is fake.

How It Works: The "Traffic Cop"

The paper describes a mathematical process called Optimal Transport. You can think of this as a traffic cop trying to direct cars (cells) from one intersection (Time A) to another (Time B).

Old Method: The cop only looks at the color of the cars. "Red car goes to Red spot." Sometimes, this sends a red sports car into a spot meant for a red school bus.
ContextFlow: The cop looks at the color AND the map of the city AND the radio chatter between drivers.
- "This red car is in the bakery district and talking to a baker. It must go to the bakery district in the next photo, not the fire station."

By adding these "rules of the road" (biological context) into the math, ContextFlow forces the AI to generate a movie where cells move in ways that actually make sense in the real world.

What They Found (The Results)

The team tested this on three real-world biological "movies":

Axolotl Brain Regeneration: Watching how a salamander's brain grows back after injury.
Mouse Organogenesis: Watching a mouse embryo develop from a tiny ball of cells into a complex body.
Liver Regeneration: Watching a mouse liver heal after damage.

The Outcome:

Fewer Mistakes: ContextFlow made far fewer "impossible" moves (like a baker turning into a firefighter) compared to previous methods.
Better Accuracy: The reconstructed "movies" of cell movement were statistically closer to the real biological data.
Efficiency: It did this without slowing down the computer significantly.

Summary

ContextFlow is a new tool that helps scientists reconstruct the life stories of cells. It stops the computer from making wild guesses by teaching it to respect the neighborhood (where the cell is) and the conversation (who the cell is talking to). This results in a much clearer, more accurate, and biologically truthful picture of how tissues grow, heal, and change.

Technical Summary: ContextFlow

Problem Statement

Inferring cellular trajectories from longitudinal, spatially resolved omics data is critical for understanding tissue dynamics in development, regeneration, and disease. However, profiling technologies are destructive, yielding only unpaired, cross-sectional snapshots at discrete time points rather than continuous trajectories. While recent flow matching (FM) frameworks have emerged as efficient generative models for learning continuous latent dynamics, state-of-the-art approaches relying on Optimal Transport (OT) couplings face a significant limitation: they determine cell-to-cell transitions solely based on transcriptomic similarity.

This reliance on transcriptomic similarity alone often leads to biologically implausible transitions. For instance, OT plans may couple cells with distinct biological functions (e.g., inhibitory and excitatory neurons) simply because their gene expression profiles are statistically similar at a given snapshot. Such errors propagate into the learned velocity field, resulting in globally inconsistent and biologically unrealistic inferred dynamics. Existing methods fail to fully exploit the contextual richness of spatial transcriptomics, such as local tissue organization and ligand-receptor communication patterns, to constrain these transitions.

Methodology: ContextFlow

The authors propose ContextFlow, a context-aware flow matching framework that integrates explicit spatial priors into the OT-coupled flow matching process. The method aims to regularize the optimal transport objective to ensure that learned trajectories are not only statistically consistent but also biologically meaningful.

1. Transitional Plausibility

ContextFlow defines a Transitional Plausibility (TP) metric to capture the likelihood of a cell evolving from one state to another across time points. This metric combines two types of spatial priors:

Spatial Smoothness (SS): Based on the observation that neighboring cells in a microenvironment respond to similar stimuli and exhibit local smoothness in expression. The SS prior quantifies the dissimilarity between the average expression profiles of the local neighborhoods of two cells.
Cell-Cell Communication Patterns (LR): Leveraging prior knowledge of ligand-receptor (LR) interactions, this prior assesses the similarity of communication patterns between the microenvironments of two cells. Cells exhibiting similar LR signaling profiles are considered more likely to transition into one another.

These priors are combined into a Transitional Plausibility Matrix (TPM), which encodes the biological constraints for cell transitions between consecutive time points.

2. Prior-Regularized Optimal Transport

The TPM is integrated into the Entropic Optimal Transport (EOT) formulation through two novel schemes:

Prior-Aware Cost Matrix (PACM): The TPM is directly incorporated into the transport cost function. A higher plausibility score (indicating a less plausible transition) increases the transport cost, thereby penalizing such couplings. This approach requires tuning a trade-off hyperparameter ( $\alpha$ ) between the original Euclidean cost and the prior-aware cost.
Prior-Aware Entropy Regularization (PAER): Instead of modifying the cost, this approach modifies the entropy regularization term. The TPM is used to define a prior joint probability matrix, which biases the entropy term to favor couplings consistent with the biological priors. This formulation avoids the need for an additional cost trade-off hyperparameter and provides a "soft filter" mechanism.

3. Training and Inference

ContextFlow utilizes the Sinkhorn algorithm to solve the prior-aware EOT problems efficiently. The resulting spatial-context-aware couplings are used to define the conditional probability paths and velocity fields for training a neural network to approximate the continuous velocity vector field. The model is trained via a regression loss that minimizes the difference between the learned velocity and the target conditional velocity derived from the context-aware couplings.

Key Contributions

Definition of Transitional Plausibility: The paper introduces a flexible mechanism to encode spatial organization and ligand-receptor communication patterns into the training of velocity fields, allowing for the explicit incorporation of domain-specific biological knowledge.
Principled Integration of Priors: The authors propose two methods (PACM and PAER) to integrate these priors into state-of-the-art OT-based flow matching. The PAER approach is highlighted for its ability to avoid additional hyperparameter tuning while maintaining efficiency via Sinkhorn optimization.
Empirical Validation: Comprehensive experiments demonstrate that ContextFlow consistently outperforms existing flow matching baselines (including MOTFM, Metric-FM, and MFM) across multiple metrics of inference accuracy and biological coherence.

Experimental Results

The authors evaluated ContextFlow on three longitudinal spatial transcriptomics datasets: Axolotl Brain Regeneration, Mouse Embryo Organogenesis, and Liver Regeneration.

Statistical Fidelity: ContextFlow, particularly the entropy-regularized variant (CTF-H), achieved superior performance in terms of 2-Wasserstein distance ( $W_2$ ), Maximum Mean Discrepancy (MMD), and Energy distance compared to baselines.
Biological Plausibility: Using a cell-type-weighted $W_2$ metric, ContextFlow generated trajectories that better preserved the distributional balance of cell types. Visual analysis confirmed a significant reduction in biologically implausible couplings (e.g., transitions between excitatory and inhibitory neurons) compared to standard EOT plans.
Generalization: The method demonstrated robust performance in both interpolation (predicting intermediate time points) and extrapolation (predicting future states beyond the training horizon) tasks. The CTF-H variant showed particular strength in Initial-Value-Problem (IVP) sampling, where error accumulation is a major challenge.
Efficiency: The proposed prior integration strategy is computationally efficient, with runtime comparable to standard minibatch OT methods and significantly faster than other spatiotemporal OT alignment tools like DeST-OT, MOSCOT, and TOAST.

Significance and Claims

The paper claims that ContextFlow offers a principled foundation for modeling cellular dynamics from spatial omics data by bridging the gap between statistical optimality and biological interpretability. By embedding spatial and functional priors directly into the transport plan, the framework generates trajectories that are consistent with known tissue organization and cell-cell communication mechanisms.

The authors position ContextFlow as a generalizable framework that improves upon existing flow matching methods without sacrificing computational efficiency. They note that while the method significantly reduces biologically implausible transitions, its effectiveness relies on the quality of the spatial priors and the availability of ligand-receptor interaction data. The work suggests that future applications could extend to modeling perturbations and disease progression, potentially offering a more interpretable alternative to purely data-driven generative models.