Title: xTrimoGene: An Efficient and Scalable Representation Learner for Single-Cell RNA-Seq Data
Venue: NeurIPS 2023
Accepted: yes
Keywords: pre-train

encoder-decoder

scRNA-seq

scalable
Abstract: Advances in high-throughput sequencing technology have led to significant progress in measuring gene expressions at the single-cell level. The amount of publicly available single-cell RNA-seq (scRNA-seq) data is already surpassing 50M records for humans with each record measuring 20,000 genes. This highlights the need for unsupervised representation learning to fully ingest these data, yet classical transformer architectures are prohibitive to train on such data in terms of both computation and memory. To address this challenge, we propose a novel asymmetric encoder-decoder transformer for scRNA-seq data, called xTrimoGene$^\alpha$ (or xTrimoGene for short), which leverages the sparse characteristic of the data to scale up the pre-training. This scalable design of xTrimoGene reduces FLOPs by one to two orders of magnitude compared to classical transformers while maintaining high accuracy, enabling us to train the largest transformer models over the largest scRNA-seq dataset today. Our experiments also show that the performance of xTrimoGene improves as we scale up the model sizes, and it also leads to SOTA performance over various downstream tasks, such as cell type annotation, perturb-seq effect prediction, and drug combination prediction. 
xTrimoGene model is now available for use as a service via the following link: https://api.biomap.com/xTrimoGene/apply.

xTrimoGene: An Efficient and Scalable Representation Learner for Single-Cell RNA-Seq Data

Abstract

Advances in high-throughput sequencing technology have led to significant progress in measuring gene expressions at the single-cell level. The amount of publicly available single-cell RNA-seq (scRNA-seq) data is already surpassing 50M records for humans with each record measuring 20,000 genes. This highlights the need for unsupervised representation learning to fully ingest these data, yet classical transformer architectures are prohibitive to train on such data in terms of both computation and memory. To address this challenge, we propose a novel asymmetric encoder-decoder transformer for scRNA-seq data, called xTrimoGene α (or xTrimoGene for short) 4 , which leverages the sparse characteristic of the data to scale up the pre-training. This scalable design of xTrimoGene reduces FLOPs by one to two orders of magnitude compared to classical transformers while maintaining high accuracy, enabling us to train the largest transformer models over the largest scRNA-seq dataset today. Our experiments also show that the performance of xTrimoGene improves as we scale up the model sizes, and it also leads to SOTA performance over various downstream tasks, such as cell type annotation, perturb-seq effect prediction, and drug combination prediction. xTrimoGene model is now available for use as a service via the following link: https://api.biomap.com/xTrimoGene/apply .

1 Introduction

Recently, Artificial Intelligence (AI) technology has demonstrated promising results for addressing scientific problems. This AI4Science paradigm witnessed diverse successful biological and pharmaceutical applications, including protein analysis [17, 20, 38, 35, 1], RNA modeling [4], and genomics modulation [27]. However, most existing AI models have predominantly focused on protein sequences, neglecting the growing volume of high-throughput experimental sequencing data in the form of gene expression values. Single-cell RNA sequencing (scRNA-seq) technology has transformed the field of cell biology and enabled us to understand cell-cell, cell-gene and gene-gene relations at the cellular level [16, 3]. This technique captures the expression levels of thousands of genes in parallel, facilitating the study of cellular heterogeneity [5, 19]. This unveiled information is crucial for understanding complex biological systems and disease progression [16, 3]. Integrating and

modeling such large-scale scRNA-seq data can reveal rich cellular information and benefit various biological task learning.

Representation learning from scRNA-seq data [9] has been an active area of research in past decades. For example, scVAE [11] and scVI [21] apply a variational autoencoder framework to derive lowdimensional cell embeddings, cscGAN [24] uses a Generative Adversarial Network (GAN) architecture to generate cell-type specific expression profiles, and SA VER-X [32] is capable of removing batch effects across datasets. Despite the success of these customized algorithms, they tend to be computationally inefficient and labor-intensive. This prompts us to explore a general-purpose model that first learns underlying knowledge from scRNA-seq data and generalizes it to different tasks in a unified manner. We draw inspiration from the pre-training and fine-tuning paradigm in Natural Language Processing (NLP), which has shown great success in improving various downstream NLP task performance [29, 12, 14]. In light of these findings, we aim to investigate the potential of applying similar approaches to representation learning in scRNA-seq data.

The first published pre-trained model for single-cell data is scBERT, which uses a low-rank transformer [36] to analyze the scRNA data. It learns the cellular representation by randomly masking a percent of non-zero gene expression values and tries to recover them. scBERT has achieved state-of-the-art results for cell-type annotation tasks. The study shows the potential of a pre-training strategy for single-cell biology research. However, scBERT has certain limitations in fully utilizing scRNA-seq data properties. These limitations include:

(1) Scalability. The large number of genes (almost 20,000) and the sparsity of scRNA-seq data, with nearly 90% of values being zero, lead to many redundant computations (e.g., self-attentions between zero tokens). It required approximately 2 . 65 × 10 19 FLOPs to train 5 million samples over 5 epochs, which equals almost 20 days of training on an A100 GPU for only an 8.9 million parameter scBERT model. (2) Limited resolution for expression values. scBERT rounds the gene expression values into integer values, which limits the model's ability to distinguish closeness and similarity between gene expression values. For instance, two close values could be mapped to separate embeddings (e.g., 1.99 and 2.01 are mapped to 1 and 2), and two distant values could be mapped to identical embeddings (e.g., 1.99 and 1.01 are mapped to 1). The strategy leads to a loss of resolution and introduces bias during model training, resulting in sub-optimal performance.

To address the challenges associated with scRNA-seq data modeling and consider the unique nature of this data (as discussed in Section 2), we present a novel and efficient framework, xTrimoGene, for pre-training large-scale scRNA-seq data. Our framework makes the following key contributions:

(1) We design an asymmetrical encoder-decoder architecture to guide the pre-training process, which enables us to learn a high-capacity model for single-cell RNA-seq data. Our model achieves an improvement in the speed of pre-training of over 3 times compared to previous encoder-only models.

(2) We illustrate that the efficiency and scalability of our model allow us to train the largest single-cell pre-trained model to date, with approximately 100 million parameters for the xTrimoGene-100M model, using a curated scRNA-seq dataset of approximately 50 billion effective gene tokens.

(3) The pre-trained model xTrimoGene achieved remarkable results in multiple downstream tasks, including cell type annotation, perturbation prediction and synergistic drug combination prediction.

