Paper: CellPLM: Pre-training of Cell Language Model Beyond Single Cells (ICLR-2024)

Reviewer comment (Theory):
The motivation and benefit of using a VAE-like architecture is not clear.

Author response:
We appreciate your inquiry regarding the use of a VAE-like architecture. The choice was motivated by the need for a robust framework capable of capturing the complex, high-dimensional and noisy nature of single-cell data. The VAE architecture offers a powerful way to learn meaningful latent representations, modeling the uncertainty from the measurement, enabling the model to effectively capture underlying biological variability and noise inherent in single-cell data. As shown in the ablation studies in Table 1 above, a proper latent distribution results in more robust cell representation and better downstream performance. Due to the same reason, VAE has been a well adopted architecture in recent single-cell analysis research on various tasks[1,2,3,4,5].

Reviewer comment (Evaluation):
Two paradigms of modeling single-cell data should both have their pros and cons but the authors do not discuss them thoroughly.

Author response:
CellPLM, with its unique approach to modeling scRNA-seq and SRT data, excels in tasks that require deep understanding of cellular heterogeneity and cell-cell relationships, such as in cell type annotation, imputation and spatial transcriptomic tasks. It also achieves impressive inference speed when generating cell embeddings. However, for tasks that explicitly requires gene-level output, CellPLM might not be capable. For example, CellPLM is not available for gene regulatory network inference, while gene-level pretrained models (e.g., scBert, scGPT, etc.) can do. Although CellPLM can implicitly capture gene interactions through gene expression embedder and encoders (as illustrated in gene embedding visualization in Appendix H.2 and gene perturbation prediction in Appendix F.3), it’s hard to query the contextualized gene representations from the model.

Reviewer comment (Experiments):
The benefit of using Gaussian mixture distribution as the prior of latent embedding is not validated.

Author response:
You concern is well-taken. As verified in the ablation studies above, the Gaussian mixture prior distribution can help model the uncertainty of the data and address the high noise inherent in transcriptomic data. It outperforms a regular Gaussian due to its capacity of handling heterogenous pre-train data. This eventually results in more robust cell representation and better downstream performance.

Reviewer comment (Evaluation):
Where do the huge speed improvements of CellPLM come from? It is because CellPLM aggregates the genes within each cell directly in the embedding layer, so it only needs to process dramatically shorter input sequences. So basically, CellPLM works at the cell level, while baselines work at the gene level. If this is true, the Claim "with 500x times higher inference speed compared to existing pre-trained models." is unfair since you are compared to a task that your model is specifically designed for. At least it should be reduced to "with 500x times higher inference speed when generating cell embeddings".

Author response:
Thank you for your great suggestion. We agree that the original claim is not rigid enough. We now change it to the suggested expression in the revised paper, i.e., “with 500x times higher inference speed on generating cell embeddings compared to existing pre-trained models”. Meanwhile, we want to emphasize that the cell embeddings are essential for various downstream tasks in single-cell analysis. For example, all the downstream tasks (clustering, denoising, cell-type annotation, imputation, gene perturbation prediction) presented in the paper rely on high-quality cell representations. Therefore, the increased inference speed in generating cell embeddings should still be considered as a highlight of CellPLM model.

Reviewer comment (Evaluation):
Why does the CellPLM model from scratch outperform most of the pre-trained baselines? This seems to be really unrealistic to me. Is it because your model is much larger than the baselines or cell-level LM is more suitable for the tasks? If the randomly initialized model already outperforms most of the baselines, then we need to re-evaluate the claims in the paper, since the performance boost may mainly result from the Transformer architecture instead of the specific model designs you mentioned.

Author response:
We appreciate your question about the performance improvement. First, we want to clarify that CellPLM does not always outperform baselines when training from scratch. For example, in spatial imputation task (both datasets) and cell type annotation task (hPancreas dataset), it cannot outperform baselines. 

