A long range foundation model for zero-shot predictions in single-cell and spatial transcriptomics data

We thank the reviewer for their helpful criticism. We address each point below:

After reading the paper I have numerous major concerns that prevent me from recommending acceptance. In short, the presented results suffer from a number of common issues in recent "single-cell foundation model" papers (e.g. lack of rigorous benchmarking against current best practice methods), and it's not clear to me that sCT represent a major advance for the analysis of single-cell data. Details below:

Thank you for your feedback. We understand your concerns regarding current single-cell foundation models. We would like to clarify that our primary goal in this paper was not to achieve a major breakthrough in single-cell data analysis. We believe such an expectation would be overly ambitious for a machine learning paper at a conference like ICLR.

Instead, we positioned our work within the literature of self-supervised models trained on single-cell and spatial transcriptomics data. Our objectives were twofold:

1. To evaluate the zero-shot capabilities of such models on tasks aligned with their pretraining objectives but applied to new data domains.
2. Motivated by the observation that state-of-the-art models like scGPT struggle in these scenarios, to propose machine learning advancements that address these limitations.

We believe this contribution is well-suited for a conference like ICLR and helps establish better practices for researchers developing such models.

Missing comparisons with non-"foundation model" baselines: A number of recent benchmarking studies (e,g, [1, 2, 3] among others) have found that recently published transformer-based single-cell foundation models perform no better than significantly simpler (e.g. linear) baseline models at a variety of downstream tasks. However, the authors' manuscript only compares against other transformer models, thus making it impossible to assess whether the authors' method indeed represents a meaningful advancement. For for the batch effect/cell type classification tasks, how do simpler baselines (e.g. logistic regression, scANVI, etc.) compare?

Thank you for raising these concerns. We have now incorporated additional bioinformatics baselines for gene imputation (MAGIC, Table 1) and cell typing (scVI, Scanpy + logistic regression, and k-nearest neighbors in Table 3). The results are reproduced in our reply to all reviewers above, as well as in the revised manuscript.

It is important to highlight that each of these baselines requires either training (sometimes referred to as "fitting") on the target data or the use of additional metadata. In contrast, we evaluate sCT in a zero-shot setting, ensuring that the evaluation data was not present in the sCT training data. Despite this significant constraint, sCT matches or outperforms these baselines. We also emphasize that comparable single-cell foundation models underperform these baselines in zero-shot scenarios.

We hope these results clarify our contributions and demonstrate the effectiveness of our approach.

Unclear real-world utility: Based on the provided experimental results, it's not clear to me that the authors' method provides additional utility beyond current standard scRNA-seq analysis workflows. For example, given that (as mentioned by the authors) <6% of protein-coding genes are expressed in an individual cell on average, I'm not sure why the results presented in Table 2 are meaningful, as most of the time all the masked values will have zero measured expression

We agree with the reviewer observation that, on average, less than 6% of protein-coding genes are expressed in an individual cell. This sparsity was a key motivator for many of our design choices and evaluation metrics. Specifically, it informed our decision to exclusively mask non-zero expressed genes in all cell imputation experiments. Thus, when we refer to masking $k\\%$ of the genes, this signifies masking $k\\%$ of the non-zero protein-coding genes. We have further emphasized this point in the manuscript.

We thank the reviewer for their helpful critique of our work. We address each comment specifically below:

First of all, current downstream tasks can not support the claim of a foundation model; it is more like an imputation model. A foundation model should deal with various tasks. Also, some concepts like "a first principles formulation of the problem" are not clearly explained.

Thank you for this insightful comment. We agree and have toned down the use of the term "foundation model" in the paper. Our goal was to investigate the recent generation of self-supervised models trained on single-cell and spatial transcriptomics data. Inspired by work in other fields, such as proteomics and genomics, we sought to understand how well these models generalize in a zero-shot setting to new data domains. We believe this capacity for generalization is essential for the practical utility of such models.

Our analysis revealed that state-of-the-art models like scGPT perform poorly when evaluated on new data domains, even for simple tasks like gene imputation, which aligns with their pretraining objective. This motivated us to further benchmark models to assess generalization capabilities and propose architectural modifications to address these limitations.

Thank you also for your comment regarding the phrase "first principles." We have removed this phrase and clarified our intended meaning. Our aim was to design an approach that minimizes assumptions when processing data to avoid introducing biases. Specifically, we sought to process the expression of all coding genes for all cells in the input without resorting to subset selection and to minimize data preprocessing. We have clarified this in Section 3 and the introduction.