2 Characteristics of Single-Cell RNA-seq Data

scRNA-seq generates a large, sparse expression matrix, where each row represents a cell (sample) and each column a gene (feature). This dataset presents several challenges and requires a specialized architecture to effectively model the data.

First, approximately 20,000 genes (columns) are shared across cells. Unlike the corpus in NLP, the genes can be arbitrarily reordered. The relation between genes depends on biological pathways rather than local contexts, where the latter shapes spatial information in Computer Vision (CV) images. Though one can roughly regard each cell (row) as a sentence or an image patch, the 20,000 genes is a vast number compared to the typical sequence length, which is mostly a few hundred and no more than a few thousand [29, 12]. Thus directly applying existing transformer architecture will not work.

Second, scRNA-seq matrices are highly sparse (90% zero in a typical dataset [15, 7]). The abundance level of RNA for each gene is measured by counting the unique molecular identifier (UMI) reads in

scRNA-seq experiments [16, 3]. However, many genes exhibit low UMI counts due to limited probing efficiency. Therefore, treating scRNA-seq data as an image and utilizing a convolutional neural network to extract features is not feasible, as it introduces a huge number of redundant computations for sparse positions.

Third, the normalized gene expression values in scRNA-seq data are continuous scalars, which typically indicate similar gene activity when they have similar values. To transform these scalars into high-dimensional tokens in the data matrix, a representation that preserves the continuous semantics is needed. Manually discretizing the gene expression values is challenging as non-optimal discretization thresholds will bias category assignment. A learned discretization approach or learnable representation, such as the one proposed in [10], is ideal for preserving the continuous semantics of the gene expression values.

Taking into the above three major features, we design a new architecture as described in the next section.

(1)

Masking

(2)

Filtering

(3)

Padding

(4)

Embedding

(9) Compute loss

1

2

3

C-2

C-1

C

...

1

2

3

19262 19263 19264

...

+

Attention

decoder

Attention

encoder

(6) Encoder

(7)

+ masked emb

+ zero emb

(8) Decoder

(5) Add

3 xTrimoGene Architecture

xTrimoGene is a highly efficient framework for pre-training large-scale single-cell RNA-seq data (illustrated in Figure 1). The training process is based on a regression-masked task, aimed at accurately recovering masked values in the expression matrix. Notably, a specific optimized asymmetrical encoder-decoder framework is employed to accelerate the learning of sparse matrices. This is achieved by feeding only the unmasked non-zero positions (less than 10% of the full length) into the encoder, while the largely masked and zero positions are input into a lightweight decoder with a reduced number of layers and attention heads. In addition, a novel auto-discretization strategy is introduced to project continuous expression values into a latent embedding space. Instead of rounding to the nearest integer, values are directly mapped to the latent space allowing for the representation of closely related values. The xTrimoGene framework consists of the following components:

Masking : A portion of the normalized gene expression matrix V is masked for prediction, including both zero and non-zero positions. c denotes cell sample size, and n denotes gene number (19,264 in our setting, see App. 1 for data collection and processing).

Filtering : The masked and zero-valued embeddings are filtered out, yielding a variable-length sequence of valuable information that is prepared for encoding.

Padding : The remaining unmasked positions are aligned with padding tokens, resulting in a much smaller unmasked-only matrix V unmasked . m denotes the maximum length of the unmasked sample. We include a scheme to illustrate the processing flow (see App. 2).

Embedding : Expression value and gene embeddings are separately projected. d denotes the dimension of the embedding. The expression embedding is calculated through an auto-discretization mapping. The gene embedding is retrieved from a randomly initialized lookup table.

Combining Expression and Gene Embeddings : The expression and gene embeddings ( E and G ) are element-wise added to form the input embedding, which is then fed into the encoder of the model.

Encoding : The sum of the embeddings is input into the encoder, which implements self-attention mechanisms using a Transformer-like architecture.

Extending masked and zero embeddings : The intermediate encoder embedding I encoder is combined with embeddings for masked and zero-value positions.

Decoding : The combined embeddings are processed by the decoder, utilizing self-attention mechanisms instead of the typical casual attention used in NLP decoders.

Loss Computation : Decoder embedding is projected to model output with an MLP layer. The mean squared error (MSE) loss is computed between the predicted masked values from the model and their corresponding ground truth values.

3.1 Encoder

The scRNA-seq data is characterized by its high sparsity, with cell information largely concentrated in the non-zero expression values. Thus, the encoder is designed to focus only on the non-zero part of the unmasked matrix, V unmasked . The encoder is based on a traditional multi-head attention transformer and takes the combination of value embedding, E , and gene embedding, G , as its input, I ∈ R c × m × d . The value and gene embeddings are similar to the word and positional embeddings in natural language modeling, respectively. The value embedding, E , is generated using the auto-discretization strategy discussed previously, while the gene embedding, G , is retrieved from a function f L that maps the gene symbols into the embedded vocabulary.

Then the encoder processes the input embeddings I and generates the high-level gene representations I encoder ∈ R b × m × d via the multi-head attention mechanism.

where f Q , f K , f V are the project functions. Trm denotes the Transformer block.

It is worth emphasizing that our encoder only operates on a subset of genes, reducing the length of the processed sequence to 1/10 of the original. This allows the full-length transformer to be used without any computational approximations.

3.2 Decoder

Unlike the encoder which focuses on the main information (non-zero expression values) in the cells, the decoder in the system performs full-length feature abstraction and extraction. The input to the decoder, I full , comprises three token types: the output from the encoder, I encoder , the genes with zero expression embs I zero , and the mask token embs I masked . Out of these tokens, genes with zero expression make up 90% of all tokens. The gene embeddings are concatenated with all of these tokens to provide the decoder with gene-specific information for the corresponding mask tokens, followed by a full connection layer.

where ⊕ represents the concatenation operation, and W p and b p are learnable parameters that project the decoder's embedding size.

The decoder in the framework is optimized for long-sequence attention calculations and employs the Performer architecture as its backbone. The decoder transforms the input I full into final gene-level embeddings, I decoder ∈ R b × n × d , and predicts the masked values through a shared linear layer, W ∈ R d × 1 , applied across all genes. The operations are expressed as follows:

The decoder has a smaller model size compared to the encoder, with a smaller embedding size, fewer attention layers, and fewer attention heads. For instance, in the largest model configuration, the layer depth ratio between the encoder and decoder is 2:1 and the head number ratio is 1.5:1 (see App. Table 2). Similarly, the principle of asymmetric encoder-decoder design has been proven powerful in masked autoencoders (MAE) [13], which is tailored for CV data pre-training. Unlike MAE, xTrimoGene utilizes the biased masking strategy to avoid the learning process being dominated by zero tokens. Though the scRNA-seq data is distinct from images, our results show that the performance gains of xTrimoGene are comparable to those of MAE, with more efficient training and better downstream task performance.

