Modeling Microenvironment Trajectories on Spatial Transcriptomics with NicheFlow

We would like to thank reviewer yBc1 for reviewing our paper. We value the criticism and feedback as an opportunity to improve the clarity of our submission, and hope that our answers will contribute to a more positive evaluation of our work.

Fit to NeurIPS

We respectfully disagree with the assessment that our manuscript does not fit NeurIPS due to its application domain. Applications of ML to computational biology have become increasingly representative at major venues, including papers on spatial transcriptomics accepted at NeurIPS, ICML and ICLR (some examples: [1-3]). Moreover, NeurIPS explicitly lists life sciences as one of the key areas under its Call for Papers, reinforcing its commitment to applied domains.

Our work does not simply apply an existing algorithm but introduces a new methodological perspective on generative Flow Matching (FM), motivated by an applied setting. Specifically, we extend Variational Flow Matching (VFM) to approximate a source-conditioned Entropic Optimal Transport (EOT) model and introduce a factorized posterior to model structured data such as attributed point clouds.

We describe these novel contributions in the introduction (L51-68).

Background organization

Our work builds on different aspects of FM that have not previously been explored together: EOT couplings across incomparable spaces, and a variational perspective on the generated space.

We believe it is important to formally introduce these concepts in a general mathematical setting. This structure allows us to tie the ideas together in a principled way later in the paper, and moves from general background to model desiderata and finally to our point-cloud-based formulation.

We introduce the mathematical formulation of spatial niches after the general background, as presenting it earlier would require framing the entire background in point cloud terms, which is specific to NicheFlow and better suited to the contribution sections.

We propose adding a new section before subsection 3.1 to introduce how biological trajectories are typically modeled using OT.

Biological introduction and scope

In the introduction, we highlight the importance of considering cells in their environments, describe static and dynamic spatial transcriptomics, identify key gaps in existing models, and state our research question explicitly (L45–46). This establishes both the biological context and our main goal. We frame the problem as a general trajectory modeling task over microenvironments, aiming to develop an algorithm applicable across diverse biological systems with spatiotemporal axes.

This broad scope motivates our submission to NeurIPS and is reflected by our experiments with generation and biological trajectory inference on development and ageing using widely available data from benchmarks [4].

We will clarify the relevance of our task in the final version with examples from tumour progression and brain development.

Biological analysis and findings

Sec. 5.2 and the Appendix present predicted dynamics and developmental hypotheses enabled by our model:

• Fig. 3, right panel: we predict that immature cell groups emigrate in the upper cranial region and that they differentiate into head mesenchymal cells, simultaneously. This demonstrates our model's combined spatial and transcriptomic insight, unlike standard single-cell RNA-seq methods.
• App. Fig. 11: We draw similar knowledge studying how facial neural crest cells evolve into jaw and teeth structures, this time predicting across multiple timepoints.
• App. Fig. 13-14: We show dynamic formation of anatomical regions during the axolotl brain development, capturing spatial and compositional changes.
• App. Fig. 15: We demonstrate our model predicts a decrease in stemness in the left dorsal pallium.

L56 - gene expression and coordinates

We refer to our previous answer for the biological use cases presented in our paper. Predicting both the location and gene expression of cells is important, as it provides an understanding of both tissue composition and organization. Different from standard scRNA-seq that only specifies how gene expression evolves across cell types, here we study how earlier cell states evolve and where they are positioned anatomically, adding a new layer of information to standard trajectory inference (e.g., prediction of both fate and migration of immature cells in Fig. 3, right). Predicting spatial and state transitions is also clinically relevant, where future model applications could, for example, study metastatic progression in cancer over time.

P2 L39-44: referenced methods

The models cited in these lines are applied to spatial transcriptomics and infer trajectories at the single-cell level. In this paragraph, we highlight how most current approaches do not consider the trajectories of entire microenvironments of cells, even though the dynamics of tissue restructuring during temporal processes are strongly determined by the coordinated evolution of cells.

