Title: Modeling Microenvironment Trajectories on Spatial Transcriptomics with NicheFlow
Venue: NeurIPS 2025
Accepted: yes
Keywords: flow matching

spatial transcriptomics

microenvironments

trajectory inference

generative modeling

point cloud
Abstract: Understanding the evolution of cellular microenvironments in spatiotemporal data is essential for deciphering tissue development and disease progression. While experimental techniques like spatial transcriptomics now enable high-resolution mapping of tissue organization across space and time, current methods that model cellular evolution operate at the single-cell level, overlooking the coordinated development of cellular states in a tissue. We introduce NicheFlow, a flow-based generative model that infers the temporal trajectory of cellular microenvironments across sequential spatial slides. By representing local cell neighborhoods as point clouds, NicheFlow jointly models the evolution of cell states and spatial coordinates using optimal transport and Variational Flow Matching. Our approach successfully recovers both global spatial architecture and local microenvironment composition across diverse spatiotemporal datasets, from embryonic to brain development.

Modeling Microenvironment Trajectories on Spatial Transcriptomics with NicheFlow

Abstract

Understanding the evolution of cellular microenvironments in spatiotemporal data is essential for deciphering tissue development and disease progression. While experimental techniques like spatial transcriptomics now enable high-resolution mapping of tissue organization across space and time, current methods that model cellular evolution operate at the single-cell level, overlooking the coordinated development of cellular states in a tissue. We introduce NicheFlow, a flow-based generative model that infers the temporal trajectory of cellular microenvironments across sequential spatial slides. By representing local cell neighborhoods as point clouds, NicheFlow jointly models the evolution of cell states and spatial coordinates using optimal transport and Variational Flow Matching. Our approach successfully recovers both global spatial architecture and local microenvironment composition across diverse spatiotemporal datasets, from embryonic to brain development 1 .

1 Introduction

Uncovering the principles governing tissue organization across space and time remains one of the most fundamental challenges in biology, with profound implications for evolutionary and developmental studies [1, 2]. While individual cells form the basic units of biological systems, they operate not in isolation but as integral parts of spatially organized microenvironments , functionally distinct neighborhoods, or niches , shaped by cell-to-cell interactions and extracellular components [3-5]. These spatial microenvironments influence crucial biological processes, from tumor progression to immune infiltration and tissue regeneration [6-8].

Spatial transcriptomics (ST) has transformed our ability to investigate these tissue architectures by providing single-cell resolution mapping of gene expression while preserving spatial context [9-12]. This technological breakthrough has enabled researchers to examine the molecular underpinnings of tissue organization with unprecedented detail. However, ST provides only static snapshots of inherently dynamic biological systems. Time-resolved spatial analysis extends ST by capturing how gene expression patterns and cellular arrangements evolve across developmental stages or experimental time [13-15]. This temporal dimension offers critical insights into the development of tissue organization in health and disease [16] that cannot be inferred from static observations alone.

Despite these technological advances, modeling trajectories on time-resolved spatial slides is complicated, as no direct correspondence exists between cells across slides due to the destructive nature of acquisition practices. Moreover, current computational methods fall short in modeling the evolution of tissue organization at the level of cellular microenvironments. Most approaches infer trajectories by modeling single-cell dynamics using velocity-based models [17-19] or optimal transport between individual cells [20, 21]. While effective at capturing cell evolution, these cell-centric methods fundamentally miss the coordinated evolution of structured niches within tissues.

How can we model the spatiotemporal evolution of cellular microenvironments preserving both local neighborhood relationships and cellular state transitions?

To address this question, we directly model the dynamics of cellular neighborhoods as cohesive units rather than focusing on isolated cell trajectories. This approach aligns naturally with tissue-scale biological processes and enables principled learning of dynamics in structured, high-dimensional, and variably sized spatial domains.

We introduce Niche Flow Matching ( NicheFlow ) (Fig. 1), a generative model for learning spatiotemporal dynamics of cellular niches from time-resolved spatial transcriptomics data. NicheFlow builds on recent advances in Flow Matching (FM) and Optimal Transport (OT) to operate over distributions of microenvironments, which we represent as point clouds. NicheFlow enables accurate modeling of global spatial architecture and local microenvironment composition within evolving tissues.

Our contributions include:

µ

θ

t

(

M

,

µ

θ

t

(

M

,

M

µ

θ

t

(

M

,

M

0

)

Source ME

Target ME

Noisy ME

M

z

∆

t

t

1

t

0

MO

口

A microenvironment-centered trajectory inference paradigm that shifts from modeling individual cells in time to modeling niches as point clouds, enabling simultaneous prediction of spatial coordinates and gene expression profiles while preserving local tissue context.

A factorized Variational Flow Matching (VFM) approach with distributional families (Laplace for spatial coordinates, Gaussian for gene expression) that jointly trains on spatial and cell state dynamics using a factorized loss, modeling spatial reconstruction and biological fidelity with tailored distributional assumptions.

A spatially-aware sampling strategy using OT between niche representations, enabling scalable training on large tissue sections while ensuring comprehensive coverage of heterogeneous regions.

Our approach consistently outperforms baselines in recovering cell-type organization and spatial structure across embryonic, brain development, and ageing datasets. NicheFlow enables principled learning of dynamics in structured, high-dimensional, and variably sized spatial domains, a challenge with parallels in other spatiotemporal modeling domains beyond biology.

2 Related Work

NicheFlow is at the interface between generative models and spatiotemporal transcriptomic data.

FMand single-cell transcriptomics. We propose a model based on FM, a framework introduced by several seminal works [22-24]. Specifically, we adopt a variational view of the FM objective, following Eijkelboom et al. [25], but extend it to mixed-factorized distributions for point cloud generation. Our method, NicheFlow, combines FM with OT, a pairing that has proven effective in modeling cellular data [26-29]. Unlike these models, however, we focus on point clouds of

<latexi sh

1_b

64="8dEoTcz2DA

uIWC

w

>

B

n

N

9LF

7Q

y

m

pZ

P

j

X

Y

+

G

R

g

5

v

f

K

<latexi sh

1_b

64="

jGPqS

Aow2

BW98

pv

C

J

c

>

H

VDL

gM

F

3

ZX

Y

Rm

r

Q

T

U

z

u

E

/

0

5

f

y

N

7

k

+

O

I

n

d

K

<latexi sh

1_b

64="

S

r

n09CIz

JN/QRFGog

Y

>A

B73

c

VDLT

E

+

p

2uy

j

H

O

w

M

f

K8

X

d

5

U

m

W

v

Z

P

k

q

<latexi sh

1_b

64="MoYcuV

CkP

/j

3I

q79

T

>A

B

n

NS8

EJ

Ur

LF

0

Q

Wy2m pZ

zX

H

+

G

O

Dd

R

g

5

v

f

K

w

<latexi sh

1_b

64="

uR

G92

0

gQycKF/

WI+

T

>A

C

X

VH

S

E

3

p

zMr7

B

d

UYZ

f

n

m

D

J

j

N58

Lq

k

v

P

O

w

o

<latexi sh

1_b

64="

uR

G92

0

gQycKF/

WI+

T

>A

C

X

VH

S

E

3

p

zMr7

B

d

UYZ

f

n

m

D

J

j

N58

Lq

k

v

P

O

w

o

<latexi sh

1_b

64="

uR

G92

0

gQycKF/

WI+

T

>A

C

X

VH

S

E

3

p

zMr7

B

d

UYZ

f

n

m

D

J

j

N58

Lq

k

v

P

O

w

o

0

)

M

0

)

spatially-resolved transcriptomic profiles. Closest to our approach is Wasserstein FM for point cloud generation [30], applied to reconstruct cellular niches. Yet, that work does not address joint generation of spatial coordinates and cellular states, nor OT-based temporal trajectory prediction, both of which are central to our contribution.

Generative models for spatial transcriptomics. Generative models have been key to spatial tasks such as gene expression prediction from histology slides [31, 32], integration with dissociated single-cell data [32], spatial imputation [33, 34], and perturbation [35]. More recently, LUNA [36] demonstrated strong performance in predicting single-cell spatial coordinates using diffusion models [37] conditioned on transcription data. While related, our model addresses the distinct task of inferring niche trajectories, enabling the simultaneous generation of coordinates and cellular states.

Trajectory inference for spatial transcriptomics. Previous work has explored learning trajectories from spatial slides. Pham et al. [38] proposed a graph-based spatiotemporal algorithm for pseudotime inference, while others leveraged tissue-resolved transcriptomics to estimate cell velocity [17-19]. Closer to our approach, Klein et al. [20] and Bryan et al. [21] use discrete OT to link cells across time and infer the evolution of cell states from spatially-resolved gene expression. Similarly, DeST-OT [39] aligns spatial slides with semi-relaxed OT couplings, preserving transcriptomic and spatial proximity between ancestor and descendant cells, while SpaTrack [19] uses Fused Gromov-Wasserstein OT [40], balancing transcriptomic and spatial differences based on spatial autocorrelation of features. Unlike our mini-batch deep learning model, these methods do not operate on entire microenvironments and rely on exact OT at the single-cell level, resulting in limitations in scalability and generalization.

3 Background

3.1 Optimal Transport with FM

FM [23] is a generative model that transforms a source density p 0 into a target density p 1 . It operates by learning a time-dependent velocity field u t ( x ) for t ∈ [0 , 1] and x ∈ R D , which generates a probability path { p t } t ∈ [0 , 1] . This path is constructed such that the marginals at time t = 0 and t = 1 match the source and target distributions, i.e., p 0 and p 1 , respectively. The velocity field induces an Ordinary Differential Equation (ODE), whose solution ϕ t ( x ) defines a flow map that transports samples from the source to the target distribution.

In practice, FM approximates u t ( x ) with a time-conditioned neural network v θ t ( x ) . While the exact marginal velocity field u t ( x ) is intractable, it can be expressed in terms of data-conditioned velocity fields and a joint distribution π ( x 0 , x 1 ) over the source and target samples x 0 ∼ p 0 and x 1 ∼ p 1 :

where p t ( x | x 0 , x 1 ) is a pre-defined interpolating probability path. Here, we consider the tractable probability path p t ( x | x 0 , x 1 ) = δ ( x -g t ( x 0 , x 1 )) , where g t ( x 0 , x 1 ) = (1 -t ) x 0 + t x 1 is a linear interpolation and δ denotes a Dirac delta function, representing a deterministic conditional path.

Lipman et al. [23] show that regressing the conditional field u t ( x | x 0 , x 1 ) is equivalent to learning the marginal field u t ( x ) in expectation. Hence, the FM objective becomes the task of learning the velocity along the conditional probability path between any pair of source and target data points. For a linear conditional probability path, the velocity u t ( x | x 0 , x 1 ) has a closed form, and the FM loss is:

In practice, the coupling π ( x 0 , x 1 ) is instantiated using sample pairs drawn from a mini-batch estimate. When one chooses π ⋆ as the OT coupling under a squared Euclidean cost between samples from p 0 and p 1 , FM approximates the dynamic OT map between source and target densities [41, 26]. Thus, the solution samples ( x 0 , x 1 ) ∼ π ⋆ from the joint distribution approximately follow:

where ϕ θ t is the solution of the ODE with velocity field v θ t .

3.2 Generative OT on incomparable source and target spaces

Klein et al. [29] generalize the OT FM formulation to settings where the source and target distributions are defined on incomparable spaces and propose an approach to generative entropic OT using FM. Given a standard normal noise distribution with samples z ∼ N ( 0 , I D ) , the authors show that the following sampling procedure:

defines a generative model that implicitly samples from an Entropic OT (EOT) coupling, where ϕ θ t ( z | x 0 ) is a FM model that maps noise to target samples, conditioned on source points. To achieve this, ϕ θ t is trained using source-target pairs ( x 0 , x 1 ) drawn from the EOT coupling π ⋆ ϵ , with ϵ denoting the entropic regularization parameter, which the model aims to approximate.

Crucially, this formulation enables OT between distinct source and target spaces, as x 0 does not flow directly into x 1 , but instead conditions the generation of target samples from noise .

3.3 Source-conditioned VFM

Consider the source-conditioned FM formulation in Sec. 3.2. Given a conditioning source x 0 and noise-based generation, the marginal field in Eq. (1) can be written as:

where we drop the conditioning on x 0 in the velocity field, as u t is entirely determined by the target x 1 when generating from noise under linear probability paths.

Since u t ( x | x 1 ) is tractable [23], one can recast FM as a variational inference problem, following Eijkelboom et al. [25], by introducing a parameterized approximation q θ t ( x 1 | x , x 0 ) to the true posterior p t ( x 1 | x , x 0 ) . Integrating the expected velocity in Eq. (5) over t ∈ [0 , 1] enables the generation of target points x 1 from noise, conditioned on x 0 .

During training, the source-conditioned Variational Flow Matching (VFM) loss is:

where π ⋆ ϵ is an entropic OT coupling modeling the joint distribution over source and target samples, and p t ( x | x 1 ) interpolates between target samples and noise. In the generation phase, one samples x 0 ∼ p 0 and noise z ∼ N ( 0 , I D ) , then simulates the marginal field in Eq. (5) starting from ϕ 0 ( x ) = z to generate a target sample x 1 .

