LangCell：用于细胞注释的语言-细胞预训练模型

细胞身份包括细胞的各种语义，包括细胞类型、pathway信息、疾病信息等。从转录组数据中了解细胞身份，例如注释细胞类型，是一项基础任务。由于语义是由人类赋予的，如果没有cell-label pair提供监督信号，AI模型很难有效地执行细胞身份识别。目前用于该任务的预训练语言模型（PLM）仅在单一模态–转录组学数据上进行训练。而且，它们必须针对标签不一致的数据集进行微调。

LangCell在预训练阶段构建单细胞数据和自然语言的统一表示，使模型能够直接纳入与细胞身份相关的见解。LangCell利用带有细胞身份信息的text来深入理解知识。

来自：LangCell: Language-Cell Pre-training for Cell Identity Understanding, ICML, 2024

背景概述

scRNA-seq数据分析中，从细胞类型、pathway信息和疾病信息等多个角度理解细胞身份至关重要。类型注释和批次整合等任务已成为该领域的基石。预训练语言模型（PLM）最近在NLP取得了成功。然而，最近的scLLM依然仅基于测序信息使用PLM分析单细胞数据。这些模型利用Transformer来学习数百万个scRNA-seq样本。尽管取得了成功，但当前的细胞表示模型面临以下挑战：

目前的模型框架仅依赖于自监督学习方法，如mask建模–用于捕捉基因共表达关系。然而，由于缺乏对人类专家知识的有效利用，模型在学习细胞表征时不能专注于理解细胞身份。这限制了模型的表示能力。
由于细胞类型其实是人类定义的，如果不对"单细胞-文本标签pair"进行微调，现有模型不可能有效地执行细胞身份识别。用于微调的数据量和质量都会显著影响模型在特定任务中的性能。存在一种困难：研究新疾病或细胞亚型通常获取标签是高成本的，因此需要小样本或零样本方法。

作者认为，高质量编码scRNA-seq数据并将其与多视角的文本描述对齐，可以显著增强文本和单细胞数据之间的有意义理解。这种整合使模型能够像人一样有语义地将知识从熟悉的类别扩展到未知的类别。LangCell是一个语言细胞预训练模型，可以将scRNA-seq数据的特征空间与文本信息无缝整合。

作者构建了一个细胞文本数据集scLibrary，包含2750万个scRNA-seq条目及其文本描述。具体而言，作者从CELLxGENE获得了原始scRNA-seq数据和相应的元数据。然后选择了细胞身份的八个关键方面，这些方面可能包含基本的见解，包括细胞类型、发育阶段和疾病信息，并从OBO Foundry获得尽可能全面的描述。

随后，作者从NLP和CV领域迁移了一些关键见解，并设计了一套在细胞文本数据上有效的多任务协同预训练方法。具体来说，作者在预训练阶段介绍了四项任务。掩蔽基因建模（MGM，mask gene model）和细胞-细胞对比学习（C-C，cell-cell contrastive learning）以增强细胞表示。此外，使用细胞-文本对比学习（C-T，cell-text contrastive learning）和细胞-文本匹配（CTM，cell-text matching）来训练模型，以识别细胞数据和文本数据之间的联系。

LangCell在零样本、少样本和全数据集场景中的一系列细胞身份理解任务上实现了SOTA性能。它解决了诸如细胞类型注释和批次整合之类的经典任务。作为第一个能够进行零样本注释的模型，LangCell在零样本场景中表现出出色的性能，在大多数情况下超过了少样本的基线方法（图1a）。

fig1

图1：a.零样本和少样本细胞类型注释。b.PBMC10K 中的实体embeding UMAP，LangCell 对齐了单细胞和文本嵌入。

方法

数据预处理