P2 Line 61: optimization

We thank the reviewer for pointing this out. Our model jointly trains on both spatial coordinates and gene expression features, but the loss function is factorized into two components, one for spatial reconstruction and another for biological fidelity, each with its likelihood and distribution (Laplace for coordinates, Gaussian for gene expression).

We will revise the sentence as follows:

"...that jointly trains on spatial and cell state dynamics using a factorized loss, modeling spatial reconstruction and biological fidelity with tailored distributional assumptions."

P2 Line 75

All these references point to Flow Matching papers that were applied to learning cellular trajectory from earlier to later time points, which is a common benchmark in OT-based generative models. All these methods predict how cell states evolve into one another, but not how they organize in the anatomical space (L75-76).

P3 Line 90: descendants

Exact-solution OT, as implemented in moscot estimates a transition matrix containing the probability of each target cell to be a descendant of each ancestor cell. We propose to make the sentence clearer as follows:

... use discrete optimal transport to link cells across time and infer the evolution of cell states from spatially-resolved gene expression.

P4 L150 tissue structure

Predicting tissue structure from a source microenvironment is biologically significant as tissue organization reflects biological function, with spatial changes indicating migration, proliferation, or reorganization during development or disease.

Additionally, spatial reorganization has been crucial for critical phenomena such as immune cell infiltration and tumour microenvironment remodeling in cancer [5], fibrosis progression and matrix remodeling [6].

In this context, our model provides a generative tool to investigate dynamics at the neighborhood level, offering insights that go beyond single-cell trajectories by modeling how spatially localized groups of cells evolve as structured units.

P5 top: type of dynamics

In scRNA-seq and its spatial variant, each cell is destroyed during sequencing. Thus, each slide or sample is an independent measurement. This is a standard setting in single-cell transcriptomics and a key reason why OT has emerged as a framework for trajectory inference across timepoints. Once an appropriate optimality criterion is defined, OT can be used to infer spatio-temporal dynamics on independent snapshots. In other words, computational methods are used to reconstruct a trajectory from successive measurements, even though longitudinal tracking of individual cells is not experimentally possible. In our work, each whole slide, as defined in Eq. (8), contains a different number of cells. We will add a sentence to make this aspect clearer in Section 4.1.

P5 eq 12: Averaging

We act on the transcriptional and spatial means to match point clouds across batches. We highlight that this defines a criterion to pair up point clouds across batches, but the dynamical model still learns to flow to all points of a target point cloud, attending all the points and coordinates of a source microenvironment. We agree that using the mean in combination with very large point clouds may mask the signal, but in a more localized setting (see App. Fig. 7), other established studies have demonstrated that cellular microenvironments are well-approximated by the statistics of a Gaussian around a central cell [7]. This supports our matching criterion.

SPFlow

SPFlow is our internal baseline, introduced in Sec. 5.1.3, which models trajectories at the single-cell level rather than the microenvironment level. While not from a published paper, it was designed to highlight the benefits of modeling point cloud dynamics over individual-cell dynamics. We will add a clarifying sentence in Sec. 5.1.3 stating that SPFlow is introduced in this work.

refs 34 and 38?

Thank you for pointing this out. We will correct the typo in our revised version.

Limitations

Discussed in Appendix B and referenced at L227.

Ref.

[1] Xie et al., "Spatially resolved gene expression prediction from histology images via bi-modal contrastive learning", NeurIPS 2023

[2] Oh et al., "Global Context-aware Representation Learning for Spatially Resolved Transcriptomics", ICML 2025

[3] Zhu et al., "Diffusion Generative Modeling for Spatially Resolved Gene Expression Inference from Histology Images", ICML 2023

[4] Halmos et al., "DeST-OT: Alignment of Spatiotemporal Transcriptomics Data", Cell Sys. 2025

[5] Zhao et al., "Spatial transcriptomics at subspot resolution with BayesSpace", Nat. Biotech. 2021