Crucially, under the assumption that u t ( x | x 1 ) is linear in x 1 , which holds when using straight-line interpolation paths, the marginal field in Eq. (5) only depends on the posterior's first moment on x 1 :

This implies that the VFM objective reduces to matching the first moment of the approximate posterior q θ t ( x 1 | x , x 0 ) to that of the true posterior p t ( x 1 | x , x 0 ) . As a result, the approximate posterior can be chosen fully factorized under a mean field assumption , since each dimension can be matched independently if the mean of the true posterior is preserved; see App. C.2 for more details.

4 NicheFlow

We introduce a flow-based generative OT model to infer the temporal evolution of spatially resolved cellular microenvironments. More specifically, given a spatial microenvironment represented as a point cloud of cell states with their coordinates, NicheFlow predicts the corresponding tissue structure at a later time point. To delineate our approach, we define a list of desiderata.

Generative model on structured data. Similar to prior work [30, 36], we consider a generative model over structured point cloud data representing cellular microenvironments. This approach implicitly accounts for spatial correlations between cells, in contrast to models that study the evolution of spatial trajectories at the single-cell level [20].

Sub-regions and variable location. Crucially, for better memory efficiency, we do not consider an entire spatial slide for trajectory inference, but instead learn the dynamics of variably located

sub-regions. This design choice enables scalability and flexibility in modeling functional regions across different parts of the screened tissue.

Changes in the number of nodes. To model the temporal evolution and densification of microenvironments, we allow the source and target regions to differ in the number of nodes. We adopt a formulation similar to Sec. 3.2, implementing OT FM between non-comparable spaces.

Flexible generative models for features and coordinates. We allow flexibility in the choice of generative models for the features and coordinates. To this end, we implement the approach described in Sec. 3.3, factorizing features and coordinates into separate posteriors from different families.

4.1 Data description and problem statement

We are given a sequence of time-resolved spatial transcriptomic measurements across biological processes such as development or ageing. For simplicity, we formulate our problem in terms of two consecutive time points indexed by s , such that s ∈ { 0 , 1 } , though the model can be extended to collections of more consecutive discrete temporal measurements. Each dataset at a time point is a full tissue slide that can be represented as an attributed point cloud :

where c s i ∈ R 2 denotes the 2D spatial coordinate of cell i at time s , and x s i ∈ R D denotes its associated feature vector, typically corresponding to its gene expression profile or a low-dimensional representation thereof. Thus, each dataset is an attributed point set in two-dimensional space. Note that each slide can have a variable number of cells, and no direct correspondence exists between single cells across subsequent slides.

To capture spatial context beyond individual cells, we define local microenvironments as fixed-radius neighborhoods. Specifically, for each cell ( c s i , x s i ) at time s , we construct a neighborhood M s i consisting of all neighboring cells within a spatial radius r :

Let {M 0 i } N 0 i =1 and {M 1 j } N 1 j =1 be collections of source and target microenvironments at consecutive time points. Our goal is to train a parameterized flow model ϕ θ t , with t ∈ [0 , 1] , that generates target microenvironments conditioned on source point clouds. Specifically, to sample a target microenvironment with k cells conditioned on a variably sized source M 0 , we define sampling as:

where M z is a point cloud composed of noisy coordinates and features, and M 1 is a generated prediction for the evolution of M 0 at the next time point. As explained in Sec. 3.1 and Sec. 3.2, we want our generative model to parameterize some notion of optimal entropic coupling π ⋆ ϵ between microenvironments across slides (see Sec. 4.2).

4.2 OT formulation

To train the flow map ϕ θ t to perform conditional EOT, we define a cost function that induces an optimal entropic coupling π ⋆ ϵ between source and target point clouds. This coupling is used to sample pairs of source and target microenvironment mini-batches during training. While there is no established notion of optimal cost in this setting, we propose to compute OT using source and target microenvironment representations based on the weighted average of features and coordinates.

Specifically, we compute a pooled representation for each microenvironment in the source and target slides by averaging spatial coordinates and gene expression features, weighted by a tunable hyperparameter λ ∈ [0 , 1] that balances spatial versus cellular state information:

where ∥ denotes concatenation, ¯ m s i ∈ R 2+ D , and s ∈ { 0 , 1 } . We then apply EOT on the sets { ¯ m 0 i } N 0 i =1 and { ¯ m 1 j } N 1 j =1 using a squared Euclidean cost and regularization parameter ϵ , yielding the

coupling π ⋆ ϵ,λ . During training, we sample matched pairs ( M 0 , M 1 ) ∼ π ⋆ ϵ,λ computed over minibatches, and use them as supervision for learning the conditional generative model. A higher value of λ prioritizes feature similarity, while lower values favor proximity in coordinates (see Fig. 12).

4.3 Mixed-factorized VFM

Once we have established a strategy for performing mini-batch OT, we proceed to describe our approach for learning the flow model ϕ θ t and simulating Eq. (11). To this end, we adopt a variant of VFM (Sec. 3.3), originally developed for graph generation, and adapt it to our point cloud setting.

In line with Sec. 3.2 and 3.3, we delineate an objective to train a parameterized, source-conditioned posterior over target point clouds q θ t ( M 1 |M , M 0 ) , where M 0 and M 1 represent source and target niches, and M denotes a noisy point cloud at interpolation time t . Importantly, our posterior comes with the following characteristics: (i) the posterior is factorized across the single points in a point cloud; (ii) the posterior is factorized across cellular features and coordinate dimensions; and (iii) the family of posteriors can be chosen differently between cellular features and coordinates.

Following (i) , we model the variational distribution over M 1 by factorizing it across individual points ( c 1 , x 1 ) ∈ M 1 . Moreover, we tackle (ii) and (iii) using the mean-field VFM assumption, modeling cellular state and positions separately (see App. C.1 and C.2 for theoretical justifications):

Here, f θ t and r θ t denote distinct approximate posterior families for cellular states ( x d 1 ) and spatial positions ( c k 1 ), respectively. We use a Laplace distribution for f θ t due to its concentration around the mean, which supports precise modeling of coordinate features, while a Gaussian distribution is used for r θ t .

As explained in Sec. 3.3, if one uses FM with straight probability paths, only the first moment of q θ t is required to simulate the generative field in Eq. (5). Therefore, the posterior is replaced by a time-dependent predictor ( ¯ f θ t , ¯ r θ t ) = µ θ t ( M , M 0 ) of the mean features ¯ r θ t and coordinates ¯ f θ t of the target point clouds, learned minimizing the following loss:

We derive the objective in App. C.4 and provide algorithms in App. E. Here, µ θ t inputs a noisy point cloud M and a source M 0 and provides a mean prediction vector for coordinates and feature dimensions, respectively indicated as ¯ f θ t and ¯ r θ t . Weimplement it as a point cloud transformer (see Sec. 4.4).

We highlight a crucial aspect about Eq. (14). While the single dimensions are fully factorized in the predictions from µ θ t , every feature's mean is a function of the whole noisy point cloud M as well as the target microenvironment M 0 . In other words, the predictions exploit structural information in the point cloud to predict the individual posterior mean of each dimension. Like most approaches, our method has modeling limitations that we outline in App. B.

4.4 Backbone architecture: Microenvironment transformer

To parameterize the conditional posterior mean µ θ t from Sec. 4.3, we use a permutation-equivariant transformer architecture designed for point clouds with variable size.

Encoder-decoder structure. The model follows an encoder-decoder layout, processing the source microenvironment M 0 via the encoder and predicting the posterior mean from a noisy target M∼ p t ( ·|M 1 ) via the decoder conditioned on the encoder's output.

Input embeddings. Each point is represented by features x ∈ R D and spatial coordinates c ∈ R 2 , embedded separately and concatenated. Time t is encoded using sinusoidal embeddings, linearly projected and broadcast across points.

Cross-attention conditioning. The encoder processes the embedded source M 0 microenvironment using self-attention. The decoder operates on the noisy target M using self-attention, followed by

cross-attention to condition on the encoder's outputs. The cross-attention mechanism allows each target point to attend to all source points.

Output projection. The decoder outputs are linearly projected to yield posterior mean estimates ( ¯ f θ t , ¯ r θ t ) for expression and coordinates.

5 Experiments

We propose quantitative and qualitative evaluations of our algorithm. Quantitatively, we test whether source-conditioned samples generated by our model preserve the biological structure and shape of future tissue states. Qualitatively, we demonstrate that our approach accurately captures compositional shifts in substructural components and developmental trajectories across time.

5.1 Quantitative evaluation

Our first research question is to assess the impact of two core modeling choices in NicheFlow: (i) learning trajectories over spatial microenvironments, rather than independently for each cell and (ii) restricting source and target point clouds to spatially co-localized neighborhoods of cells, instead of sampling them randomly across the slide.

We use NicheFlow and baseline FM approaches that do not incorporate (i) and (ii) to simulate spatial trajectories conditioned on early time point observations. Assuming that spatial arrangements and the biological composition at later stages evolve from earlier slides, the global correspondence between predicted and true slides indicates the quality of the generative trajectory.

5.1.1 Training setup

Datasets. We assess model performance across three spatiotemporal datasets: (i) Mouse embryogenesis [20, 8] and (ii) the axolotl brain development [42], two Stereo-seq datasets profiling the spatially-resolved cellular development of a mouse embryo and axolotl brain across three (E9.5, E10.5 and E11.5) and five time points, respectively. We also consider the (iii) mouse brain ageing dataset [43], profiled with MERFISH [44] across twenty time points (see Apps. F.1 and F.2).

Dataset construction. For each dataset and time point s ∈ S , we construct a set of cellular microenvironments by applying the fixed-radius neighborhood definition introduced in Sec. 4.1. Each microenvironment M s i is centered at cell i and contains all cells within a fixed radius r . This results in:

where N s is the number of cells in the tissue at time s , and each M s i is an attributed point cloud encoding both spatial and gene expression information. We standardize coordinates for cross-time comparability and reduce the normalized gene expression to its top 50 Principal Components (PC).

Batching. We train NicheFlow with mini-batches of source and target cellular point clouds. To ensure spatial diversity during training, we sample individual batches uniformly from within discrete

Target (10.5)

SPFlow CFM

SPFlow GVFM

SPFlow GLVFM

NicheFlow CFM

NicheFlow GVFM

NicheFlow GLVFM

Source (9.5)

Target (11.5)

SPFlow CFM

SPFlow GVFM

SPFlow GLVFM

NicheFlow CFM

NicheFlow GVFM

NicheFlow GLVFM

Source (10.5)

regions of the slides computed with K -Means clustering over the 2D coordinates (see Fig. 13 for a visualization with different K values). From these regions, we collect M source and target microenvironments and resample N ≤ M matching pairs from the entropic OT coupling ( M 0 , M 1 ) ∼ π ⋆ ϵ,λ as described in Sec. 3.1, where M 0 and M 1 denote the sampled sets (see Sec. 4.2 for details on our OT coupling).

Evaluation data. For consistent and reproducible evaluation, we discretize each tissue into a fixed 2D grid and define evaluation microenvironments as fixed-radius neighborhoods around the nearest cells to each grid point. This guarantees full spatial coverage and ensures deterministic comparison across methods. See App. F.5 for details.

Multiple time-point. NicheFlow predicts piecewise trajectories between subsequent time points. Instead of learning one flow for each couple of subsequent slides, we train a single model with additional conditioning on source and target labels (see App. F.6).

5.1.2 Quantitative evaluation metrics

Spatial structure. We quantify coordinate generation accuracy using two asymmetric distance metrics. The point-to-shape distance (PSD) measures how far predicted coordinates deviate from the true structure by computing the mean squared distance from each generated point to its nearest ground truth counterpart. In contrast, the shape-to-point distance (SPD) evaluates how well the generated points cover the target region by averaging the squared distance from each ground truth point to its nearest generated point (see App. F.3 for a mathematical formulation of the metrics).

Cell-type organization. To assess how well the model reconstructs the spatial organization of different cell types, we use a 1-nearest-neighbor (1NN) classification setup. Since the model generates only gene expression profiles and spatial coordinates, we assign cell type labels to generated cells using a classifier trained on ground truth gene expression data (see App. F.4). Each predicted cell is then matched to its nearest real cell, and we report the weighted F1 score (1NN-F1).

5.1.3 Models and results

Baselines. We compare against what we call SPFlow (Single-Point Flow), a standard FM-based model that predicts temporal trajectories across slides at a single-cell level using an MLP-based velocity field. We also consider RPCFlow (Random Point Cloud Flow), which has the same backbone as NicheFlow, but conditions on randomly sampled point clouds instead of radius neighborhoods. Additionally, we include LUNA [36], a diffusion model for spatial reconstruction from dissociated cells. Note that LUNA does not model temporal dynamics and only generates coordinates from noise with their respective biological annotations. Therefore, we use it as a reference for spatial generation accuracy via the 1NN-F1 metric rather than a proper baseline.

Ablations. We assess different training objectives by comparing standard Conditional Flow Matching (CFM) [26] with two variational formulations modeling posteriors over coordinates and features: Gaussian-only (GVFM) and Gaussian-Laplace (GLVFM). The former uses Gaussian posteriors for coordinates and features. The latter uses the factorized formulation in Sec. 4.3. For the point-cloudbased methods, we use a fixed value of λ = 0 . 1 and sampled batches of 64 regions chosen from

