SToFM: a Multi-scale Foundation Model for Spatial Transcriptomics

Dear Reviewer tnYZ:

Thanks for your appreciation and detailed review! We try our best to response your questions. Due to space limitations, we include tables and references in the anonymous link https://anonymous.4open.science/api/repo/stofm-rebuttal/file/rebuttal-tnYZ.pdf?v=af96a9f5, and refer to them in the following as Rebuttal-Table X and [X], separately.

Q1: About SToCorpus-88M.

• Is it the largest? To our best knowledge, it is the cell-resolution ST pretraining corpus that contains the largest number of both slices and data points.
• Will it be released? Yes, it will be publicly released.
• Recently-published ST datasets. Thank you! We will add citations to the revised manuscript.

Q2: Lack of clarity and details.

Apologies for the lack of clarity. We promise to add details to the revised version and release detailed code.

• Domain-adaptation Domain adaptation performs an epoch in random order on all 88M data points. The loss had converged at the late stage of the first epoch.
• How are the sub-slice generated exactly? First, we divide the slice into rectangles of the same size. Then, for each rectangle, we merge or split it according to the number of cells it contains. The code will be released for reproducing.
• Number of morphological classes? They are multi-classifacation tasks. We have provided experimental details in the Rebuttal-Table 1.
• How is the ST imputation done? We trained a fully connected neural network with cell representations from SToFM to obtain 50 outputs, as a multi-target regression task.

Q3: Lack of ablations

We will add more ablation studies in the revised version.

• Data efficiency Due to the high computational cost, we pretrain the multi-scale ST representation learning phase using 12.5% and 50% of the data during rebuttal, as shown in Rebuttal-Table 2. The results show that reduction in data volume led to a significant decrease in model performance. Considering that the SToCorpus-88M consists of approximately 2,000 ST slices, we believe that reducing the amount of data may reduce data diversity and limit the model's transferability. Additionally, we would like to clarify that, as mentioned in Section 4.1, we conducted 1 epoch and 3 epochs in each of the two training stages.
• SE(2) Transformer We conduct ablation studies on the PDR loss and the spatial distance matrix, as shown in Rebuttal-Table 3. For more analysis, please refer to the response Q1 to Reviewer LBxH. In addition, the effectiveness and efficiency of the SE(2) Transformer are discussed in detail in Uni-Mol[1]. And when incorporating spatial information in ST data, we primarily focus on the interactions between neighboring cells. Therefore, using the distance-based architecture is in line with the data characteristics.
• Masking ratio Masking ratios are set according to the specific data. Each point in ST data contains high-dimensional features, which increase the data complexity. We had attempted to mask 20% of the cell expressions and perturb another 20% of the cell positions, but found training difficult to converge. In addition, [2] has demonstrated that models are usually robust to mask probabilities when they can converge properly.
• Geneformer initialization If randomly initialized, it will incur expensive computational overhead to train a cell encoder from scratch. And due to the lower quality of gene expressions in ST data compared to scRNA-seq data, our pre-experiments indicate that it is difficult to converge in early training. Considering that Geneformer is one of the SOTA models, and domain adaptation was performed using a large amount of data, we believe there will not be a significant difference if using other SOTA models such as scGPT [3].

Q4: Batch effect correction.

Please refer to the response Q5 to Reviewer Rgcj.

Q5: Catastrophic forgetting.

Catastrophic forgetting may cause the model's suitability for scRNA-seq data to decrease while becoming more suitable for ST data.

• SToFM is a model specialized for ST data, and we do not recommend users to apply SToFM on scRNA-seq data scenarios.
• We conduct an experiment on scRNA-seq data. As shown in Rebuttal-Table 4, there is almost no drop in performance. This may be because the distribution of gene expressions in ST data and scRNA-seq data is similar.

Q6: scGPT in Table 1.

Apologies for the confusion. The table header is "ST Pretraining", and what we want to express is that scGPT did not use gene expressions from ST data during pretraining.

Q7: Leiden cannot control the number of clusters.

After clustering, the number of clusters will be checked. If it does not fall within 20-100, the resolution of Leiden will be adjusted to re-cluster until the cluster number meets the requirement.

---

Thank you again for your insights which are invaluable in solidifying our work. Should our responses address your queries, we would deeply appreciate your support.

Dear Reviewer LBxH:

Thanks for your appreciation and detailed review. We try our best to response the questions below:

Q1: The ablation study is not sufficient to support the effectiveness of the fine-tuning with the incorporation of distance information.

Thank you for your suggestions! We conduct an ablation study on the spatial distance matrix. As shown in the Table below, the results demonstrate that removing the spatial distance matrix significantly decreases model performance. Essentially, the SE(2) Transformer relies on the spatial distance matrix to establish relationships between cells. If spatial distance matrix is removed, the ability of the model to perform intercellular information interactions will be compromised.

Q2: For the Pairwise Distance Recovery (PDR), it is not sure that the cell embedding contains sufficient distance information to make the distance recovery mechanism feasible.