[6] Raredon et al., "Single-cell connectomic analysis of adult mammalian lungs", Science Advances 2021

[7] Haviv et al., "The covariance environment defines cellular niches for spatial inference", Nat. Biotechnology 2025

We thank Reviewer fydQ for their constructive feedback and positive assessment of our method and manuscript. Below, we address the reviewer’s suggestions and remaining concerns point by point.

Dynamic with conditional model

Our approach uses source-conditioned CFM to approximate a continuous Entropic Optimal Transport (EOT) coupling between source and target distributions [1]. This makes our method a principled predictor for the next time point, provided the OT assumption holds.

Unlike previous scRNA-seq methods that evolve individual cells, we model the joint evolution of spatially defined point clouds, treating each microenvironment as a rigid object. Since spatial slides can differ in cell density, matching source and target point clouds becomes an OT problem across incomparable spaces with different cell counts.

A standard ODE model like the cited ones would evolve each point individually and would be inadequate in the point cloud setting. Instead, we use the continuous EOT coupling of a source-conditioned CFM model to approximate how a source point cloud evolves over time, capturing both compositional and spatial changes in cellular neighborhoods. As shown in Fig. 13-15 (Appendix), our model generates spatially resolved point clouds from an initial state at $t=0$.

Functional similarity to static OT approaches.

Our method is intrinsically more flexible than moscot and WaddingtonOT, as it provides a generative model approximating a neural Monge map. NicheFlow predicts how a cellular point cloud evolves from $t$ to $t+1$ in both space and gene expression. In contrast, static OT methods only compute the probability that each earlier cell evolves into each later cell, and do not generate later time points.

For example, see Fig. 13-15 in the Appendix, where we predict the evolution of a source spatial region to a target point cloud across time, using the generated neighbourhood at $t+1$ as a source to predict the dynamics until $t+2$. All the colored points are purely generated. In other words, our model can be used for discrete time point interpolation and cell state generation, which cannot be achieved by static OT.

The reason for our comparison with moscot is that it is the most established method performing spatiotemporal trajectory predictions across time, and we believe it is important to show the community how our model represents a promising alternative in an applied setting. With an extra page of results, we will add some results on generative interpolations from the Appendix to the main text.

Comparison with references.

We thank the reviewer for highlighting these relevant works and will cite them in our revised manuscript. Below, we briefly summarize the key differences between our approach and the suggested models, DeST-OT and SpaTrack.

DeST-OT aligns spatial slides using semi-relaxed optimal transport, incorporating a loss that preserves transcriptomic and spatial proximity between ancestors and descendants. It also introduces a growth vector to estimate the expected number of descendants per source cell.

SpaTrack uses Fused Gromov-Wasserstein OT, balancing transcriptomic and spatial differences based on spatial autocorrelation of features.

Our work is comparable to SpaTrack and DeST-OT in terms of biological settings: they all aim to capture temporal information in time-resolved spatial transcriptomics. However, there are crucial differences:

• Both suggested methods use discrete OT to probabilistically assign each source cell to an existing target cell, similar to moscot. In contrast, our approach parameterizes a continuouos map in the Monge sense (as in [1-4]), predicting how cell distributions (or point clouds, in our case) evolve over time via a flow.
• Our model trains an approximate coupling in batches, making NicheFlow scalable to large spatial slides and datasets (e.g., the 4.2 million-cell brain ageing dataset we used). Moreover, parameterized deep models also enable future generalization to large treatment studies or cross-species tasks.
• The main difference is that NicheFlow models the evolution of spatial microenvironments (neighborhoods), not individual cells. It generates how a spatial region changes across slides, while prior methods focus on discovering single-cell descendants from ancestors. By adjusting neighborhood size and shape, NicheFlow enables both localized and large-scale tissue analysis.

In summary, with NicheFlow we propose two new perspectives on spatio-temporal modeling: the use of a continuous flow-based generative model to learn spatial trajectories and the formulation of the trajectory problem at the level of point clouds.