Mapping spinal cord E10.5-E11.5

Mapping neural crest cells E9.5-E10.5

Source

NicheFlow mapping probability

Source

NicheFlow mapping probability

入 =0.1

入=0.5

Frequency

0.3

Frequency

0.4

·Spinal cord

0.2

0.2

0.1

Neural crest

0.0

0.0

jion

me

gangli

Cavity

root

urface

xe

Expected target

moscot mapping probability

Expected target

moscot mapping probability

●Dermomyotome

Off-target

Q=0.9

Off-target

0.6

Q=0.5

Sclerotome

Frequency

Surface ectoderm

Dorsal root ganglion

0.1

·Head mesenchyme

0.0

0.0

Cavity

Sclerotome

the ablations in App. D.5 and D.6. We present results with a single value for λ as this experiment compares different models on reconstructing fixed tissue structures over time, for which we enforce spatial preservation (see Sec. 4.2). However, we also include model comparisons across multiple values of λ in Tab. 3 for completion.

Results. Our quantitative evaluation (Tab. 1) demonstrates that NicheFlow trained with the GLVFM objective consistently achieves strong performance across both spatial and semantic metrics. It outperforms all baselines in reconstructing spatial structure (PSD, SPD) and cell-type organization (1NN-F1) on developmental datasets, while remaining competitive on ageing data. These results highlight the importance of structured microenvironment modeling for capturing the spatiotemporal dynamics of complex tissues. We complement our quantitative results visually in Fig. 2 and Fig. 6 and 10 in the Appendix, where we show that SPFlow fails to capture tissue-level organization, producing blurry and spatially incoherent samples. In contrast, NicheFlow generates predictions that preserve spatial structure and cell-type organization, despite learning only from local microenvironments. In App. D.4 and App. D.7, we additionally demonstrate that NicheFlow produces conditionally consistent outputs with the source, while RPCFlow generates very diffused mappings across the slide, failing to preserve spatial consistency across time.

5.2 Qualitative evaluation and biological analysis

We explore the capabilities of NicheFlow on the spatial trajectory inference task through qualitative and biological assessments. Specifically, we focus on validating whether our model captures compositional changes within fixed spatial structures and developmental trajectories.

Experimental setup. We train the model as described in Sec. 5.1.1. Using the mouse embryonic dataset [8], we select specific microenvironments as source niches for which we want to study the trajectory over time. Specifically, we propose two scenarios for the application of NicheFlow depending on the choice of the OT parameter λ (see Sec. 4.2):

Compositional changes in fixed structures across time . We choose the evolution of the spinal cord of the embryo from E10.5 to E11.5 as an example and set a low value of λ to prioritize the preservation of the spatial location in the trajectory.

Spatial and cellular development of immature cells . Developing cells may displace to different areas of the embryo, requiring a higher value of λ to account for gene expression. As a case study, we inspect how neural crest cells in the head evolve into mesenchymal and cranial structures.

Baseline. We compare NicheFlow with the spatiotemporal framework in moscot [20], which models spatial trajectories at the single-cell level. In contrast to NicheFlow, moscot integrates spatial

coordinates directly into the OT formulation using a Fused Gromov-Wasserstein cost [45], where the hyperparameter α controls the trade-off between spatial and feature-based distances (see App. F.8). A high α increases the influence of spatial distances, whereas in our framework this role is played by the hyperparameter λ (see Sec. 4.2). We provide more details on the selection of the hyperparameter α for our experiments in App. D.2. Notably, moscot relies on exact OT and learns a transition matrix between source and target samples. As such, it is not a generative model over point clouds like NicheFlow. However, given the overlap in downstream tasks, we consider the comparison relevant.

Evaluation and results. For both scenarios (1) and (2), we have prior knowledge of the ground truth regions that the source microenvironments are expected to occupy at later developmental stages, as well as their corresponding biological compositions. For both moscot and NicheFlow, we assess whether the transported mass of source samples concentrates within the correct anatomical region at the target time point, and whether the predicted descendant cell types are biologically consistent (see App. F.8). Results are summarized in Fig. 3. When modeling the evolution of the spinal cord, moscot assigns considerable mass to unrelated regions such as the urogenital ridge and branchial arches, whereas NicheFlow correctly maps source niches to the maturing spinal cord. Similarly, NicheFlow captures the differentiation of neural crest cells into mesenchymal and cranial tissues within the head region, while moscot exhibits substantial off-target leakage towards lower regions.

6 Conclusions and Discussion

We introduce NicheFlow, a point-cloud-based generative model designed to capture the spatiotemporal dynamics of cellular niches in time-resolved spatial transcriptomics data. Unlike methods that model single-cell trajectories independently, NicheFlow implicitly captures spatial correlations between cells by learning trajectories on variably sized local neighborhoods. To this end, we combine OT with a new version of VFM that factorizes features and coordinates into distinct posteriors from different distribution families. We showed that NicheFlow outperforms standard FM approaches at reconstructing spatial context from previous time points and improves the mapping of biological structures in time over established exact OT approaches. With the expected increase in the volume and quality of spatial data, modeling coordinated cellular state translations with generative models is a promising avenue for generalization beyond spatiotemporal inference. We envision that models like NicheFlow will enable spatial perturbation prediction and modality translation tasks requiring principled parameterized maps beyond discrete OT to extrapolate and drive biological hypotheses.

Acknowledgments

A.P. is supported by the Helmholtz Association under the joint research school Munich School for Data Science (MUDS). A.P., F.G., S.G. and F.J.T. also acknowledge support from the German Federal Ministry of Education and Research (BMFTR) through grant numbers 031L0289A and 031L0289C. F.J.T. acknowledges support from the Helmholtz Association's Initiative and Networking Fund via the CausalCellDynamics project (grant number Interlabs-0029) and the European Union (ERC, DeepCell, grant number 101054957). The authors of this work take full responsibility for its content.

References

Urs Mayr, Denise Serra, Prisca Liberali, Allon Klein, and Barbara Treutlein. Exploring single cells in space and time during tissue development, homeostasis and regeneration. Development , 146(12), 2019.

Marietta Zinner, Ilya Lukonin, and Prisca Liberali. Design principles of tissue organisation: How single cells coordinate across scales. Current opinion in cell biology , 67:37-45, 2020.

Yanming Ren, Zongyao Huang, Lingling Zhou, Peng Xiao, Junwei Song, Ping He, Chuanxing Xie, Ran Zhou, Menghan Li, Xiangqun Dong, Qing Mao, Chao You, Jianguo Xu, Yanhui Liu, Zhigang Lan, Tiejun Zhang, Qi Gan, Yuan Yang, Tengyun Chen, Bowen Huang, Xiang Yang, Anqi Xiao, Yun Ou, Zhengzheng Su, Lu Chen, Yan Zhang, Yan Ju, Yuekang Zhang, and Yuan Wang. Spatial transcriptomics reveals niche-specific enrichment and vulnerabilities of radial glial stem-like cells in malignant gliomas. Nature Communications , 14(1), February 2023.

Anna Christina Schaar, Alejandro Tejada-Lapuerta, Giovanni Palla, Robert Gutgesell, Lennard Halle, Mariia Minaeva, Larsen Vornholz, Leander Dony, Francesca Drummer, Mojtaba Bahrami, and Fabian J Theis. Nicheformer: a foundation model for single-cell and spatial omics. bioRxiv , pages 2024-04, April 2024.

Song Liu, Xiaoyan Li, Zhiyue Gu, Jiayu Wu, Shuangzheng Jia, Jinghua Shi, Yi Dai, Yushi Wu, Hailan Yan, Jing Zhang, Yan You, Xiaowei Xue, Lulu Liu, Jinghe Lang, Xiaoyue Wang, and Jinhua Leng. Single-cell and spatial transcriptomic profiling revealed niche interactions sustaining growth of endometriotic lesions. Cell Genomics , 5(1):100737, January 2025.

Reuben Moncada, Dalia Barkley, Florian Wagner, Marta Chiodin, Joseph C. Devlin, Maayan Baron, Cristina H. Hajdu, Diane M. Simeone, and Itai Yanai. Integrating microarray-based spatial transcriptomics and single-cell RNA-seq reveals tissue architecture in pancreatic ductal adenocarcinomas. Nature Biotechnology , 38(3):333-342, January 2020.

Ludvig Larsson, Jonas Frisén, and Joakim Lundeberg. Spatially resolved transcriptomics adds a new dimension to genomics. Nature Methods , 18(1):15-18, January 2021.

Ao Chen, Sha Liao, Mengnan Cheng, Kailong Ma, Liang Wu, Yiwei Lai, Xiaojie Qiu, Jin Yang, Jiangshan Xu, Shijie Hao, Xin Wang, Huifang Lu, Xi Chen, Xing Liu, Xin Huang, Zhao Li, Yan Hong, Yujia Jiang, Jian Peng, Shuai Liu, Mengzhe Shen, Chuanyu Liu, Quanshui Li, Yue Yuan, Xiaoyu Wei, Huiwen Zheng, Weimin Feng, Zhifeng Wang, Yang Liu, Zhaohui Wang, Yunzhi Yang, Haitao Xiang, Lei Han, Baoming Qin, Pengcheng Guo, Guangyao Lai, Pura Muñoz-Cánoves, Patrick H. Maxwell, Jean Paul Thiery, Qing-Feng Wu, Fuxiang Zhao, Bichao Chen, Mei Li, Xi Dai, Shuai Wang, Haoyan Kuang, Junhou Hui, Liqun Wang, JiFeng Fei, Ou Wang, Xiaofeng Wei, Haorong Lu, Bo Wang, Shiping Liu, Ying Gu, Ming Ni, Wenwei Zhang, Feng Mu, Ye Yin, Huanming Yang, Michael Lisby, Richard J. Cornall, Jan Mulder, Mathias Uhlén, Miguel A. Esteban, Yuxiang Li, Longqi Liu, Xun Xu, and Jian Wang. Spatiotemporal transcriptomic atlas of mouse organogenesis using DNA nanoball-patterned arrays. Cell , 185(10):1777-1792.e21, May 2022.

Patrik L. Ståhl, Fredrik Salmén, Sanja Vickovic, Anna Lundmark, José Fernández Navarro, Jens Magnusson, Stefania Giacomello, Michaela Asp, Jakub O. Westholm, Mikael Huss, Annelie Mollbrink, Sten Linnarsson, Simone Codeluppi, Åke Borg, Fredrik Pontén, Paul Igor Costea, Pelin Sahlén, Jan Mulder, Olaf Bergmann, Joakim Lundeberg, and Jonas Frisén. Visualization and analysis of gene expression in tissue sections by spatial transcriptomics. Science , 353(6294): 78-82, July 2016.

Samuel G. Rodriques, Robert R. Stickels, Aleksandrina Goeva, Carly A. Martin, Evan Murray, Charles R. Vanderburg, Joshua Welch, Linlin M. Chen, Fei Chen, and Evan Z. Macosko. Slideseq: A scalable technology for measuring genome-wide expression at high spatial resolution. Science , 363(6434):1463-1467, March 2019.

Chee-Huat Linus Eng, Michael Lawson, Qian Zhu, Ruben Dries, Noushin Koulena, Yodai Takei, Jina Yun, Christopher Cronin, Christoph Karp, Guo-Cheng Yuan, and Long Cai. Transcriptomescale super-resolved imaging in tissues by RNA seqFISH+. Nature , 568(7751):235-239, March 2019.

Lambda Moses and Lior Pachter. Museum of spatial transcriptomics. Nature Methods , 19(5): 534-546, March 2022.

James A. Briggs, Caleb Weinreb, Daniel E. Wagner, Sean Megason, Leonid Peshkin, Marc W. Kirschner, and Allon M. Klein. The dynamics of gene expression in vertebrate embryogenesis at single-cell resolution. Science , 360(6392), June 2018.

Daniel E. Wagner, Caleb Weinreb, Zach M. Collins, James A. Briggs, Sean G. Megason, and Allon M. Klein. Single-cell mapping of gene expression landscapes and lineage in the zebrafish embryo. Science , 360(6392), June 2018.

Jingyi Ren, Haowen Zhou, Hu Zeng, Connie Kangni Wang, Jiahao Huang, Xiaojie Qiu, Xin Sui, Qiang Li, Xunwei Wu, Zuwan Lin, Jennifer A. Lo, Kamal Maher, Yichun He, Xin Tang, Judson Lam, Hongyu Chen, Brian Li, David E. Fisher, Jia Liu, and Xiao Wang. Spatiotemporally

resolved transcriptomics reveals the subcellular RNA kinetic landscape. Nature Methods , 20(5): 695-705, April 2023.

Michael Papanicolaou, Amelia L Parker, Michelle Yam, Elysse C Filipe, Sunny Z Wu, Jessica L Chitty, Kaitlin Wyllie, Emmi Tran, Ellie Mok, Audrey Nadalini, et al. Temporal profiling of the breast tumour microenvironment reveals collagen xii as a driver of metastasis. Nature communications , 13(1):4587, 2022.

Tamim Abdelaal, Laurens M Grossouw, R Jeroen Pasterkamp, Boudewijn PF Lelieveldt, Marcel JT Reinders, and Ahmed Mahfouz. Sirv: Spatial inference of rna velocity at the single-cell resolution. NAR genomics and bioinformatics , 6(3):lqae100, 2024.