To further illustrate how CellPLM benefits from each component, we add an ablation study on the from-scratch version on MS dataset of cell type annotation task, presented in Table 1 below. We find that both the bi-level masking module and the transformer encoder positively contribute to the from-scratch performance. When removing these modules, CellPLM performance is closer to the non-pretrain baselines (e.g., CellTypist). Since CellPLM is the first method that applies bi-level masking and cell-level transformer encoders to these tasks, this can be considered as our contribution. Lastly, the gap between pre-trained and from-scratch is still considerably larger than the gap between from-scratch versions. This sufficiently demonstrates the benefits brought by pre-training and transfer, proving the effectiveness of the pre-training framework we proposed and laying a foundation for future downstream tasks. This is indeed our most essential contribution.

|  | Macro F1 | Macro Precision |
| --- | --- | --- |
| CellPLM (pre-trained) | 0.766 $\pm$ 0.007 | 0.803 $\pm$ 0.008 |
| CellPLM (from scratch) | 0.709 $\pm$ 0.007 | 0.732 $\pm$ 0.015 |
|    w/o masks | 0.683 $\pm$ 0.013 | 0.703 $\pm$ 0.027 |
|    w/o transformer | 0.685 $\pm$ 0.012 | 0.701 $\pm$ 0.028 |
| CellTypist | 0.667 $\pm$ 0.002 | 0.693 $\pm$ 0.001 |

Table 1. Ablation studies of mask technique and transformer encoder on MS dataset, in training from scratch scenario.

Reviewer comment (Experiments):
There is a lack of ablation study to show the effectiveness of each proposed component.

Author response:
We appreciate your advice. We now add three new ablation studies on two representative tasks to demonstrate the effectiveness of proposed latent distribution and transformer encoder. In each setting, we change one component in the model architecture and go through the whole pre-train and fine-tune pipeline to get the downstream performance.  The results are presented in Table 2 below.

- First, when we replace the proposed mixture of gaussian prior distribution with a gaussian prior distribution (noted as “w/o Mixture of Gaussian”, commonly used in previous methods like scVI), the performance significantly drops on all datasets, indicating that an unsuitable prior distribution can greatly hurt the performance. A regular Gaussian distribution cannot accommodate the highly heterogeneous data present in the pre-train dataset, which were collected from different people, organs, and sequencing platforms.
- Second, we removed the latent distribution in its entirety, noted as “w/o latent distribution”, i.e., we converted from a VAE-like probabilistic generative model to a deterministic masked auto-encoder. The performance consistently falls between the original one and the first ablation. On one hand, this supports our motivation of using probabilistic models with Gaussian mixture prior distribution. The latent distribution helps model the uncertainty of the data and address the high noise inherent in transcriptomic data. This results in a robust cell representation. On the other hand, the selection of prior distribution is very important because an improper prior (e.g., regular Gaussian) can be even worse than no latent distribution.
- Lastly, we replace the transformer encoder with an MLP encoder, keeping the same number of layers and hidden dimension (the total parameters reduce from 85M to 50M). The performance significantly drops on spatial imputation task, while the gap is relatively small on cell-type classification task. This aligns with our intuition, as spatial transcriptomic data provide spatial location information, enabling the model to better identify and utilize the relationships between cells. In contrast, the cell type annotation dataset does not provide spatial location information, which makes the benefits gained from the transformer encoder more limited.

Overall, through a series of ablation studies, we have verified that our CellPLM model can capture the relationships between cells via the transformer encoder and enhance the performance of downstream tasks, generating more robust and useful cell representations through appropriate prior distributions.