W1 and 2 metrics

We agree that incorporating gene-level and spatial assessments using Wasserstein distances can offer additional insight into the generative capabilities of the models. In response, we have now computed $\mathcal{W}_1$ and $\mathcal{W}_2$ distances for both spatial coordinates (“Pos.”) and gene expression vectors (“Genes”) for all models presented in Table 1.

To ensure scalability, we compute Wasserstein distances per cell type and report the average across types. Shown below are the best-performing configurations for each model on the mouse embryonic development dataset; additional results will be provided in the appendix due to space constraints.

These results reinforce our findings: NicheFlow with the GLVFM objective outperforms other models across all metrics, including the new Wasserstein evaluations, confirming its ability to generate realistic spatial and transcriptional patterns.

Proposed baselines: MIOFlow, TrajectoryNet and DeepRUOT?

These methods are designed for single-cell transcriptomic trajectory inference and do not model spatial structure or microenvironment-level dynamics. TrajectoryNet, MIOFlow, and DeepRUOT learn continuous trajectories in gene expression space (or PCA reductions thereof) rather than generating spatially resolved point clouds conditioned on tissue context.

Our work addresses a different problem: modeling how spatial neighborhoods of cells evolve jointly in both space and gene expression composition. Adapting these models to predict coordinates cell state simultaneously, as well as generating an increased number of cells in time, would alter their original design and put them at a disadvantage. Intuitively, the lack of spatial information in the prediction prevents the computation of the SPD and PSD spatial metrics on the outputs of the proposed models.

We will clarify these aspects in the revised manuscript and cite the papers.

The use of unbalanced OT.

While unbalanced OT is a promising direction, integrating it into NicheFlow is non-trivial. Our model tracks the evolution of cell groups as point clouds, not individual cells. Unlike single-cell frameworks, where growth reflects mass conservation, growth at the point cloud level is harder to define. Since our OT formulation models how cell neighborhoods evolve jointly in space, it does not rely on explicit growth modeling or enforce mass conservation at the individual-cell level.

Instead, our model captures changes in cell density by using a conditional generative approach to approximate the EOT coupling, allowing mappings between source and target clouds of differing sizes. The number of sampled cells at each time point defines the noisy cloud size during training, enabling the model to learn spatially resolved density shifts, as proliferation in development or expansion in disease.

Interpolation capabilities.

Great question. In principle, NicheFlow can interpolate unseen time points. Given a densely sampled developmental process, one could condition generation on a continuous target time and simulate evolution from a source. If the samples sufficiently approximate a continuum, this extends to unseen slides.

Unfortunately, this approach faces key data limitations:

• Current spatial datasets are sparsely sampled in time and do not approximate a continuum. Intermediate slides are often anatomically distinct from both their predecessors and successors, making interpolation a fully out-of-distribution (OOD) task.

• One could test interpolation using replicates (e.g., training on one embryo and predicting held-out slides from another sampled at the same timepoints), but to our knowledge, such replicated spatiotemporal datasets are not available.

• Crucially, some intermediate slides contain unique cell types or anatomical structures not observed at other time points. In spatial transcriptomics, where transcriptional profiles are constrained by spatial context, this can lead to biologically infeasible predictions.

In other words, while interpolation is algorithmically possible, its biological feasibility depends on denser temporal sampling or the availability of replicates, which are currently lacking in existing datasets.

Ref

[1] Klein et al. "GENOT: Entropic (Gromov) Wasserstein flow matching with applications to single-cell genomics." Advances in Neural Information Processing Systems 37 (2024): 103897-103944.

[2] Tong et al. "Conditional flow matching: Simulation-free dynamic optimal transport." TMLR (2023).

[3] Atanackovic et al. "Meta flow matching: Integrating vector fields on the wasserstein manifold." ICLR (2024).

[4] Wang et al. "Joint Velocity-Growth Flow Matching for Single-Cell Dynamics Modeling." arXiv preprint arXiv:2505.13413 (2025).