While new convolutional architecture is interesting, the authors did not verify its effectiveness by comparing it with simple MLP, and also did not clearly explain the reason behind the design.

Thank you for raising this point. Unfortunately, using a standard multi-layer perceptron (MLP) in this context is computationally infeasible.

Our input consists of $k \times n$ gene expression tokens, where $k$ represents the number of cells $(50)$ and n represents the number of coding genes $(20,000)$. This results in $10^6$ tokens. As noted in the manuscript, processing this input directly with a Transformer model is impractical due to the quadratic scaling of self-attention with respect to sequence length. To address this, we employ a convolutional tower to reduce the input sequence length from $10^6$ tokens to $4096$ tokens, a size manageable for processing on a single high-performance GPU.

Achieving this reduction with an MLP would require transforming an input of shape (10^6, token_embed_dim) to (4096, transformer_embed_dim). Even a single-layer MLP would necessitate 4 x 10^9 x token_embed_dim x transformer_embed_dim parameters. Even with minimal embedding dimensions of 1, this would result in 4 x 10^9 parameters in the MLP alone, which would be roughly 50 times more than the current total number of parameters (70 million) in the model. We have clarified this constraint in the manuscript.

Conversely, convolutional layers provide a mechanism for gradually reducing the sequence length while increasing the embedding dimension from token_embed_dim to transformer_embed_dim. This is achieved with a manageable number of parameters due to weight sharing within the convolutions. Transpose convolutions then enable us to reverse this process, expanding the sequence length back to its original size while reducing the embedding dimension for gene expression reconstruction. This architecture also allows for skip connections between the downsampling and upsampling layers.

While we acknowledge that using convolutions for unordered data like ours might seem counterintuitive, our primary motivation was dimensionality reduction. We were pleased to find that this architecture not only effectively downsamples and upsamples the sequence with a controlled number of parameters but also demonstrates strong empirical performance. We believe this approach can serve as valuable guidance for others developing similar models. We have clarified the rationale behind this architectural choice in the manuscript.

Dear reviewer PjDS,

Thank you for constructive feedback on our paper. We appreciate your time and consideration towards reviewing our work. In response to your comments, we summarize our rebuttal below and highlight the changes in our manuscript.

Terminology clarifications: We replaced "foundation model" with more precise descriptions and clarified our goal of improving generalization in self-supervised models for single-cell RNA-seq data. We also modified section 3 in our manuscript to clarify our contributions.

Clarifications on Architecture: We justified the use of convolutional layers over MLPs for dimensionality reduction due to computational constraints. We also highlighted the efficiency and empirical performance of the convolutional design.

Imputation and Baseline Methods: We have added references to gimVI and scFoundation. We have also added additional baselines for gene imputation and cell-typing studies.

Experimental Design and Data Leakage: We confirmed no data leakage in held-out FOVs, and clarified questions about the finetuning on spatial data.

As the discussion deadline approaches, we kindly request any further feedback from you to help further refine our paper, and resolve any remaining concerns. We appreciate your insights and suggestions and thank you for your consideration.

Best regards, The authors

We thank the reviewer for their helpful criticism. Below we address them point by point:

The proposed improvements in the paper are rather straightforward. Learned Gene Normalization refers to the layer normalization, Positional Embeddings refer to standard sinusoidal position embeddings. The authors stack 50 cells into a single sequence, which from my perspective, creates sudo bulks and lowers the granularity instead of modeling intercellular dependencies. It would be great if the authors could elaborate more on how the proposed strategy can help with modeling intercellular dependencies.

Thank you for your insightful comments. We appreciate you bringing these points to our attention and have revised Section 4 of the paper to address them.

Shared Gene Normalization: To avoid confusion, we have renamed "Learned Gene Normalization" to "Shared Gene Normalization." This emphasizes that the normalization is applied across all cells in the input sequence. While based on standard layer normalization, it is crucial to note that this normalization is repeated for each of the N cells in the sequence, ensuring consistent gene normalization across all cells. This process significantly impacts performance, as demonstrated in Table 4.

Positional Embeddings: We have clarified the section on positional embeddings. Our approach differs from standard sinusoidal positional embeddings. We have removed any periodic 1D embedding, as there is no inherent order in the cells or genes. Instead, we use a constant embedding for each cell in single-cell data to enable cell differentiation. For spatial transcriptomics data, we utilize a 2D-aware embedding that captures the relative positions of cells within the field of view (FOV). This is detailed in Equation 1.

Cell Stacking: Our approach to stacking cells is similar to CellPLM (ICLR 2024), which demonstrated performance gains on various tasks. Our motivation aligns with CellPLM's:

1. Single-cell data: Stacking significantly improves performance, as shown in our ablation study in Table 4.
2. Cross-modal applicability: It enables the same model to process both spatial and single-cell data, facilitating transfer learning between modalities.

Furthermore, our whole-cell masking experiments with spatial data show that the model can predict entire cell gene expression profiles based on neighboring cell information. This strongly suggests that sCT has learned intercellular interactions. We acknowledge the reviewer's suggestion to further explore intercellular dependencies and plan to investigate this in future work. We believe these clarifications and the supporting ablation studies validate our design choices (see Table 4). These choices provide a robust foundation for developing future deep learning models for single-cell and spatial transcriptomics data analysis.

We thank you again for your additional response to our revised manuscript. Below we hope to clarify some of our positions:

The authors state that "cell stacking enables the same model to process both spatial and single-cell data, facilitating transfer learning between modalities." Could you please elaborate on this further? While the cell stacking strategy may capture intercellular relationships, it remains unclear how this design directly enables cross-modal transferability. A more detailed explanation or empirical evidence to support this claim would be beneficial.

Thank you for highlighting the need for further clarification on this point. In our architecture, cell stacking provides a mechanism for incorporating gene expression profiles of cell neighborhoods in spatial transcriptomics. While single-cell data lacks explicit neighborhood information, we simulate this by sampling cells from the same study, allowing sCT to learn intercellular coexpression patterns. This shared input setup enables cross-modal transfer, as a model pre-trained on single-cell data can be fine-tuned on spatial data. Empirical evidence for this is presented in Table 1, where we show that sCT pre-trained on single-cell and spatial data achieves improved performance. We also add comparisons with sCT trained only on spatial transcriptomics data (sCT (ST only) ) which has been reproduced below. It is worth noting that sCT pretrained on single-cell data and fine-tuned on spatial data consistently outperforms the model trained solely on spatial data. This performance advantage is particularly evident at higher masking ratios, highlighting the effective transfer of knowledge from single-cell data to spatial data. We will add the additional results to the next version of the manuscript.

My primary concern with this manuscript remains the limited application scenario. While the proposed method introduces some improvements over existing approaches, these advancements appear to be relatively straightforward and do not provide significant methodological novelty. Additionally, the downstream tasks are restricted to imputation and cell typing. I would encourage the authors to explore more biologically meaningful applications, such as imputing missing genes in spatial transcriptomics data, rather than masking some cells and predicting them. Expanding the scope of downstream tasks would strengthen the manuscript's impact.

We acknowledge the reviewer’s concerns about the scope and novelty of our contributions. While we do not claim significant methodological breakthroughs, we believe our integration of multiple architectural improvements (detailed in Table 4) represents a meaningful advancement. Our results (Tables 1, 2, and 7) demonstrate that these design choices collectively improve model performance, providing a robust foundation for transcriptional data modeling. Regarding the suggestion to explore more biologically meaningful tasks, we have incorporated experiments simulating the prediction of missing gene expression values (Table 7, Appendix Section C.1, and Section 4). These results decompose performance by gene sparsity levels, showing that sCT outperforms CellPLM. We believe this addresses the reviewer's suggestion while demonstrating sCT's potential for biologically relevant applications.

Thank you again for your insightful feedback and constructive suggestions. We hope these clarifications address your concerns and highlight the value of our work.

Best regards,

The authors

I appreciated the experiment in Table 3. But I think it would add value to it to add how the two models perform on gene predictions in the three sparsity levels separately. Also, Table 3’s caption does not describe the depicted evaluation metric.

Thank you for your comment and for pointing this out. We have edited the caption to clarify this, and have added more detailed results on the different sparsity levels. We have also moved the table to the appendix due to space constraints. We summarize our results here:

Note that sCT outperforms cellPLM on two out of three datasets on all classes.

I believe the presentation of results in Figure 3 is suboptimal. I would add titles to the single plots and definitely include a UMAP plot from the real cell gene expression to show how the embedding compares with the gene expression space. As of now, saying that sCT “corrects gene expression” sounds unsupported, as correction implies visualising and comparing with the data when the batch effect is still present. I also think that clustering requires some simpler baseline (like scVI or even the plain PCA embedding of the data).

Thank you for this comment. We now have replaced Fig. 3 to show UMAPs of the cell embeddings for several algorithms. We also have toned down our claims regarding batch effects. In addition, we have updated Table 3 to include additional baselines like scVI, and Scanpy + basic ML algorithms that need to be trained and measure clustering metrics where appropriate. We observe that sCT in zero-shot mode is comparable to strong bioinformatic baselines which require training, or access to query data.