||Cell-type Classification||||Spatial Imputation||||
|---|---|---|---|---|---|---|---|---|
||MS||hPancreas||Lung||Liver||
||Macro F1|Macro Precision|Macro F1|Macro Precision|corr|cosine|corr|cosine|
|CellPLM|0.766$\pm$0.007|0.803$\pm$0.008|0.749$\pm$0.010|0.753$\pm$0.010|0.318$\pm$0.015|0.481$\pm$0.011|0.328$\pm$0.011|0.481$\pm$0.010|
|w/o Mixture of Gaussian|0.737$\pm$0.042|0.766$\pm$0.069|0.711$\pm$0.025|0.701$\pm$0.025|0.258$\pm$0.011|0.449$\pm$0.005|0.232$\pm$0.013|0.433$\pm$0.008|
|w/o Latent Distribution|0.750$\pm$0.024|0.809$\pm$0.032|0.733$\pm$0.034|0.731$\pm$0.033|0.262$\pm$0.011|0.449$\pm$0.008|0.246$\pm$0.017|0.428$\pm$0.012|
|w/o Transformer Encoder|0.750$\pm$0.050|0.794$\pm$0.074|0.751$\pm$0.010|0.750$\pm$0.012|0.244$\pm$0.016|0.443$\pm$0.008|0.250$\pm$0.032|0.440$\pm$0.021|
---
Table 2. Ablation studies of latent distribution and transformer encoder on cell-type classification and spatial imputation task.

Reviewer comment (Evaluation):
Although the improvement in experimental results is very significant, further analysis is needed to determine whether it is due to the relationships between cells.

Author response:
Thank you for the suggestion. As verified in the ablation studies in Table 1 above, the model does leverage the cell-cell relations via the transformer encoder. With the help of transformer encoder, on cell classification tasks, the F1 score increases from 0.75 to 0.766 on MS dataset, and on spatial imputation tasks, the correlation increases from 0.244 to 0.318 on Lung dataset.

Reviewer comment (Experiments):
I'm not sure if the Gaussian distribution can achieve the expected effect in this method, this point also needs to be validated through experiments.

Author response:
Thank you for your suggestion. As verified in the ablation studies in Table 1 above, the Gaussian mixture prior distribution is more suitable than a regular Gaussian prior distribution, which can help model the uncertainty of the data and address the high noise and heterogeneity inherent in transcriptomic data. This eventually results in more robust cell representation and better downstream performance. In the experiments, when pre-training with Gaussian mixture prior distribution, the F1 score increases from 0.75 to 0.766 on MS dataset on cell classification task, and the correlation increases from 0.262 to 0.318 on Lung dataset on spatial imputation tasks.

Reviewer comment (Evaluation):
The source of the pre-training gains is unclear, given the current presentation.

Author response:
Thank you for highlighting the need for clarity in our data processing. We appreciate the opportunity to provide more details. In our study, the subset of 13,500 genes was resulted from the intersection of all scRNA-seq pre-train datasets, based on their prevalence and relevance in the existing literature. Most of these genes are protein-coding genes. A more detailed list of datasets is now included in Appendix E.2.  All data are pre-filtered (i.e., quality control) by the original authors and are collected in AnnData objects. We further perform normalization and log transformation following the convention in Seurat[1]. This important detail is now added to Section 3.2 and Appendix E.2.

Reviewer comment (Experiments):
The authors restrict the gene set to a subset of 13,500 genes, but it is not specified which genes are included or whether they are protein-coding or non-protein-coding.

Author response:
Your question about cell masking is insightful. In our approach, cell masking refers to the process of masking at the gene level within each cell, same as the “bi-level masking technique” proposed in [2]. The core idea is to mask most genes (e.g., 75%) in a subset of cells (e.g., 25%), forcing the model to leverage information from other cells to impute the missing information. Therefore, your understanding of “other cell genes have the necessary information to impute the cell” is totally correct.
    
**Our method is not deemed “cheating” as it uses the same amount of input as other models and does not use any label information from the test set.**  In fact, leveraging intercellular information is an essential idea of our proposed “cell language modeling”. In addition to utilizing intracellular partially observed gene expression data, our model is also trained to identify and exploit similar cells within the same tissue, cellular communication, and spatial neighborhood composition to impute indeed information. This capability is essential for better denoising and spatial-resolved transcriptomic applications, and previous single-cell pre-trained models are not capable of doing this.

