Regularized Optimal Transport for Single-Cell Temporal Trajectory Analysis

Major

• Although I believe the authors when they claim good representation performance of their autoencoder, I feel claiming its “superiority” is rather unsupported. Probably it is a mischoice in the terminology, but there is no proper benchmark demonstrating in what way the masked autoencoder is better than traditional approaches (e.g. scVI [1]) for the proposed task. Also, masked autoencoders for single-cell have now been proposed in multiple settings, which makes their use a nice asset, but a source of novelty. I think the work would benefit from proof of the reported superiority in the representation, or, conversely, the authors may want to tone the representation aspect a little bit in the presentation.
• In my opinion, referring to the task of “recasting learning trajectories” as a machine learning problem overlooks the fact that ML for trajectory inference is now a standard in the computational biology field. I would rather phrase the contribution as an additional method in this direction. This shortcoming is reflected in the Related Works section, where only traditional methods (not based on deep learning) are discussed. I think such a section would benefit from an update including how machine learning has already been applied to OT for single-cell trajectory inference. 
• I wonder why Mouse-IEP was not included in the benchmarks comparing Tarot with existing models.
• Though I generally value attention to biology, I believe the biological details on page 4 may be shortened a little bit (and integrated with details in the appendix) to dedicate more time to the methods part, as some aspects (like the masked autoencoder part) felt a bit rushed. 
• In my opinion, one of the major drawbacks of the method is the operation at the cluster level. I see the benefit of such an approach on the dataset at hand since Figure 1 demonstrates low variability within clusters. Nevertheless, some datasets display larger heterogeneity per biological cluster [4], which would make the representation choice suboptimal. I think acceptance at ICLR would require a more comprehensive dataset choice, including more challenging settings (e.g. the MEF dataset in [5]). The authors claim that individual cellular representations are also an option, but no evidence was provided in these regards, so I cannot validate the claim properly. 
• Though I see the value of the way the authors perform continuous trajectory learning, other methods have been released based on dynamic optimal transport that learns continuous vector fields (and so interpolations). An example is the standard OT conditional flow matching [6]. The lack of a comparison to such methods should be discussed. Comparison with at least a couple of them would be very beneficial for the contribution. 
• I wonder if limiting the evaluation of GPT-G and GPT-L to monotonically varying genes is comprehensive enough. In other more complex datasets, this would potentially impair modelling more complex regulatory processes. It would be useful for the authors to elaborate more on this. 
• In my opinion, the paper lacks a little bit of elaboration on some aspects of the experimental process. How did you run the several methods to make them compliant with the setting you propose? Did you simply run them on averaged cluster representations? Were other adjustments needed?
• In my opinion, the way the GSEA analysis is presented is a bit unclear and the biological value of the finding is not discussed. It would be beneficial to at least disclose the GO terms in the main text and explain why they make sense in the considered system. 
• The change in trajectory between the perturbed and unperturbed states in Figure 9 is not discussed in terms of biological relevance. Are the changes in cluster transitions biologically relevant? 
• I think the last result on real-world experiments would have benefitted from some visualisation or table.

Minor

• The authors refer to cells with the term “gene expressions”. In my opinion, it would be more appropriate and readable to address them as “sequenced cells” or “gene expression profiles of n cells”. The way it is phrased now may be confused with the concept of “number of genes”.
• “Specifically, the number of nearest neighbors was chosen to be 50, according to the rich experiences of biology scientists”. I would personally remove this sentence unless supported by a solid citation. Packages like scanpy use a default lower number of neighbours and that seems to work properly. It is fine if more neighbours worked better for the datasets in this study, but I would phrase this sentence in this direction.
• I would make sure indexes are not bolded in the notations (e.g. line 245).

As of know, I unfortunately believe that the reasons to reject the current version of the paper outweigh the reasons to accept it. But I am happy to discuss the author’s views on my comments.

[1] Lopez, Romain, et al. "Deep generative modeling for single-cell transcriptomics." Nature methods 15.12 (2018): 1053-1058.

[2] Fang, Zhaoyu, Ruiqing Zheng, and Min Li. "scMAE: a masked autoencoder for single-cell RNA-seq clustering." Bioinformatics 40.1 (2024): btae020.

[3] Kim, Jaesik, et al. "Single-cell Masked Autoencoder: An Accurate and Interpretable Automated Immunophenotyper." NeurIPS 2023 AI for Science Workshop. 2024.

