跳转至

CellMSA: Context Modeling for Single-Cell Representation Learning

会议: NeurIPS2026
arXiv: 2609.38908
代码: https://github.com/PharMolix/CellMSA
领域: 计算生物学
关键词: 单细胞表示学习、多细胞上下文、基因对表示、批次整合、细胞状态分类

一句话总结

CellMSA 将跨批次与相关类型细胞按基因对齐，用低维上下文提取基因对依赖并引导目标细胞编码，在有标签辅助检索的整合评测、无测试标签检索的分类评测和结合 STATE-ST 的扰动预测中取得较强结果。

研究背景与动机

单细胞 RNA 测序给每个细胞提供一份高维表达向量，但一个零值既可能是真正不表达，也可能来自测序深度不足或随机缺失。 scGPT、Geneformer 等基础模型主要独立编码单个细胞，因此只能从这一次稀疏观测中推断细胞身份与基因关系；跨供体、平台和批次时，技术差异还会掩盖真实生物学变化。 CellPLM、STATE、Stack 开始引入多个细胞，但若先将细胞压成向量或基因模块，细粒度基因关系容易丢失；若只看同批次同类型细胞，上下文又往往主要提供局部去噪信息。

问题不在于简单增加邻居数量，而在于哪些差异值得比较、比较结果如何进入目标细胞模型。 同类型跨批次细胞有助于暴露稳定的身份信号，相关但不同类型细胞则提供辨别功能差异的背景；模型需要同时保留这些来源的差异以及每个基因的位置。 蛋白质 MSA 提供了一个建模类比：多个相关样本中一致与协同变化的位点，可以汇聚成成对关系，再用于处理单个目标样本。 这里对齐的是基因身份而非进化同源序列，因此借用的是结构化上下文归纳偏置，并不是把表达相关性解释为蛋白共进化或因果调控。

核心 idea：先在基因级对齐的多细胞上下文中提取状态依赖的基因对关系，再将这些关系作为注意力偏置注入目标细胞编码，使邻居信息改变“如何理解目标细胞”，而不只是平均邻居的细胞向量。

方法详解

整体框架

输入是目标细胞及其检索到的相关细胞，输出是目标细胞的 512 维表示，以及可用于分析的基因对表示。 整个流程依次经过关系分组检索、CellMSA-Module 和 GenePairformer：低维模块处理所有上下文行，高维编码器只处理目标行。 三目标预训练提供训练监督；实际提取表示时不需要重建标签或计算训练损失，但仍需要检索上下文。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["目标细胞与候选库"] --> B["关系分组检索"]
    B --> C["CellMSA-Module"]
    C -->|基因对表示| D["GenePairformer"]
    A -->|目标基因嵌入| D
    D --> E["目标细胞表示"]
    D -.->|仅训练：基因预测| F["三目标预训练"]
    E -.->|仅训练：重建与对比| F
    E -->|推理| G["分类或 STATE-ST"]

细胞表达被离散成 11 个取值类别，零表达单列，非零值按细胞内分位数分箱；最大输入长度为 2,048，基因词表有 61,982 个条目。 目标与邻居必须共享对应基因列，不能分别按各自表达高低排列后把同一列当作同一个基因。 基因身份嵌入与表达值嵌入相加，邻居行再加入它与目标的关系类型嵌入；<cls> 用来汇聚最终细胞表示。 基因对状态由两个基因的身份嵌入初始化，随后由上下文和目标细胞逐层修正，并非外部给定的一张固定调控图。

关键设计

1. 关系分组检索：让上下文同时提供局部重复、跨批次稳定性与相关类型差异

预训练时，每个目标检索 16 个同批次同类型细胞、16 个跨批次同类型细胞、8 个相关但不同类型细胞，共 40 个邻居；加上目标形成 41 行输入。 第一组帮助辨认局部共享表达，第二组帮助区分生物学一致性与批次特征，第三组提供相近细胞身份之间的差异背景。 三类来源通过可学习关系嵌入明确标记，因此模型不必把所有邻居都视为等价的正样本。

相关类型并非任意选取：作者对各类型的表达进行跨批次汇总，构造类型均值向量，通过余弦距离与 Ward 层次聚类形成 28 个簇，每个类型最多保留同簇中最相似的 50 个相关类型。 附录报告簇内与簇间平均余弦相似度分别为 0.845 和 0.581，说明这套上下文有表达相似性依据，但不等于经实验确认的谱系关系。 这一步预训练利用 CELLxGENE 的类型和批次元数据，不能把整个方案描述成完全无标签自监督学习。