We thank Reviewer yXEk for the thoughtful review. We appreciate the positive assessment and the insightful suggestions. Below, we address the remaining questions and points of criticism.

Global tissue coordination.

While we recognize the value of global coordination, our work focuses on individual functional regions in a tissue, enabling users to study the migratory or compositional development of custom biological areas over time. This aligns with the locality paradigm in graph and point-cloud methods for cell communication, which we extend to the spatiotemporal setting.

In future work, we could incorporate global context by conditioning velocity fields on global niche embeddings from long-range graph or point cloud representation models. This allows biasing the model by the influence of the global cellular context. Still, scaling point-cloud methods like NicheFlow to whole-slide data remains a computational challenge, especially at higher resolutions, and is left for future exploration.

Tradeoff parameter.

We note that the role of $\lambda$ is discussed and explored throughout the main text and Appendix (L314–323):

• Sec. 5.2. In our comparison with moscot, we test two scenarios:

  • For studying compositional changes in fixed spatial structures, we use a low value ($\lambda=0.1$) to preserve spatial location over time.
  • For modeling migratory processes (e.g., neural crest to head mesenchyme), we use $\lambda=0.5$, showing that adding cell-state similarity aids fate mapping across space.

• App. D.3. The value $\lambda=0.1$ is used for the point cloud generation task in Tab. 1. We also provide:

  • An ablation across $\lambda$ values (Tab. 2),
  • A visualization of its effect on coupling distance (Fig. 7, ref. in L199).

• App. D.6.1 and D.6.2. We use $\lambda=0.5$ for additional biological experiments, such as mapping contained organs over time and inferring anatomical changes during axolotl brain development.

Pointwise metrics

PSD and SPD are standard metrics in point cloud generation and correspond to the two directions of the Chamfer distance, quantifying fidelity and coverage [1]. Importantly, we compute them over many aggregated microenvironments to reconstruct the entire tissue, so they indirectly reflect global morphology.

To directly address this concern, we include Gromov–Wasserstein (GW) and Fused GW (FGW) distances [2, 3] between generated and real point clouds to compare different training paradigms for NicheFlow. These metrics compare intra-point-cloud pairwise distance matrices between source and target, capturing internal structure (GW) and gene expression similarity (FGW). Values are averaged over 5 generation runs.

As computing GW/FGW on entire slides is computationally infeasible due to the number of points, we evaluate them at the generated niche level, where structural comparisons are tractable. A direct comparison with SPFlow or RPCFlow is not appropriate: SPFlow predicts single points and lacks structure, while RPCFlow produces random point clouds that vary across runs.

MED

ABD

MBA

NicheFlow with the GLVFM objective consistently achieves the lowest or competitive GW and FGW distances, confirming the value of our novel mixed-factorized posterior at preserving geometric constraints.

Cellular interactions

While modeling cell-cell interactions is important, spatial proximity alone can reveal meaningful biology, e.g., how brain regions evolve with disease and treatment, or how processes like fibrosis and inflammation alter tissue structure beyond direct cell-cell signaling. Our goal is to provide a framework for modeling combined spatial and molecular shifts as a response to time.

Additionally, our qualitative analyses (Fig. 3, App. Figs. 11–14) show that NicheFlow generalizes to biologically annotated regions in the mouse embryo and axolotl brain, demonstrating its flexibility to handle more sophisticated definitions of source clouds based on labeling and prior biological knowledge.

Adapt $\lambda$

$\lambda$ controls how source and target batches are paired during training, influencing the modeled dynamics.

In principle, $\lambda$ can vary over time (e.g., higher in early development to capture large-scale migrations, lower later to emphasize local changes) or space (e.g., focusing long-range cell migration only in specific regions). However, spatial variation requires biological knowledge and annotation to identify areas where such dynamics are expected, making this approach less easily automatable.