Reviewer comment (Experiments):
**Cell masking:** The authors aggregate cell genes using their embeddings and pass these tokens to the LLM. What does it mean to mask a cell? In the diagram, the cells are masked, but in the equations, the genes are masked. Assuming the masking is at the gene level and samples are taken from the same batch, wouldn't other cell genes have the necessary information to impute the cell? In essence, is the model cheating?

Author response:
For a single cell, the model indeed operates by feeding a single cell token into the model, utilizing its latent representation (the performance is shown in Table 1 below, in our response to “W4 Sequencing Length”). This process is similar to handling larger samples but adapted to maximize information extraction from a single-cell context. However, our model is more suitable for batch processing of data, which is in line with the characteristics of high-throughput sequencing data.

Reviewer comment (Reproducibility):
**Single-cell operation:** How would the model operate if you only have a single-cell sample? Would you feed a single cell token to the LLM and use the latent representation?

Author response:
For a single cell, the model indeed operates by feeding a single cell token into the model, utilizing its latent representation (the performance is shown in Table 1 below, in our response to “W4 Sequencing Length”). This process is similar to handling larger samples but adapted to maximize information extraction from a single-cell context. However, our model is more suitable for batch processing of data, which is in line with the characteristics of high-throughput sequencing data.

Reviewer comment (Experiments):
**Sequence length:** How would the model performance change if you vary the "sequence length"? For example, in your batch, do you have the same number of genes for each cell?

Author response:
The question of varying sequence length is pertinent. In our framework, the number of genes per cell can be variable.  This is managed by the gene expression embedder, which is implemented by a sparse aggregation operation, and thus is flexible to any input gene sets (within the pretrain gene set). 

However, we want to emphasize that the “sequence length” in our transformer model does not refer to the number of genes but the number of cells. We treat cells as tokens and use the transformer model to process a batch of cells. Therefore, the “sequence length” in CellPLM is equivalent to “batch size” in traditional deep learning models. To evaluate the impact of “sequence length” (a.k.a, “batch size”), we separate the input dataset (MS dataset for cell-type annotation task) into smaller batches and run a same fine-tuned cell-type annotation model for inference. The result is presented in the table 1 below. From the experiments, the performance of the model is not sensitive to batch size, but too small a batch size can lead to slight performance degradation. 

|Batch Size|32768|2048|128|32|16|8|4|2|1|
|---|---|---|---|---|---|---|---|---|---|
|Macro F1| 0.773 | 0.773 | 0.773 | 0.773 | 0.770 | 0.768 | 0.769 | 0.767 | 0.768 |
---
Table 1. Parameter analysis of batch size on MS dataset for cell-type annotation task.

Reviewer comment (Experiments):
**Spatial information:** How is the spatial information leveraged here? Is it through the positional encoding? What would happen if this positional encoding is removed?

Author response:
Spatial information in our model is indeed leveraged through positional encoding, with the help of transformer encoders. We appreciate the need to clarify its impact. To address this, we conducted additional experiments to fine-tune the same pre-trained model with and without positional encoding on spatial imputation task, the results of which are presented in the Table 2 below.

||Lung||Liver||
|---|---|---|---|---|
||corr|cosine|corr|cosine|
|CellPLM|0.318$\pm$0.015|0.481$\pm$0.011|0.328$\pm$0.011|0.481$\pm$0.010|
|w/o Positional Encoding|0.237$\pm$0.027|0.442$\pm$0.007|0.215$\pm$0.029|0.424$\pm$0.016|
---
Table 2. Ablation study of positional encoding on spatial imputation task.

This ablation experiment demonstrates the significance of positional encoding in enhancing the performance on spatial transcriptomic data, indicates that the model can leverage the spatial information from positional encoding.

