xTrimoGene: An Efficient and Scalable Representation Learner for Single-Cell RNA-Seq Data
Authors: Jing Gong
Minsheng Hao
Xingyi Cheng
Xin Zeng
Chiming Liu
Jianzhu Ma
Xuegong Zhang
Taifeng Wang
Le Song
Venue: NeurIPS 2023 poster (2023)
Keywords: pre-train
encoder-decoder
scRNA-seq
scalable
TL;DR: A novel and efficient Single-Cell RNA-Seq pre-training model.

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
Jing Gong
Minsheng Hao
Xingyi Cheng Xin Zeng'
Chiming Liu
Jianzhu Ma²
Xuegong Zhang² Taifeng Wang' BioMap Research 2 Tsinghua University
Le Song!
Mohamed bin Zayed University of Artificial Intelligence
(gongjing, minsheng 2022, zingyi, zengzin, chiming, taifeng, songle)@biomap.comm {zhangzg, majianzzu))tainngua
www.:///tttnnnnnuuu..u.. edu.
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 surpass- ing 50M records for humans with each record measuring 20,000 genes. This highlights the need for unsupervised representation learning 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 pro- pose novel asymmetric encoder-decoder transformer for scRNA-sed data, called XTrimoGene" (or xTrimoGene for short)*, which leverages the sparse character- istic of the data to scale up the pre- pre-training. This scalable design of x TrimoGene reduces FLOPs by one two orders of magnitude compared classical transform- ers while maintaining high accuracy, enabling us to train the largest transformer models over the largest scRNA-seg dataset today. Our experiments also show that the performance of x XTimmoGnne improves 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 service via the following link: https://api. biomap. com/ xTx www.ioTeeeeaaaaaaa /apply
Introduction
Recently, Artificial Intelligence (AI) technology has demonstrated promising results for address- ing scientific problems. This A14Science paradigm witnessed diverse successful biological and pharmaceutical applications, including protein analysis [17 20 38 35 1. RNA modeling (4), and genomies modulation (27). However, most existing Al models have predominantly focused protein sequences, neglecting the growing volume of high-throughn
experimental sequencing data the form 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 [6 3 This technique captures the expression levels of thousands of genes in parallel, facilitating the study of cellular heterogeneity [ 19] This unveiled information is crucial for understanding complex biological systems and disease progression [16 3. Integrating and
"qqual contribution. Mingsheng Hao conducted this work during his internship BioMap. Correspondence to: xingyi @ biomap.com, songle @ @wiimap.cco........aaa 'xTrimoGene member of BioMap's large-scale Al engine "xTrimo" series. "a" denotes the academic version, distinguish from the commercial product TrimoGGnne.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

modeling such large-scale scRNA-seq data can reveal rich cellular information and benefit various biological task learning.
Representation learning from scRNA-seg data 9 has been an active area of research in past decades. For example, ScVAE [T]] and VT [21 apply variational autoencoder framework to derive low- dimensional cell embeddings, escGAN [24] uses a Generative Adversarial Network (GAN) archi- tecture to generate cell-type specific expression profiles, and SAVER-X [32] is capable of removing batch effects across datasets. Despite the success of these customized algorithms, they tend to be computationally inefficient and adbor-intensive This prompts us explore general-purpose purpose model that first learns underlying knowledge from scRNA-seq data and generalizes it to different tasks in 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 low-rank trans- former (36) to analyze the scRNA data. It learns the cellular representation by randomly masking a percent of non-zero gene expression values and tries 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 seRNA-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 FLOPSs 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 I and 2), and two distant values could be mapped to identical embeddings (e.g.. 1.99 and 1.01 are mapped 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 curated scRNA-seq dataset of approximately 50 billion effective gene tokens.
(3) The pre-trained model x' 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 large, sparse expression matrix, where each row represents cell (sample) and each column 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 roughly regard each cell (row) as sentence or image patch, the 20,000 genes is vast number compared to the typical sequence length, which is mostly a few hundred and no more than few thousand 22 [2). Thus directly applying existing transformer architecture will not work.
Second, scRNA-seq matrices are highly sparse (90% zero in a typical dataset [55 7). The abundance level of RNA for each gene is measured by counting the unique molecular identifier (UMI) reads in

---

However, many genes exhibit low UMI counts due to limited probing efficiency. Therefore, treating seRNA-seq data as an image and utilizing convolutional neural network to extract features is not feasible, as it introduces huge number of redundant computations for sparse positions.
scRNA-seq experiments [].3].
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, 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 [T0), ideal for preserving the continuous semantics of the gene expression values.
Taking into the above three major features, we design a architecture described in the next section.
(30 Pacting
(4) Embedding
Masking
Expression matttk
Unnuskod--an matia Varwnial e gr
Expression emb E groone
Gene omb 6 peeess
Compute
(5)
(7) masked zoro
(8) Decoder
(8) Encoder
Attntion decoder
Attertion enooder
Decodar
Fullienggh emb
Encoder emb
Inpul emb
Faecser Reve
Jnu gersd
fexcader geooodd
Figure The XTrimoGene Framework: (1) Random positions (including both zero and non- zero values) masked for prediction. (2) Masked and zero-valued positions filtered out. (3) Remaining unmasked positions are aligned with padding tokens (grey) to ensure maximum length consistency within batch. (4) Gene expression values and gene embeddings
separately projected into embeddings. (5) These two embeddings are element-wise added. (6) The resulting input is fed into the encoder. (7) The intermediate encoder embedding is combined with embeddings for masked positions and zero embeddings.
(8) This combined representation is then fed into the decoder. (9) Decoder embedding is projected model output with an MLP layer. The MSE loss is calculated between the model output and ground truth values for the masked positions.
XTrimoGene Architecture
XTrimoGene is highly efficient framework for pre-training large-scale single-cell RNA-seq data (illustrated in Figure The training process is based on a regression masked task, aimed at accurately recovering masked values in the expression matrix. Notably, specific optimized asymmetrical encoder-decoder framework is employed to accelerate the lcarning 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 reduced number of layers and attention heads. In addition, 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 XTimmoenne 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. e denotes cell sample size, and denotes gene number (19,264 in setting, see App. 1 for data collection and processing).
3

---

Filtering: The masked and zero-valued embeddings are filtered out, yielding a variable-length sequence valuable information that is prepared for encoding.
The remaining unmasked positions are aligned with padding tokens, resulting in a much smaller unmasked-only matrix Vunmasked- denotes the maximum length of the unmasked sample. We include a scheme to illustrate the processing flow (see App. 2).
Padding:
Embedding: Expression value and gene embeddings separately projected. d denotes the di- mension of the embedding. The expression embedding calculated through an auto-discretization mapping. The gene embedding is retrieved from 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-lik architecture.
Extending masked and zero embeddings: The intermediate encoder embedding Lencoder is com- bined with embeddings for masked and zero-value positions.
Decoding: The combined embeddings are processed by the decoder, utilizing self-attention mecha- nisms instead of the typical casual attention used in NLP decoders.
Loss Computation: Decoder embedding is projected to model output with an MIP 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 focus only on the non-zero part of the unmasked matrix, V Vunmasked- The encoder is based on a traditional multi-head attention transformer and takes the combination of value embedding, E, and gene embedding. G, its input, / E Rexmxd 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 -uto-discretizztion strategy discussed previously, while the gene embedding, G. is retrieved from a function fL that maps the gene symbols into the embedded vocabulary.
E = Autobin Vunmasked @ Mnonzero): G = fo (genes), I= E G
(1)
Then the encoder processes the input embeddings I and generates the high-level gene representations Iencoder E mxd via the multi-head attention mechanism.
Iencoder = Trm(fq(D), (I).
(2)
where fo. fK, fv are the project functions. Trm denotes the Transformer block.
is worth emphasizing that our encoder only operates on 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-zere expression values) in the cells, the decoder in the system performs full-length feature abstraction and extraction. The input to the decoder, Tfull, comprises three token types: the output from the encoder, Tenenders the genes with zero expression embs Izero, and the mask token embs T 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.
T W Wp.Teneader @ Izero ① Imasked) bp
(3)

---

where b represents the concatenation operation, and Wp and by 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 its backbone. The decoder transforms the input full into final gene-level embeddings, decoder bxxxx and predicts the masked values through shared linear layer, W E Rdxi, applied across all genes. The operations are expressed as follows:
decoder = Trm (fo((((((rut))) fK ( fv (v(Lsu))),
Leecoder W
The decoder has smaller model size compared to the encoder, with 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 1.5:1 (see App. Table 2). Similarly, the principle of asymmetric encoder-decoder design has been proven powerful in masked autoencoders (MAE) [3] 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 x TrimmoGnne 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 e into a hidden embedding, denoted as The transforma- tion is achieved using an auto-discretization block. This auto-diseretization process involves random look-up table 7' defined in Rdxb In this representation, d refers to the embedding dimension, while b the number of bins with a default value of 100. The transformation starts by applying a linear layer to the expression value, given by 21 v W1, where Wi represents the weight vector. The resulting V1 is then subjected to a leaky ReLU activation, resulting in U/2 = Leaky ReLU(v)). Subsequently, a cross-layer projection is applied, represented by V/ u2 D2 + U2. Here, W2 denotes the weight vector, and a is a scaling mixture factor. Next, the bin weights of t3 are normalized using the softmax function, resulting in ita softmax(Us) Finally, the transformed value is represented combination of individual embeddings from the look-up table, given by = 14. It's important to note that the weights in this combination serve learnable parameters.
weighted
To validate the effectiveness 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 nutto-iscrretiztionn strategy effectively represents continuous values with high resolution while preserving relatively rich meaning.
We also compared the performance of the proposed auto-discretizztion 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 calculated by value distribution and frequency. We evaluated the different strategies on standard cell clustering task (see App. 4) and found that the proposed ditto-discrttization strategy outperformed the others (as shown in Figurd2 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. for acceleration strategy.
4.1 Regression masked task
The traditional masked language task is a multi-class classification problem, where the predicting target is single token with limited, naturally distinct categories. contrast, the normalized gene expression value is continuous scalar. To fit the data property, we modify the pre-trained learning
5

---

(A)
Auto discretization
Up bin two/
(B)
Round bin (w/ 0)
Equal bin
Regression
Classification
0.8
0.6
0.6
Value
Value
0.4
ARI
NMI
HOMO
CP
Metric
Metric
Figure 2: Pre-training strategy ablation study. (A) Performance comparison between auto discretiza- tion strategy and other binning methods for expression value projection. The cell clustering task is evaluated and five metrics are displayed. ARI for Adjusted Rand index, NMI for Normalized Mutual Information, HOMO for Homogeneity, CP for Completeness and SIL for Silhouette Coefficient. (B) Performance of pre-trained models with different task modes, including regression and classification settings. The cell clustering task is evaluated. See the main text for details.
objective regression task, aimed recovering the absolute value of the masked positions. The loss function employed is the MSE between the ground truth and the predicted values:
Loss
(n m) C E(V. - VV,)²
(5)
where n represents the number of all genes, m represents the maximum length of the unmasked posi- tions sample, and e 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 providing evidence of the benefits of learning more fitted representation.
4.2 Masking strategy
We mask both non-zeros and 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 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 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 the entire sentence still relatively redundant and encoded by the unmasked tokens. We investigated the density information needed for the scRNA-seq regression task by training models with different masking ratios (for non-zero values, the ratio was 10 times higher than for zero values) ranging from 15% 90% with 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 peak (App. Figure 3). We also found current biased masking optimal (App. Figure 4) and the percentage of [MASK] tokens agrees well with NL. tasks (App. Figure 5). These results suggest that the scRNA-seq expression vector contains more redundant information than sentence and highlight the role of hidden regulations between genes in constraining the inference of expression values.

---

Experiments
Next we will explain our experimental settings and results. The dataset description can be referred to App. I.
5.1
Computational efficiency
We quantitatively compared the training cost of XTrimoGene with other two encoder-only models, including full-length attention Transformer and bernel-based approximation Performer (scBERT). For approximately 10 million trainable parameters
apple-to-apple comparison, three models are 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 T). Notably, XTrimoGenes runs 3-times faster than Performer. The results validate the efficiency xTrimoGene, which is readily adapted for large-scale data pre- tre-training.
Table 1: Computational efficiency comparison between different algorithms. The resource column is normalized by the Transformer row.
Model name
Parameter
Forward backward (FLOPs/saaple)
Total train (FLOPs)
Resource
(M)
Transformer
11.3
9.86E+12
2.46E+20
100%
Performer
8.9
1.06E+12
2.65E+19
10.8%
xTrimoGene
9.8
3.35E+11
8.38E+18
3.4%
5.2
Scalability
The Deep Learning community has shown significant interest in the scalability of proposed models Vanilla Transformer models are challenging to scale due to their computational time and resource requirements, which increase quadratically with model size. Varieties attention mechanisms have been proposed to accelerate training speed, critical factor for model scaling.
18 21.
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 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 x TrimoGene-31 model. The tendency is consistent across different data sizes. The results suggest XTrimoGene framework 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-seg data often exhibit varying levels of sparsity, thus it's necessary to assess whether xTrimo- Gene 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 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 A). However, the correlation remains above 0.8 even when the sparsity level reaches 96% (Figure A), indicating the robustness of XTrmmoGene. We also compared XTrimoGene's performance with Performer and found that x XTimooGnne consistently achieves higher correlation across most subgroups (Figure B B). 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 NL.P domain, with the former demonstrating effectiveness in language comprehension and the latter context generation. Apart from comparison on masked value recovery, we further evaluated xTrimoGene against encoder-only Performer on the cell clustering task. The results
7

---

[4]
1.00
4.,62++1.31
0.35
sparse Jeve
0.8
Peurson comnaatiiinn cxeftiinnt
value 0.04
0.35
(pradict (ruhh!
0.30-
Valuus
0.8
XTiimaGene
0.4
0.50
0.2
0.75
0.75
0.88
0.90
0.92
6.94
0.96
0.98
0.75
0.80
0.85
0.50
0.95
AR1
KX/1
HOMO
CP
SIL
Sparse
Performer
Metric
Figure 3: Comparison of performance for different sparse level data. (A) XTrimoGene performance for recovering masked values at different sparse levels. Each dot represents a subset defined by cell type. Sparse level is calculated the ratio between zero value percentages. Pearson correlation coefficient metric is calculated on masked positions. (B) Performance comparison xTrimoGene and Performer while recovering masked values different sparse levels. Dot has the dot size is proportional to the sparse level. Both the x and y axis denotes the Pearson correlation coefficient metric for particular algorithm. (C) Comparison of performance for XTrimoGene framework and encoder-only framework. Cell clustering task is evaluated.
meaning as (A) but the
demonstrate that x 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 x XrimmoGnne on these single-cell tasks. Additionally, we explored the potential application on bulk RNA-sequeneing data, with 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
and Segerstolpe [31] dataset, which has been widely benchmarked. We compared the x xTrimoGene against other several methods, including scBERT [36), ACTINN (23, Scanpy (34), CellTypist (6), scVI [21] and singleCelINet (37]. For the XrimmoGnne model, added max-pooling layer and 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 Fl score, surpassing all the other methods (Table The results indicated that XTrimoGene learns a well-represented cellular embedding (visualized App. Figure 7) by simply aggregating contextual gene embedding.
Table 2: The cell annotation performance on the Zheng68K and Segerstolpe dataset. x xTrimoGene is evaluated with 100M parameter model.
Method Name
Zheng68K
Segerstolpe
Precision
Fl score
Precision
Fl score
XTimmoGene
0.7335 + 0.0226
0.7354 + 0.0189
0.8112 + 0.0009
0.8140 + 0.0008
ScBERT
0.7029 + 0.0115
0.6695 + 0.0077
0.6818 + 0.0736
0.6703 + 0.0653
ACTINN
0.6720 + 0.0021
0.6486 + 0.0041
0.7545 0.0018
0.7219 + 0.0073
Scanpy
0.6111 + 0.0017
0.5474 + 0.0085
0.6274
+ 0.0000
0.5398 +: 0.0000
CellTypist
0.7454 = 0.0009
0.7151
0.0038
0.7923
0.0003
0.8117 0.0001
scVI
0.4883 + 0.0005
0.4843 = 0.0008
0.5101
+ 0.0022
0.5208 + 0.0016
singleCelINet
0.6452 + 0.0013
0.5982 0.0027
0.7551 0.0096
0.8055 + 0.0076

---

5.4.2 Perturbation response prediction
Recently, perturb-seq technology was established screen gene expression response given pooled perturbations at single-cell level [8]. Several algorithms have also been developed to predict per- turbation effects [30, 22 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 XTrimoGnne.
The normal state (before perturbation) gene expression profile is fed into xTrimoGene and we obtained the context embedding. which replaces 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 levels. As shown in TigureAA. GEARS with XTrimoGene embedding scores 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 cireumstances, including post-perturbation.. The application further proved the efficacy and potential of Xrimooenne to boost scRNA-seq based tasks.
different perturbation
(A)
Gears
xTrimoGene
(B)
Total
0.8
True Positive Rate
1-gene
0.6
2-9ene(seen0)
0.4
2-gene(seen1)
0.2
DeepDDS
2-gene(seen2)
DeepSynergy
Random Forest
0.0
0.0
0.1
0.2
0.0
0.2
0.4
0.6
0.8
Top20
Positive
Figure 4: (A) The MSE of the top 20 deferentially expressed (DE) genes given by different models on perturbation response prediction. The top 20 DE genes are calculated between the before and post-perturbation expression profiles. "Total" denotes evaluating all test perturbation sets. "I-gene" denotes evaluation on the single gene perturbation subset, where the perturbed target is not seen in the training set. "2-gene" represents the sub-test set for perturbing two genes simultaneously. "seen0", "seenl" " and "seen2" denotes zero, one or two perturbed targets are not seen in the training set, respectively. The black line denotes a 95% confidence interval. (B) ROC curve of different models on drug combination synergy prediction task. xTrimoGene denotes replacing the raw expression profile with context embeddings in the DeepDDS framework and others remain unchanged. Refer App. 8.3 for more details.
5.4.3 Synergistic drug combinations prediction
The drug synergistic task evaluates how patients or cells respond to a drug combination interven- tion [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 syner- gistic 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 XTimmoGnne is able to generate good latent embedding for this bulk expression data.
Similar 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 utilizing embedding from the xTrimoGene model outperforms all other models. The result proved that xTrimoGene can accurately capture cell-level representation, even for bulk sequencing data. This also opens the avenue for xTrimoGine be applied across other biological modeling tasks, especially where bulk-level transcriptome data is available.

---

6 Conclusion
xTrimoGene is 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 estab- lishes the projection strategy of continuous values with higher resolution. The results show that x 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 vre-trainnng 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, x 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 cell. The memory consumption for inference with the x XTrimoGene- 100M model is approximately 50GB, whose hardware requirement (Nvidia A100 80G GPU) beyond some academic labs, thus computational or memory-efficiet engineering techniques tend to advance the model pre-training and application.
xTrimoGene has been integrated into BioMap's single-cell analysis platform, functioning 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 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 xTrimo- Gene framework, nuto-discretization and unsupervised objectives, contributing concrete ideas and pseudocode, along with code review. Jing Gong and Mingsheng Hao (Research Intern at BioMap) 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 x XTrmmoenne 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 FLOOPs 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 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, 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 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 2002YFF1200004)) the Beijing Municipal Science & Technology Commission and the Administrative Commission of Zhongguancun Science Park/222111000032222222 And this work was also funded by BioMap.
References
[1] Nadav Brandes, Dan Ofer, Yam Peleg, Nadav Rappoport, and Michal Linial. ProteinBERT: universal deep-learning model of protein sequence and function. Bioinformatics, 38(8):2102- 21 10, 02 2022.
[2] 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, 13:187771901,, 2020.
10

---

[3] Geng Chen, Baitang Ning, and Tieliu Shi. Single-cell ma-seq technologies and related compu- tational data analysis. Frontiers in genetics, page 317, 2019.
[4] 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.
[5] Sijie Chen, Yanting Luo, Haoxiang Gao, Fanhong Li, Yixin Chen, Jiaqi Li, Renke You, Min- sheng Hao, Haiyang Bian, Xi Xi, et al. heca: The cell-centric assembly of a cell atlas. Iscience, 2545):004318, 2022.
[6] 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, 37666594)eaab15197, 2022.
[7] 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 Alex K. Shalek, Alexandra- Chhléé Villani, Aviv Regev, and Joshua Z. Levin. Systematic comparison of single-cell and single-nucleus rna-sequencing methods. Nature Biotechnology, (8(6))737-744 Jun 2020.
[8] Atray Dixit, Oren Parnas, Biyu Li, Jenny Chen, Charles P Fulco, Livnat Jerby-Amon, Ne- manja 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. 2016.
cell,
[9] Mario Flores, Zhentao Liu, Tinghe Zhang, Md Musaddaqui Hasib, Yu-Chiao Chiu, Zhenqing Ye, Karla Paniagua, Sumin Jo, Jianqiu Zhang. Shou-Jiang Gao, al. Deep learning tackles single- cell analysis- survey of deep learning for serna-seq analysis. Briefings in bioinformatics, 23(1) 23(1):bbab531, 2022.
[10] 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.
[11] Christopher Heje Grønbech, Maximillian Fornitz Vording, Pascal N Timshel, Casper Kaac Sønderby, Tune H Pers, and Ole Winther. VAE: variational auto-encoders for single-cell gene expression data. Bioinformatics, 36(16):4415 4422, 05 2020.
[12] Xu Han, Zhengyan Zhang, Ning Ding, Yuxian Gu, Xiao Liu, Yugi 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. Al Open, 2:225-250, 2021.
[13] 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.
[14] Dan Hendrycks, Kimin Lee, and Mantas Mazeika. Using pre-training can improve model robust- ness and uncertainty. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Leurning, volume 97 of Proceedings of Machine Learning Research, pages 2712-2721. PMLR, 09-15 Jun 2019.
zero-inflation controversy about serna-seq data. Genome Biology, 23(1):31, Jan 2022. [15] Ruochen Jiang, Tianyi Sun, Dongyuan Song, and Jingyi Jessica Li. Statistics or biology: the
[16] Dragomirka Jovic, Xue Liang, Hua Zeng, Lin Lin, Fengping Xu, and Yonglun Luo. Single-cell ma sequencing technologies and applications: brief overview. Clinical and Translational Medicine, 12(3):6694, 2022.
11

---

[17] John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ron- neberger. Kathryn Tunyasuvunakool, Russ Bates, Augustin Zídek. Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583-58)) 2021.
[18] 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 wxXv:2001.08361. 2020.
[19] 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):D59--D002 2022.
[20] 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 language model. Science, 379(6637) 1130, 2023.
[21] Romain Lopez, Jeffrey Regier, Michael B Cole, Michael Jordan, and Nir Yosef. Deep 2018.
generative modeling for single-cell transcriptomics. Nature methods,
[22] 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. Paz. Compositional perturbation autoencoder for single-cell response modeling. BioRxiv, 2021.
[23] Feiyang Ma and Matteo Pellegrini. Actinn: automated identification of cell types in single cell ma sequencing. Bioinformatics, 36(2):533-538, 2020.
[24] 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 ma-seq data using generative adversarial networks. Nature Communications, 11(1)(166)7786 793, 2020.
[25] 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):3802,, 2017.
[26] 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(64555)7786-793, 2019.
[27] Ahmad Pesaranghader, Stan Matwin, Marina Sokolova, lean-CChrsttophr Grenier, Robert G Beiko, and Julie Hussin. deepSimDEF: deep neural embeddings of gene products and gene ontology terms for functional analysis of genes. Bioinformatics, (8(11))3051-0611 05 2022.
(28) Kristina Preuer, Richard P Lewis. Sepp Hochreiter, Andreas Bender, Krishna C Bulusu, and Günter Klambauer. Dyeeryn: predicting anti-cancer drug synergy with Deep Learning. Bioinformatics, 34(9): 1538 1546-1546, 12 2017.
[29] 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-1197, 2020.
[30] Yusuf Roohani, Kexin Huang, and Jure Leskovec. Gears: Predicting transcriptional outcomes of novel multi-gene perturbations. BioRxiy, pages 2022-07, 07, 2022.
[31] Asa 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 tran- scriptome profiling of human pancreatic islets health and type diabetes. Cell Metab., 24(4):593-607, October 2016.
[32] 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.
12

---

[33] Jinxian Wang, Xuejun Liu, Siyuan Shen, Lei Deng, and Hui Liu. Deepdds: deep graph neural network with attention mechanism predict synergistic drug combinations. Briefings in Bioinformatics, 23(1 23(1))bbab390, 2022.
[34] F Alexander Wolf, Philipp Angerer, and Fabian J Theis. Scanpy: large-scale single-cell gene expression data analysis. Genome biology, 19:1-5, 2018.
[35] 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.
[36] Fan Yang. Wenchuan Wang, Fang Wang, Yuan Fang, Duyu Tang, Junzhou Huang. Hui Lu, and Jianhua Yao. scbert as large-scale pretrained deep language model for cell type annotation of single-cell raasseq data. Nature Machine Intelligence, 4(10):852-866, 2022.
[37] Patrick Cahan Yuqi Tan. Singlecellnet: A computational tool to elassify single cell ma-seq data across platforms and across species. Cell Systems, 19(2):207-213, 2019.
[38] 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.
[39] 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):1404) 2017.
13


Reviews:
[{"review_id": "VN8cojcjca", "reviewer": "Reviewer_9yec", "summary": "Here 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.", "questions": "See \"Weaknesses\" section.", "limitations": "I 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.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* **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).", "weaknesses": "The 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the constructive comments. We have shown additional discussions and experiments to further strengthen our work. The point-by-point responses to the comments are as follows. If our response does not fully address your concerns, please post additional questions and we will be happy to have further discussions.\n\n> **Q1: Was data leakage avoided?**\n\n**A1**: We did not specifically introduce a precaution strategy during the pre-training data preparation.\n\nHowever, we suppose this has no significant influence on the downstream task evaluation for the following two reasons:\n\n1. Nature of Training:\n   - The pre-training process is self-supervised, focused solely on gene expression values, while other meta-information like cell type, tissue type, sequencing platform, etc., being disregarded.\n   - Conversely, the downstream tasks (such as perturbation prediction) are supervised, employing distinct cell label information or relationship of pre-post perturbed cells.\n\n2. Learning Focus:\n   - During pre-training, the model is primed to understand intrinsic gene relationships without any exposure to label information. While for the downstream task, the model learns the relationship between gene expressions and cellular labels.\n\nHence, even if there were overlaps between pre-trained and downstream evaluation data, the lack of exposure to label information during pre-training should safeguard against any significant data leakage effects.\n\n> **Q2: Proper modeling of scRNA-seq count distributions.**\n\n**A2**: The modeling of xTrimoGene is different from DCA. DCA uses the observed gene expression to predict the parameters in the ZINB distribution. In contrast, xTrimoGene is a masked autoencoder [1], i.e., recovering the unobserved gene expression values depends on other genes within a cell. MSE loss is generally a default choice. However, the pre-training data distribution of raw counts is very skewed, where most values are small but a few are very large, using MSE may cause the model to focus too much on those large values. Therefore, we use the normalized expression value.\n\nFollowing the suggestion, we explored an ablation study to compare the raw counts versus normalized expression values in our modeling. Concretely, we trained a 10M parameter model with raw counts as input, while all the other configurations keep consistent as the xTrimoGene-10M model. Then, the pre-trained model is evaluated on the PBMC3K clustering task, in which the raw count matrix is fed into the model to obtain the context embedding for subsequent analysis. The results show that pre-training with normalized gene expression value achieves a better performance than raw counts, which is consistent with our aforementioned assumption.\n\n|Input| ARI| NMI| HOMO| CP|SIL|\n|-|-|-|-|-|-|\n| Normalized value | 0.7767|0.7810|0.7841|0.7778 | 0.1406 |\n| Raw count| 0.6575| 0.7318|0.7459| 0.7183|0.1086|\n\n[1] He et al. Masked Autoencoders Are Scalable Vision Learners. 2021\n\n> **Q3: How to handle batch effects?**\n\n**A3**: For xTrimoGene, two strategies could handle batch effects:\n\n1. Fine-Tuning with Batch Information: xTrimoGene can be fine-tuned with user-specific datasets incorporating batch information. We achieve this by: \n(1) Converting the batch ID to embeddings via a lookup table.\n(2) Summing these batch embeddings with the original value and gene embeddings. This aggregated input is then processed by the transformers. During fine-tuning, gene expression embeddings from different datasets are enriched with their corresponding batch embeddings.\n(3) At the inference stage, to harmonize cells from different batches, a consistent batch embedding can be utilized, ensuring the output remains consistent across batches.\n\n2. Integration with Other Batch Correction Techniques: One could dump embeddings from xTrimoGene and feed them into other methods.\nFor this strategy, we used the pancreas dataset (comprising 8 batches and is a benchmark for batch integration) as an example. We subsampled this dataset to 3k cells for efficiency.\nAfter getting embeddings from xTrimoGene, we employed a lightweight method, BBKNN(Polański, et al. Bioinformatics 2020) to correct for batches. As depicted in Figure R1, the results of xTrimoGene+BBKNN combination mitigated batch effects and maintained cell-type variations.\n\n> **Q4: Rare cell types**\n\n**A4**: To investigate how xTrimoGene behaves on these unseen data, we conduct the following evaluation analysis.\n\n1. We first collected the scRNA-seq data from a previous study (Ji et al., 2020, Cell), which profiles human skin squamous cell carcinoma landscape with scRNA-seq and spatial transcriptomics technology simultaneously. The scRNA-seq data is not present during xTrimoGene training process. Notably, the authors found that clustering the scRNA-seq data (Figure R2, left panel) yields a novel cell subgroup named TSK (Tumor-Specific Keratinocyte), which is clearly in the tumor sample  but not the normal sample. \n\n2. We employed the tumor scRNA-seq data to explore whether xTrimoGene is robust to distinguish the cell subpopulations from others. The expression matrix is fed into xTrimoGene and the dumped context embedding is used for subsequent UMAP visualization. The results show that the TSK subgroup is clearly separated (Figure R2, right panel). More importantly, the two TSK subgroups are merged with xTrimoGene, demonstrating its generalization ability to generate good cell-specific embeddings for unseen data.\n\n> **Q5: Memory/hardware requirements?**\n\n**A5**: For the pre-trained xTrimoGene models, the memory consumption for inference with a sample of approximately 2000 non-zero expressed genes is approximately 50GB for the xTrimoGene-100M model and around 18GB for the xTrimoGene-10M model. It's worth noting that, in line with our pre-training settings, we conducted our tests using bf16 mode on an Nvidia A100 80G GPU. We hope this provides clarity, and we're here to offer further details if necessary."}, {"review_id": "oQBSdInkZt", "reviewer": "Reviewer_MmMn", "summary": "The 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.", "questions": "- 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? ", "limitations": "Limitations 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.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The 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.", "weaknesses": "\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\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**WA1**: We are grateful for all your comments and will revise our manuscript to make more discussion and emphasize the connection. In our scenario, this asymmetric encoder-decoder architecture is not only efficient but also designed for handling the high sparsity in scRNA-seq data. Specifically, we feed the non-zero expressed genes into the encoder and let the zero expressed genes only be processed by the decoder.\n\n**WA2**：\n1. Using the count values to query embeddings can cause the model to overly capture noise since the MSE loss is sensitive to the data scale. This noise makes it difficult for model optimization. Therefore, the normalized data is more applicable for model training and a \"coarser-grained\" representation strategy- discretization or quantization, is needed to make the continuous values be mapped into similar embeddings.\n2. Our method is more aptly described as a form of soft discretization. This contrasts with traditional methods of hard discretization, which can lead to problems like Similar value But Dissimilar embedding (SBD) as highlighted in ll53-ll55. The soft discretization, allows for flexibility, effectively addressing the SBD challenge. Additionally, our Auto-discretization method is differentiable, supporting end-to-end training. Also, the embedding methods for numerous or continuous values can be observed in other fields[1,2]. Appendix Figure 1 shows bin distribution for low-expression values. This is because the majority of gene expressions in our training data fall within the 1-2 range. Those exceeding a value of 4 are relatively rare. We are sorry for the confusion and the shapes for all the matrices are added in the paragraph.[1] AutoEmb: Automated ... Recommendations \n[2] AutoDis: An Embedding ...  Prediction  \n\n3. Both the look-up table entries and the weights, w1 and w2, are learnable. The process can be viewed as a soft discretization, and while the outcome is a weighted sum of the embedding table (refer to Q2), we label it as a \"look-up\" due to its meaning of mapping raw data to its respective embedding, we are sorry for this confusion. The low expression values are not strictly discrete since the data distribution is not uniform.\n\n**WA3**: \nIn total, 1,140 (5.7\\% of all genes) entries are masked, which on average includes 600 non-zero entries (30\\% of total non-zero value genes) and 540 zero entries (3\\% of total zero value genes).\n\nWe ablated the ratio of the zero under the same total masking ratios/numbers (5.7\\% of all genes, n=1,140) as in the public comments experiment no.4. For each setting, we trained a 10M parameter model over 5 million data for 50,000 steps. Then, we evaluated the model on PBMC3K downstream cell clustering task. As the zeros masking ratio increases, the performance tends to be better and then achieves a plateau (1.3\\% - 3.0\\%). The current zeros masking ratio lies in the plateau interval.\n\nFor appendix Figure 4, we aim to investigate if all the masked values should be replaced with a [MASK] token. The three percentages (80\\%,10\\%,10\\%) represent the probability of masked values being replaced with a [MASK] token, a random expression token, and the original token, respectively.\n\n**WA4**:\nThank you for the proposed logic to organize the work, we have revised the main text accordingly. \n\n**WA5**:\nYes, the comparison can also be applied to other tasks with produced cell embedding. Here we investigate whether xTrimoGene is advantageous over scBERT on the perturbation effect prediction task. Similar to xTrimoGene, the generated embedding from scBERT is in conjunction with GEARS. As shown in the public comments experiment no.5, xTrimoGene achieves a lower MSE value than scBERT across all metrics, demonstrating its superior efficiency in generating cellular context embedding under different biological conditions.\n\n**Minor**: For minor comments, we revised the manuscript accordingly. And for datasets, we used two datasets named PBMC and Zheng68K with ~2k and 68k cells, respectively.\n\n**Q2**: The numbers in brackets are the absolute depth/head value between the encoder and decoder. We deleted the numbers and added the reference to Appendix Table 2 for clarification.\n\n**Q3**: Equation 5 defines the loss at masked positions (including zero and non-zero entries), whose range and prefactor vary across different samples.\n\n**Q4**: The ablation studies across the paper are all conducted with model context embedding. Concretely, we fed the PBMC3K expression matrix into the evaluated model and dumped the context embedding. Then, we utilized the embedding to cluster all the cells and calculated the clustering metrics. In the scenario of classification pre-training mode (Figure 2B), the expression value is rounded into an integer (each integer as a token) and the pre-training objective is to predict the token (which is discrete). After the pre-training stage is finished, the model is evaluated on the ablation cell clustering task as aforementioned with the exception that the PBMC3K expression matrix are rounded integers.\n\n**Q5**: Sorry for the typo, the result is from xTrimoGene-100M model.\n\n**Q6**:  We can fine-tune the pre-trained model in mice as the scRNA data is comparable with humans. Upon fine-tuning over a large amount of data, the model is possible to capture the intrinsic species-specific regulations and learn good representations for species-specific genes.\n\n**Q7**:  Yes, the gene embeddings are changed along with genes instead of the cell. For Fig 1 upper right corner part, we are retrieving the expression and gene embeddings for the unmasked-only matrix. The matrix has been processed with filtering and padding steps. In the former step, we filtered out the zeros and masked entries, which positions/orders are not consistent across different cells. Thus, we utilized a color gradient to denote the gene differences.\n\n**Q8**: The former, they are also converted to embedding similar to unmasked entries, instead of pre-processed value.\n"}, {"review_id": "gveM8Oikvz", "reviewer": "Reviewer_TD4Y", "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.", "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.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The training process is accelerated by utilizing sparse input.\n\n2. A novel auto-discretization strategy is introduced to map continuous expression values into a latent embedding space.\n\n3. Considering its high accuracy and efficiency, the proposed large model xTrimoGene is a valuable contribution to the single-cell community.\n\n4. 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.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the constructive comments. Here we showed additional analysis and discussions to further strengthen our work, with a focus on the advantage over CellTypist. If our response does not fully address your concerns, please post additional questions and we will be happy to have further discussions.\n\nWe have further explored the performance of xTrimoGene in comparison to CellTypist and demonstrated the following advantages and potential biological insights:\n\n1. xTrimoGene is more robust to identify rare cell types than CellTypist. \n\nFor the Zheng68K annotation task, xTrimoGene achieves a marginal improvement over CellTypist (F1-score: 0.7354 versus 0.7151). However, the boosting performance across individual cell type varies. Among all the profiled 11 different cell types, xTrimoGene gains a large margin (F1 score: 0.21 versus 0.0) for CD4+ T Helper2 cell type, which is the smallest subgroup in the total population. The observation suggests xTrimoGene is superior over CellTypist to detect and distinguish rare cell types.\n\n|Cell type|xTrimoGene-Precision | xTrimoGene-Recall | xTrimoGene-F1 | CellTypist-Precision | CellTypist-Recall | CellTypist-F1|\n| :--- | :---: | :---: | :---: | :---: | :---: | ---: |\n| CD14+ Monocyte (195) | 0.86 | 0.81 | 0.84 | 0.86 | 0.85 | 0.85 |\n| CD19+ B (558) | 0.96 | 0.81 | 0.88 | 0.90 | 0.84 | 0.87 |\n| CD34+ (19) | 0.90 | 0.95 | 0.92 | 1.00 | 0.84 | 0.91 |\n| CD4+ T Helper2 (9) | **0.20** | **0.22** | **0.21** | 0.00 | 0.00 | 0.00 |\n| CD4+/CD25 T Reg (612) | 0.78 | 0.68 | 0.73 | 0.72 | 0.69 | 0.71 |\n| CD4+/CD45RA+/CD25- Naive T (185) | 0.53 | 0.71 | 0.61 | 0.66 | 0.54 | 0.59 |\n| CD4+/CD45RO+ Memory (303) | 0.65 | 0.64 | 0.64 | 0.70 | 0.47 | 0.56 |\n| CD56+ NK (853) | 0.93 | 0.90 | 0.91 | 0.93 | 0.92 | 0.92 |\n| CD8+ Cytotoxic T (2031) | 0.87 | 0.88 | 0.87 | 0.86 | 0.83 | 0.84 |\n| CD8+/CD45RA+ Naive Cytotoxic (1636) | 0.84 | 0.91 | 0.88 | 0.80 | 0.94 | 0.87 |\n| Dendritic (194) | 0.80 | 0.84 | 0.82 | 0.84 | 0.83 | 0.83 |\n\n\n\n\n2. xTrimoGene embedding reveals potential cell type specific gene-gene networks.\n\nxTrimoGene, as a pre-trained model, can generate unique embeddings which encapsulate intrinsic gene-gene relationships. In the Zheng68K dataset, after cell-type annotation, we conducted the differential gene expression analysis between B and non-B cell groups. And then we retrieved the top 10 genes specific to B cells and retrieved the context embedding for these 10 genes from xTrimoGene. These embeddings are used to construct a gene-gene network (Figure R3), where the edge represents gene embedding similarity. From the gene network, we found HLA genes family (HLA-) have higher similarities within each other, while gene CD74 has less similarity with others. The analysis illustrates that xTrimoGene embedding can be utilized to decipher cell type specific gene networks, which is not within the scope of CellTypist.\n\nIn summary, xTrimoGene not only outshines in handling rare cell types but also offers the potential to decode gene-gene relationships, setting it apart from existing methods. "}, {"review_id": "O1KAANnFeU", "reviewer": "Reviewer_N2nG", "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.", "questions": "Please refer to the weaknesses.", "limitations": "No significant limitations of the method are found.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "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.\n2. 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.\n3. The embedding of the proposed method and existing works could be visualized to provided an intuitive understanding of the performance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the constructive comments. Here we showed additional analysis and discussions to further strengthen our work. If our response does not fully address your concerns, please post additional questions and we will be happy to have further discussions.\n\n**A1**: Introducing padding tokens within a batch is major compatible with parallel matrix operation. However, computes related to the [PAD] placeholder are useless as the corresponding attention is all masked out in implementation. In this scenario, diminishing computes on [PAD] token will further boost the training efficiency. \n\nOne widely used strategy to avoid padding is sequence packing [1], which theoretically concatenates all the variable-length sequences into one sequence and then splits it into individual fragments with a fixed length. The resulting batches have a consistent shape and no further padding is needed. The method has been utilized in natural language modeling and protein sequence pre-training. However, the design is not applicable to xTrimoGene as gene expression pattern is cell specific, cross-concatenation will introduce noise signal and collapse the intrinsic cell representations.\n\nAnother strategy called bucketing [2] tends to fit xTrimoGene more closely. While preparing the training set, one can sort the samples by expression pattern distribution. Specifically, the single-cell samples are sorted by dropout ratio and thus cells with similar non-zero length are divided into the same buckets. This will greatly reduce the needed [PAD] tokens within a batch and potentially accelerate the overall training process.\n\n[1] M Kosec et al, Packing: Towards 2x NLP BERT Acceleration. 2021\n\n[2] Tom Kocmi et al, Curriculum Learning and Minibatch Bucketing in Neural Machine Translation. 2017\n\n**A2**: Though individual component of the architecture has been explored before, the overall design is predominately motivated to align scRNA-seq data characteristics. Regarding the novelty, we would like to highlight the following contributions and advancements over previous methods.\n\n1. xTrimoGene is the first asymmetrical encoder-decoder architecture to guide single-cell RNA-seq data pre-training. The scRNA-seq is highly sparse, thus using the encoder-only design tends to introduce huge amounts of redundant computations. The encoder in xTrimoGene concentrates on capturing intrinsic features from the most informative non-zero values, while the decoder additionally integrates zero-value genes to further tweak the gene-gene interactions. This strategy empowers the model to learn gene representations both efficiently and comprehensively.\n\n2. Projection of expression value needs to maintain the continuous feature. Different from language tokens, the 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. We verified the effectiveness of the proposed auto-discretization strategy.\n\n3. Training strategy setting is comprehensively ablated. We have conducted a series of ablation studies to validate the optimal training configurations. Concretely, we proved regression task (objective) is superior to the conventional classification one (Figure 2A). Instead of masking all positions randomly, we need to separately mask zero and non-zero values with a trade-off ratio (Appendix Figure 2,3,4).\n\n4. We scaled the pre-trained model and achieved remarkable performances on multiple downstream tasks, including cell type annotation, perturbation prediction and synergistic drug combination prediction. The results demonstrate the generalization ability of xTrimoGene, which is expected to drive more advancements in other downstream task learning.\n\nIn summary, xTrimoGene is developed to pre-train large-scale scRNA-seq data efficiently, where the underlying design is motivated and optimized to align scRNA-seq data characteristics. We envision the established framework is meaningful for further algorithm improvement.\n\n**A3**: For all the compared methods related to cell type annotation tasks, we investigated the embedding distributions with Zheng68K data set and uncovered a potential connection to model performance. \n\nSpecifically, CellTypist and scVI are included for comparison versus xTrimoGene. We projected the embedding in a UMAP plot, where the cells are colored by either ground truth (upper row) or model predicted (bottom row) cell type labels (Figure R4). The model performance is correlated with the consistency level between the two plots.\n\nIn contrast with CellTypist, xTrimoGene achieves a better performance in predicting the CD19+ B cell type (Figure R4, 1-1 v.s. 1-2). Meanwhile, the CellTypist tends to identify some CD8+ Cytotoxic T cells (2-1) as CD19+ B cell type (2-2). The CD8+ Cytotoxic T cell is the largest sub-population across all 11 cell types, potentially leading to the inaccurate assignment. scVI performs much worse than xTrimoGene, as a proportion of CD8+ Cytotoxic T cells (3-1) are detected as CD56+ NK cell types (3-2). It seems that two batches of Cytotoxic T cells are present and scVI is inferior to discriminate the batch effects. Notably, a smaller subgroup CD34+ cells (4-1) are nearly assigned as Dendritic (4-2) cell type incorrectly. This may suggest scVI has a limited resolution to separate rare cells.\n\nVisualization analysis of embedding is useful to interpret model behavior, especially when it remains challenging to establish explainability for deep models. The paradigm provides a convenient and intuitive manner to validate the performance and will be further utilized to decode the limitations and advantages of xTrimoGene."}]

Author rebuttal:
We extend our sincere gratitude to all the esteemed reviewers for dedicating their time and expertise to meticulously evaluate our work. Their valuable feedback has significantly contributed to enhancing the quality and depth of our research. We have conscientiously delved into every facet of the comments.

Guided by the insightful suggestions, we have made several noteworthy adjustments to the content. These revisions encompass elucidating intricate concepts for improved clarity, bolstering our textual content, and incorporating additional analyses. To elucidate, we have conducted a comprehensive array of ablation studies and experimental investigations, which in turn have shed light on the architectural intricacies, merits, and limitations of xTrimoGene.

Here we briefly summarize the added experiments:

1. **Ablation study on normalized value versus raw count**. The results show that pre-training with normalized gene expression value achieves a better performance than raw counts. The details are responded to Reviewer 9yec.

2. **Application on handling batch effects**. We discussed the potential usage with xTrimoGene to remove batch effects. Concretely, we provided an analysis (Figure R1) and demonstrated xTrimoGene's capability on the issue. The details are responded to Reviewer 9yec.

3. **Rare cell types detection**. We explored the ability of xTrimoGene to identify rare cell types from a large population (Figure R2). The details are responded to Reviewer 9yec.

4. **Ablation of the relative masking frequency of zeros**. Apart from the presented ablation study on masking strategy, we provided an additional investigation on zero masking frequency. The results show that the current configuration is within the optimal interval. The details are as below and responded to Reviewer MmMn.

| Value    | Masked1     | Masked2    | Masked3   | Masked4    | Masked5    | Total   |
|----------|-------------|------------|-----------|------------|------------|---------|
| $\neq$ 0 | 100(5%)    | 300(15%)  | 600(30%) | 900(45%)  | 1,100(55%) | 2,000   |
| = 0      | 1,040(5.8%) | 840(4.7%) | 540(3%)  | 240(1.3%) | 40(0.2%)  | 18,000  |
| Sum      | 1,140       | 1,140      | 1,140     | 1,140       | 1,140       | 20,000  |
| Sum ratio        | (5.7%)     | (5.7%)    | (5.7%)   | (5.7%)    | (5.7%)    | (100%) |
| ARI      | 0.6654      | 0.5043     | 0.7767    | 0.7817     | 0.5048     |         |
| NMI      | 0.7170      | 0.6884     | 0.7810    | 0.7833     | 0.6348     |         |
| HOMO     | 0.7285      | 0.7481     | 0.7841    | 0.78226    | 0.6826     |         |
| CP       | 0.7058      | 0.6375     | 0.7778    | 0.7843     | 0.5932     |         |
| SIL      | 0.1418      | 0.1522     | 0.1406    | 0.1675     | 0.1247     |         |


5. **Comparison to other models on downstream task**. We also included scBERT for comparison on the perturbation effection prediction task. The results indicate that xTrimoGene is superior over scBERT to capture the intrinsic context embedding under perturbed conditions. The details are as below and responded to Reviewer MmMn.

| Pre-trained model | Total  | 1-gene | 2-gene(seen0) | 2-gene(seen1) | 2-gene(seen2) |
|-------------------|--------|--------|---------------|---------------|---------------|
| xTrimoGene    | 0.1983 | 0.1930 | 0.2385        | 0.2100        | 0.1286        |
| scBERT            | 0.2231 | 0.2116 | 0.2581        | 0.2386        | 0.1522        |


6. **In-depth comparison with CellTypist**. For the cell type annotation task, we performed further analysis and showed advantages in two aspects: (1) xTrimoGene is more robust to identify rare cell types than CellTypist; (2) xTrimoGene embedding reveals potential cell type specific gene-gene networks (Figure R3). The details are responded to Reviewer TD4Y.

7. **Embedding visualization analysis**. We also utilized embedding visualization to interpret the behavior of three models on the cell type annotation task (Figure R4). The details are responded to Reviewer N2nG.

We take this opportunity to reiterate our gratitude to each reviewer, as their incisive comments have propelled our work. If our responses have not entirely addressed your concerns, we cordially invite you to share additional queries. We stand ready to engage in further discussions and provide any necessary clarifications.


Meta-review: All reviewers liked the paper and both authors and reviewers engaged in a good discussion that has helped improve the paper. Acceptance is recommended.
Decision: Accept (poster)