3.3 Auto-discretization strategy

Our aim is to transform an expression value v into a hidden embedding, denoted as e . The transformation is achieved using an auto-discretization block. This auto-discretization process involves a random look-up table T defined in R d × b . In this representation, d refers to the embedding dimension, while b is the number of bins with a default value of 100. The transformation starts by applying a linear layer to the expression value, given by v 1 = v · w 1 , where w 1 represents the weight vector. The resulting v 1 is then subjected to a leaky ReLU activation, resulting in v 2 = Leaky_ReLU ( v 1 ) . Subsequently, a cross-layer projection is applied, represented by v 3 = w 2 · v 2 + α · v 2 . Here, w 2 denotes the weight vector, and α is a scaling mixture factor. Next, the bin weights of v 3 are normalized using the softmax function, resulting in v 4 = softmax ( v 3 ) . Finally, the transformed value is represented as a weighted combination of individual embeddings from the look-up table, given by e = T · v 4 . It's important to note that the weights in this combination serve as learnable parameters.

To validate the effectiveness of the expression value projection, we conducted an analysis of viewing the weight distribution pattern for continuous values. Our results showed that the normalized weight distribution of the close values exhibited smooth transitions and that of the distant values being clearly distinguishable (App. section 3 Figure 1). This supports the conclusion that the auto-discretization strategy effectively represents continuous values with high resolution while preserving relatively rich meaning.

We also compared the performance of the proposed auto-discretization strategy with three other discretization methods: (1) Round bin with zero, in which values are rounded to the nearest integer, and zeros are kept as it is, (2) Up bin without zero. Values greater than zero are converted to the nearest ceiling integer, while zero is represented as individual 0. (3) Equal bin. All the values fall into a fixed percentage interval, which is calculated by value distribution and frequency. We evaluated the different strategies on a standard cell clustering task (see App. 4) and found that the proposed auto-discretization strategy outperformed the others (as shown in Figure2 A), demonstrating the importance of high-resolution projections in handling expression values.

4 Training Strategy

We now explain the strategy used to train the asymmetric encoder-decoder transformer. Pre-trained task and masking strategy are outlined, see App. 5 for acceleration strategy.

4.1 Regression masked task

The traditional masked language task is a multi-class classification problem, where the predicting target is a single token with limited, naturally distinct categories. In contrast, the normalized gene expression value is a continuous scalar. To fit the data property, we modify the pre-trained learning

(A)

(B)

objective to a regression task, aimed at recovering the absolute value of the masked positions. The loss function employed is the MSE between the ground truth and the predicted values:

where n represents the number of all genes, m represents the maximum length of the unmasked positions in a sample, and c represents the number of cells. To evaluate the efficacy of this modification, we compared the regression setting with the classification setting on the cell clustering task. The results indicate that the regression model outperforms the classification model (Figure 2B), providing evidence of the benefits of learning a more fitted representation.

4.2 Masking strategy

We mask both non-zeros and zeros positions though the scRNA-seq expression matrix is highly sparse (where zero percentage is usually over 90%). As the zero positions percentage is much higher than non-zero positions, the masked ratio can't be the same for the two types. Otherwise, the model tends to predict all zeros and still obtains a low error level. We propose to mask an almost equal number of positions for zero and non-zeros positions (see App. section 6 Table 1). The setting enforces the model to learn embeddings for all values and not to be dominated by zero representation. We found zero values supervision is necessary to boost the performance (App. Figure 2), which demonstrates that some zeros represent the true extremely low expression level. This type of zeros is informative to illustrate how the gene abundant behaves inside the cell.

The recovery of masked tokens in NLP is challenging due to the fact that word comprehension relies heavily on long-range interactions rather than local context. Accurate inference of the missing tokens can be achieved at low masking ratios (15%) where the information in the entire sentence is still relatively redundant and encoded by the unmasked tokens. We investigated the density of information needed for the scRNA-seq regression task by training models with different masking ratios (for non-zero values, the ratio was set 10 times higher than for zero values) ranging from 15% to 90% with a 15% interval. The models were then evaluated on the cell clustering task, with the results showing that performance improved first and then degraded as the masking ratio increased. When the masking ratio was close to 30%, the majority of metrics reached a peak (App. Figure 3). We also found current biased masking is optimal (App. Figure 4) and the percentage of [MASK] tokens agrees well with NLP tasks (App. Figure 5). These results suggest that the scRNA-seq expression vector contains more redundant information than a sentence and highlight the role of hidden regulations between genes in constraining the inference of expression values.

5 Experiments

Next we will explain our experimental settings and results. The dataset description can be referred to App. 1.

5.1 Computational efficiency

We quantitatively compared the training cost of xTrimoGene with other two encoder-only models, including full-length attention Transformer and kernel-based approximation Performer (scBERT). For an apple-to-apple comparison, three models are set to approximately 10 million trainable parameters and trained on 5 million samples over 5 epochs. We calculated the corresponding FLOPs, where matrix multiplication operations are considered only. We observed that total FLOPs for Performer (scBERT) decreased to 10% of native Transformer (see Table 1). Notably, xTrimoGenes runs 3-times faster than Performer. The results validate the efficiency of xTrimoGene, which is readily adapted for large-scale data pre-training.

5.2 Scalability

The Deep Learning community has shown significant interest in the scalability of proposed models [18, 2]. Vanilla Transformer models are challenging to scale due to their computational time and resource requirements, which increase quadratically with model size. Varieties of attention mechanisms have been proposed to accelerate training speed, a critical factor for model scaling.

To test the scale-up ability of xTrimoGene, we pre-trained three models across multiple compute regions and scales (e.g., from 3M to 100M parameters). The detailed hyperparameter setting is displayed in the App. Table 2. The training curve clearly shows all models are steadily down to a lower loss when training steps increase (App. Figure 6). More importantly, the xTrimoGene-100M model obtains a significant improvement over the xTrimoGene-10M model, which is also superior to the xTrimoGene-3M model. The tendency is consistent across different data sizes. The results suggest xTrimoGene framework is robust to scale-up, making it possible and convenient to pre-train larger models with more data.

5.3 Robustness on high sparse data