Reviewer comment (Experiments):
**Positional Encoding:** It is not clear if this had a significant improvement in the quality of representation learnt.

Author response:
Spatial information in our model is indeed leveraged through positional encoding, with the help of transformer encoders. We appreciate the need to clarify its impact. To address this, we conducted additional experiments to fine-tune the same pre-trained model with and without positional encoding on spatial imputation task, the results of which are presented in the Table 2 below.

||Lung||Liver||
|---|---|---|---|---|
||corr|cosine|corr|cosine|
|CellPLM|0.318$\pm$0.015|0.481$\pm$0.011|0.328$\pm$0.011|0.481$\pm$0.010|
|w/o Positional Encoding|0.237$\pm$0.027|0.442$\pm$0.007|0.215$\pm$0.029|0.424$\pm$0.016|
---
Table 2. Ablation study of positional encoding on spatial imputation task.

This ablation experiment demonstrates the significance of positional encoding in enhancing the performance on spatial transcriptomic data, indicates that the model can leverage the spatial information from positional encoding.

Reviewer comment (Experiments):
**Comparison with previous models:** Although the authors compare with previous single-cell models, these models only operated on gene expression data and fed a single cell at a time (i.e., the sequence dimension is the genes within the cell as opposed to other cells, as done in this work).

Author response:
We thank you for pointing out the distinction between CellPLM and existing models. As we introduced in Section 2, the existing pre-trained models all focus solely on modeling gene relationships within individual cells, neglecting the intercellular information in an organism. CellPLM is the first pre-trained model that can leverage this important information inherent in the high-throughput sequencing data and spatially-resolved data. We highlighted this difference in Figure 1 in the paper. 

In addition, cell-cell relations have been leveraged in non-pretrained baselines. For example, graph-based models (e.g., GraphSCI and scGNN, stPlus, SpaGE) in denoising and imputation tasks explicitly construct **cell-cell nearest-neighbor graphs** to leverage such information. Compared to those methods, CellPLM adopts a transformer to implicitly model the cell-cell relations.

Reviewer comment (Experiments):
**Benefit of pre-training:** The benefit of which part of the pre-training is the result of the improvements shown in the result section is not clear from the current presentation. Also, the authors don't compare with scVI (although they do cite the method) or even a simpler baseline like HVG.

Author response:
Clarifying the specific benefits of pre-training is crucial, and we thank you for emphasizing this. The direct evidence of the benefit of **knowledge transfer** from pre-training is the gap between from-scratch performance and fine-tune performance as we presented in Table 2, 3 and 4. On all tasks and all datasets, pre-train fine-tune outperforms training CellPLM models directly on downstream tasks (a.k.a, “from scratch”). 

To further verify the contribution of each **architectural component** in CellPLM model, we have now added three new ablation studies on two representative tasks to examine the effectiveness of proposed latent distribution and transformer encoder. In each setting, we change one component in the model architecture and go through the whole pre-train and fine-tune pipeline to get the downstream performance. The results are presented in Table 3 below and added into Appendix I.

- First, when we replace the proposed mixture of gaussian prior distribution with a gaussian prior distribution (noted as “w/o Mixture of Gaussian”, commonly used in previous methods like scVI), the performance significantly drops on all datasets, indicating that an unsuitable prior distribution can greatly hurt the performance. A regular Gaussian distribution cannot accommodate the highly heterogeneous data present in the pre-train dataset, which were collected from different people, organs, and sequencing platforms.
- Second, we removed the latent distribution in its entirety, noted as “w/o latent distribution”, i.e., we converted from a VAE-like probabilistic generative model to a deterministic masked auto-encoder. The performance consistently falls between the original one and the first ablation. On one hand, this supports our motivation of using probabilistic models with Gaussian mixture prior distribution. The latent distribution helps model the uncertainty of the data and address the high noise inherent in transcriptomic data, which results in a robust cell representation. On the other hand, the selection of prior distribution is very important because an improper prior (e.g., regular Gaussian) can be even worse than no latent distribution.
- Lastly, we replace the transformer encoder with an MLP encoder, keeping the same number of layers and hidden dimension (the total parameters reduce from 85M to 50M). The performance significantly drops on spatial imputation task, while the gap is relatively small on cell-type classification task. This aligns with our intuition, as spatial transcriptomic data provide spatial location information, enabling the model to better identify and utilize the relationships between cells. In contrast, the cell type annotation dataset does not provide spatial location information, which makes the benefits gained from the transformer encoder more limited.