Learning $\lambda$ is technically feasible, but it requires additional supervision. Simply including it in the gradients of Flow Matching lacks biological interpretability. Biologically motivated proxies (e.g., stemness) could guide learning, but at the cost of added complexity. Instead, we treat $\lambda$ as a biologically informed prior that shapes learning in an interpretable way.

Sensitivity to $r$

$r$ defines the spatial resolution at which NicheFlow learns tissue dynamics and constraints its generated outputs to it. According to our experiments, the choice of $r$ has an influence on the resulting microenvironments, which highlights the importance of selecting a biologically meaningful radius that matches the desired level of granularity (i.e., resolution of the anatomical regions to study). In practice, $r$ can be chosen to emphasize either coarse or fine structures, and it could even vary across timepoints to reflect biological contexts.

Alternative neighborhoods

We refer the reviewer to our response above on "Microenvironments and cellular interactions", where we discuss the flexibility of NicheFlow in handling non-trivial neighborhoods. As noted there, one could define point clouds using alternative logics, such as pruning or weighting edges based on cell–cell interactions, rather than using a fixed spatial radius. While we opted for radius-based neighborhoods due to their interpretability and prevalence in spatial transcriptomics [4], our framework is not limited to this choice and could integrate graph-based or learned groupings in future work.

Continuous-time/Multi-timepoint

NicheFlow enables multi-timepoint generation by sequentially predicting future states: we predict $t{+}1$ from $t$, then use the output as a source to generate $t{+}2$, and so on, without relying on ground truth between steps (see App. D.6). This is demonstrated in App. Figs. 11, 13–15 across various developmental processes in mouse (e.g., neural crest cell propagation) and axolotl (e.g., tracking hemisphere and cavity formation).

A true continuous formulation would require aligning the model’s internal flow with biological time, for instance via a Neural ODE–style vector field. This extension is challenging due to point cloud density changes and the issue of incomparable source and target spaces, but it represents an interesting direction for future work. More crucially, such a formulation would require a more densely sampled temporal dataset to approximate a continuum. This is a technological limitation of nowadays' spatiotemporal data acquisition techniques, which focus on a coarse view of a dynamical process.

Limitations

Our representation of attributed point clouds enables flexible modeling of continuous cellular states and can be combined with predefined or learned embeddings. In this work, we focused on datasets with a single sample per time point, where time is the dominant source of variation in both gene expression and spatial organization. We thus avoided scVI or similar approaches that remove sample-specific batch effects, as they risk discarding meaningful temporal signals. When replicates are available, VAE-based models can help remove technical variation, but in our setting, training a VAE would likely capture similar biological structure as PCA, albeit via a non-linear encoder.

We also highlight the generative nature of NicheFlow: it uses Conditional Flow Matching to transform noise into spatially organized cell neighborhoods, conditioned on earlier time points. The model integrates spatial geometry via a point cloud transformer, capturing both expression and spatial context for generating spatial neighborhoods.

Lastly, unlike pseudotime methods (e.g., DPT) that infer latent trajectories from expression data only, our approach models how spatial cell assemblies evolve over physical time. Our method does not attempt to infer latent cell orderings from expression data, but instead focuses on modeling how local tissue organization changes over time.

While manifold learning is relevant, spatially constrained versions thereof are an interesting future direction, but beyond the scope of this work.

Ref

[1] Fan, Haoqiang, et al.. "A point set generation network for 3D object reconstruction from a single image." CVPR (2017).

[2] Mémoli, Facundo. "Gromov–Wasserstein distances and the metric approach to object matching." Foundations of Computational Mathematics 11.4 (2011): 417–487.

[3] Vayer, Titouan, et al. "Fused Gromov-Wasserstein Distance for Structured Objects." Algorithms (2020)

[4] Haviv, Doron, et al. "The covariance environment defines cellular niches for spatial inference." Nature Biotechnology 43 (2025): 269–280.

I thank the authors for their complete response. Their responses address my concerns, therefore I would like to increase my ratings from 5 to 6.