scRNA-seq data often exhibit varying levels of sparsity, thus it's necessary to assess whether xTrimoGene is robust in handling different sparse data. To verify the robustness, we divided the test samples into subgroups based on cell type and calculated the sparsity level (i.e., percentage of zero values in the expression matrix) and Pearson correlation coefficient between the predicted and actual values. Our results reveal that the correlation gradually decreases as the sparsity level increases, as expected (Figure 3A). However, the correlation remains above 0.8 even when the sparsity level reaches 96% (Figure 3A), indicating the robustness of xTrimoGene. We also compared xTrimoGene's performance with Performer and found that xTrimoGene consistently achieves a higher correlation across most subgroups (Figure 3B). These findings demonstrate that xTrimoGene is robust in handling highly sparse data and outperforms encoder-only architectures.

The performance of the encoder-decoder and encoder-only architectures have been comparatively analyzed in the NLP domain, with the former demonstrating effectiveness in language comprehension and the latter in context generation. Apart from comparison on masked value recovery, we further evaluated xTrimoGene against encoder-only Performer on the cell clustering task. The results

(A)

(B)

(C)

demonstrate that xTrimoGene achieves superior performance, reaffirming its proficiency in latent embedding extraction (Figure 3C).

5.4 Evaluation on downstream tasks

Currently, multiple tasks have been established to evaluate different models, including well-defined cell type annotation and recently developed perturbation response prediction tasks. We first assessed the performance of xTrimoGene on these single-cell tasks. Additionally, we explored the potential application on bulk RNA-sequencing data, with a focus on synergistic drug combination prediction.

5.4.1 Cell type annotation

First, we evaluated xTrimoGene's performance on cell type annotation task with Zheng68K [39] and Segerstolpe [31] dataset, which has been widely benchmarked. We compared the xTrimoGene against other several methods, including scBERT [36], ACTINN [23], Scanpy [34], CellTypist [6], scVI [21] and singleCellNet [37]. For the xTrimoGene model, we added a max-pooling layer and a linear layer to predict cell type labels with fine-tuning mode (see App. 8.1). For other methods, we followed their instruction with the default parameter setting. We observed that xTrimoGene achieves a high Precision and F1 score, surpassing all the other methods (Table 2). The results indicated that xTrimoGene learns a well-represented cellular embedding (visualized in App. Figure 7) by simply aggregating contextual gene embedding.

5.4.2 Perturbation response prediction

Recently, perturb-seq technology was established to screen gene expression response given pooled perturbations at single-cell level [8]. Several algorithms have also been developed to predict perturbation effects [30, 22] at the single cell level, i.e., what is the expression value of genes after perturbation? We compared the native GEARS[30] model with and without incorporating embeddings from xTrimoGene.

The normal state (before perturbation) gene expression profile is fed into xTrimoGene and we obtained the context embedding, which replaces raw expression value input in the GEARS model (App. 8.2). All the other settings remain unchanged. The evaluated dataset (Norman et al. [26]) contains both single and double gene perturbation and we thus assess the performance across different perturbation levels. As shown in Figure 4A, GEARS with xTrimoGene embedding scores a lower MSE (decreased 14.8%) for top20 differential expressed genes across all perturbation scenarios. Notably, the tendency is consistent across different perturbation levels, regardless the perturbed target is seen or not. We also compared against the scBERT embedding and observed a similar trend, where xTrimoGene achieves better results (App. Table 3). The results demonstrated that the pre-training strategy empowers xTrimoGene to capture constraints under various circumstances, including post-perturbations. The application further proved the efficacy and potential of xTrimoGene to boost scRNA-seq based tasks.

(B)

(A)

5.4.3 Synergistic drug combinations prediction

The drug synergistic task evaluates how patients or cells respond to a drug combination intervention [25]. However, the generated wet-lab experimental data only covers a tiny search space of possible drug combinations. Multiple models have been proposed to accelerate predicting the synergistic landscape of drugs [28, 33]. For instance, DeepDDS integrates genomic expression profiles and drug chemical information, greatly improving the prediction performance. We further explored whether xTrimoGene is able to generate good latent embedding for this bulk expression data.

Similar to the perturbation prediction test, we adapted xTrimoGene to DeepDDS with the intermediate context embedding (see App. 8.3). We also included DeepSynergy and Random Forest for comparison. As illustrated in Figure 4B, utilizing embedding from the xTrimoGene model outperforms all the other models. The result proved that xTrimoGene can accurately capture cell-level representation, even for bulk sequencing data. This also opens the avenue for xTrimoGene to be applied across other biological modeling tasks, especially where bulk-level transcriptome data is available.

6 Conclusion

xTrimoGene is a new, efficient framework for learning scRNA-seq data. It proposes an asymmetric encoder-decoder framework that takes advantage of the sparse gene expression matrix and establishes the projection strategy of continuous values with a higher resolution. The results show that xTrimoGene is scalable and performs well on tasks like cell type annotation, perturbation response prediction, and synergistic drug combination prediction. The experiments demonstrate the efficacy of pre-training in single-cell biology. xTrimoGene is potentially adapted to other types of cell modeling analysis, including rare cell detection (App. 8.4), batch effect removal and regulatory network construction.

Certain limitations exist for xTrimoGene and further work is desired to advance the design. At present, xTrimoGene major utilizes gene expression values during the pre-training stage, overlooking varieties of other related meta-information like sample condition (health/disease), cell type, tissue type, sequencing platform, etc. These rich annotations are biologically meaningful and highly correlated with the expression pattern within a cell. The memory consumption for inference with the xTrimoGene-100M model is approximately 50GB, whose hardware requirement (Nvidia A100 80G GPU) is beyond some academic labs, thus computational or memory-efficient engineering techniques tend to advance the model pre-training and application.

xTrimoGene has been integrated into BioMap's single-cell analysis platform, functioning as a fundamental and essential model (as depicted in the App. Figure 9). The pre-trained model services have been publicly available. In the future, with the increase of data, larger pre-trained models are expected to drive more advancements in various downstream task learning.

Author Contributions and Acknowledgments

