[论文解读] 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 transcriptomics被引用 12

一句话总结

该论文提出 CellLM，一种通过新颖的分而治之对比学习方法在大规模细胞语言模型中实现将批量大小与 GPU 内存解耦，在多个单细胞任务上达到最先进的结果。

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).

研究动机与目标

• 促进对 scRNA-seq 数据的细胞表征学习的改进，并克服类似 BERT 的模型带来的非各向同性问题。
• 开发一个可扩展的预训练框架，使批量大小与 GPU 内存解耦。
• 提出在正常细胞和癌细胞上训练的大规模细胞语言模型（CellLM）。
• 结合自监督任务以增强表征并转移到下游任务。

提出的方法

• 提出分而治之的对比学习，以在固定内存下实现大有效批量大小。
• 使用一个三部分的模型架构，将表达编码、基因编码和基于 Performer 的编码器结合。
• 以 MLM、细胞类型判别和分而治之对比学习进行预训练。
• 通过选择性基因表达表征和基因嵌入的图信息来降低计算负载。
• 证明分而治之方法等价于端到端对比学习（附录 A）。

实验结果

研究问题

• RQ1 分而治之的对比学习在保持端到端对比学习的好处的同时，是否能将批量大小与 GPU 内存解耦？
• RQ2 在正常和癌性 scRNA-seq 数据上训练的大规模 CellLM 是否比基于 BERT 的 scRNA-seq 模型在下游任务上有改进？
• RQ3 CellLM 是否在单细胞和细胞系情境下的细胞类型注释和药物敏感性预测方面达到SOTA性能？
• RQ4 将癌症数据纳入对细胞系药物敏感性预测的转移有何影响？

主要发现

• CellLM 在 Zheng68K 的宏 F1 上达到 71.8，并在 Baron 上达到 80.1 宏 F1，达到SOTA。
• 在少-shot 单细胞药物敏感性任务中，CellLM 的 F1 平均为 88.7–89.0，优于 scBERT。
• CellLM 在单组学细胞系药物敏感性预测中以 0.934 的 Pearson 相关性（暖启动）优于 scBERT 的 0.872。
• CellLM 在暖启动下实现单组学细胞系药物敏感性预测的 0.934 Pearson 相关性，超过基线。
• 分而治之对比学习提供了在不产生内存溢出的情况下的巨大有效批量大小，并且与端到端对比学习的性能相匹配。