检索的信息权限随下游任务变化：主整合实验用已有类型注释，而分类验证与测试使用高变基因（HVG）、PCA 和 KNN，不以待预测的类型或状态标签挑邻居。 因此部署时需要明确候选库、可用元数据和检索规则；并不存在一个对所有任务都通用且不依赖数据组织方式的“40 个最佳邻居”。

2. CellMSA-Module：把跨细胞证据压入基因对，而不是先压掉基因分辨率

所有行的基因嵌入先降到 128 维，再交替执行外积均值（Outer-product Mean）和基因对加权平均（Pair-weighted Averaging），共 4 层。 前者对每个基因对，先对两列基因特征分别线性投影，再在每一行计算外积，最后沿细胞行平均并写回基因对状态。 若某组基因在多个邻居中呈现一致或共同变化，成对表征就能积累这种跨细胞证据；它不是直接计算原始表达的 Pearson 相关系数。

\[ \mathbf{p}_{uv}^{l+1}=\mathbf{p}_{uv}^{l}+W_p^l\operatorname{Flatten}\left(\frac{1}{S}\sum_{s=1}^{S}\mathbf{a}_{su}^{l}\otimes\mathbf{b}_{sv}^{l}\right). \]

其中 \(S\) 包括目标与上下文行，两个投影向量来自对应行中基因 \(u\)、\(v\) 的低维特征。 外积结果展平后投影回基因对空间，基因对嵌入维度为 8；这保留了成对结构，也限制了上下文计算的通道宽度。

后者把更新后的基因对状态投影成各注意力头的基因间权重，沿另一基因位置做 softmax，再用这些权重聚合每一行的 value 特征。 每行各基因自己的 sigmoid 门控制写回强度；不同细胞共享由上下文得出的成对权重，但保留自己的表达特征。 这样“从行特征提取关系”与“用关系修正行特征”构成迭代，而不是做一次邻居平均后便停止。