Le Song led the project by designing its scope, conceptualizing ideas, integrating resources, and making decisions on techniques. Xingyi Cheng played a key role in the development of the xTrimoGene framework, auto-discretization and unsupervised objectives, contributing concrete ideas and pseudocode, along with code review. Jing Gong and Mingsheng Hao (Research Intern at BioMap) were primarily responsible for conducting pre-training and downstream experiments, serving as the first authors of the paper. Their work covered areas such as model scaling, cell type annotation, perturbation response prediction, and synergistic drug combinations prediction. Xin Zeng made significant contributions to the code of the xTrimoGene framework, worked with an early performer version, and conducted initial downstream experiments. Chiming Liu oversaw the engineering aspects of the project, including the implementation of the data pipeline and FLOPs computation. Jianzhu Ma, Xuegong Zhang, Taifeng Wang, and Le Song conceived the project and provided invaluable guidance for the project and contributed their expertise in computational biology knowledge. Taifeng Wang also played a pivotal role in pushing for the model's service implementation. Finally, Jing Gong, Mingsheng Hao, Xingyi Cheng, and Le Song collectively contributed to writing this paper.

In addition, we would like to express our gratitude to the individuals at BioMap who have made contributions to our project. Chenrui Xu and Yucheng Guo played roles in data preprocessing and integration. Zhaoren He's expertise in data analysis and application greatly enhanced our work, and we deeply appreciate his contributions.

This work was supported by the Ministry of Science and Technology of the People's Republic of China (2022YFF1203004), the Beijing Municipal Science & Technology Commission and the Administrative Commission of Zhongguancun Science Park(Z221100003522022). And this work was also funded by BioMap.

References

Nadav Brandes, Dan Ofer, Yam Peleg, Nadav Rappoport, and Michal Linial. ProteinBERT: a universal deep-learning model of protein sequence and function. Bioinformatics , 38(8):21022110, 02 2022.

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems , 33:1877-1901, 2020.

Geng Chen, Baitang Ning, and Tieliu Shi. Single-cell rna-seq technologies and related computational data analysis. Frontiers in genetics , page 317, 2019.

Jiayang Chen, Zhihang Hu, Siqi Sun, Qingxiong Tan, Yixuan Wang, Qinze Yu, Licheng Zong, Liang Hong, Jin Xiao, Tao Shen, et al. Interpretable rna foundation model from unannotated data for highly accurate rna structure and function predictions. bioRxiv , pages 2022-08, 2022.

Sijie Chen, Yanting Luo, Haoxiang Gao, Fanhong Li, Yixin Chen, Jiaqi Li, Renke You, Minsheng Hao, Haiyang Bian, Xi Xi, et al. heca: The cell-centric assembly of a cell atlas. Iscience , 25(5):104318, 2022.

C. Domínguez Conde, C. Xu, L. B. Jarvis, D. B. Rainbow, S. B. Wells, T. Gomes, S. K. Howlett, O. Suchanek, K. Polanski, H. W. King, L. Mamanova, N. Huang, P. A. Szabo, L. Richardson, L. Bolt, E. S. Fasouli, K. T. Mahbubani, M. Prete, L. Tuck, N. Richoz, Z. K. Tuong, L. Campos, H. S. Mousa, E. J. Needham, S. Pritchard, T. Li, R. Elmentaite, J. Park, E. Rahmani, D. Chen, D. K. Menon, O. A. Bayraktar, L. K. James, K. B. Meyer, N. Yosef, M. R. Clatworthy, P. A. Sims, D. L. Farber, K. Saeb-Parsy, J. L. Jones, and S. A. Teichmann. Cross-tissue immune cell analysis reveals tissue-specific features in humans. Science , 376(6594):eabl5197, 2022.

Jiarui Ding, Xian Adiconis, Sean K. Simmons, Monika S. Kowalczyk, Cynthia C. Hession, Nemanja D. Marjanovic, Travis K. Hughes, Marc H. Wadsworth, Tyler Burks, Lan T. Nguyen, John Y. H. Kwon, Boaz Barak, William Ge, Amanda J. Kedaigle, Shaina Carroll, Shuqiang Li, Nir Hacohen, Orit Rozenblatt-Rosen, Alex K. Shalek, Alexandra-Chloé Villani, Aviv Regev, and Joshua Z. Levin. Systematic comparison of single-cell and single-nucleus rna-sequencing methods. Nature Biotechnology , 38(6):737-746, Jun 2020.

Atray Dixit, Oren Parnas, Biyu Li, Jenny Chen, Charles P Fulco, Livnat Jerby-Arnon, Nemanja D Marjanovic, Danielle Dionne, Tyler Burks, Raktima Raychowdhury, et al. Perturb-seq: dissecting molecular circuits with scalable single-cell rna profiling of pooled genetic screens. cell , 167(7):1853-1866, 2016.

Mario Flores, Zhentao Liu, Tinghe Zhang, Md Musaddaqui Hasib, Yu-Chiao Chiu, Zhenqing Ye, Karla Paniagua, Sumin Jo, Jianqiu Zhang, Shou-Jiang Gao, et al. Deep learning tackles singlecell analysis-a survey of deep learning for scrna-seq analysis. Briefings in bioinformatics , 23(1):bbab531, 2022.

Yury Gorishniy, Ivan Rubachev, and Artem Babenko. On embeddings for numerical features in tabular deep learning. Advances in Neural Information Processing Systems , 35:24991-25004, 2022.

Christopher Heje Grønbech, Maximillian Fornitz Vording, Pascal N Timshel, Casper Kaae Sønderby, Tune H Pers, and Ole Winther. scVAE: variational auto-encoders for single-cell gene expression data. Bioinformatics , 36(16):4415-4422, 05 2020.

Xu Han, Zhengyan Zhang, Ning Ding, Yuxian Gu, Xiao Liu, Yuqi Huo, Jiezhong Qiu, Yuan Yao, Ao Zhang, Liang Zhang, Wentao Han, Minlie Huang, Qin Jin, Yanyan Lan, Yang Liu, Zhiyuan Liu, Zhiwu Lu, Xipeng Qiu, Ruihua Song, Jie Tang, Ji-Rong Wen, Jinhui Yuan, Wayne Xin Zhao, and Jun Zhu. Pre-trained models: Past, present and future. AI Open , 2:225-250, 2021.

Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 16000-16009, 2022.

Dan Hendrycks, Kimin Lee, and Mantas Mazeika. Using pre-training can improve model robustness and uncertainty. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning , volume 97 of Proceedings of Machine Learning Research , pages 2712-2721. PMLR, 09-15 Jun 2019.

Ruochen Jiang, Tianyi Sun, Dongyuan Song, and Jingyi Jessica Li. Statistics or biology: the zero-inflation controversy about scrna-seq data. Genome Biology , 23(1):31, Jan 2022.

Dragomirka Jovic, Xue Liang, Hua Zeng, Lin Lin, Fengping Xu, and Yonglun Luo. Single-cell rna sequencing technologies and applications: A brief overview. Clinical and Translational Medicine , 12(3):e694, 2022.