||Cell-type Classification||||Spatial Imputation||||
|---|---|---|---|---|---|---|---|---|
||MS||hPancreas||Lung||Liver||
||Macro F1|Macro Precision|Macro F1|Macro Precision|corr|cosine|corr|cosine|
|CellPLM|0.766$\pm$0.007|0.803$\pm$0.008|0.749$\pm$0.010|0.753$\pm$0.010|0.318$\pm$0.015|0.481$\pm$0.011|0.328$\pm$0.011|0.481$\pm$0.010|
|w/o Mixture of Gaussian|0.737$\pm$0.042|0.766$\pm$0.069|0.711$\pm$0.025|0.701$\pm$0.025|0.258$\pm$0.011|0.449$\pm$0.005|0.232$\pm$0.013|0.433$\pm$0.008|
|w/o Latent Distribution|0.750$\pm$0.024|0.809$\pm$0.032|0.733$\pm$0.034|0.731$\pm$0.033|0.262$\pm$0.011|0.449$\pm$0.008|0.246$\pm$0.017|0.428$\pm$0.012|
|w/o Transformer Encoder|0.750$\pm$0.050|0.794$\pm$0.074|0.751$\pm$0.010|0.750$\pm$0.012|0.244$\pm$0.016|0.443$\pm$0.008|0.250$\pm$0.032|0.440$\pm$0.021|
---
Table 3. Ablation studies of latent distribution and transformer encoder on cell-type classification and spatial imputation task.

In addition, the result of **scVI** is now added in the denoising task (Table 2 in Section 4.2) using `get_normalized_expression` function of scVI package and the performance of scVI is very close to DCA. Meanwhile, we want to clarify that the PCA baseline in zero-shot clustering experiments (in Section 4.1) indeed come from **HVG**. As we mentioned in Section 4.1, the PCA results are obtained from 4,500 highly variable genes (HVG). We also include a new clustering result with **scVI** on the same dataset. It achieves ARI=0.843, NMI=0.823, slightly outperforming HVG+PCA baseline but still worse than CellPLM (zero-shot). The visualization is added in the Appendix H.1.

Reviewer comment (Evaluation):
**Gene interactions:** Does the model learn gene level interactions? By feeding cell-level representations, how can this be assessed?

Author response:
Thank you for the great question about gene-level interactions. Our model, despite feeding cell-level representations, is designed to capture gene-level interactions through its inherent architecture. The gene functionality is first captured by gene embeddings in the expression embedder module, and then interacted through the feed-forward networks in the transformer encoder, coupled with cell-cell interaction.
    
To evaluate whether the model learns meaningful gene representations, we provide a new visualization of gene embeddings in appendix H.2. From the visualization, the gene embeddings maintain some latent structures. As a case study, we highlight a specific family of genes, HLA genes. There are multiple classes of HLA genes[3]. For example, HLA class I genes (e.g., HLA-A, -B, and -C) present endogenous peptides to responding CD8+ T Cells while the class II (e.g., HLA-DR, -DP, and –DQ) process exogenous peptides for presentation to CD4+ helper T Cells. From the UMAP visualization, HLA gene embeddings present clusters that perfectly match the functionality and characteristics of those genes. To further assess whether the model implicitly encode gene-gene interactions, we conduct experiments of gene perturbation prediction (presented in appendix F.3), which demonstrates that our model is able to identify the knock-out effect of specific genes.