Wenxin Long, Tianyu Liu, Lingzhou Xue, and Hongyu Zhao. spvelo: Rna velocity inference for multi-batch spatial transcriptomics data. bioRxiv , pages 2025-03, 2025.

Xunan Shen, Lulu Zuo, Zhongfei Ye, Zhongyang Yuan, Ke Huang, Zeyu Li, Qichao Yu, Xuanxuan Zou, Xiaoyu Wei, Ping Xu, et al. Inferring cell trajectories of spatial transcriptomics via optimal transport analysis. Cell Systems , 16(2), 2025.

Dominik Klein, Giovanni Palla, Marius Lange, Michal Klein, Zoe Piran, Manuel Gander, Laetitia Meng-Papaxanthos, Michael Sterr, Lama Saber, Changying Jing, et al. Mapping cells through time and space with moscot. Nature , pages 1-11, 2025.

John Peterson Bryan, Samouil L Farhi, and Brian Cleary. Accurate trajectory inference in time-series spatial transcriptomics with structurally-constrained optimal transport. bioRxiv , pages 2025-03, 2025.

Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow. In The Eleventh International Conference on Learning Representations , 2023.

Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Eleventh International Conference on Learning Representations , 2023.

Michael Samuel Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants. In The Eleventh International Conference on Learning Representations , 2023.

Floor Eijkelboom, Grigory Bartosh, Christian A. Naesseth, Max Welling, and Jan-Willem van de Meent. Variational flow matching for graph generation. In The Thirty-eighth Annual Conference on Neural Information Processing Systems , 2024.

Alexander Tong, Kilian FATRAS, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research , 2024.

Alexander Y. Tong, Nikolay Malkin, Kilian Fatras, Lazar Atanackovic, Yanlei Zhang, Guillaume Huguet, Guy Wolf, and Yoshua Bengio. Simulation-free Schrödinger bridges via score and flow matching. In Proceedings of The 27th International Conference on Artificial Intelligence and Statistics , volume 238 of Proceedings of Machine Learning Research , pages 1279-1287. PMLR, 02-04 May 2024.

Lazar Atanackovic, Xi Zhang, Brandon Amos, Mathieu Blanchette, Leo J Lee, Yoshua Bengio, Alexander Tong, and Kirill Neklyudov. Meta flow matching: Integrating vector fields on the wasserstein manifold. In The Thirteenth International Conference on Learning Representations , 2025.

Dominik Klein, Théo Uscidda, Fabian Theis, and Marco Cuturi. Genot: Entropic (gromov) wasserstein flow matching with applications to single-cell genomics. Advances in Neural Information Processing Systems , 37:103897-103944, 2024.

Doron Haviv, Aram-Alexandre Pooladian, Dana Pe'er, and Brandon Amos. Wasserstein flow matching: Generative modeling over families of distributions. arXiv preprint arXiv:2411.00698 , 2024.

Sichen Zhu, Yuchen Zhu, Molei Tao, and Peng Qiu. Diffusion generative modeling for spatially resolved gene expression inference from histology images. In The Thirteenth International Conference on Learning Representations , 2025.

Xiaomeng Wan, Jiashun Xiao, Sindy Sing Ting Tam, Mingxuan Cai, Ryohichi Sugimura, Yang Wang, Xiang Wan, Zhixiang Lin, Angela Ruohao Wu, and Can Yang. Integrating spatial and single-cell transcriptomics data using deep generative models with spatialscope. Nature Communications , 14(1):7848, 2023.

Doron Haviv, Ján Remšík, Mohamed Gatie, Catherine Snopkowski, Meril Takizawa, Nathan Pereira, John Bashkin, Stevan Jovanovich, Tal Nawy, Ronan Chaligne, et al. The covariance environment defines cellular niches for spatial inference. Nature Biotechnology , 43(2):269-280, 2025.

Kongming Li, Jiahao Li, Yuhao Tao, and Fei Wang. stdiff: a diffusion model for imputing spatial transcriptomics through single-cell transcriptomics. Briefings in Bioinformatics , 25(3): bbae171, 2024.

Amir Akbarnejad, Lloyd Steele, Daniyal J Jafree, Sebastian Birk, Marta Rosa Sallese, Koen Rademaker, Adam Boxall, Benjamin Rumney, Catherine Tudor, Minal Patel, et al. Mapping and reprogramming human tissue microenvironments with mintflow. bioRxiv , pages 2025-06, 2025.

Tingyang Yu, Chanakya Ekbote, Nikita Morozov, Jiashuo Fan, Pascal Frossard, Stéphane d'Ascoli, and Maria Brbi´ c. Tissue reassembly with generative ai. bioRxiv , pages 2025-02, 2025.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems , 33:6840-6851, 2020.

Duy Pham, Xiao Tan, Brad Balderson, Jun Xu, Laura F Grice, Sohye Yoon, Emily F Willis, Minh Tran, Pui Yeng Lam, Arti Raghubar, et al. Robust mapping of spatiotemporal trajectories and cell-cell interactions in healthy and diseased tissues. Nature communications , 14(1):7739, 2023.

Peter Halmos, Xinhao Liu, Julian Gold, Feng Chen, Li Ding, and Benjamin J Raphael. Dest-ot: Alignment of spatiotemporal transcriptomics data. Cell Systems , 16(2), 2025.

Titouan Vayer, Laetitia Chapel, Rémi Flamary, Romain Tavenard, and Nicolas Courty. Fused gromov-wasserstein distance for structured objects. Algorithms , 13(9):212, 2020.

Aram-Alexandre Pooladian, Heli Ben-Hamu, Carles Domingo-Enrich, Brandon Amos, Yaron Lipman, and Ricky T. Q. Chen. Multisample flow matching: Straightening flows with minibatch couplings. In Proceedings of the 40th International Conference on Machine Learning , volume 202 of Proceedings of Machine Learning Research , pages 28100-28127. PMLR, 23-29 Jul 2023.

Xiaoyu Wei, Sulei Fu, Hanbo Li, Yang Liu, Shuai Wang, Weimin Feng, Yunzhi Yang, Xiawei Liu, Yan-Yun Zeng, Mengnan Cheng, Yiwei Lai, Xiaojie Qiu, Liang Wu, Nannan Zhang, Yujia Jiang, Jiangshan Xu, Xiaoshan Su, Cheng Peng, Lei Han, Wilson Pak-Kin Lou, Chuanyu Liu, Yue Yuan, Kailong Ma, Tao Yang, Xiangyu Pan, Shang Gao, Ao Chen, Miguel A. Esteban, Huanming Yang, Jian Wang, Guangyi Fan, Longqi Liu, Liang Chen, Xun Xu, Ji-Feng Fei, and Ying Gu. Single-cell stereo-seq reveals induced progenitor cells involved in axolotl brain regeneration. Science , 377(6610):eabp9444, 2022.

Eric D. Sun, Olivia Y. Zhou, Max Hauptschein, Nimrod Rappoport, Lucy Xu, Paloma Navarro Negredo, Ling Liu, Thomas A. Rando, James Zou, and Anne Brunet. Spatial transcriptomic clocks reveal cell proximity effects in brain ageing. Nature , 638(8049):160-171, Feb 2025.

Kok Hao Chen, Alistair N Boettiger, Jeffrey R Moffitt, Siyuan Wang, and Xiaowei Zhuang. Spatially resolved, highly multiplexed rna profiling in single cells. Science , 348(6233):aaa6090, 2015.

Titouan Vayer, Nicolas Courty, Romain Tavenard, and Rémi Flamary. Optimal Transport for Structured Data with Application on Graphs. In Proceedings of the 36th International Conference on Machine Learning , volume 97 of Proceedings of Machine Learning Research , pages 6275-6284. PMLR, 2019.

Haoqiang Fan, Hao Su, and Leonidas Guibas. A point set generation network for 3d object reconstruction from a single image. In Computer Vision and Pattern Recognition , 2016.

Facundo Mémoli. Gromov-wasserstein distances and the metric approach to object matching. Foundations of Computational Mathematics , 11(4):417-487, Aug 2011.

Alexander Tong, Jessie Huang, Guy Wolf, David Van Dijk, and Smita Krishnaswamy. TrajectoryNet: A dynamic optimal transport network for modeling cellular dynamics. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning , volume 119 of Proceedings of Machine Learning Research , pages 9526-9536. PMLR, 2020.

Guillaume Huguet, Daniel Sumner Magruder, Alexander Tong, Oluwadamilola Fasina, Manik Kuchroo, Guy Wolf, and Smita Krishnaswamy. Manifold interpolating optimal-transport flows for trajectory inference. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022.

Zhenyi Zhang, Tiejun Li, and Peijie Zhou. Learning stochastic dynamics from snapshots through regularized unbalanced optimal transport. In The Thirteenth International Conference on Learning Representations , 2025.

A Broader impacts

This work addresses fundamental challenges in spatial transcriptomics by modeling complex spatial and compositional changes in developing tissues. We demonstrate how efficient representations of high-dimensional spatial cellular data can advance the understanding of developmental trajectories and microenvironment dynamics. We anticipate releasing NicheFlow as an open-source, user-friendly tool to enable broad application in spatial biology studies. Given its use with biological data, NicheFlow may also be applied in sensitive contexts involving clinical or patient information.

B Limitations

Our approach relies on fixed OT feature weighting by a parameter λ during training (see Sec. 4.2), limiting flexibility at inference and potentially constraining certain biological analyses. Moreover, radius-based niche definitions may also be sub-optimal for small or irregularly shaped microenvironments, where the radius captures excessive spatial context and does not allow fine-grained modeling of the functional region's evolution.

The model assumes that spatial slides can be aligned with respect to each other in time and requires normalization-based pre-processing. Future work will be directed towards rotational and translational invariant spatial constraints and the incorporation of cell-to-cell communication priors in the neighborhood definition. While the learned flow models cell population dynamics, it does not explicitly capture biological events such as division or death.

Finally, this study focuses on in-distribution testing and does not consider generalization to unseen slides or full anatomical regions excluded from the training process. To achieve prediction in unseen settings, we foresee the need for technical replicates of the same slide across time points, as the model cannot extrapolate spatial arrangements without prior exposure to the associated region. We leave these analyses to future work when spatio-temporal measurements across multiple replicates become increasingly available.

C Mixed-factorized Variational Flow Matching

C.1 Theoretical aspects of Variational Flow Matching

Variational Flow Matching (VFM) [25] relies on the observation that one can write the time-resolved marginal vector field u t ( x ) in FM as the expected conditional field u t ( x | x 1 ) under the posterior p t ( x 1 | x ) as:

Since u t ( x | x 1 ) has a closed form and u t ( x ) is all that we need to generate the probability path p t from noise to data, this opens the door to a new interpretation of the objective as a variational inference problem, where we approximate p t ( x 1 | x ) with a variational posterior q θ t ( x 1 | x ) . In other words, one can optimize the following objective:

where p 1 ( x ) is the data distribution and p t ( x | x 1 ) a straight probability path. When u t ( x | x 1 ) is linear in x 1 , this model formulation acquires convenient properties listed below.

Mean parameterization. The expected conditional field under the posterior only depends on the posterior mean:

suggesting that it is sufficient to parameterize the posterior mean to simulate data under the marginal flow. The posterior mean can be regressed against real samples x 1 during training.

Equivalence between posterior and approximate posterior formulation. From the previous point, it follows that the expectation of the conditional field is the same under the true and approximate posterior, as long as their first moments match.

Efficient simulation. Given a parameterized posterior mean µ θ t , simulating the generative field in Eq. (15) is efficient under the linearity condition. For example, in the standard FM setting with

straight paths [23], the marginal generative field becomes:

which can be easily simulated in the range t ∈ [0 , 1] .

C.2 Factorized posterior

Similar to Eijkelboom et al. [25], in our work, we use a fully factorized posterior, where individual dimensions can follow different families of distributions with finite moments (see Sec. 4.3). Notably, a factorized approximate posterior over x 1 is allowed as a choice for q θ t , since the only requirement to simulate u t ( x ) is for q θ t ( x 1 | x ) to match the expectation of p t ( x 1 | x ) over x 1 , irrespectively of higher moments or correlations between factors.

In this regard, it is useful to consider the following proposition.

Proposition 1. Let x 1 ∈ R D be a D -dimensional target data point, p t ( x 1 | x ) the posterior probability path conditioned on a noisy point x ∼ p t ( x ) , and u t ( x | x 1 ) the conditional velocity field. Assume that u t ( x | x 1 ) is linear in x 1 . Then, for any dimension d ∈ { 1 , . . . , D } , the following holds:

where x d refers to the d th scalar dimension of the vector x .

Proof. We begin by proving Eq. (19) using marginalization:

Next, we prove Eq. (20). Under the assumption that the conditional velocity field u t ( x | x 1 ) is linear in x 1 , we have:

Here, step (1) follows from the linearity assumption, which ensures that the conditional velocity at x d depends only on x d 1 .

In other words, the expected value under the posterior at an individual feature d does not depend on the other features and has an influence only on the d -th dimension of the conditional vector field. This flexibility allows each dimension's approximate posterior to be chosen from a potentially different distributional family, as long as the first moment exists and is correctly parameterized.