• Intuitively, the spatial autocorrelation of cellular expression profiles [1], as well as the ligand-receptor intercellular signaling pathways of neighboring cells can help to recover the original distances from the noisy data.
• Experimentally, the PDR loss can be converged well during our pre-training process. Moreover, following the discussion in Uni-Mol [2] about recoverability of distances, we only add limited noise to the cell coordinates, which does not obscure all spatial location information. This makes this pre-training task more feasible.

Q3: The proposed model only trained on cell resolution spatial transcriptomics data, and the feasibility of extending to low resolution ST remains unknown.

• SToFM can be applied to low-resolution ST data without domain adaptation. In the DLPFC experiment of Section 4.2, we specifically chose the low-resolution 10x Visium data, and SToFM performed exceptionally well in this task, demonstrating that the model can be scaled to low-resolution ST data.
• Advancements in biotechnology in recent years have continuously improved the resolution of ST data, making the analysis of high-resolution ST data a more valuable direction in the field of bioinformatics [3].

Q4: Not very sure the effectiveness of the way to formulate the virtual cell, i.e., simple by averaging of the representation and location.

• Intuitively, the main purpose of constructing virtual cells is to provide a summary of information from global tissue sections. After clustering by combining expression embedding and location information, each cluster should contain a set of cells that are close in location and have similar expression. Thus, by averaging the embedding and position of the cells within a cluster, it is possible to say "there is a cluster of cells with similar expression at this location". We believe this is methodologically sound. The use of average or sum for simple yet effective pooling is also widely applied in fields such as graph representation learning [4].
• Experimentally, our ablation experiments demonstrate that incorporating macro-scale information through virtual cells can effectively improve model performance. In addition, we calculate the similarity of expression embeddings and spatial positions of some virtual cells with each cell in the cluster, as shown in the Table below. |Sample ID|Pearson correlation of expression embeddings|Cosine Similarity of positions| |-|:-:|:-:| |1|0.863|0.804| |2|0.871|0.833| |3|0.893|0.853|

The results indicate a high similarity between the expression embeddings and positions of the virtual cell and individual cells within the cluster.

---

Thank you again for your detailed review! Your insights have been invaluable in aiding us to enhance and solidify our work. Should our responses satisfactorily address your queries, we would deeply appreciate your support for our work.

Refs:

[1] Mapping the transcriptome: Realizing the full potential of spatial data analysis

[2] Uni-Mol: a universal 3D molecular representation learning framework

[3] Methods and applications for single-cell and spatial multi-omics

[4] Graph pooling in graph neural networks: methods and their applications in omics studies

Dear Reviewer Rgcj:

Thanks for your appreciation and detailed review. We try our best to response the questions below:

Q1: Ablation experiment on micro-scale components. & Q2: The effect of the spatial distance matrix. & Q4: The effect of incorporating spatial information.

Thank you for your suggestions for refining the ablation experiments!

• For Q1, We ablate the micro-scale by limiting the scale of sub-slices to 1, to ensure that the cells could only interact informationally with the virtual cells that represent macro-scale information, but not with other cells in the micro-environment. Experimental results are shown in the Table below. A significant decrease in performance is observed, which demonstrates the effectiveness of incorporating the micro-scale information.
• For Q2, we conducted an ablation study on the spatial distance matrix. The spatial information is very underutilized in this case, only used to construct virtual cells and divide sub-slices. As shown in the Table below, the results demonstrate that removing the spatial distance matrix significantly decreases model performance.
• For Q4, we had presented the results of completely ablating spatial information in the second row "Cell Encoder w/ DA" of Table 5 of our paper. We repeat the results in the Table below. Essentially, we rely on spatial information to establish relationships between multiple cells, with both micro- and macro-scale information being part of spatial information. If spatial information is removed, the model will not be able to interact between cells at all, essentially degrading to independently encoding each single cell.

Q3: The details on the second-time cell encoding mentioned in Sec. 3.2 should be written. If a sample rate is involved in the second-time cell encoding, the sensitivity of sample rate should be explored.

• We will present more details in the revised paper, and will release detailed code to help readers better understand and reproduce our approach.
• We will add experiments and discussion on the sample rate. The purpose of the second time cell encoding is to enable L_MCM and L_PDR to optimize the cell encoder through backpropagation. To balance the training cost and model performance, we use only a small number of cells for this computation, which is similar to selecting a smaller batch size to update the cell encoder. For the selection of the sampling number, we determined the number of samples to be 12, given that the original Geneformer paper gave a training batch size of 12.
• Considering the computational cost, we test the impact of the sample number on a small amount of data (1/8 of the SToCorpus-88M), as shown in the Table below. The results show that the model has some robustness to this hyperparameter, just as the batch size often only affects the convergence speed rather than the model performance. However, setting the sample size to 0, i.e. freezing the cell encoder, will lead to a decrease in model performance. |Sample number|Embryo2-F1| EmbryoCross-F1| |-|:-:|:-:| |0|0.722|0.417| |4|0.754|0.424| |12|0.758|0.423|

