[論文レビュー] 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メモリを分離する新しいdivide-and-conquer contrastive learningアプローチを用いた大規模セル言語モデルで、複数の単一細胞タスクで最先端の結果を達成します。

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）を提案する。
• 表現を強化し下流タスクへの転移を促進する自己教師付きタスクを組み込む。

提案手法

• 固定メモリ下で大きな有効バッチサイズを実現するためのdivide-and-conquer対照学習を提案する。
• 発現エンコード、遺伝子エンコード、そしてPerformerベースのエンコーダを組み合わせた3部構成のモデルアーキテクチャを使用する。
• MLM、細胞タイプ識別、およびdivide-and-conquer対照学習を用いて事前学習する。
• 選択的な遺伝子発現表現とグラフに基づく遺伝子埋め込みによって計算負荷を削減する。
• divide-and-conquer法がエンドツーエンド対照学習と同等であることを証明する（付録A）。

実験結果

リサーチクエスチョン

• RQ1 divide-and-conquer対照学習はエンドツーエンド対照学習の利点を保持しつつ、バッチサイズをGPUメモリからデカップリングできるか？
• RQ2 正常細胞およびがん細胞のscRNA-seqデータで訓練された大規模CellLMは、BERTベースのscRNA-seqモデルと比較して下流タスクを改善するか？
• RQ3 CellLMは単一細胞および細胞系の文脈で細胞タイプ注釈と薬物感受性予測においてSOTA性能を達成できるか？
• RQ4 がんデータを含めることが細胞系列薬物感受性予測への転移に与える影響は何か？

主な発見

モデル	Zheng68KマクロF1	Zheng68K加重F1	Zheng68K精度	BaronマクロF1	Baron加重F1	Baron精度
Scanpy	49.5 ±2.2	60.6 ±2.3	63.5 ±2.6	61.5 ±2.0	90.1 ±1.0	91.7 ±1.0
scBERT	68.8 ±0.6	77.5 ±0.8	77.9 ±0.8	84.4 ±4.1	97.0 ±0.3	97.2 ±0.3
CellLM w/o CL	66.4 ±0.9	76.1 ±1.1	76.6 ±1.2	85.5 ±3.1	97.6 ±0.3	97.8 ±0.3
CellLM MoCo	67.0 ±1.2	75.2 ±1.5	75.6 ±1.6	89.0 ±2.4	97.3 ±0.2	97.4 ±0.2
CellLM	71.8 ±1.0	80.1 ±1.1	81.0 ±0.9	90.1 ±2.3	98.1 ±0.1	98.2 ±0.2

• CellLMはZheng68Kで71.8のマクロF1、Baronで80.1のマクロF1でSOTAの細胞タイプ注釈を達成。
• Few-shotの単一細胞薬物感受性では、CellLMが平均で88.7–89.0のF1を達成し、scBERTを上回る。
• CellLMは単一オミクス細胞系薬物感受性予測を0.934のピアソン相関係数（ウォームスタート）で改善、scBERTの0.872を上回る。
• ウォームスタート時の単一オミクス細胞系薬物感受性予測で93.4のピアソン相関を達成し、ベースラインを上回る。
• Divide-and-conquer対照学習はメモリオーバーフローなしに大きな有効バッチサイズを提供し、エンドツーエンドCLの性能に匹敵する。