We thank Reviewer 1m53 for the constructive and thorough feedback. We appreciate the positive remarks on our model's structure and core idea, as well as the critical comments. Below, we address the remaining concerns.

Notation

As noted, $\mathbf{x}_0$ and $\mathbf{x}_1$ explicitly refer to source and target samples. When source and target spaces are incomparable, $\mathbf{x}_0$ serves as a conditioning input, not as the ODE’s initial state, and our notation reflects it.

We intentionally avoid using $\mathbf{x}_t$ for the ODE solution, instead denoting it as $\phi_t^\theta(\cdot)$ throughout. Using $\mathbf{x}_t$ would suggest $\mathbf{x}_0 = \mathbf{x}\_{t=0}$, which is ambiguous when source and target domains are incomparable. By using $\phi_t^\theta$, we clarify that $\phi_0^\theta(\mathbf{z} \mid \mathbf{x}_0) = \mathbf{z}$, with $\mathbf{x}_0$ as a condition rather than the starting point of the trajectory. This is made explicit in L138 to L140. The same convention applies to interpolations, where we introduce $g_t$ (e.g., Eq. (2), L105).

Hold-out evaluation

We agree that hold-out evaluation is generally desirable, and this could be implemented by leaving out either an anatomical region or an entire slide. However, both options are challenging due to data constraints.

Predicting the trajectory of a region on a new slide requires prior observation of the same region in another sample at the same time point. Without this, the task becomes fully out-of-distribution, which remains difficult for current generative models [1]. Existing datasets lack replicated slides per time point, and to our knowledge, none provide repeated spatiotemporal measurements. In addition, current datasets span only coarse developmental stages and cover a few time points due to experimental costs. Holding out a slide would therefore introduce large temporal gaps, making it hard for the model to infer transitions.

Instead, we showcase the model as follows:

• ML perspective: We show that our approach improves over baselines at enforcing spatial constraints for spatiotemporal prediction (Fig. 2 and Tab. 1), evaluated on a grid of test point clouds (Sec. 5.1.1, App. F.5).
• Biological perspective: We highlight its use for hypothesis generation, in contrast to exact OT methods, by analyzing spatial and transcriptional cell fate and the compositional evolution of selected anatomical regions (Sec. 5.2, App. D.6).

Our goal is to introduce a scalable deep generative framework for niche-based spatiotemporal analysis, addressing the limitations of exact OT methods, which do not scale to large slides. With advances in dataset generation, we expect neural OT maps to play a key role in enabling generalization and modeling in spatiotemporal and perturbation contexts.

Flow function/residual

In NicheFlow, $r_t$ and $f_t$ are neural networks predicting posterior means over denoised target coordinates and cellular features (L213–222). These means are used to infer the velocity field on noisy coordinates and cell states, conditioned on a source point cloud during generation (Sec. 3.3 and Sec. 4.3). Thus, $r_t$ is not a residual, but a posterior mean predictor that enables computing the generative velocity on the coordinates.

$r_t$ and $f_t$ are outputs of a point cloud transformer that attends to all nodes and positions in both the source and noisy point clouds with cross-attention, and predicts a velocity for each sample. Samples from $\mathcal{M}$ are passed through the same global function, but this is not a shared MLP over individual points, but a permutation-invariant point cloud transformer. This architecture computes the velocities of samples and coordinates, integrating the intra-cloud context as well as the conditioning source point cloud (Sec. 4.4 and App. F.6 for details).

Justification for Tab. 1

The metrics in Tab. 1 provide a technical evaluation of the source-conditioned point cloud generation task, supporting our architectural and modeling choices over baselines. Since we assess how well the global point cloud structure is recovered by a flow model and training strategy, this comparison was performed using a fixed $\lambda$ to enforce spatial preservation and the same biological semantics across models. For the rebuttal, we followed the suggestion and evaluated RPCFlow across varying $\lambda$ values (see below).

$\lambda$-dependent improvement