John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature , 596(7873):583-589, 2021.

Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361 , 2020.

Mengwei Li, Xiaomeng Zhang, Kok Siong Ang, Jingjing Ling, Raman Sethi, Nicole Yee Shin Lee, Florent Ginhoux, and Jinmiao Chen. Disco: a database of deeply integrated human single-cell omics data. Nucleic acids research , 50(D1):D596-D602, 2022.

Zeming Lin, Halil Akin, Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Nikita Smetanin, Robert Verkuil, Ori Kabeli, Yaniv Shmueli, et al. Evolutionary-scale prediction of atomic-level protein structure with a language model. Science , 379(6637):1123-1130, 2023.

Romain Lopez, Jeffrey Regier, Michael B Cole, Michael I Jordan, and Nir Yosef. Deep generative modeling for single-cell transcriptomics. Nature methods , 15(12):1053-1058, 2018.

Mohammad Lotfollahi, Anna Klimovskaia Susmelj, Carlo De Donno, Yuge Ji, Ignacio L Ibarra, F Alexander Wolf, Nafissa Yakubova, Fabian J Theis, and David Lopez-Paz. Compositional perturbation autoencoder for single-cell response modeling. BioRxiv , 2021.

Feiyang Ma and Matteo Pellegrini. Actinn: automated identification of cell types in single cell rna sequencing. Bioinformatics , 36(2):533-538, 2020.

Mohamed Marouf, Pierre Machart, Vikas Bansal, Christoph Kilian, Daniel S. Magruder, and Stefan Bonn Christian F. Krebs. Realistic in silico generation and augmentation of single-cell rna-seq data using generative adversarial networks. Nature Communications , 11(1)(166):786793, 2020.

Reza Bayat Mokhtari, Tina S Homayouni, Narges Baluch, Evgeniya Morgatskaya, Sushil Kumar, Bikul Das, and Herman Yeger. Combination therapy in combating cancer. Oncotarget , 8(23):38022, 2017.

Thomas M. Norman, Max A. Horlbeck, Joseph M. Replogle, Alex Y. Ge, Albert Xu, Marco Jost, Luke A. Gilbert, and Jonathan S. Weissman. Exploring genetic interaction manifolds constructed from rich single-cell phenotypes. Science , 365(6455):786-793, 2019.

Ahmad Pesaranghader, Stan Matwin, Marina Sokolova, Jean-Christophe Grenier, Robert G Beiko, and Julie Hussin. deepSimDEF: deep neural embeddings of gene products and gene ontology terms for functional analysis of genes. Bioinformatics , 38(11):3051-3061, 05 2022.

Kristina Preuer, Richard P I Lewis, Sepp Hochreiter, Andreas Bender, Krishna C Bulusu, and Günter Klambauer. DeepSynergy: predicting anti-cancer drug synergy with Deep Learning. Bioinformatics , 34(9):1538-1546, 12 2017.

Xipeng Qiu, Tianxiang Sun, Yige Xu, Yunfan Shao, Ning Dai, and Xuanjing Huang. Pre-trained models for natural language processing: A survey. Science China Technological Sciences , 63(10):1872-1897, 2020.

Yusuf Roohani, Kexin Huang, and Jure Leskovec. Gears: Predicting transcriptional outcomes of novel multi-gene perturbations. BioRxiv , pages 2022-07, 2022.

Åsa Segerstolpe, Athanasia Palasantza, Pernilla Eliasson, Eva-Marie Andersson, Anne-Christine Andréasson, Xiaoyan Sun, Simone Picelli, Alan Sabirsh, Maryam Clausen, Magnus K Bjursell, David M Smith, Maria Kasper, Carina Ämmälä, and Rickard Sandberg. Single-cell transcriptome profiling of human pancreatic islets in health and type 2 diabetes. Cell Metab. , 24(4):593-607, October 2016.

Jingshu Wang, Divyansh Agarwal, Mo Huang, Gang Hu, Zilu Zhou, Chengzhong Ye, and Nancy R Zhang. Data denoising with transfer learning in single-cell transcriptomics. Nature methods , 16(9):875-878, 2019.

Jinxian Wang, Xuejun Liu, Siyuan Shen, Lei Deng, and Hui Liu. Deepdds: deep graph neural network with attention mechanism to predict synergistic drug combinations. Briefings in Bioinformatics , 23(1):bbab390, 2022.

F Alexander Wolf, Philipp Angerer, and Fabian J Theis. Scanpy: large-scale single-cell gene expression data analysis. Genome biology , 19:1-5, 2018.

Yijia Xiao, Jiezhong Qiu, Ziang Li, Chang-Yu Hsieh, and Jie Tang. Modeling protein using large-scale pretrain language model. arXiv preprint arXiv:2108.07435 , 2021.

Fan Yang, Wenchuan Wang, Fang Wang, Yuan Fang, Duyu Tang, Junzhou Huang, Hui Lu, and Jianhua Yao. scbert as a large-scale pretrained deep language model for cell type annotation of single-cell rna-seq data. Nature Machine Intelligence , 4(10):852-866, 2022.

Patrick Cahan Yuqi Tan. Singlecellnet: A computational tool to classify single cell rna-seq data across platforms and across species. Cell Systems , 19(2):207-213, 2019.

Ningyu Zhang, Zhen Bi, Xiaozhuan Liang, Siyuan Cheng, Haosen Hong, Shumin Deng, Jiazhang Lian, Qiang Zhang, and Huajun Chen. Ontoprotein: Protein pretraining with gene ontology embedding. arXiv preprint arXiv:2201.11147 , 2022.

Grace XY Zheng, Jessica M Terry, Phillip Belgrader, Paul Ryvkin, Zachary W Bent, Ryan Wilson, Solongo B Ziraldo, Tobias D Wheeler, Geoff P McDermott, Junjie Zhu, et al. Massively parallel digital transcriptional profiling of single cells. Nature communications , 8(1):14049, 2017.

Peer reviews:
comments: SUMMARY:
This paper proposes a scRNA-seq pretraining method, overcoming the scalability and resolution weaknesses of previous works. Experiments on various downstream tasks proves the efficiency and effectiveness of the proposed method.

STRENGTHS:
The paper is well structured and easy to follow in general. The challenges it aims to tackle are also clear, namely, scalability and resolution. The proposed method is technically sound. Its performance and efficiency are evaluated on several downstream tasks. The experiments are thoroughly conducted to investigate the influence of different mask strategies and self-supervised objectives.