C.3 Marginal field derivation in source-conditioned VFM

When applying source conditioning to VFM, the marginal conditional vector field given a source x 0 is:

where the p t ( x | x 1 ) is a probability path interpolating observations x 1 with noise. Note that we omit x 0 from the probability path and conditional velocity as they are fully determined by x 1 under linear conditional probability paths. Furthermore, we can rewrite the marginal as an expectation:

where we used that p t ( x | x 1 ) = p t ( x | x 0 , x 1 ) .

C.4 Gaussian and Laplacian Hybrid VFM

We define a hybrid Variational Flow Matching (VFM) model using a fully factorized variational distribution over individual points in the target microenvironment M 1 . Following the mean-field assumption, the variational distribution factorizes over spatial and feature dimensions:

In line with Eijkelboom et al. [25], using FM with straight paths enables us to efficiently simulate the marginal generating field using the first moment of the posterior distribution. In other words, for a fully factorized posterior, we only need to parameterize a mean predictor. In our setting, the mean prediction is a neural network µ θ t as a time-condition function of a noisy microenvironment M and a source M 0 with outputs:

where ¯ f θ,k t and ¯ r θ,d t are the expected values for the k th coordinate and d th cell feature. Then, we choose a parameterization for the variational factors at time t ∈ [0 , 1] as follows:

Substituting into the negative log-likelihood yields:

This results in a loss consisting of an ℓ 1 error on spatial coordinates and a mean squared error on gene expression features, consistent with the hybrid variational design.

D Additional results

D.1 Additional comparisons with moscot on embryonic development

We propose a similar analysis as presented in Sec. 5.2.

Mapping anterior neural crest cells E9.5 → E10.5 → E11.5

Source

Predicted intermediate

Expected target

E9.5

E10.5

E11.5

Jaw and tooth

Neuralcrest

Neural crest

Cavity

push

push

NicheFlow mapping probability - E10.5

moscot mapping probability - E10.5

入=0.5

Q=0.5

0.6

0.8

Frequency

0.4

0.2

0.2

0.0

0.0

NicheFlow mapping probability-E11.5

moscot mapping probability - E11.5

入=0.5

Q=0.5

0.3

0.0

0.0

Brain

Cavity

Brain

Cavity

wand tooth

eectoderm

Jaw and tooth

Surface

Surfacee

In Fig. 4, we compare the ability of NicheFlow to predict an entire spatial structure trajectory by pushing an initial source cloud through all the developmental stages. To this end, the trajectory of an initial point cloud is first pushed to the next time point, and the model's prediction is used as a source for predicting the subsequent time point. An accurate niche trajectory reconstruction signifies that our model can be used to sequentially predict microenvironment evolution by treating its intermediate predictions as inputs, corroborating their accurate reflection of real point clouds.

In Fig. 4, we show that pushing anterior neural crest cells twice from E9.5 to E11.5 through the flow generates realistic target point clouds with a cell composition reflecting the expected cranial structure, mostly made of cavity cells, jaw and teeth (arising at E11.5 for the first time) and surface ectoderm. Doing the same with moscot oversamples regions outside of the cranial structure, thereby incorrectly mapping most of the neural crest density to brain cells.

In Fig. 5, we also show that NicheFlow is more accurate than moscot at transporting mass from defined organs like the liver across development. More specifically, while density leaks from the liver to the GI tract in the mapping produced by moscot, the prediction computed by NicheFlow more accurately retrieves the liver structure at the later time point. Together with previous evidence, our results underscore the importance of accounting for spatial correlations between cells during OT-based trajectory inference to buffer out the noise resulting from single-cell-based predictions.

Mapping liver E10.5-E11.5

Source

NicheFlow mapping probability

入=0.5

0.6

Frequency

0.4

0.2

0.0

G

Expected target

moscot mapping probability

Q=0.5

0.4

Frequency

Liver

0.1

0.0

tissue

D.2 α parameter sweep in moscot

The comparison with moscot assesses how well each model captures compositional changes in anatomical structures or migratory patterns, depending on the use case. Qualitatively, we found that the spatial component in the OT problem considered by moscot and regulated by the hyperparameter α is comparable to our spatial term α across the 0.1-0.9 range in the Fig. 3 analysis.

In Tab. 2, we report the proportion of source density mapped to the correct cell type across α values, to support our choice for the results in Sec. 5.2. For the spinal cord (Fig. 3, left), values above 0.75 yielded better qualitative and quantitative results. For neural crest cells (Fig. 3, right), α = 0 . 5 performs best.

More specifically, when mapping fixed structures over time, values below 0.75 caused excessive density dispersion outside the anatomical region. For migration, no value led to generally accurate transitions, though α = 0 . 5 mapped the highest density to the expected cell type.

D.3 Additional experiments on the axolotl brain development and aging datasets

We provide additional visualizations of the generated samples on the axolotl brain development dataset, presented in Fig. 6. As can be seen from the figure, NicheFlow correctly retrieves the spatial and anatomical characteristics of the brain, including hemisphere formation and cavity.

Target (St. 54)

SPFlow CFM

SPFlow GVFM

SPFlow GLVFM

NicheFlow CFM

NicheFlow GVFM

NicheFlow GLVFM

Source (St. 44)

Target (St. 57)

SPFlow CFM

SPFlow GVFM

SPFlow GLVFM

NicheFlow CFM

NicheFlow GVFM

NicheFlow GLVFM

Source (St. 54)

Target (Juv.)

SPFlow CFM

SPFlow GVFM

SPFlow GLVFM

NicheFlow CFM

NicheFlow GVFM

NicheFlow GLVFM

Source (St. 57)

Target (Adult)

SPFlow CFM

SPFlow GVFM

SPFlow GLVFM

NicheFlow CFM

NicheFlow GVFM

NicheFlow GLVFM

Source (Juv.)

Target (Meta)

SPFlow CFM

SPFlow GVFM

SPFlow GLVFM

NicheFlow CFM

NicheFlow GVFM

NicheFlow GLVFM

Source (Adult)

CMPN

Immature cckIN

MCG

VLMC

dNBL2

dpEX

npyIN

sfrpEGC

CP

Immature dpEX

·

MSN

·

CckIN

dNBL3

mpEX

ntng1IN

·

sstIN

Immature CMPN

Immature mpEX

Oligo

dEGC

dNBL4

mpIN

ribEGC

·

tINBL

Immature DMIN

Immature nptxEX

Unknown

dNBL5

scgnIN

wntEGC

Immature MSN

To further support our result, we propose a more in-depth qualitative analysis of the NicheFlow application on the axolotl brain development dataset. More specifically, in Fig. 7 we demonstrate that our model predicts the formation of crucial anatomical structures like the left and right lobes both spatially and compositionally. This vouches for flexibility in NichFlow's performance, which extends to non-trivial topology changes and simultaneously accounts for accurate cell state and coordinate generation in time. Similar results can be observed in Fig. 8, where we showcase the correct prediction of the formation of a left lobe cavity, predicting trajectories from an immature brain region.

Moreover, in Fig. 9 we predict the compositional and structural time evolution of the left dorsal pallium in the axolotl brain development. Following Wei et al. [42], we know that early time points populate the dorsal pallium of immature cell types like ependymoglial cells (EGC), neuroblasts (NBL), and immature neurons. In the left dorsal pallium, these disappear at the juvenile stages (the 3 rd ) and lead to differentiation into mature neurons (nptxEX) and later EGC (WntEGC, sfrpEGC). This fixed structural development was accurately predicted by NicheFlow when pushing left dorsal pallium cells forward across the trajectory.

Finally, similar to Fig. 2 and Fig. 6, in Fig. 10, we qualitatively show that NicheFlow with microenvironment sampling strategy and mixed-factorized VFM is the best approach for reconstructing mouse brain trajectories in time.

Source

Ground Truth

CP

VLMC

dEGC

0.1

0.0

Predicted frequencies

Generated niche

Generated region

St.54

Reference

CP

0.3

Generated

Immature cckIN

Immature dpEX

0.2

Immature mpEX

0.1

VLMC

0.0

dEGC

dNBL2

SStIN

wntEGC

St.57

CP

0.4

Reference

Immature dpEX

Frequency

Generated

Immature mpEX

0.2

VLMC

dEGC

dNBL2

0.0

dpEX

sfrpEGC

SStIN

wntEGC

Juv.

CMPN

sstIN

Reference

CP

tINBL

0.3

Generated

MSN

wntEGC

Unknown

VLMC

0.1

dpEX

mpEX

0.0

npyIN

scgnIN

Adult

CP

0.25

Reference

Generated

Unknown

VLMC

dpEX

mpEX

mpIN

0.05

0.00

scgnIN

SSTINHINB

sstIN

wntEGC

Meta.

CMPN

mpIN

Reference

CP

0.2

Generated

MCG

npyIN

MSN

scgnIN

Oligo

sfrpEGC

Unknown

sstIN

VLMC

tINBL

0.0

dpEX

wntEGC

mpEX

Source

80

dEGC

dNBL2

0.2

Predicted frequencies

Generated niche

Generated region

St.54

Reference

CP

Immature cckIN

Generated

Immature mpEX

VLMC

dEGC

0.0

dNBL2

dNBL4

sfrpEGC

wntEGC

St.57

CP

Reference

Immature mpEX

0.2

Generated

Immature nptxEX

VLMC

dEGC

dNBL2

dNBL4

sfrpEGC

wntEGC

Juv.

CP

sstIN

MSN

tINBL

Reference

Generated

Unknown·

wntEGC

VLMC

CCKIN

0.1

mpEX

0.0

npyIN

ribEGC

sfrpEGC

Adult

CP

SStIN

Reference

MSN

tINBL

0.25

Generated

Unknown

wntEGC

VLMC

CCKIN

mpEX

0.05

0.00

npyIN

ribEGC

sfrpEGC

CMPN

sfrpEGC

Meta.

CP

SstIN

MCG

tINBL

0.3

Reference

Generated

MSN

Oligo

wntEGC

0.2

Unknown

VLMC

dpEX

mpEX

D.0

npyIN

ScgnIN

Compositional evolution left dorsal pallium

Composition

Time point = 0

Time point = 1

Time point = 2

Time point = 3

Time point = 4

Time point = 5

0.8

1

Frequency

0.6

Immature

Mature

0.4

0.2

0.0

Celltype

Source

Generatedregion

Target (3.8m.)

SPFlow CFM

SPFlow GVFM

SPFlow GLVFM

NicheFlow CFM

NicheFlow GVFM

NicheFlow GLVFM

Source (3.4m.)

Target (4.3m.)

SPFlow CFM

SPFlow GVFM

SPFlow GLVFM

NicheFlow CFM

NicheFlow GVFM

NicheFlow GLVFM

Source (3.8m.)

Target (5.4m.)

SPFlow CFM

SPFlow GVFM

SPFlow GLVFM

NicheFlow CFM

NicheFlow GVFM

NicheFlow GLVFM

Source (4.3m.)

Target (6.6m.)

SPFlow CFM

SPFlow GVFM

SPFlow GLVFM

NicheFlow CFM

NicheFlow GVFM

NicheFlow GLVFM

Source (5.4m..)

Target (9.8m.)

SPFlow CFM

SPFlow GVFM

SPFlow GLVFM

NicheFlow CFM

NicheFlow GVFM

NicheFlow GLVFM

Source (6.6m.)

Ependymal

NSC

Neuron-Inhibitory

OPC

Tcell

B cell

Macrophage

·

Neuroblast

Neuron-MSN

VLMC

Endothelial

Microglia

Neuron-Excitatory

Neutrophil

Pericyte

VSMC

D.4 Conditional generation

Source (t1)

Target(t2)

Predictions (50x)

Sample Likelihood

0.8

0.6

Likelihood

0.4

0.2

0.0

Source (t1)

Target (t2)

Predictions (50x)

Sample Likelihood

0.8

0.6

Likelihood

0.4

0.2

0.0

Source (t1)

Target (t2)

Predictions (50x)

SampleLikelihood

0.8

0.6

Likelihood

0.4

0.2

0.0

Source (t1)

Target (t2)

Predictions (50x)

Sample Likelihood

0.8

0.6

Likelihood

0.4

0.2

0.0

Source (t1)

Target (t2)

Predictions (50x)

Sample Likelihood

0.8

0.6

Likelihood

0.4

0.2

0.0

Source (t1)

Target (t2)

Predictions (50x)

Sample Likelihood

0.8

0.6

Likelihood

0.4

0.2

0.0

To assess whether the model accurately conditions on input microenvironments, we visualize the spatial distribution of generated samples given a fixed source microenvironment. Figure 11 illustrates such cases on the embryonic development dataset, showing the input region at time t 1 , the corresponding target at time t 2 , and 50 independently generated samples from NicheFlow. Directly computing the likelihood of ground truth cell coordinates under our generative model is intractable, as it would require evaluating the density of an implicitly defined distribution over point clouds. Instead, we approximate the spatial likelihood using kernel density estimation (KDE) over Monte Carlo samples drawn from the model. Given a set of generated coordinates { ˆ c i } N i =1 , we estimate the likelihood at a ground truth location c as:

where σ is a fixed bandwidth parameter. The resulting KDE heatmap visualizes the spatial concentration of samples, allowing us to qualitatively assess whether the model produces consistent predictions conditioned on the source microenvironment.

