SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Table 1: Computational efficiency comparison between different algorithms. The resource column is normalized by the Transformer row.
Model name	Parameter (M)	Forward + backward (FLOPs/sample)	Total train (FLOPs)	Resource
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 [18, 2]. Vanilla Transformer models are challenging to scale due to their computational time and resource requirements, which increase quadratically with model size. Varieties of attention mechanisms have been proposed to accelerate training speed, a critical factor for model scaling.
To test the scale-up ability of xTrimoGene, we pre-trained three models across multiple compute regions and scales (e.g., from 3M to 100M parameters). The detailed hyperparameter setting is displayed in the App. Table 2. The training curve clearly shows all models are steadily down to a lower loss when training steps increase (App. Figure 6). More importantly, the xTrimoGene-100M model obtains a significant improvement over the xTrimoGene-10M model, which is also superior to the xTrimoGene-3M model. The tendency is consistent across different data sizes. The results suggest xTrimoGene framework is robust to scale-up, making it possible and convenient to pre-train larger models with more data.
5.3 Robustness on high sparse data
scRNA-seq data often exhibit varying levels of sparsity, thus it's necessary to assess whether xTrimoGene is robust in handling different sparse data. To verify the robustness, we divided the test samples into subgroups based on cell type and calculated the sparsity level (i.e., percentage of zero values in the expression matrix) and Pearson correlation coefficient between the predicted and actual values. Our results reveal that the correlation gradually decreases as the sparsity level increases, as expected (Figure 3A). However, the correlation remains above 0.8 even when the sparsity level reaches  $96\%$  (Figure 3A), indicating the robustness of xTrimoGene. We also compared xTrimoGene's performance with Performer and found that xTrimoGene consistently achieves a higher correlation across most subgroups (Figure 3B). These findings demonstrate that xTrimoGene is robust in handling highly sparse data and outperforms encoder-only architectures.
The performance of the encoder-decoder and encoder-only architectures have been comparatively analyzed in the NLP domain, with the former demonstrating effectiveness in language comprehension and the latter in context generation. Apart from comparison on masked value recovery, we further evaluated xTrimoGene against encoder-only Performer on the cell clustering task. The results
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 as the ratio between zero value percentages. Pearson correlation coefficient metric is calculated on masked positions. (B) Performance comparison of xTrimoGene and Performer while recovering masked values at different sparse levels. Dot has the same meaning as (A) but the dot size is proportional to the sparse level. Both the x and y axis denotes the Pearson correlation coefficient metric for a particular algorithm. (C) Comparison of performance for xTrimoGene framework and encoder-only framework. Cell clustering task is evaluated.
demonstrate that xTrimoGene achieves superior performance, reaffirming its proficiency in latent embedding extraction (Figure 3C).
5.4 Evaluation on downstream tasks
Currently, multiple tasks have been established to evaluate different models, including well-defined cell type annotation and recently developed perturbation response prediction tasks. We first assessed the performance of xTrimoGene on these single-cell tasks. Additionally, we explored the potential application on bulk RNA-sequencing data, with a focus on synergistic drug combination prediction.
5.4.1 Cell type annotation
First, we evaluated xTrimoGene's performance on cell type annotation task with Zheng68K [39] and Segerstolpe [31] dataset, which has been widely benchmarked. We compared the xTrimoGene against other several methods, including scBERT [36], ACTINN [23], Scany [34], CellTypist [6], scVI [21] and singleCellNet [37]. For the xTrimoGene model, we added a max-pooling layer and a linear layer to predict cell type labels with fine-tuning mode (see App. 8.1). For other methods, we followed their instruction with the default parameter setting. We observed that xTrimoGene achieves a high Precision and F1 score, surpassing all the other methods (Table 2). The results indicated that xTrimoGene learns a well-represented cellular embedding (visualized in App. Figure 7) by simply aggregating contextual gene embedding.
Table 2: The cell annotation performance on the Zheng68K and Segerstolpe dataset. xTrimoGene is evaluated with 100M parameter model.
Method Name	Zheng68K	Segerstolpe
Precision	F1 score	Precision	F1 score
xTrimoGene	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
singleCellNet	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 to screen gene expression response given pooled perturbations at single-cell level [8]. Several algorithms have also been developed to predict perturbation effects [30, 22] at the single cell level, i.e., what is the expression value of genes after perturbation? We compared the native GEARS[30] model with and without incorporating embeddings from xTrimoGene.
The normal state (before perturbation) gene expression profile is fed into xTrimoGene and we obtained the context embedding, which replaces raw expression value input in the GEARS model (App. 8.2). All the other settings remain unchanged. The evaluated dataset (Norman et al. [26]) contains both single and double gene perturbation and we thus assess the performance across different perturbation levels. As shown in Figure 4A, GEARS with xTrimoGene embedding scores a lower MSE (decreased  $14.8\%$  ) for top20 differential expressed genes across all perturbation scenarios. Notably, the tendency is consistent across different perturbation levels, regardless the perturbed target is seen or not. We also compared against the scBERT embedding and observed a similar trend, where xTrimoGene achieves better results (App. Table 3). The results demonstrated that the pre-training strategy empowers xTrimoGene to capture constraints under various circumstances, including post-perturbations. The application further proved the efficacy and potential of xTrimoGene to boost scRNA-seq based tasks.
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. "1-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", "seen1" 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 to 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 intervention [25]. However, the generated wet-lab experimental data only covers a tiny search space of possible drug combinations. Multiple models have been proposed to accelerate predicting the synergistic landscape of drugs [28, 33]. For instance, DeepDDS integrates genomic expression profiles and drug chemical information, greatly improving the prediction performance. We further explored whether xTrimoGene is able to generate good latent embedding for this bulk expression data.
Similar to the perturbation prediction test, we adapted xTrimoGene to DeepDDS with the intermediate context embedding (see App. 8.3). We also included DeepSynergy and Random Forest for comparison. As illustrated in Figure 4B, utilizing embedding from the xTrimoGene model outperforms all the other models. The result proved that xTrimoGene can accurately capture cell-level representation, even for bulk sequencing data. This also opens the avenue for xTrimoGene to be applied across other biological modeling tasks, especially where bulk-level transcriptome data is available.
xTrimoGene is a new, efficient framework for learning scRNA-seq data. It proposes an asymmetric encoder-decoder framework that takes advantage of the sparse gene expression matrix and establishes the projection strategy of continuous values with a higher resolution. The results show that xTrimoGene is scalable and performs well on tasks like cell type annotation, perturbation response prediction, and synergistic drug combination prediction. The experiments demonstrate the efficacy of pre-training in single-cell biology. xTrimoGene is potentially adapted to other types of cell modeling analysis, including rare cell detection (App. 8.4), batch effect removal and regulatory network construction.
Certain limitations exist for xTrimoGene and further work is desired to advance the design. At present, xTrimoGene major utilizes gene expression values during the pre-training stage, overlooking varieties of other related meta-information like sample condition (health/disease), cell type, tissue type, sequencing platform, etc. These rich annotations are biologically meaningful and highly correlated with the expression pattern within a cell. The memory consumption for inference with the xTrimoGene-100M model is approximately 50GB, whose hardware requirement (Nvidia A100 80G GPU) is beyond some academic labs, thus computational or memory-efficient engineering techniques tend to advance the model pre-training and application.
xTrimoGene has been integrated into BioMap's single-cell analysis platform, functioning as a fundamental and essential model (as depicted in the App. Figure 9). The pre-trained model services have been publicly available. In the future, with the increase of data, larger pre-trained models are expected to drive more advancements in various downstream task learning.
Author Contributions and Acknowledgments
Le Song led the project by designing its scope, conceptualizing ideas, integrating resources, and making decisions on techniques. Xingyi Cheng played a key role in the development of the xTrimoGene framework, auto-discretization and unsupervised objectives, contributing concrete ideas and pseudocode, along with code review. Jing Gong and Mingsheng Hao (Research Intern at BioMap) were primarily responsible for conducting pre-training and downstream experiments, serving as the first authors of the paper. Their work covered areas such as model scaling, cell type annotation, perturbation response prediction, and synergistic drug combinations prediction. Xin Zeng made significant contributions to the code of the xTrimoGene framework, worked with an early performer version, and conducted initial downstream experiments. Chiming Liu oversaw the engineering aspects of the project, including the implementation of the data pipeline and FLOPs computation. Jianzhu Ma, Xuegong Zhang, Taifeng Wang, and Le Song conceived the project and provided invaluable guidance for the project and contributed their expertise in computational biology knowledge. Taifeng Wang also played a pivotal role in pushing for the model's service implementation. Finally, Jing Gong, Mingsheng Hao, Xingyi Cheng, and Le Song collectively contributed to writing this paper.
In addition, we would like to express our gratitude to the individuals at BioMap who have made contributions to our project. Chenrui Xu and Yucheng Guo played roles in data preprocessing and integration. Zhaoren He's expertise in data analysis and application greatly enhanced our work, and we deeply appreciate his contributions.
This work was supported by the Ministry of Science and Technology of the People's Republic of China (2022YFF1203004), the Beijing Municipal Science & Technology Commission and the Administrative Commission of Zhongguancun Science Park(Z221100003522022). And this work was also funded by BioMap.
[1] Nadav Brandes, Dan Ofer, Yam Peleg, Nadav Rappoport, and Michal Linial. ProteinBERT: a universal deep-learning model of protein sequence and function. Bioinformatics, 38(8):2102-2110, 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, 33:1877-1901, 2020.
[33] Jinxian Wang, Xuejun Liu, Siyuan Shen, Lei Deng, and Hui Liu. Deepdds: deep graph neural network with attention mechanism to predict synergistic drug combinations. *Briefings in Bioinformatics*, 23(1):bbab390, 2022.
[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 a large-scale pretrained deep language model for cell type annotation of single-cell rna-seq data. Nature Machine Intelligence, 4(10):852-866, 2022.
[37] Patrick Cahan Yuqi Tan. Singlecellnet: A computational tool to classify single cell rna-seq data across platforms and across species. Cell Systems, 19(2):207-213, 2019.
[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):14049, 2017.