首先提供原始scRNA-seq数据作为计数矩阵。假设有 $m$ 个细胞和 $n$ 个基因，则计数矩阵表示为 $A∈\N^{m×n}$ 。作者采用了rank-value编码方法，用于将计数矩阵转换为类似于自然语言的序列数据。首先，分别对每个细胞的基因表达进行归一化，以削弱测序深度的影响，得到 $A^{'}$ 。然后，找到 $A^{'}$ 每列的非零中值 $β∈\R^{n}$ 作为每个基因的中值表达，并用 $β$ 归一化 $A^{'}$ 每一列得到 $A^{''}$ 。即：
eq1

与常规预处理的不同之处：与 $A^{'}$ 相比， $A^{''}$ 消除了整体表达水平带来的差异，其值可以反映一个基因在一个细胞中表达以及在所有细胞中的相对水平。例如，一些housekeeping基因可能很容易具有更高的绝对表达，但这并不一定表明该基因在该细胞中具有特别值得注意的高表达。基于 $A^{''}$ ，根据每个细胞中表达的基因的相对表达对其进行排序，再获得该细胞的基因序列。值得注意的是，在对大规模预训练数据集进行统计时，可以获得更通用的 $β$ ，并在所有后续模型应用中使用该 $β$ 。

模型架构

fig2

图2：LangCell架构。对于每个包含scRNA-seq数据和元数据pair的输入：（1）在预处理过程中，将scRNA-seq数据转换为按相对表达水平降序排列的基因序列，并使用OBO Foundry从元数据中获得细胞的多视角文本描述。（2）使用细胞编码器 f 和文本编码器的单模态 g1 来获得细胞和文本的嵌入，并且使用文本编码器的多模态 g2 来计算matching得分 $p_{c,t}$ 。（3）通过四个损失函数的联合优化进行预训练。其中，掩蔽基因建模（MGM）和细胞-细胞对比学习（C-C）旨在增强单细胞表征。相反，细胞-文本对比学习（C-T）和细胞-文本匹配（CTM）旨在训练模型理解单细胞和文本数据之间的联系。

LangCell包含细胞编码器和文本编码器：

细胞编码器：使用预训练的Geneformer来初始化细胞编码器，该编码器将连续的细胞输入编码为embedding序列。值得注意的是，作者在序列的开头添加了一个[CLS] token，其嵌入通过线性投影为细胞嵌入。
文本编码器：该编码器有两种编码模式，单模态和多模态。对于单模态文本编码，它相当于BERT。对于多模态编码，作者在注意力层中的每个自注意模块之后添加一个可插入的交叉注意模块，以通过线性层计算联合嵌入和细胞文本匹配概率。权重使用PubMedBERT进行初始化，这被证明是生物医学领域中最好的预训练BERT之一。

将细胞编码器定义为 $f$ ，用于从单细胞数据 $c$ 导出嵌入 $z_c$ 。将文本编码器的单模态定义为 $g_1$ ，其负责从文本数据 $t$ 生成嵌入 $z_t$ 。将文本编码器的多模态定义为 $g_2$ ，用于计算单细胞和文本数据之间的匹配概率 $p_{c,t}$ 。这三种编码方法阐述如下： $z_{c}=f(c)\\ z_{t}=g_{1}(t)\\ p_{c,t}=g_{2}(z_{c},t)=g_{2}(f(c),t)$

预训练

LangCell旨在将scRNA-seq数据和文本映射到共享的潜在空间，并利用自然语言中包含的非结构化知识来优化细胞表征。为此，在预训练过程中，作者联合优化了四个目标损失函数，包括掩蔽基因建模、模态内和模态间对比学习以及细胞文本匹配。

掩蔽基因建模

作者随机mask细胞输入序列中的一些基因，并使用模型在这些位置的输出embedding来预测原始输入的重建。使用交叉熵损失函数作为损失函数： $L_{MGM}=\frac{1}{N}\sum_{i=1}^{N}H(v_{ij},\widehat{v}_{ij})$ 其中， $N$ 是mask基因的数量， $v_{ij}$ 和 $\widehat{v}_{ij}$ 分别表示第 $i$ 个mask位置被识别为第 $j$ 个基因的label和预测概率。

细胞-细胞对比学习

作者引入了细胞-细胞对比学习。在scRNA-seq数据中，每个基因表达水平都具有独特的含义，而人工数据增强方法，如对输入数据的搅乱和干扰，可能会破坏基因表达语义。作者认为，特征水平的扰动更适合于scRNA-seq数据中的数据增强。因此，使用两个应用于同一单细胞的标准dropout实例来构建正样本，而同一批次中的其他单细胞作为负样本，这在自然语言研究中已被证明是有效的。损失函数的使用如下： $L_{C-C}=-\frac{1}{T}\sum_{i=1}^{T}log\frac{e^{sim(z_{c}^{(i)},z_{c}^{(i)+})/\tau}}{\sum_{j=1}^{T}e^{sim(z_{c}^{(i)},z_{c}^{(j)+})/\tau}}$ 其中 $T$ 是batch size，sim是余弦相似度， $\tau$ 是温度参数， $z_{c}^{(i)}$ 是细胞 $i$ 的embedding， $z_{c}^{(j)+}$ 为细胞 $j$ 的正样本。

细胞-文本对比学习

作者通过细胞-文本对比学习将细胞和文本投影到相同的嵌入空间中。文本编码器采用单模态编码。损失函数为：
eq2

细胞-文本matching

在计算这种损失时，文本编码器采用多模态编码模式，在每个自注意力层后进行细胞嵌入的交叉注意力计算，最终输出用于二进制分类，以预测细胞是否与文本匹配。该任务旨在探索分辨率更高的细胞和文本之间的匹配关系，选择与细胞-文本对的正样本尽可能相似的细胞-文本来形成负样本。损失函数是二进制交叉熵： $L_{CTM}=H(y,p_{c,t})$ 其中 $y$ 表示指示细胞是否与文本匹配的标签

下游任务

零样本分类

基于细胞数据和文本的对齐表示空间，LangCell可用于零样本细胞注释（图3）。具体而言，对于给定的单细胞 $c$ 和 $N$ 个候选文本描述 $\left\{t^{(i)}\right\}^N_{i=1}$ ，我们通过比较它们在共享嵌入空间中的余弦距离来获得 $logits_1$ ，并通过比较细胞文本匹配模块给出的分数来获得 $logits_2$ 。在实际应用中，由于 $logits_2$ 的计算速度较慢，我们只为 $logits_1$ 得分较高的候选者计算 $logits_2$ 。然后两者相加。

微调

还可以在细胞编码器后添加分类或者回归head，用于在下游任务微调，这样的模型统称为LangCell-CE。

fig3

图3：LangCell在零样本细胞注释中的应用。LangCell利用细胞和文本数据的共享嵌入空间获得相似性得分，通过匹配模块获得细胞-文本匹配得分，并综合考虑这些得分，得到最终的分类logits。

数据集构建与预训练

作者建立了scLibrary，这是一个全面的数据集，包括大约2750万对scRNA-seq数据和相关的文本描述。该数据集来源于CELLxGENE数据库，作者在该数据库中获得了原始计数矩阵格式的scRNA-seq数据和相应的元数据。接下来，从元数据中选择了8个可能包含重要见解的信息，并使用这些条目生成OBO Foundry多视图文本描述。具体而言，所选条目包括化验、细胞类型、发育阶段、组织信息、器官信息、疾病信息以及捐赠者的性别和种族。

预训分为两个阶段。在第一阶段，使用Geneformer初始化细胞编码器参数，仅使用 $L_{MGM}$ 和 $L_{C−C}$ 损失函数进行单模态训练。在第二阶段，使用PubMedBERT初始化文本编码器参数，并使用所有四个损失函数进行多模态训练。两个阶段分别训练三个epochs。

比较LangCell和Cell2Sentence

LangCell和Cell2Sentence各自具有功能上的特点。Cell2Sentence还不能支持零样本场景，虽然C2S是设计为元数据+细胞基因排序构成LLM微调模板，但其在实验中覆盖的数据集不是最广泛的，所以在面对新细胞类型时不能识别出来。当然，Cell2Sentence最大的特点是自然语言可交互，这是LLM+指令微调的必然结果。

LangCell不能做到自然语言交互，但是其架构类似CLIP，变成了一个基于深度学习的检索任务，来自AIR的另一个方法是DrugCLIP。从目前发展看来，包括过去的开放词汇表目标检测等任务均基于CLIP，这是实现零样本的一类方法。因此LangCell的最大特点是支持零样本分类。注意在不同的数据集上，为了零样本分类，我们依然需要提供当前数据集下每个细胞类型的文本以供检索。