D.5 OT Ablation Study

We begin by emphasizing that the optimal choice of the OT feature-weighting parameter λ in Eq. (12) depends on the downstream application. This parameter determines the relative importance assigned to gene expression versus spatial coordinates during OT. For developmental processes such as organogenesis or regeneration, higher values of λ prioritize transcriptional similarity, facilitating the reconstruction of continuous differentiation trajectories and capturing fate-driven transitions, which may span large spatial distances. In contrast, for applications that aim to monitor changes within a fixed spatial region, such as shifts in cell-type composition over time, a lower λ is more appropriate. In such cases, emphasizing spatial locality helps avoid spurious long-range transport assignments caused by molecular noise.

In Tab. 3, we sweep the value of λ across multiple models and training strategies on the same task as in Sec. 5.1, showing consistent results across settings as reported in Tab. 1.

Furthermore, Figure 12 visualizes the impact of varying λ on the OT plans described in Sec. 4.2, where intermediate values yield more coherent and biologically plausible couplings.

入=0.00

入 = 0.10

入 = 0.25

入 = 0.50

入 = 0.75

入 = 0.90

入 = 1.00

D.6 K -Means regions ablation study

To ensure diverse and spatially distributed sampling during training, we partition each tissue section into K spatial regions using K -Means clustering over the 2D cell coordinates (see Sec. 5.1.1). At each training step, microenvironments are sampled uniformly from within these regions, encouraging broad spatial coverage and preventing oversampling of densely populated areas. In this ablation study, we investigate how varying the number of spatial regions, K , affects model performance.

As shown in Figure 13, increasing K leads to increasingly fine-grained spatial partitions. While moderate values of K help improve spatial resolution, excessively high values (e.g., K = 128 or 256 ) result in overly small and fragmented regions. This can cause significant overlap between sampled microenvironments and reduce sampling diversity. Moreover, in sparsely populated tissue sections, high K values may yield regions with insufficient cells, degrading both representativeness and stability.

T = E9.5

T = E10.5

T = E11.5

= 4

K

8=

K = 16

K = 32

K = 64

K =128

K=256

Conversely, using too few regions (e.g., K = 4 ) results in broad spatial partitions that may encompass multiple heterogeneous tissue compartments. This undermines the locality assumptions of our model and increases intra-region variability, which can impair the model's ability to learn.

Given the trade-offs outlined above, we focus our evaluation on K ∈ { 8 , 16 , 32 , 64 } , which spans a range of granularities that preserve both spatial interpretability and sampling robustness. Tab. 4 sum-

marizes the quantitative results for this ablation across three datasets: mouse embryonic development (MED), axolotl brain development (ABD), and mouse brain aging (MBA). We observe that using K = 64 consistently yields strong performance, achieving the highest 1NN-F1 scores on both the MED and MBA datasets, while also performing competitively on ABD. These findings indicate that K = 64 offers an effective balance between spatial resolution and stability, and we adopt it as the default configuration in our main experiments.

D.7 Justifying the choice of NicheFlow over RPCFlow

Although RPCFlow achieves competitive performance on spatial metrics such as PSD and SPD, it lacks the essential capability of meaningful conditional generation. In RPCFlow, conditioning is performed using randomly sampled point clouds from the spatial regions without explicit microenvironment structure. As shown in Fig. 14, when conditioned on such random sources, RPCFlow tends to reconstruct the entire target tissue, rather than capturing local dynamics driven by the source input. This undermines its ability to model spatiotemporal evolution in a biologically grounded manner.

Source (t1)

Target(t2)

Predictions(50x)

Sample Likelihood

0.045

0.040

0.035

0.030

Likelihood

0.025

0.020

0.015

0.010

Source (t1)

Target (t2)

Predictions(50x)

Sample Likelihood

0.04

0.03

Likelihood

0.02

0.01

In contrast, NicheFlow is explicitly designed to push localized microenvironments through time. By conditioning on fixed-radius neighborhoods centered around specific spatial positions, the model learns how cellular contexts evolve, preserving both spatial coherence and transcriptional identity. Figure 15 visualizes this distinction: while NicheFlow consistently generates well-localized predictions aligned with the input microenvironment, RPCFlow fails to retain spatial specificity, often diffusing the prediction across broader regions.

From a biological perspective, predicting the fate of a local tissue region over time is far more relevant than mapping random point sets. Microenvironments encode structured cellular contexts, such as signaling interactions or niche-specific cell states, that are crucial for downstream analysis (e.g., lineage fate prediction, microenvironment-based intervention simulation). Because RPCFlow lacks this interpretability and fails to enable such downstream tasks, it cannot serve as a practical generative model in biological settings.

In summary, while RPCFlow may appear performant under some aggregate spatial metrics, only NicheFlow enables localized, conditionally consistent generation of evolving tissue regions. This makes it not only superior for evaluation but also for practical use in biological modeling.

D.8 Structure-aware evaluation

While we refer to our primary evaluation metrics as PSD and SPD, they correspond to the two asymmetric directions of the Chamfer Distance (CD) - a widely used metric in point cloud generation [46]. PSD measures how closely the predicted points adhere to the ground truth structure (fidelity), while SPD captures how well the prediction covers the full extent of the target (coverage). We

(a) NicheFlow - Sample 1

(b) NicheFlow - Sample 2

(c) RPCFlow - Sample 1

(d) RPCFlow - Sample 2

Source(t1)

Target (t2)

Predictions(50x)

Sample Likelihood

0.8

0.6

Likelihood

0.4

0.2

0.0

Source (t1)

Target(t2)

Predictions(50x)

Sample Likelihood

0.8

0.6

Likelihood

0.4

0.2

0.0

Source(t1)

Target (t2)

Predictions (50x)

Sample Likelihood

0.20

0.15

Likelihood

0.10

0.05

Source(t1)

Target(t2)

Predictions(50x)

Sample Likelihood

0.175

0.150

0.125

Likelihood

0.100

0.075

0.050

0.025

compute both metrics over collections of microenvironments tiled across the tissue, providing an indirect assessment of global morphological accuracy.

To directly assess the preservation of internal structure within generated microenvironments, we additionally evaluate Gromov-Wasserstein (GW) and Fused Gromov-Wasserstein (FGW) distances [47, 40]. These metrics compare the intra-point pairwise distances in predicted and ground truth point clouds, capturing latent structural relationships. FGW further incorporates transcriptomic similarity, offering a holistic measure of structural and feature-based alignment.

Given the computational cost of GW and FGW, we restrict their evaluation to microenvironment-level generations. For each source niche, we sample predictions from the trained model and compute GW/FGW against the corresponding real microenvironment, averaging results across multiple runs.

The source-target pairs are obtained by computing OT between all microenvironments at time t and t +1 , yielding a one-to-one coupling used for conditioning and evaluation.

As SPFlow generates individual points and RPCFlow produces randomly sampled, unstructured clouds, these models lack coherent internal structure and are not amenable to structural comparison via GWor FGW. We therefore report these metrics only for models that generate full microenvironments.

The results in Tab. 5 show that NicheFlow, particularly when trained with the GLVFM objective, consistently achieves lower GW and FGW distances compared to baseline models. These findings underscore the structural fidelity of our predictions and further validate the modeling advantages of microenvironment-aware generative training.

D.9 Wasserstein metrics

To complement the local structural assessments with GW and FGW, we additionally include global evaluation metrics that are widely adopted in spatial transcriptomics [48-50]. Specifically, we report the 1-Wasserstein ( W 1 ) and 2-Wasserstein ( W 2 ) distances between predicted and real cell distributions, computed separately over spatial coordinates and gene expression features. These metrics are calculated at the level of individual cell types and averaged across all timepoints and generated samples. This provides a more global perspective on how well the model reconstructs tissue-wide spatial and transcriptional distributions.

We evaluate all model variants and training objectives across the three datasets (MED, ABD, MBA) and report the results in Tabs. 6 to 8. NicheFlow with the GLVFM objective consistently achieves the lowest Wasserstein distances across nearly all dimensions in the embryonic development dataset (MED), outperforming both baseline models and alternative objectives.

In the ABD and MBA datasets, which exhibit more gradual spatial evolution or less distinct cell-type boundaries, we observe that RPCFlow sometimes matches or marginally outperforms NicheFlow in isolated metrics. However, as we discussed in App. D.7, RPCFlow is not a conditionally consistent model: it does not produce meaningful trajectories given a source microenvironment and lacks

biological interpretability. Therefore, even when RPCFlow achieves slightly lower Wasserstein distances in aggregate, these gains are not actionable in practice, as the model cannot be used to infer biologically meaningful transitions or trace the evolution of specific spatial regions.

Taken together, Wasserstein metrics offer a complementary, global validation of the structural and semantic accuracy of NicheFlow. They reinforce the quantitative evidence already provided by Chamfer-based metrics (PSD/SPD) and Gromov-based structural evaluations (GW/FGW), confirming that our model accurately captures the tissue-scale organization of both spatial and gene expression patterns.

E Algorithms

Here, we present the algorithmic procedures underlying NicheFlow. To streamline the exposition, we first define compact notation for representing noisy microenvironments and their interpolations.

We denote a single noisy microenvironment (simplifying Eq. (11)) as:

where k denotes the number of spatial points per microenvironment, c z i are 2D cell coordinates, and z i are associated feature vectors in R D . For a collection of N microenvironments, we define the corresponding set of noisy microenvironments

Given a noisy sample M z and a corresponding ground-truth microenvironment M 1 , we define their linear interpolation at time t ∈ [0 , 1] as: M t = (1 -t ) M z + t M 1 , where the interpolation is

performed element-wise across the matrix rows. For batched data, this generalizes to:

Algorithm 1 SAMPLEANDINTERPOLATE

E.1 OT Conditional Flow Matching

The OT Conditional Flow Matching (OT-CFM) algorithm consists of a training and a generation phase. During training, the model learns a time-dependent vector field u θ t conditioned on a source microenvironment M 0 , which maps a noisy microenvironment M z -sampled from Gaussian noise-to its corresponding target microenvironment M 1 . The supervision is provided via source-target pairs ( M 0 , M 1 ) obtained through EOT over pooled microenvironment representations. At generation time, the learned vector field is integrated starting from M z , conditioned on a given source M 0 , to generate the predicted target microenvironment M 1 .

We optimize the following loss:

The full pseudocode for both phases is provided in Algs. 2 and 3.

Algorithm 2 OT CFM - Training

Input: Number of samples N , feature dimension D , microenvironment size k , OT plan π ∗ ϵ,λ , conditional velocity field u θ t Output: Trained parameters θ of u θ t

1: ( M 0 , M 1 , M t , M z ) ← SAMPLEANDINTERPOLATE ( N,D,k,π ∗ ϵ,λ ) ▷ Microenvironments (Alg. 1) 2: θ ←∇ θ L ( M 0 , M 1 , M t , M z ; θ ) ▷ Compute loss (Eq. (36)) & update parameters θ

Algorithm 3 OT CFM - Generation

E.2 OT Gaussian Variational Flow Matching

The OT Gaussian Variational Flow Matching (OT-GVFM) algorithm adopts a variational perspective on Flow Matching. Instead of directly learning a time-dependent conditional velocity field, the model learns a factorized variational posterior q θ t ( M 1 | M t , M 0 ) over target microenvironments M 1 , conditioned on an interpolated microenvironment M t and a source microenvironment M 0 . The predicted velocity field is then computed as the difference between the posterior mean µ θ t ( M t , M 0 ) and the current state M 1 .

The training objective minimizes the discrepancy between ground-truth targets and the predicted posterior means ( ¯ f θ t , ¯ r θ t ) :

At generation time, trajectories are generated by integrating the vector field µ θ t ( M t , M 0 ) -M t , starting from a noise-sampled microenvironment M z and conditioned on a given source M 0 . To ensure stability near t = 1 , the vector field is scaled by a time-dependent denominator, as in prior VFM formulations.

The pseudocode for both phases is provided in Algs. 4 and 5.

Algorithm 4 OT-GVFM - Training

Input: Number of samples N , feature dimension D , microenvironment size k , OT plan π ∗ ϵ,λ , source-conditioned posterior mean predictor µ θ t

Output: Trained parameters θ of µ θ t

1: ( M 0 , M 1 , M t , M z ) ← SAMPLEANDINTERPOLATE ( N,D,k,π ∗ ϵ,λ ) ▷ Microenvironments (Alg. 1) 2: θ ←∇ θ L ( M 0 , M 1 , M t ; µ θ t ) ▷ Compute loss (Eq. (37)) & update parameters θ

Algorithm 5 OT-GVFM - Generation

Input: Source microenvironment M 0 , learned source-conditioned posterior mean predictor µ θ t Output: Generated microenvironment M 1

1: M z ∼ N ( 0 , I ) 1 ×

D +2 k ▷ Sample noisy sample 0 1 - t + ϵ t ▷ Solve ODE

E.3 NicheFlow: OT Gaussian-Laplace Variational Flow Matching

NicheFlow extends OT-GVFM by modifying the variational posterior: it assumes a Gaussian distribution for gene expression features and a Laplace distribution for spatial coordinates. This change leads to a hybrid loss that combines an L 2 loss on gene expression and an L 1 loss on spatial locations:

At generation time, the model integrates the velocity field defined by the difference between the predicted mean µ θ t ( M t , M 0 ) and the current state M 1 , starting from noise and conditioned on the source M 0 , identical to the OT-GVFM procedure.

F Experimental setup

F.1 Dataset description

We evaluate our model on three publicly available, time-resolved spatial transcriptomics datasets spanning development and aging processes. Each dataset provides single-cell resolution profiles with matched spatial coordinates and curated cell type annotations. A detailed summary of the organism, tissue, number of time points, cell types, total number of cells, and acquisition technology for each dataset is provided in Table 9.

F.2 Dataset and microenvironments preprocessing

Dataset preprocessing. All datasets used in our study underwent a preprocessing procedure appropriate for spatial transcriptomic analysis, involving total count normalization, logarithmic

transformation, and principal component analysis (PCA). For the mouse embryonic development [20] and axolotl brain development [42] datasets, total count normalization, and logarithmic transformation had already been applied; we, therefore, performed PCA ourselves on the transformed data. For the mouse brain aging dataset [43], we applied all three steps: we first normalized raw count matrices so that each cell had the same total expression, followed by a natural logarithm transformation of the form log( X +1) to stabilize variance and mitigate the influence of large values. We then computed PCA on the log-transformed data.

To reduce computational overhead due to high cell counts in the aging dataset, we subsampled the data by a factor of 0.2.

Finally, we standardized PCA components across all cells to ensure consistent scaling across time points. Spatial coordinates were standardized independently for each time point by subtracting the per-time-point mean and dividing by the standard deviation. This standardization preserves the relative spatial configuration while accounting for scale and position differences over developmental time.

Microenvironments preprocessing. To facilitate efficient training and enable consistent microenvironment construction, we precompute all fixed-radius neighborhoods for each dataset using a radius r chosen based on spatial resolution. To reduce variability in the number of neighbors and improve batching efficiency, we fix the number of nodes per microenvironment to the most frequent neighbor count observed within each slide. This standardization ensures structural comparability across microenvironments while significantly reducing computational overhead during training, as costly radius or nearest-neighbor queries are avoided at runtime.

F.3 Quantitative evaluation metrics

We assess spatial fidelity using two complementary asymmetric distance measures. The point-toshape distance (PSD) captures how much predicted cell positions diverge from the actual tissue layout, computed as the average squared distance from each simulated point to its nearest neighbor in the ground truth. Conversely, the shape-to-point distance (SPD) quantifies how comprehensively the predicted distribution spans the target structure by averaging the squared distance from each ground truth point to its closest generated counterpart.

Let G t denote the set of generated coordinates and R t the set of ground truth coordinates at time t . Define NN t ref ( c i ) as the nearest neighbor of c i ∈ G t in R t , and NN t gen ( c i ) as the nearest neighbor of c j ∈ R t in G t . Then, the two metrics are given by:

where G := ∪ t ∈T G t and R := ∪ t ∈T R t .

F.4 Cell-type classification for evaluation

To evaluate cell-type fidelity of generated microenvironments, we train a supervised classifier to assign cell type labels based on gene expression profiles. We apply the same preprocessing steps

used for training our generative models: total count normalization, log-transformation, and PCA (see App. F.2). The resulting low-dimensional embeddings are used as input features for a simple multilayer perceptron (MLP), trained to predict discrete cell type labels.

The classifier consists of a two-layer feedforward network with ReLU activations and a final linear projection to the number of cell types. It is trained using cross-entropy loss and optimized with the AdamW optimizer. We report performance using the weighted F1-score.

We obtain strong classification results across all datasets. On the mouse embryonic development dataset, the classifier achieves a weighted F1-score of 0.85; on axolotl brain development, 0.80; and on the aging dataset, 0.97. These results correlate with the number of input genes and the variance retained in the PCA-reduced space. The aging dataset contains only 300 genes, and 50 principal components explain sufficient variance to accurately distinguish most cell types. In contrast, the embryonic development dataset contains approximately 2,000 genes, and the axolotl brain development dataset includes over 12,700 genes, making classification more challenging due to higher gene expression variability.

F.5 Discretized microenvironments

To ensure consistent and reproducible evaluation across methods and datasets, we construct a fixed set of evaluation microenvironments by discretizing the spatial domain of each tissue section. For each time point, we define a regular 2D grid over the tissue and select the closest cell to each grid point as the centroid of a microenvironment. Each centroid is then used to construct a fixed-radius neighborhood, following the microenvironment definition in Section 4.1. This procedure ensures full spatial coverage by verifying that every cell belongs to at least one microenvironment.

E9.5

E10.5

E11.5

Figure 16 illustrates this discretization for the mouse embryonic development dataset. Each blue dot corresponds to a selected centroid. In cases where the number of grid-based centroids falls below a target threshold, additional centroids are randomly sampled to match a fixed total count per slide. This augmentation ensures that all slides contribute equally to the evaluation and prevents bias from sparse regions.

We apply the same discretization procedure to all three datasets used in our experiments: mouse embryogenesis, axolotl brain development, and mouse brain aging. By standardizing the evaluation regions spatially and deterministically, we eliminate the need for stochastic region sampling during evaluation, which would otherwise lead to nondeterministic and irreproducible results.

F.6 Microenvironment Transformer

To model the spatiotemporal evolution of local cellular neighborhoods, we design a permutationequivariant transformer-based architecture tailored to structured point cloud data. Our Microenvironment Transformer processes local microenvironments-sets of cells with gene expression features and spatial coordinates-and predicts time-dependent outputs such as velocity fields or future states.

The model operates on a source M 0 and noisy M z microenvironments with per-cell features x i ∈ R D and 2D coordinates c i ∈ R 2 . The architecture consists of the following components:

Input Embeddings:

(a) Feature Embedding: A linear transformation is applied to the input features x i of each cell.

(b) Coordinate Embedding: Spatial coordinates c i are linearly projected and concatenated to the feature embedding.

(c) Time Embedding: For the noisy microenvironment only, time t ∈ [0 , 1] is encoded using sinusoidal functions cos( ωt ) and sin( ωt ) , followed by a linear projection and concatenation with the input embedding.

2. Transformer Encoder (Source Microenvironment):

(a) Self-Attention: A stack of transformer encoder blocks with multi-head self-attention processes the embedded source microenvironment.

(b) No Time Embedding: Time information is not provided to the encoder, as it encodes the source M 0 .

(c) Residual Feedforward: Each block contains a feedforward subnetwork with LeakyReLU activation and a residual connection.

(d) Layer Normalization: Applied after both the attention and feedforward layers.

(e) Masking: Binary masks are used to ignore padding in variable-length microenvironments.

3. Transformer Decoder (Noisy Microenvironment):

(a) Time Embedding: Temporal context is injected into the decoder by embedding the time t and concatenating it to the target point embedding.

(b) Cross-Attention: Decoder layers apply cross-attention between the noisy microenvironment and the encoded source microenvironment.

(c) Self-Attention and Feedforward: Each decoder block includes standard self-attention and residual feedforward layers.

(d) Layer Normalization and Masking: As with the encoder, normalization, and masking are applied throughout.

Final Output Projection:

(a) Prediction Head: A linear layer maps the decoder outputs to the desired dimensionality.

This architecture allows for flexible and expressive modeling of temporal dynamics in cellular point clouds. By encoding only the source and decoding the temporally conditioned target, the model supports variational and flow-based training objectives with explicit temporal conditioning.

F.7 Hyperparameters and Computational Costs

Model hyperparameters. For all experiments, we use the same configuration for the Microenvironment Transformer architecture. The full set of hyperparameters is as follows:

Input feature dimension: 50 PCA-based gene expression features concatenated with a one-hot encoding of the time-point, resulting in a total dimensionality of 50 + |T | , where |T | is the number of slides (timepoints) in the dataset.

Input coordinate dimension: 2

Embedding dimension: 128

MLP hidden dimension: 256

Number of attention heads: 4

Number of encoder layers: 2

Number of decoder layers: 2

Dropout rate: 0.1

Output dimension: 52 (gene expressions features + coordinates)

OT and mini-batching. To ensure spatial diversity and computational tractability, we uniformly sample 256 source-target microenvironment pairs from the K spatial clusters obtained via K -Means. We then compute the entropic OT plan between these sampled pairs and resample 64 source-target pairs from this plan to define a single training instance. During training, we process 16 such instances per batch.

Optimization. All models are trained using the AdamW optimizer with a learning rate of 2 · 10 -4 and a weight decay of 1 · 10 -5 . We train each model until convergence.

Computational cost. All models were trained on a single NVIDIA GeForce GTX 1080 Ti GPU with 11GB of memory. Depending on the dataset and training objective (e.g., CFM or VFM), training takes approximately 12-16 hours per model.

F.8 Comparison with moscot

Here, we describe how we conduct the comparisons with moscot, as shown in Fig. 3, Fig. 4, and Fig. 5. For both NicheFlow and moscot, we select a source microenvironment that we want to track over time. In the case of NicheFlow, this corresponds to an aggregate of point clouds. For moscot, it refers to a group of single cells spatially located within the region of interest. The same set of cells is used for both methods.

NicheFlow. To generate contour plots over the spatial slide, we push forward the selected region and assign the generated points to their nearest real neighbors based on spatial coordinates. We then compute a probability value for each real position by normalizing the number of assigned generated points. In other words, the more generated points that are close to a given real point, the higher the probability assigned to that location in the contour plot. Cell type proportions are computed as the aggregated frequencies of the generated cell types across the slide. Each plot considers 10 independent generation runs from the same niche.

moscot. This baseline is not a generative model, but rather a standard discrete OT framework using a Fused Gromov-Wasserstein cost. As such, it does not generate new features or coordinates. Instead, it outputs a transition matrix that assigns matching probabilities between each source slide cell and each target slide cell. We use these transition probabilities to compute contour plots over the target slide and to aggregate cell type probabilities for the bar plots. For the latter, moscot provides a custom method called cell_transition() .

For the Appendix figures Fig. 4,Fig. 7, Fig. 8, and Fig. 9, we propagate the initial source region across multiple time steps. In NicheFlow, this is achieved by using the simulated point cloud from step t to predict the next state at t +1 , and then feeding this output as the input for the following trajectory step from t +1 to t +2 , and so on. Ground truth points are not used as intermediate sources during this process. For moscot, pushing the source across time points can be automatically done by setting non-subsequent time points in the cell_transition() function.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The abstract and introduction state the core contributions. These claims are substantiated through targeted experiments and ablation studies.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We discuss the limitations in App. B.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory assumptions and proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: We introduce the Mixed-Factorized VFM formulation alongside NicheFlow in Sec. 4, and provide complete theoretical derivations in App. C.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental result reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: The experimental setup is described in Sec. 5.1.1, with additional details provided in Apps. E and F. We outline the dataset preprocessing, evaluation discretization procedure, backbone architecture, and all hyperparameter and optimizer settings. The appendices also include training and inference algorithms for OT-CFM and OT-VFM under both GVFM and GLVFM objectives.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: A link to the code is included in our project page, which is reported in the footnotes of page 1.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental setting/details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We report the values used in our experiments in Sec. 5.1.1 and provide additional justification through ablations in Apps. D.5 and D.6 and App. F.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment statistical significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We report the mean and standard deviation across five independent evaluation runs for all experiments and evaluation metrics.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments compute resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We provide details on the hardware that has been used in this research in App. F.7.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code of ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The research presented in this work conforms with the NeurIPS Code of Ethics. Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: We discuss the broader impacts of our work in App. A.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The research discussed in this paper does not require safeguards to be put in place.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: Where relevant, we acknowledge and cite the original authors of prior work, including associated codebases and datasets, and specify the licenses under which these resources were released.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: We use only publicly available datasets, which are cited appropriately. The complete codebase, including documentation and data preprocessing scripts, is public.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and research with human subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: This work did not conduct research on human subjects or crowdsourcing experiments.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional review board (IRB) approvals or equivalent for research with human subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: This work did not conduct experiments where human subjects were involved and therefore does not require IRB approvals.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

16. Declaration of LLM usage

Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.

Answer: [NA]

Justification: LLMs were not used for the development of our core method.

Guidelines:

The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.