Q5: This paper migrates the single-cell model through a domain adaptation strategy. However, there may be relatively large differences between different ST datasets, such as different technical platforms. How should we further improve adaptability?

• Research such as scGPT [1] and LangCell [2] has already proven that large-scale pretraining is one of the best ways to remove batch effects in scRNA-seq data. Nicheformer [3] has also demonstrated that by pretraining, the model can gain modeling capabilities across different ST technology platforms. Indeed, despite the gap between different datasets, the underlying gene co-expression, intercellular signaling pathways, and other information should be largely uniform, which makes them have similar distributions that can be captured by the pre-trained model.
• In the DLPFC experiments in Section 4.2, we purposely chose 10x Visium data that was not used in the pretraining. The excellent performance of SToFM on this task also proves that the model has gained the ability to transfer across technology platforms from the pretraining. We believe that our model and dataset can be of great help for future research in the application scenario of batch integration.

---

Thank you again for your detailed review! Your insights have been invaluable in aiding us to solidify our work. Should our responses address your queries, we would deeply appreciate your support for our work.

Refs:

[1] scGPT: toward building a foundation model for single-cell multi-omics using generative AI

[2] LangCell: language-cell pre-training for cell identity understanding

[3] Nicheformer: a foundation model for single-cell and spatial omics

Dear Reviewer TZDp:

Thanks for your appreciation and detailed review. We try our best to response the questions below:

Q1: Bridging the gap between ST and scRNA-seq data through transfer learning.

ST data consists of two parts: spatial location and gene expression values. As SToFM is a model for ST data, the purpose of domain adaptation is to utilize well-trained scRNA-seq model to better encode the gene expression values of ST data. The gene expression values of ST data have a similar distribution to scRNA-seq data and follow the same underlying gene co-expression patterns. Many well-established bioinformatics methods have also proven the effectiveness of transferring knowledge between scRNA-seq and ST data, such as Tangram[1]. Therefore, we believe it is reasonable to apply transfer learning between scRNA-seq data and the gene expression values of ST data.

Q2: Since this model requires pretraining of the cell encoder, its scalability is limited.

Research like scGPT[2] has already shown that Transformer-based cell encoder models have good scalability. We believe that a well-pretrained cell encoder can help improve the model's scalability.

Q3: Computational costs. Lack of running time experiments and time complexity analysis.

• The cell encoder is based on the Transformer architecture. And the time complexity of the SE(2) Transformer is similar to that of a standard Transformer [3]. Therefore, the time complexity of both parts of SToFM is the same as the standard Transformer. Specifically, the time complexity is $O(N *n^2+M *m^2)$, where $N$ and $n$ are the number of cells and genes, and $M$ and $m$ are the number and the scale of the sub-slices, respectively.
• Section 4.2 of our paper provides details on the pretraining time cost. We have also added runtime experiments for inference, as shown in the Table. Specifically, running inference on an ST slice containing tens of thousands of cells typically takes 1–5 minutes, which we consider acceptable for practical applications.

Q4: Model's sensitivity to different alpha values.

Thank you for your suggestion! We have added experiments to show how different alpha values affect the model’s performance, as shown in the Table below:

The model has a certain robustness in alpha, and we believe this may be because cells that are closer in location are more likely to have similar gene expressions [4]. The alpha=0.8 that we chose is essentially the optimal setting.

Q5: The impact of the combination ratio of the two loss functions.

We will introduce further analysis on the loss ratio in the revised version of the paper. These two loss are relatively close in scale, and in our experiments, combining them in a 1:1 ratio allows both to converge normally. Considering the computational cost, we test the impact of the ratio of the two losses on the speed of convergence and the performance of the model on a small amount of data (1/8 of the SToCorpus-88M), as shown in Rebuttal-Fig.1 (https://anonymous.4open.science/api/repo/stofm-rebuttal/file/rebuttal-TZDp.pdf?v=5b2ab9f4) and the Table below. ($\gamma$ is the loss ratio in $L=\gamma*L_{MCM} + (1- \gamma) * L_{PDR}$)

The results show that the model's performance is robust to the ratio of loss functions. However, removing one of the loss functions will lead to a decrease in model performance. The $\gamma$=0.5 that we chose is essentially the optimal setting.

Q6: The provided code does not run.

We had uploaded the code of the model and core algorithms along with the paper to help better understand the method. The executable code has been organized now, but anonymous links containing code are not allowed during rebuttal. We promise to release the executable code on github after the paper is published.

---

Thank you again for your detailed review! Your insights are invaluable in aiding us to solidify our work. Should our responses satisfactorily address your queries, we would deeply appreciate your support for our work.

Refs:

[1] Deep learning and alignment of spatially resolved single-cell transcriptomes with Tangram

[2] scGPT: toward building a foundation model for single-cell multi-omics using generative AI

[3] Uni-Mol: a universal 3D molecular representation learning framework

[4] Mapping the transcriptome: Realizing the full potential of spatial data analysis