WEAKNESSES:
1. Would it be possible to further improve the efficiency by optimizing the padding operation? It seems that the padded zeros are also not useful in the encoding process.
2. The masking and reconstruction strategy, the auto-discretization module, and optimization objective seem to be a simple combination of previous works. This could slightly lower the novelty of the proposed method.
3. The embedding of the proposed method and existing works could be visualized to provided an intuitive understanding of the performance.

QUESTIONS:
Please refer to the weaknesses.

LIMITATIONS:
No significant limitations of the method are found.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-06-10

comments: SUMMARY:
In this study, the authors propose an asymmetric encoder-decoder transformer for scRNA-seq data, called xTrimoGene, for large scale dataset pre-training. Quantitative comparison with various SOTA approaches shows the advantage in efficiency and scalability. In addition, a series of downstream analysis further validate the performance of xTrimoGene.

STRENGTHS:
1. The training process is accelerated by utilizing sparse input.

2. A novel auto-discretization strategy is introduced to map continuous expression values into a latent embedding space.

3. Considering its high accuracy and efficiency, the proposed large model xTrimoGene is a valuable contribution to the single-cell community.

4. The authors conducted validation of xTrimoGene using multiple downstream analysis tasks.

WEAKNESSES:
In table 2, compared with CellTypist, the propsoed method only obtained a marginal improvement. In this case, it will be better if the author can provide more analysis to show the advantage of xTrimoGene on this task. See questions for detailed suggestions.

QUESTIONS:
Compared with CellTypist, is there any biological insights can only be found by xTrimoGene from these datasets? Is there any specific circumstance in which xTrimoGene consistently demonstrates superior performance？