Please refer to our LLM policy ( https://neurips.cc/Conferences/2025/LLM ) for what should or should not be described.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this manuscript the authors propose a flow matching framework to model the temporal evolution of spatially resolved microenvironments from time-resolved spatial transcriptomics data. \n\nThey develop a method denoted as NicheFlow and apply it to several spatiotemporal datasets involving embryonic and brain development.\n\nSTRENGTHS AND WEAKNESSES:\nThe main strength of the manuscript is the development of NicheFlow, a method to model the temporal evolution not of individual cells,\nbut of niches.  \n\nThat said, the manuscript has several major weaknesses, as I outline below. \n\nClarity of Exposition and Fit to NeurIPS: The manuscript is not very accessible to a non-expert, and not sufficiently self contained. \nTo me it seems that the target audience of this manuscript are researchers in computational biology. \nAs such, in my opinion, the manuscript would be more suitable for publication in such venues rather than in a ML conference. \nAlso, the organization of the paper could be improved. For example, section 3 with a background on flow matching is a bit out of context\ngiven that the precise formulation of what is measured and what are the tasks of interest appear only later. \n\nQuality and Significance: While I am not an expert in this field, to the best of my understanding, the reason for collecting\nspatial-temporal biological datasets is to allow to gain novel insights into temporal biological processes, such as development\nand disease progression. \nIn this manuscript, the authors present their method of modeling the evolution of Niches.  \nHowever, there is no sufficient biological background, no concrete biological questions considered. \nThe manuscript would benefit if it included an in-depth example where new insight is gained using the proposed model \non a concrete biological problem of major importance. \n\nOriginality and Novelty: The model proposed is to the best of my knowledge novel.  \n\nThe evaluation section would benefit from a specific example where the authors would illustrate how NicheFlow allows\ngaining novel insight into some dynamical biological process. \n\nAdditional comments: \nP2 L56 prediction of spatial coordinates and gene expression... It may be good to provide some concrete biological examples\nwhere this would be useful / important. \n\nP2 L66 NicheFlow enables principled learning of dynamics... Again, it may be good present a concrete biological example\nwith specific scientific findings enabled by NicheFlow.\n\nQUESTIONS:\nP2 Lines 39-44 Are the methods mentioned here applied to spatio-temporal datasets, or just to spatial datasets ? \n\nP2 Line 61: independently optimizes.... why is it important to perform this optimization independently ? \n\nP2 Line 75: in modeling cellular data [25-28]. How are references [25-27] relevant to this task ? \n\nP3 L90 identify likely descendants.. sentence is unclear. \n\nP4 L150 NicheFlow predicts the corresponding tissue structure... Why is this interesting from a biological / scientific perspective?  \n\nP5 top: it is a bit unclear what dynamics can be learned by looking at only two consecutive time points. It is also unclear if\nthe two time points one measures exactly the same cells (did the cells not move / where they tracked / no cell divisions occurred) ? \n\nP5 Eq. 12 - it seems a bit simplistic to simply average the gene expression of neighborhood cells? What are typical processes\nin the dynamical evolution of niches and are these easily captured by such averaging ? \n\nP8 Which reference is SPFlow ? \n\nP12 are refs 34 and 38 the same?\n\nLIMITATIONS:\nThe authors do not discuss the limitations of their approach in the main text, but they do mention several of them in the supplementary, appendix B.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "CLARITY": 1, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces NichFLOW, a flow-based generative model designed for modeling time-series spatial transcriptomics data. The authors propose incorporating microenvironments, variational flow matching, and conditional generative modeling to learn correspondences between spatially adjacent cell slices in terms of spatial position and gene expression. The method is validated on several datasets. Overall, the idea is interesting, as while static optimal transport (OT) methods have been widely used in this context, flow matching-based methods are less considered.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\n1. The paper is generally well-written, includes extensive experiments, and presents a comprehensive study.  \n\n2. The work builds on some concepts from prior studies, such as GENOT and VFM, and introduces the concept of microenvironments to handle time-series spatial transcriptomics data. The integration of these ideas in this context is new.\n\n3.   Reconstructing dynamics from time-series spatial transcriptomics data is an important problem and has gained increasing attention.\n\nWeakness:\n\n1. The paper employs flow matching as a conditional generative model, which appears limited in capturing the continuous dynamics of cells across time points; however, this can be achieved by many dynamical-based optimal transport (OT) methods. Specifically, the method seems less suited for interpolation tasks, such as predicting data at time point 1 given data at times 0 and 2, unless relying on simple linear interpolation. A major weakness to me is that NichFLOW appears to achieve functionality similar to static OT methods (such as WaddingOT, Moscot, etc). The paper does not clearly show whether NichFLOW can accomplish tasks that static OT methods cannot, limiting its broader applications in modeling complex spatiotemporal cellular dynamics.\n\n\n2. The paper lacks a comprehensive discussion comparing the proposed method with other works working on the spatio-temporal ST data. Including such a comparison, e.g., with [1,2], could better clarify the similarities and differences between NichFLOW and other methods. \n\n\n3. While the evaluation includes metrics for cell type prediction, it lacks detailed gene-level assessments, such as Wasserstein-1 (W1) and Wasserstein-2 (W2) distances between generated and true gene expression vectors, spatial positions, or their combined distributions. These metrics could better elucidate the method’s generative performance. Additionally, the paper does not include neural ODE-based methods (e.g., TrajectoryNet [3], DeepRUOT[4]) as baselines for comparison. \n\n\n4. The method does not account for cell growth or death, as acknowledged in the limitations. Given that some tested datasets involve developmental processes, neglecting cell proliferation effects may lead to inaccurate dynamics. The paper could benefit from exploring unbalanced OT (UOT) frameworks, such as those used in WaddingOT/MOSCOT, which could be a straightforward way to address this issue.\n\nReference:\n\n[1] DeST-OT: Alignment of Spatiotemporal Transcriptomics Data, Peter Halmos, Xinhao Liu, Julian Gold, Feng Chen, Li Ding, Benjamin J. Raphael\n\n[2] Inferring cell trajectories of spatial transcriptomics via optimal transport analysis, Xunan Shen, Lulu Zuo, Zhongfei Ye, et. al.\n\n[3] TrajectoryNet: A Dynamic Optimal Transport Network for Modeling Cellular Dynamics, Alexander Tong, Jessie Huang, Guy Wolf, David van Dijk, Smita Krishnaswamy\n\n[4] Learning stochastic dynamics from snapshots through regularized unbalanced optimal transport, Zhenyi Zhang, Tiejun Li, Peijie Zhou\n\nQUESTIONS:\n1.  Can the authors clarify whether NichFLOW can perform interpolation (e.g., interpolate time point 1 from time points 0 and 2) or similar tasks? If so, how were these capabilities evaluated, and could additional experiments demonstrate these advantages?  \n \n\n2.  Could the authors provide insights into the key differences or advantages of NichFLOW over other methods to better position its contribution?  \n\n3. The evaluation metrics focus on cell type performance but lack gene-level assessments, such as W1 or W2 distances for gene expression vectors, spatial positions, or combined metrics.  Would adding such metrics provide a clearer understanding of the method’s generative capabilities? Additionally, why were neural ODE-based methods like TrajectoryNet, MIOFlow, DeepRUOT not included as baselines? \n\n4. The method does not account for cell growth, death, which may impact the accuracy of dynamics, especially in developmental datasets. Why was an unbalanced OT (UOT) or unbalanced Fused Gromov-Wasserstein OT framework, such as those in MOSCOT, not considered to model these effects, and could its inclusion improve the method’s applicability to such datasets?\n\nOverall, I view this work positively, and if the authors can address the weaknesses outlined, incorporate additional experiments, baselines (e.g., TrajectoryNet, MIOFlow, DeepRUOT), gene-level metrics (e.g., W1, W2 distances), and comprehensive discussions, I will consider raising my score.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-25"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces NicheFlow, a generative model designed to capture the spatiotemporal evolution of cellular microenvironments in spatial transcriptomics (ST) data. Unlike existing cell-centric approachStes, NicheFlow models local neighborhoods (“niches”) as point clouds of cells, each comprising gene expression and spatial coordinates, and learns their dynamics over time using a combination of: Optimal Transport (define soft couplings between niches across timepoints), Variational Flow Matching (simulate niche trajectories via learned velocity fields), factorized probabilistic model over expression and spatial modalities, and a Transformer-based point cloud architecture. Empirical evaluations on three datasets (mouse embryonic development, axolotl brain development, and mouse brain aging) demonstrate that NicheFlow outperforms both traditional FM-based models and the widely used MOSCOT framework in terms of spatial and semantic accuracy. Qualitative results highlight NicheFlow’s ability to capture known biological differentiation trajectories.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n- Shifts from modeling single-cell trajectories to niche-level dynamics, which is more aligned with tissue-scale processes like differentiation, migration, and patterning.\n- Uses point cloud representations of microenvironments, which are permutation-invariant, variable in size, and expressive of spatial structure.\n- Integrates EOT and source-conditioned VFM, extending prior FM formulations to handle unpaired, noisy, and variable-sized spatiotemporal point clouds.\n\nWeaknesses:\n- NicheFlow is trained on local microenvironments only — there’s no mechanism to enforce global tissue-wide coordination, which is often important in developmental processes governed by morphogen gradients or topological constraints.\n- The trade-off parameter $\\lambda$ (controlling spatial vs gene similarity in the OT cost) is fixed at 0.1. There’s no exploration of $\\lambda$ sensitivity or adaptation, which is important for different biological regimes (e.g., morphogenesis vs migration).\n- PSD and SPD capture spatial proximity but not structural or topological similarity (e.g., niche boundaries, internal configuration, or inter-cell distances). All metrics are pointwise — no graph or structure-aware evaluation is used.\n- While using the microenvironments concept is good, the key part of utilizing this concept is missing -- the cells actually interact with their neighbors. Without considering cell-cell interactions, grouping cells solely based on their spatial location has limited biological significance.\n\nQUESTIONS:\n- Can $\\lambda$ be learned or adapted across time or tissue regions? This seems essential given its major role in balancing spatial vs transcriptomic cues.\n- How sensitive is the model to the choice of neighborhood radius $r$? Could learned neighborhoods or graph-based neighborhood definitions be more effective? \n- Could this model be extended to continuous-time predictions or interpolation across multiple timepoints?\n\nLIMITATIONS:\nWhile NicheFlow models gene expression trajectories over time, it does so directly in the observed (or PCA-projected) expression space. However, gene expression data is known to lie on a lower-dimensional manifold shaped by biological processes like differentiation or zonation. Many successful models in scRNA-seq (e.g., scVI, DPT, diffusion models) explicitly exploit this structure for denoising, interpolation, and trajectory learning. By ignoring the data manifold, NicheFlow may be susceptible to noise and struggle to distinguish technical variability from biological transitions. Incorporating latent space modeling or manifold-aware trajectories could improve robustness and interpretability.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 5, "CLARITY": 4, "IMPACT": 4, "ORIGINALITY": 4, "DATE": "2025-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a framework, NicheFlow, for predicting both the physical migration of cells within the microenvironment and the associated changes in gene expression. The framework combines variational flow matching with optimal transport-based conditioning, and is evaluated on three time-resolved spatial transcriptomics datasets to assess its representational capacity and fitting accuracy.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\n* The problem of jointly estimating niche migration and gene expression changes is well-motivated and biologically relevant.\n* The proposed model integrates recent advances, including the variational formulation of flow matching and OT-based alignment of unpaired source and target spaces, into a cutting-edge framework.\n* The paper conducts component-wise validation of the model, including ablation studies, which strengthens the credibility of the empirical findings.\n* The manuscript is well-structured, and the main idea is clearly explained.\n\n**Weaknesses**\n1. Ambiguous notation: The notation is somewhat unclear and potentially confusing. The paper refers to $x_0$ and $x_1$ as source and target variables, respectively, which implies that they live in different, unaligned domains. However, this creates confusion because $x_{t=0}$ and $x_0$ may refer to different variables.\n2. Lack of hold-out evaluation: The most critical issue is the lack of explanation regarding how the hold-out test set was constructed. Throughout the paper, it appears that the model is evaluated directly on the source and target slides used for training, thereby assessing fitting performance rather than generalization. Since the proposed model is generative and intended to handle unseen states, it is essential to evaluate its performance on held-out samples or slides. Without such validation, the model’s generalizability remains unclear.\n3. Model design: It is not evident that the current formulation of NicheFlow is effective in capturing heterogeneous niche dynamics. As described around eq. (13), the model appears to share the same flow functions $f_t$ and residual terms $r_t$ across all samples in $\\mathcal{M}^1$, which could lead to excessive averaging and loss of niche-specific information. Although Table 1 shows some improvements of NicheFlow over RPCFlow, the performance likely depends on proper tuning $\\lambda$, whose selection process is not discussed. This omission makes it difficult to judge whether the observed improvements are robust or incidental.\n4.  Limited evaluation validity: The evaluation in Fig. 3 appears to measure how well the model fits the training data, not its predictive performance on unseen inputs. This makes it difficult to assess the true modeling capability. Moreover, the comparison of hyperparameters such as $\\alpha$ (in MOSCOT) and $1 - \\lambda$ may not be fair due to differences in underlying loss formulations. A comparison under optimal settings for each method is needed.\n\nQUESTIONS:\n* Appendix D.3 explores how changing $\\lambda$ affects prediction and metrics, but the paper does not explain how to select $\\lambda$ on data not seen during training. Without a principled way to choose $\\lambda$, the method may have limited practical utility.\n* In Fig. 10, NicheFlow appears to better fit niche migration than RPCFlow, but could this difference be due to hyperparameter choices or distribution assumptions (e.g., Gaussian vs. Laplace), rather than model architecture itself?\n* In Fig. 3, NicheFlow seems to outperform Moscot, but this might simply be due to better parameter tuning. While both $\\alpha$ (in MOSCOT) and $1-\\lambda$ (in NicheFlow) control the spatial weighting, their formulations differ (e.g., GW loss vs. distance), so direct comparison at the same values is not necessarily fair.\n\nMinor Comments\n* What do the names RPCFlow and SPFlow mean? Please provide explicit definitions for these models.\n* Table 1 is not referenced in the main text.\n\nI would be open to raising my score if the authors are able to address the critical concerns raised above.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-06"}