To address the reviewer's concern, we conducted an extensive ablation study where we varied $\lambda \in \\{0.1, 0.25, 0.5, 0.75 \\}$ for both RPCFlow and NicheFlow across all three datasets and across all three training objectives. This resulted in 72 evaluations, each reported as the mean ± std. over 5 generation runs.

Due to the limited space, we present only the best-performing method for each metric on each dataset. We will add the full table to the Appendix. When NicheFlow outperforms RPCFlow on all metrics, we report the overall best RPCFlow configuration.

MED

ABD

MBA

We conclude that RPCFlow benefits from tuning the $\lambda$ parameter. However, even with its optimal configurations, it does not outperform NicheFlow with different patterns than when $\lambda=0.1$. Biologically, different values of $\lambda$ have distinct meanings (see below).

Biology of $\lambda$

While Tab. 1 evaluates point cloud generation, the choice of $\lambda$ carries strong biological significance, which we explore throughout the paper. Specifically, $\lambda$ controls the trade-off between spatial coherence and feature similarity in the OT coupling and should be selected based on the underlying question (L318–323):

1. High $\lambda$: suited for modeling cell type migration across anatomical regions (Fig. 3, right).
2. Low $\lambda$: appropriate for studying how fixed structures evolve compositionally (Fig. 3, left).

Since the optimal value depends on the biological context, we emphasize a qualitative analysis of $\lambda$’s impact, rather than relying solely on quantitative metrics. We show examples where a high or low value is necessary for meaningful biological interpretation.

Notably, RPCFlow is excluded from our qualitative evaluation, as Sec. D.5 (Fig. 10) shows it fails to condition predictions on the source. Despite strong generative performance, it ignores the source microenvironment and yields biologically implausible mappings.

Comparison with moscot

Moscot solves an exact OT problem, i.e. it computes transition probabilities between all source and target point pairs. Since its transition matrix is fixed, the model cannot generalize to test samples. We therefore compare our model to moscot on its native task of mapping embryonic structures across developmental time. Our results show that the proposed spatially constrained generative model offers a promising alternative to this established method. Additionally, App. D.6.2 presents tasks that exact OT models like moscot cannot handle, such as generating niche-specific anatomical changes and compositional remodeling.

$\alpha$ in moscot

The comparison with moscot assesses how well each model captures compositional changes in anatomical structures or migratory patterns. Qualitatively, we found that the spatial component in the FGW problem is comparable to our spatial term across the 0.1–0.9 $\alpha$ range in the Fig. 3 analysis. When mapping fixed structures over time, values below 0.75 caused excessive density dispersion outside the anatomical region. For migration, no $\alpha$ value led to generally accurate transitions, though $\alpha = 0.5$ mapped the highest density to the expected cell type.

Following 1m53's suggestion, we report the proportion of source density mapped to the correct cell type across $\alpha$ values, to support our choice. For the spinal cord (Fig. 3 left), values above 0.75 yielded better qualitative and quantitative results. For neural crest cells (Fig. 3 right), $\alpha = 0.5$ performs best.

RPCFlow vs. NicheFlow

The difference between RPCFlow and NicheFlow is that the former is trained on point clouds sampled across the slide, while the latter uses strictly defined spatial neighborhoods. Both methods are evaluated on the same test niches, and across the same model backbones, OT strategies and posterior losses (Gassian vs. Laplacian). Moreover, we showed above that NicheFlow overcomes RPCFlow for different values of $\lambda$, making our quantitative results robust to it.

RPCFlow and SPFlow?

RPCFlow refers to a Random Point Cloud Flow model, where point clouds are sampled randomly across the tissue. SPFlow stands for Single Point Flow, which models trajectories at the single-cell level rather than on microenvironments. While described in Sec. 5.1.3, we will clarify them further in the revision.

Tab. 1 ref.

We reference Tab. 1 in L299.

[1] Nalisnick et al. Do deep generative models know what they don't know?, ICLR (2019).