SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jing Gong $^{1*}$  Minsheng Hao $^{12*}$  Xingyi Cheng $^{1\dagger}$  Xin Zeng $^{1}$   
Chiming Liu $^{1}$  Jianzhu Ma $^{2}$  Xuegong Zhang $^{2}$  Taifeng Wang $^{1}$  Le Song $^{13\dagger}$ $^{1}$  BioMap Research  $^{2}$  Tsinghua University  
 $^{3}$  Mohamed bin Zayed University of Artificial Intelligence
{gongjing, minsheng_2022, xingyi, zengxin, chiming, taifeng, song1e}@biomap.com, {zhangxg, majianzhu}@tsinghua.edu.cn
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) $^{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.
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
*Equal contribution. Mingsheng Hao conducted this work during his internship at BioMap.
$\dagger$  Correspondence to: xingyi@biomap.com, songle@biomap.com.
$^4$  xTrimoGene is a member of BioMap's large-scale AI engine "xTrimo" series. "α" denotes the academic version, to distinguish it from the commercial product xTrimoGene.
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 low-dimensional cell embeddings, cscGAN [24] uses a Generative Adversarial Network (GAN) architecture 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 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-attention between zero tokens). It required approximately  $2.65 \times 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.
Figure 1: The xTrimoGene Framework: (1) Random positions (including both zero and nonzero values) are masked for prediction. (2) Masked and zero-valued positions are filtered out. (3) Remaining unmasked positions are aligned with padding tokens (grey) to ensure maximum length consistency within a batch. (4) Gene expression values and gene embeddings are 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 to model output with an MLP layer. The MSE loss is calculated between the model output and ground truth values for the masked positions.
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_{\text{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.
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_{\text{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 \in \mathbb{R}^{c \times m \times 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.
$$
E = \operatorname {A u t o b i n} \left(V _ {\text {u n m a s k e d}} \odot M _ {\text {n o n z e r o}}\right), G = f _ {L} (\text {g e n e s}), I = E + G \tag {1}
$$
Then the encoder processes the input embeddings  $I$  and generates the high-level gene representations  $I_{encoder} \in \mathbb{R}^{b \times m \times d}$  via the multi-head attention mechanism.
$$
I _ {\text {e n c o d e r}} = \operatorname {T r m} \left(f _ {Q} (I), f _ {K} (I), f _ {V} (I)\right) \tag {2}
$$
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.
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.
$$
I _ {f u l l} = W _ {p} \left(I _ {\text {e n c o d e r}} \oplus I _ {\text {z e r o}} \oplus I _ {\text {m a s k e d}}\right) + b _ {p} \tag {3}
$$
where  $\oplus$  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} \in \mathbb{R}^{b \times n \times d}$ , and predicts the masked values through a shared linear layer,  $W \in \mathbb{R}^{d \times 1}$ , applied across all genes. The operations are expressed as follows:
$$
I _ {\text {d e c o d e r}} = \operatorname {T r m} \left(\left(f _ {Q} \left(I _ {\text {f u l l}}\right), f _ {K} \left(I _ {\text {f u l l}}\right), f _ {V} \left(I _ {\text {f u l l}}\right)\right), \tilde {V} = I _ {\text {d e c o d e r}} \cdot W \right. \tag {4}
$$
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  $\mathbb{R}^{d \times 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 \cdot 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} = \text{Leaky\_ReLU}(v_{1})$ . Subsequently, a cross-layer projection is applied, represented by  $v_{3} = w_{2} \cdot v_{2} + \alpha \cdot v_{2}$ . Here,  $w_{2}$  denotes the weight vector, and  $\alpha$  is a scaling mixture factor. Next, the bin weights of  $v_{3}$  are normalized using the softmax function, resulting in  $v_{4} = \text{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 \cdot 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
Figure 2: Pre-training strategy ablation study. (A) Performance comparison between auto discretization 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 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:
$$
\operatorname {L o s s} = \frac {1}{(n - m) * c} \sum \left(V _ {i, j} - \tilde {V} _ {i, j}\right) ^ {2} \tag {5}
$$
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.
Next we will explain our experimental settings and results. The dataset description can be referred to App. 1.
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.