[논문 리뷰] Large-Scale Cell Representation Learning via Divide-and-Conquer Contrastive Learning

Suyuan Zhao, Jiahuan Zhang외 1명|arXiv (Cornell University)|2023. 06. 07.

Single-cell and spatial transcriptomics인용 수 12

한 줄 요약

논문은 대형 셀 언어 모델 CellLM을 소개하며, 배치 크기를 GPU 메모리와 분리하는 새로운 Divide-and-Conquer contrastive 학습 접근법으로 단일 세포 작업에서 최첨단(SOTA) 결과를 달성합니다.

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 계열 모델로부터의 비등방성(anisotropy)을 극복합니다.
• 배치 크기를 GPU 메모리와 분리하는 확장 가능한 사전 학습 프레임워크를 개발합니다.
• 정상 세포와 암세포를 대상으로 학습된 대규모 셀 언어 모델(CellLM)을 제안합니다.
• 표현 학습 및 다운스트림 작업으로의 전달을 향상시키기 위해 자기지도(Task)를 포함합니다.]
• method:[

제안 방법

• 고정 메모리 하에서 큰 유효 배치 크기를 가능하게 하는 divide-and-conquer contrastive 학습을 제안합니다.
• 표현 인코딩, 유전자 인코딩, Performer 기반 인코더를 결합한 3파트 모델 아키텍처를 사용합니다.
• MLM, 세포 타입 판별, divide-and-conquer contrastive 학습으로 사전 학습합니다.
• 선택적 유전자 발현 표현 및 그래프 정보 기반 유전자 임베딩으로 계산 부담을 줄입니다.
• divide-and-conquer 방법이 엔드-투-엔드 대조학습과 동등하다는 것을 입증합니다(Appendix A).

실험 결과

연구 질문

• RQ1 divide-and-conquer contrastive 학습이 엔드-투-엔드 대조학습의 이점을 보존하면서 GPU 메모리에서 배치 크기를 분리해 낼 수 있습니까?
• RQ2 정상 및 암세포의 scRNA-seq 데이터로 학습된 대형 CellLM이 BERT 기반 scRNA-seq 모델보다 다운스트림 작업을 개선합니까?
• RQ3 CellLM이 단일 세포 및 세포주 맥락에서 셀 타입 주석 및 약물 민감도 예측에서 SOTA 성능을 달성할 수 있습니까?
• RQ4 암 데이터의 포함이 세포주 약물 민감도 예측으로의Transfer에 미치는 영향은 무엇입니까?

주요 결과

• CellLM은 Zheng68K에서 71.8 매크로 F1, Baron에서 80.1 매크로 F1로 SOTA 셀 타입 주석을 달성합니다.
• 몇 샷 단일 세포 약물 민감도에서 CellLM은 88.7–89.0 F1을 평균으로 기록하며 scBERT를 능가합니다.
• CellLM은 0.934 강온 시작(warm start)에서 단일 오믹스 셀룰라인 약물 민감도 예측을 향상시키며 scBERT의 0.872 대비 우수합니다.
• CellLM은 warm-start에서 단일 오믹스 세포주 약물 민감도에서 93.4 피어슨 상관계수를 달성하며 베이스라인을 능가합니다.
• Divide-and-conquer 대조학습은 메모리 초과 없이 큰 유효 배치 크기를 제공하고 엔드-투-엔드 CL 성능과 일치합니다.