代价仍然随基因数平方增长：外积均值约为 \(O(SG^2d_p^2)\)，基因对加权平均约为 \(O(G^2d_p+SG^2d')\)。 降维节省的是每个成对操作的宽度，不会把稠密基因对矩阵变成线性复杂度；基因截断因此既是工程预算，也是生物学信息覆盖边界。

3. GenePairformer：让上下文关系进入目标注意力，并被目标表达继续更新

完成上下文提炼后，模型将基因对状态投影到注意力头空间，同时以目标行原始基因嵌入开始 6 层目标编码。 这一分工很重要：上下文负责提供成对先验，最终表示仍来自目标细胞，而不是把 41 行全部送入同等宽度的高维 Transformer。

每层目标的 query–key 相似度加到已有成对状态，再对更新后的状态做 softmax 来聚合 value。 因此所谓 pair bias 不只是一次性附加项，而是会吸收目标序列的注意力 logits，逐层反映当前目标表达与跨细胞证据的共同作用。

\[ R_{uv}^{L+1,H}=R_{uv}^{L,H}+\frac{Q_u^{L,H}(K_v^{L,H})^\top}{\sqrt{d_H}}. \]

\(H\) 是注意力头，\(d_H\) 是头维度；更新后的 \(R\) 决定基因间注意力权重，输出 <cls> 再投影成细胞表示。 这与普通 Transformer 的区别，不是输出端多接一张关系图，而是内部每层基因交流已经受到上下文关系约束。

4. 三目标预训练：避免邻居复制，同时兼顾类型身份与细胞特异变化

掩码基因建模（MGM）预测被遮挡位置的原始表达分箱，且目标与所有上下文行在相同基因位置同时遮挡。 只遮目标会允许模型直接读取邻居相同位置的表达；同步遮挡迫使它依赖其他基因及其成对关系，减少这种捷径。 表达重建则从最终细胞向量和基因身份恢复随机选取基因的表达，要求细胞向量保留全局转录信息，而非只记住类型名称。

细胞对比学习（CCE）选同类型细胞为正例，与同批次候选细胞进行 InfoNCE 对比，使用余弦相似度和温度 0.05。 类型元数据在这里也是弱监督；其作用是帮助跨噪声建立生物学可区分表示，但过强会抹掉同类型内部的疾病状态差异。 附录中的目标权重分析显示，强调重建也可能保留更多输入噪声，所以作者没有只用类型对比或只用表达重建。

一个完整示例

以一个近端小管（PT）细胞为例，预训练先收集 16 个同批次同类型邻居、16 个其他批次同类型邻居及 8 个相关类型邻居。 同一基因在 41 行中始终占据相同列，模型可以比较哪些表达特征跨批次仍稳定，以及哪些关系只在目标附近的细胞状态中显著。 CellMSA-Module 将这些比较变成成对状态，GenePairformer 再利用该状态解释目标表达，最后 <cls> 提供分类器所需的细胞向量。

如果正在做 PT 状态分类的测试，邻居改由 HVG–PCA–KNN 选择，不能先使用真实 aPT、dPT、dPT/DTL 标签再声称预测成功。 上述例子说明机制，并非一份额外实验结果；疾病相关基因对激活也只能作为状态关联线索，不能直接叫作因果调控边。

损失函数 / 训练策略

三目标的加权关系为：

\[ \mathcal{L}=\mathcal{L}_{\mathrm{MGM}}+0.01\mathcal{L}_{\mathrm{Rec}}+10\mathcal{L}_{\mathrm{CCE}}. \]

MGM 和重建均使用表达分箱交叉熵；MGM 选取概率为 0.15，重建基因选取概率为 0.1。 模型有 47.13M 参数，使用 AdamW、1,000 步线性 warmup 后维持学习率 \(10^{-5}\)，有效 batch size 为 32。 在 4 张 A800 上训练一轮约需 20 天，预训练语料来自 2,090 个数据集、818 个细胞类型，并在训练前排除下游评测数据集。

“109M”指细胞观测条目，其中 65.6M 是 primary observations，而不是 109M 个独立生物学细胞。 作者保留部分重复聚合或图谱整合中的非 primary 条目，因为同一细胞可对应不同上下文；这不能消除重复细胞或图谱结构被过度加权的风险。

实验关键数据

主实验

下表保留能回答不同问题的核心结果，分类数字是 5 次随机供体划分的均值与标准差，扰动预测则只有一次固定划分。 整合总分按 \(0.6S_{\mathrm{bio}}+0.4S_{\mathrm{batch}}\) 计算，表中各类指标均越高越好，但不同任务的数值不能互相比大小。

数据集 / 任务 指标 CellMSA 对照 比较边界
Tabula Sapiens 整合 Total 0.736 Stack 0.663 两者均有类型标签辅助检索
Tabula Sapiens 整合 Bio / Batch 0.850 / 0.566 Stack 0.740 / 0.547 不代表每个批次指标都领先
Blood 类型分类 Macro-F1 \(0.912\pm0.005\) CellPLM \(0.906\pm0.010\) 验证与测试检索不用真实类型标签
Kidney Atlas PT 状态分类 Macro-F1 \(0.931\pm0.006\) STATE-SE \(0.912\pm0.009\) 同一谱系内的 3 类状态
Replogle + STATE-ST Pearson \(\Delta\) 0.433 Expression + STATE-ST 0.398 同一状态转移框架，单次固定划分
Replogle + STATE-ST PRAUC / DE Overlap 0.334 / 0.215 STATE-SE + STATE-ST 0.287 / 0.180 比较差异表达基因的恢复情况

主整合评测中，scGPT、Geneformer、scVI 等基线提取表示时不用类型标签，因此不能把所有方法的主表排序都当作同一信息预算的零标签比较。 附录另给出仅 Bladder 的无标签检索实验：CellMSA 的 Batch / Bio / Total 为 0.509 / 0.807 / 0.688，Stack 为 0.451 / 0.699 / 0.600。 该实验支持优势并非完全来自类型标签，但不能替代所有组织上的无标签复验，也不能单独归因于某个模块。

分类按 donor ID 以 7:1:2 划分，冻结基础模型提取表示后训练三层 MLP；Blood 只保留占比至少 1% 的类型。 PT 的 aPT、dPT、dPT/DTL 分别对应正常样、损伤、严重损伤并伴转录身份漂移，CellMSA 的两个分类任务 Accuracy 分别为 \(0.962\pm0.002\)、\(0.958\pm0.004\)。 这不是对任意罕见细胞类型或任意临床疾病的诊断结论。

扰动数据包含 4 个细胞系、100 个共同且样本较多的扰动，连同对照约 132k 细胞；HepG2 的 45% 扰动作测试、5% 作验证，其余与其他细胞系用于训练。 上下文限制在同细胞系，跨类型在这里指不同扰动状态，不是本体细胞类型；验证与测试的扰动后表达不能作为上下文。 Pearson \(\Delta\) 是预测与真实“扰动均值减对照均值”向量在基因维度上的 Pearson 相关，0.433 相对 0.398 的增益约为 8.8%。 CellMSA + STATE-ST 的 Spearman-FC 为 0.431，对照 Expression + STATE-ST 为 0.404；这评估真实显著差异表达基因上的 fold-change 排序一致性。

消融实验

配置 PT Accuracy PT Macro-F1 扰动 Pearson \(\Delta\) 扰动 Spearman-FC
去掉 CellMSA-Module 与成对表示 0.854 0.811 0.355 0.366
保留架构但不输入上下文 0.878 0.817 0.389 0.367
只使用同批次上下文 0.929 0.893 0.400 0.397
完整模型 0.958 0.931 0.433 0.431

相比只用同批次邻居，完整模型的 PT Macro-F1 增加 0.038、扰动 Pearson \(\Delta\) 增加 0.033，支持跨批次与相关类型上下文的附加价值。 删除整个模块与成对状态时，PT Macro-F1 降低 0.120，但该消融同时改变多个因素，不能把全部差值归给某一个算子。 上下文数量超过 40 后收益趋于饱和；原文没有在正文给出曲线各点的精确数值，因此不补造逐点表。

关键发现

  • 基因对表示能提供状态相关的解释线索：PT 疾病分析中，头 6、7 更接近疾病标记，头 4、5 更接近稳态标记，但头编号不是固定的因果功能标签。
  • 外部 STRING 比较中，KDR 的前 50 个配对基因有 26 个得到功能关联支持，即 52%；没有随机背景显著性检验，不能称为显著富集或 GRN 恢复精度。
  • 随机移除 Blood 细胞 50% 的非零基因后，原始与损坏表示的平均 Pearson 相关仍为 0.972，中位数 0.984；这是表示稳定性，而非损坏输入下分类准确率。

亮点与洞察

  • 上下文被用来估计“基因如何共同参与一个状态”，而不是作为最终向量的均值。这使多细胞信息在目标内部注意力层面发挥作用。
  • 同位置同步掩码是机制成立的重要配套：不给邻居提供被预测位置的答案，才能更有力地要求模型学习跨基因依赖。
  • 低维多行、高维单行的计算分工值得迁移到其他组学数据，但应保留特征身份对齐，不能把任意邻居集合都视作有效 MSA。

局限与展望

  • 成对激活是统计依赖，可能受共同上游因素、细胞组成和技术偏差影响；即使矩阵有方向，也不意味着因果调控方向。
  • 主整合实验的信息权限不齐，完整无标签证据仅覆盖 Bladder；需要扩展到更多组织和注释不可靠的候选库。
  • 表达分箱、2,048 长度截断和稠密平方级成对计算限制了表达精度与基因覆盖；低秩或稀疏成对表示需要额外验证是否损失状态辨别力。
  • 语料的非 primary 观测可能重复，扰动实验只有单次划分，不能从当前结果推断跨随机种子、跨物种或临床泛化稳定性。
  • 未来可用独立扰动实验检验成对关系的功能意义；当前输出更适合提出研究假设，不应作为直接临床决策依据。

相关工作与启发

  • vs scGPT / Geneformer：独立细胞编码可直接提取表示，CellMSA 则需要检索并利用上下文关系；收益伴随候选库和计算预算的额外要求。
  • vs CellPLM / Stack：共同点是超越单细胞输入，差异是 CellMSA 维持基因列对齐并显式提炼基因对；Blood Macro-F1 对 CellPLM 的优势较小，不能只引用较弱基线营造普遍巨大提升。
  • vs AlphaFold 风格 MSA：迁移的是“多样本一致性与变化形成 pair representation”的建模逻辑，而不是残基接触预测的生物学含义。
  • vs STATE：本文在扰动实验中使用 STATE-ST 作为状态转移模型，因此结论是表示增强能帮助该框架，而不是 CellMSA 单独替代完整扰动预测系统。

评分

  • 新颖性: 4/5。将基因对结构、来源可区分的上下文与目标细胞编码有效连接，但核心算子借鉴成熟架构。
  • 实验充分度: 4/5。覆盖整合、分类、扰动与解释分析，信息权限和单次扰动划分仍限制比较强度。
  • 写作质量: 4/5。主流程与附录设置较清楚，需谨慎区分相关性解释和因果机制。
  • 价值: 4/5。为上下文增强单细胞表示提供可复用路线，真实部署还需检索质量与独立生物学验证。