[4] He, Zhisong, et al. "An integrated transcriptomic cell atlas of human neural organoids." bioRxiv (2023): 2023-10.

[5] Schiebinger, Geoffrey, et al. "Optimal-transport analysis of single-cell gene expression identifies developmental trajectories in reprogramming." Cell 176.4 (2019): 928-943.

[6] Tong, Alexander, et al. "Improving and generalizing flow-based generative models with minibatch optimal transport." arXiv preprint arXiv:2302.00482 (2023).

• Limited novelty: $`TAROT`$ is presented as a novel framework however all three components have been previously suggested for single cell data analysis. Specifically, both WOT (Schiebinger et al. 2019) and moscot (Klein et al., 2023) utilize OT for single cell trajectory analysis, rely on lower dimensional representations for the mapping, and allow incorporation of biological priors via the source and marginal distributions as well as un-balanced mapping. Of note, while the authors relate to these as "preliminary" works they are well accepted and used by the community (e.g. WOT has ~770 citations and moscot's github has ~100 stars).
• Reliance on cell clustering: As remarked by the authors (in Appendix C.3), cell clustering is a critical step in data preprocessing for $`TAROT`$ application. This can be a limitation when reliable clusters are unavailable and generating them may require "experience of biology scientists" (as used by the authors for the $`Mouse-RGC`$ data, Sec. 3.2). The ablations show that clustering quality as well is simply removing the cost results in poor performance, it will be valuable to relate to this limitation.
• Performance assessment: Qualitative results rely on metrics which are designed to validate the developmental and functional priors induced by $`TAROT`$. Though gene patterns assessment rely on genes not included in the prior, $`TAROT`$ has a clear advantage in these tasks, as it is designed to capture these patterns. Next, the biological contribution is presented only for $`TAROT`$ (Figure 7) and knockout analysis, whereas it may be that other methods capture the same trends, or compared to least performing method (Figure 8).
• Data availability: The curation of $`Mouse-RGC`$ is presented by the authors as a major contribution however the data is not made available.

Despite the strengths discussed above, there still exists some concerns:

• Methodological novelty : Although the paper presents a unified framework, the components of this framework—such as cell representation learning via masked autoencoders, regularized OT for trajectory inference, and the use of splines for continuous trajectory generation—are not novel. Each of these techniques has been explored in previous works (e.g., WaddingOT [1], Moscot [2], Slingshot, Monocle). The combination of these existing methods may come across as a straightforward integration. The paper may need to further address how this combination leads to new insights or capabilities that were not achievable with existing techniques, with more systematic benchmarkings.

• Limited to Cluster-Level Resolution: This paper produces continuous trajectories at the cluster level rather than achieving single-cell resolution. Recent advancements, particularly in dynamic optimal transport, can directly generate continuous single-cell trajectories, and simulation-free approaches such as conditional flow matching (see [3,4,5]) can scale to high-dimensional gene spaces. The paper may need to adequately compare TAROT against these more recent methods and explain the advantages TAROT provides over approaches that can inherently handle single-cell granularity.

• Generalization: While the authors demonstrate TAROT’s effectiveness on Mouse-RGC, the paper does not provide sufficient discussion on the generalizability of the approach. For example on cross-dataset beyond Mouse-RGC, whether the framework is broadly applicable or tailored to a specific type of data. In the abstract , it is stated that "....Existing approaches show a limited generalization of our datasets*". The readers may naturally question what is the performance of TAROT on other datasets where existing methods demonstrated good performance.

• Cell Division and Death: The current framework does not consider cell division or death (e.g., Moscot), which is crucial in biological processes that affect trajectory inference.

• Data Availability: The data availability statements seemed not found.

Based on the discussions above, I suggest the work could be further revised substantially, and I would be very happy to increase the score if the concerns are adequately addressed by the authors.

[1] Schiebinger, Geoffrey, et al. "Optimal-transport analysis of single-cell gene expression identifies developmental trajectories in reprogramming." [2] Klein, Dominik, et al. "Mapping cells through time and space with moscot." [3] Tong, Alexander, et al. "Simulation-free schrödinger bridges via score and flow matching." [4] Lipman, Yaron, et al. "Flow matching for generative modeling." [5] Tong, Alexander, et al. "Improving and generalizing flow-based generative models with minibatch optimal transport."