LIMITATIONS:
The limitations are not presented in the paper. Consider discussing the potential limitations or challenges associated with the proposed method. Acknowledging and addressing these factors would further enhance the study's robustness and reliability.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-06-12

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe manuscript proposes an adaptation of the key contributions of [1] to the scRNA-seq data modality, which is dubbed xTrimoGene. Additionally, the authors introduce a new embedding encoder, their \"auto-discretization strategy\". Several ablation experiments motivate specific choices in the model. Finally, xTrimoGene is tested on three downstream tasks.\n\nSTRENGTHS:\nThe model is described (for the most part) clearly, so that the paper is easy to follow. Utilizing the sparsity of scRNA-seq data is important in the context of transformers and, therefore, adapting the ideas from [1] to this modality seems like a very good match. I also appreciated the ablation experiments because they justify specific choices for the model and make the develoment process more transparent.\n\nWEAKNESSES:\n**Connection to [1]:**  \nThe paper is for the most part a (useful) adaptation of the ideas of [1] to the scRNA-seq domain. This is mentioned in one sentence on page 5 (l 156), which, in my opinion, under emphasizes the connection. I appreciate that additional steps, such as determining the correct size of the mask and reducing the frequency of masking zeros are necessary to transfer the ideas to scRNA-seq data. Nevertheless, I strongly suggest putting references to [1] more prominently and earlier in the paper, e.g., into the abstract, the introduction and the beginning of Section 3 (at least beginning of section 3.1).\n\n\n**Exposition of the auto-discretization strategy:**  \nWhile the paper is overall clearly written, I did not understand several aspects about the auto-discretization strategy.  \nFirst, why is discretization necessary in the first place; unnormalized count data would already be discrete, so why not use this? Also, why is discreteness necessary?  \nSecond, does the proposed method actually discretize? According to ll168 ff it is just a weighted sum of vectors and Appendix Figure 1 confirms this empirically at least for low expression values.  \nThird, I did not get the description in the first paragraph of 3.3. My main question is what the shapes of the various w's and v's are in that paragraph. My best guess is that $v_1, v_2, v_3, v_4 \\in \\mathbb{R}^{c\\times b}$. But then the product $EXP_{lookup} \\cdot v_4$ is the wrong way around and would only produce output in $\\mathbb{R}^{c\\times d}$, while from Fig 1 I understand that the output should be in $\\mathbb{R}^{c \\times m \\times d}$, which makes sense as m is the sequence length, crucial for the transformer. Also, I wonder why the two products of $v_2$ are not collapsed into a single one in line 167. Should perhaps one of the two $v_2$'s be a $v_1$?  \nWhen the authors write in ll 168 ff \"The final output is a weighted combination of individual embeddings from the look-up table [...], where the weights are learnable parameters.\", I wonder whether the entries of the look-up table are themselves learnable, or just the weights $w_1, w_2$. Also, why is $EXP_{lookup}$ called a look-up table, when actually a weighted sum of its rows is returned? I assume that the softmax performs a soft version of discritization, with a winner-takes-(almost-)all effect. But at least for low expression values the output really is not discrete.\n\n**Ablation of the relative masking frequency of zeros:**  \nI agree that zeros need to be masked much less frequently than non-zero entries of the gene expression matrix. The authors' choice of masking an \"almost equal number of positions for zero and non-zero positions\" is plausible (Why \"almost\" here? What is the precise masking strategy?). I would be curious as to how the downstream results change for other ratios. I.e. I suggest a plot like in the Fig 3 of the appendix, in which the total masking ratio is kept at 30% but the frequency of masking zeros and non-zeros is changed. Perhaps something similar is depicted in Fig 4 of the appendix, which I did not understand (Could you please explain the three percentages in more detail?).\n\n**Order of subsections in section 5:**  \nThe main experimental results appear in subsections 5.4, which also describes the downstream tasks for the first time and compares to competitors. Sections 5.1-5.3 are valuable in that they analyze the model in more depth and perform ablations. Some ablations are even based on the task described only later in 5.4. Therefore, I would suggest to move the most important subsection, 5.4., to the beginning of section 5. In addition, one might also put the experiments of Fig 2 into section 5 (after the downstream experiments have been described) because 2A also relies on the downstream task. This way, the reader would be familiar with the downstream tasks before they are used to compare models.\n\n**Comparison to other models:**  \nFor the cell type annotation task many methods were used to produce a representation which was subsequently used for clustering. I guess that nearly all of these methods (for sure scBERT, scVI) can also be used in conjunction with GEARS and DeepDDS in much that same way as xTrimoGene is used for perturbation prediction and drug prediction, respectively. While the current experiments for these two task show that xTrimoGene representations offer benefits compared to raw gene expressions, it would strengthen the paper if xTrimoGene also outperformed, say, a combination of scBERT representations and GEARS. \n\n\n\n**Minor:**  \n- In line 15 of the abstract the downstream task is described as \"cell classification\". This sounds as if it was supervised. But according to the appendix unsupervised Leiden clustering is used. Please change this to \"cell type annotation\" or \"cell clustering\". \n- In line 19 the term AI4Science is introduced but never used. Perhaps it can be omitted?\n- In l 84 the authors say that gene expression values in scRNA-seq data are continuous scalars. But the raw data is counts, i.e., discrete. This confused me when reading the paper the first time. From l 103 onwards, they speak of *normalized* gene expressions, which indeed are real numbers and not necessarily natural numbers. Please add \"normalized\" or \"pre-processed\" to l 84. The pre-processing is discussed in the appendix, but not linked from the paper. Perhaps it would be useful to include a cross-reference at the beginning of section 3, as a first step of the pipeline. \n- Line 131 refers to the auto-discretization strategy \"discussed previously\", but the auto-discretization strategy is only discussed later in section 3.3.\n- It would be useful to mention the key aspects of the large pre-training dataset at the beginning of section 5 (size, number of genes, that it was scraped from GEO) at the beginning of section 5. \n- It would be great if Table 1 was extended by the runtime and memory consumption of each model.\n- Line 271 and Table 2 speak of Zheng68K, which has 68K cells, while appendix l 43 claims the PBMC dataset had only ~2k cells.\n- Line 162: should it be $V\\in \\mathbb{R}^{c\\times m}$ rather than $V\\in\\mathbb{R}^{c\\times n}$?\n- Fig 3 panel A's caption and main caption \"Sparse level\" --> \"Sparsity level\"\n- First line of Fig 2: Typo in \"Performance\" and missing hyphen in \"auto-discretization\".\n- Line 227: word order \"other two\" --> \"two other\"\n- Line 228 \"comparison, three models\" --> \"comparison, all three models\"\n\n**Reference:**    \n[1] He, Kaiming, et al. \"Masked autoencoders are scalable vision learners.\" Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022.\n\nQUESTIONS:\n- How is the name \"xTrimoGene\" motivated? Please include an explanation of this non-obvious name.\n- What are the numbers in brackets in line 154?\n- What is the range of the sum in Eq 5? Just the masked entries of the gene experession matrix (similar to [1]) or all entries? The normalization prefactor does not match either. \n- How is classification in the classification task mode of Fig 2B) performed? I understand that the output is continuous and not discrete, so that I do not see what the classes would be. For this reason, I have accepted the use of MSE as loss immediately. Nevertheless, performing the ablation is appreciated.\n- Why was the comparison in Table 2 only performed for the 10M parameter model and not the 100M parameter model? The latter has lower validation loss, so it should perform even better, right?\n- The model was trained on a very large dataset of human gene expressions. Is there any hope for it to transfer to other species?\n- I understood that the gene embedding was akin to positional embeddings in NLP. Therefore, would have assumed that it only changes with the gene, not the cell. Is this correct? If so, why are the rows in Fig 1 upper right corner not all the same? \n- Are $I_{masked}$ and $I_{zero}$ in eq (3) also the sums of the auto-discretization encoder and the gene encoder? Or are the pre-processed gene expression levels directly used here?\n\nLIMITATIONS:\nLimitations are not explicitly discussed. Code is not part of the submission, but is promised to be released on GitHub alongside the pre-trained model, which is particularly interesting.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nHere the authors propose xTrimoGene, a scalable transformer-based model for learning representations of scRNA-seq data. The authors demonstrate that their proposed method is more computationally efficient than alternatives for training transformers on scRNA-seq data, and they also validate their model on cell type annotation and perturbation response prediction tasks.\n\nSTRENGTHS:\n* **Clarity**: I found the manuscript very well organized and the writing easy to follow. Well done!\n* **Significance**: There has been much recent excitement about the potential utility of applying transformer architectures to scRNA-seq data. Here the authors demonstrate a method for making the training of such models more efficient, and also they demonstrate that the embeddings from such models are useful for multiple downstream tasks. Thus, I believe this work is significant.\n* **Originality**: The authors' proposed strategy for training xTrimoGene is, to my knowledge, novel (though I note that I am not an expert in transformer training strategies).\n\nWEAKNESSES:\nThe results presented by the authors are promising, and I indeed enjoyed reading the paper. However, I do have some minor issues that I would like to see clarified during the rebuttal period before I can give a recommendation of acceptance. If the authors are able to address my concerns, I would be happy to raise my score:\n\n* **Was data leakage avoided?**: The authors briefly describe their data collection pipeline in Section 1 of the Appendix. However, this description is a bit sparse (just mentioning that data was collected from the GEO). Were any precautions taken to avoid data leakage issues in the evaluation of xTrimoGene on downstream tasks? That is, did the authors ensure that any test cells from train/test splits e.g. on the perturbation prediction task were not present in the data used to pretrain xTrimoGene?\n* **Proper modeling of scRNA-seq count distributions**: Previous works (e.g. the Deep Count Autoencoder https://www.nature.com/articles/s41467-018-07931-2) have found significantly better preservation of biological variation in latent representations of scRNA-seq data by directly modeling scRNA-seq counts (e.g. using a negative binomial or zero-inflated negative binomial distribution) compared to minimizing MSE loss for normalized counts. However, xTrimoGene opts for the latter. I think it would be valuable to know how this choice affects the model's latent representation. An ablation study comparing modeling the raw counts versus normalized counts would be very interesting here. However, if the authors are unable to perform such an experiment in the rebuttal period, performing an (easier) experiment similar to that presented in Figure 2 of the Deep Count Autoencoder paper (depicting how increasing levels of noise can potentially result in low-quality latent representations) would also provide valuable results.\n* **How to handle batch effects?**: A classic problem in scRNA-seq analysis is integrating datasets from multiple batches that may have systematic differences unrelated to any underlying biology. Given that xTrimoGene's pretraining procedure does not account for batch effects, could the authors discuss how would one analyze data from different batches with xTrimoGene?\n* **Rare cell types**: A potential issue with pretrained models like xTrimoGene is that they may misbehave when applied to new datasets with phenomena not present in the training set (e.g. new cell types or tissues not seen in the training data). Did the authors explore xTrimoGene's behavior in such a scenario? A new experiment investigating this behavior would be useful.\n* **Memory/hardware requirements?**: Could the authors provide additional details on the hardware/memory requirements for loading and performing inference with a pretrained xTrimoGene model?\n\nQUESTIONS:\nSee \"Weaknesses\" section.\n\nLIMITATIONS:\nI would like to see the authors clarify some potential limitations with their method (see \"Weaknesses\" for specific questions). I do not forsee any negative societal impacts resulting from this work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}