We thank the reviewer for their insightful comments and suggestions. We address their questions below:

In lines 59-60, I find that the connection between modelling continuity and the lack of accounting for over-dispersion is not very clearly motivated. I think this is a very important point when modelling single-cell RNAseq and I am happy the authors brought it up. In general, though, it could be stated clearly why modelling continuous raw counts is a problem when trying to recover the over-dispersion in scRNA-seq.

Thank you for recognizing the strengths of our approach and for highlighting the importance of addressing over-dispersion in single-cell RNAseq data.

We model gene expression levels as discrete tokens, rather than continuous values, for several reasons:

1. Improved Loss Function and Optimization: This allows us to predict probability distributions over tokens and utilize a cross-entropy loss, which offers a smoother loss landscape and is less sensitive to the scale of input values compared to a mean squared error loss used for continuous value prediction.
2. Robust Evaluation Metrics: Using discrete values enables the use of classification metrics like Matthews Correlation Coefficient (MCC), which are more suitable for evaluating performance with heavily imbalanced data, as is common in this domain. Reporting MCC, in addition to continuous metrics, provides a more comprehensive assessment of model performance. We believe this practice should be widely adopted in the field.
3. Stratified Masking Strategy: Discretization facilitates our stratified masking procedure, where we apply different masking rates for zero tokens and non-zero tokens. This strategy significantly impacts performance, as demonstrated in our ablation study (Table 4). While we acknowledge that discretizing gene expression is not a novel concept, as seen in scBERT and scGPT, the points above highlight the advantages of our approach and the best practices we advocate for. We have further clarified these aspects in the revised manuscript.

From a terminology perspective, I would not define individual cell expression vectors as “gene expressions”. I would replace the term with something like “expression values” or so

We thank the reviewer for the suggestion, and have appropriately modified the manuscript to reflect this.

I find myself disagreeing with the sentence “Unlike scRNA-seq, spatial transcriptomics is a relatively new technology and publicly available data is still scarce.” Nowadays there are a lot of public spatial transcriptomics datasets out there and technologies to produce them, which makes the field so exciting. I would probably tune the claim down a bit.

While we agree that spatial data is getting more prolific, the volume of openly available data is still very limited when compared to single-cell RNA-seq. We have updated the manuscript to clarify this point in Section 3.

The architecture in 3.2 is understandable and clear. However, I feel it is a bit confusing for a less expert audience what a sequence here actually is (i.e. a stack of cells), especially because this is what makes sCT different from foundation models processing one cell at a time. You describe this in Learning intercellular dependencies but I believe it would make the flow better if you define this formalisation from before.

Thank you for this insightful comment. We have edited Section 3.2 accordingly to clarify that a sequence here represents a stack of cells as well as to clarify how sCT differs from existing models.

I think the structure of the paper could be improved. Especially, Table 1 is referenced multiple times in the methods section, but it contains a metric (MCC) that goes undefined till much later in the paper as a representation of optimality. Probably I would either refer to Table 1 only in the results or give some elements of understanding before.

Thank you for pointing this out. We have edited the table to improve clarity as well as the overall structure of the paper. We moved table 1 and the ablation study at the end of the experimental section.

I would not discuss results in the Table caption (e.g. in Table 2) but only in the text.

We have addressed this in the revised manuscript, and improved the table captions to make them more succinct and clear. We now discuss results in more detail in the experimental section.

We sincerely thank the reviewer for their thoughtful comments and for carefully reviewing our rebuttal. One of our other reviewers asked about the importance of multi-modality in pre-training, so we are reproducing the results of an additional experiment here, and would be grateful to hear your thoughts:

Performance of sCT Trained Solely on Spatial Data:

We have included an additional comparison in the gene imputation experiment (reproduced below for reference). The results show that the sCT variant pretrained on single-cell data and fine-tuned on spatial data outperforms the variant trained solely on spatial data, as well as the variant trained only on single-cell data. This supports our rationale that integrating single-cell data enhances the performance of spatial imputation tasks, especially as the severity of masking of test data increases.

Best regards, The authors

Below we reproduce tables for fixed gene masking, with results shown by sparsity level, as well as cell-typing results with additional baselines.

New detailed fixed gene masking

New celltype results with additional bioinformatic baselines

Below are new results reproduced here for convenience and also present in the revised manuscript:

New gene imputation baselines

• adds MAGIC as a bioinformatic baseline.