[Paper Review] Large-Scale Cell Representation Learning via Divide-and-Conquer Contrastive Learning

Suyuan Zhao, Jiahuan Zhang+1 more|arXiv (Cornell University)|Jun 7, 2023

Single-cell and spatial transcriptomics12 citations

TL;DR

The paper introduces CellLM, a large-scale cell language model trained with a novel divide-and-conquer contrastive learning approach to decouple batch size from GPU memory, achieving state-of-the-art results on multiple single-cell tasks.

ABSTRACT

Single-cell RNA sequencing (scRNA-seq) data is a potent tool for comprehending the "language of life" and can provide insights into various downstream biomedical tasks. Large-scale language models (LLMs) are starting to be used for cell representation learning. However, current LLM-based cell representation learning methods depend solely on the BERT architecture, causing an anisotropic embedding space that leads to inefficient semantic representation. Contrastive learning alleviates this problem by distributing the embeddings uniformly. As a larger batch size in contrastive learning results in better representation, the practical application of contrastive learning in cell representation learning is hampered by the high dimensionality of scRNA-seq data and the large parameter volume of LLMs. To address the batch size limitation, we propose a novel divide-and-conquer contrastive learning approach to decouple the batch size from the GPU memory size for cell representation learning. Based on our divide-and-conquer contrastive learning approach, we introduce Single-Cell Language Model CellLM, a large-scale cell representation learning model to handle high-dimensional scRNA-seq data with tens of thousands of genes. CellLM has over 50 million parameters trained with 2 million scRNA-seq data and makes the first attempt to learn cell language models from both normal cells and cancer cells. CellLM achieves new state-of-the-art (SOTA) results in all evaluated downstream tasks: including a 71.8 F_1-score for cell type annotation (a 3.0% absolute improvement over scBERT), an average F_1-score of 88.9 for single-cell drug sensitivity prediction in a few-shot scenario (an 8.3% absolute improvement), and a 93.4 Pearson's correlation for single-omics cell line drug sensitivity prediction (a 6.2% absolute improvement).

Motivation & Objective

• Motivate improved cell representation learning for scRNA-seq data and overcome anisotropy from BERT-like models.
• Develop a scalable pre-training framework that decouples batch size from GPU memory.
• Propose a large-scale cell language model (CellLM) trained on normal and cancer cells.
• Incorporate self-supervised tasks to enhance representation and transfer to downstream tasks.

Proposed method

• Propose divide-and-conquer contrastive learning to enable large effective batch sizes under fixed memory.
• Use a three-part model architecture combining expression encoding, gene encoding, and a Performer-based encoder.
• Pre-train with MLM, cell type discrimination, and divide-and-conquer contrastive learning.
• Reduce computational load by selective gene expression representations and graph-informed gene embeddings.
• Prove the divide-and-conquer method is equivalent to end-to-end contrastive learning (Appendix A).

Experimental results

Research questions

• RQ1 Can divide-and-conquer contrastive learning decouple batch size from GPU memory while preserving end-to-end contrastive learning benefits?
• RQ2 Does a large-scale CellLM trained on normal and cancer scRNA-seq data improve downstream tasks compared to BERT-based scRNA-seq models?
• RQ3 Can CellLM achieve SOTA performance on cell type annotation and drug sensitivity prediction in single-cell and cell line contexts?
• RQ4 What is the impact of including cancer data on transfer to cell line drug sensitivity prediction?

Key findings

• CellLM achieves SOTA cell type annotation with 71.8 macro F1 on Zheng68K and 80.1 macro F1 on Baron.
• In few-shot single-cell drug sensitivity, CellLM averages 88.7–89.0 F1, outperforming scBERT.
• CellLM improves single-omics cell line drug sensitivity prediction with 0.934 Pearson correlation (warm start) vs. 0.872 for scBERT.
• CellLM attains 93.4 Pearson’s correlation for single-omics cell line drug sensitivity in warm-start, outperforming baselines.
• Divide-and-conquer contrastive learning provides large effective batch sizes without memory overflow and matches end-to-end CL performance.