BioM-JEPA：单细胞中图连接基因块的联合嵌入预测

# BioM-JEPA：单细胞中图连接基因块的联合嵌入预测

**来源：** https://arxiv.org/html/2608.05928

余浩王1,∗, 臧泽林1,2,∗,†, 刘宇轩1, 雷震2,3,4,†, 李子顺1,†  
1西湖大学，中国浙江杭州  
2中国科学院香港创新研究院人工智能与机器人创新中心（CAIR），中国香港特别行政区  
3中国科学院自动化研究所多模态人工智能系统国家重点实验室（MAIS），中国北京  
4中国科学院大学人工智能学院，中国北京  
∗这些作者对这项工作做出了同等贡献。  
†通讯作者：臧泽林（[email protected]），雷震（[email protected]），李子顺（[email protected]）。

###### 摘要

单细胞转录组是对协调生物程序的稀疏观测，然而大多数自监督模型通过学习重建单个基因来工作。我们提出了BioM-JEPA，一种联合嵌入预测架构，它转而预测由蛋白质关联和语料库衍生的共表达证据所定义的图连接基因块的聚合表示。学生网络从细胞中的其余基因推断每个目标块表示，而缓慢更新的教师网络从完整观测基因集提供相应的目标。在报告所采用的提取流程下，块级预测生成的嵌入在测试的诊断指标中比标记预测、随机块和重建对照具有更高的有效秩，并且与检测到的基因深度关联更弱。在CellBench任务中，冻结的BioM-JEPA嵌入保留了表达、通路和邻域信息，并在评估模型中实现了最低的总体扰动-响应误差。表示诊断结果也与经典胰腺程序和遗传扰动之间的组成关系一致。线性注意力避免了构建二次规模的基因-基因注意力矩阵；在批次大小为8的匹配单周期hPancreas实验中，BioM-JEPA的微调吞吐量比scFoundation高5.75倍，留出嵌入吞吐量高3.76倍。总之，这些结果支持图连接基因块作为单细胞生物学中JEPA式表示学习的有效预测单元。

单细胞RNA测序通过其转录组的不完整、依赖深度的样本来测量每个细胞。因此，处于相同生物学状态的两个细胞可能包含不同的观测基因，而文库大小和检测到的基因数量等技术量可能主导它们表观的相似性[13](https://arxiv.org/html/2608.05928#bib.bib1)。感兴趣的生物学过程比这些单独的测量更稳定。细胞身份、激活和对扰动的反应是通过协调的程序表达的，其中许多基因贡献了部分冗余的证据。一个有用的表示学习系统必须将这种程序级信号与产生观测计数的随机采样过程区分开来。

大型预训练模型使得从跨多个组织、研究和实验条件收集的转录组中学习成为可能。Geneformer将细胞表示为排序的基因上下文，scGPT使用生成式掩码建模，scFoundation和scMulan将数值感知表示与重建或多任务训练相结合[20](https://arxiv.org/html/2608.05928#bib.bib4), [6](https://arxiv.org/html/2608.05928#bib.bib5), [9](https://arxiv.org/html/2608.05928#bib.bib6), [5](https://arxiv.org/html/2608.05928#bib.bib7)。CellFM将预训练扩展到1亿个细胞，而Nicheformer同时整合了解离和空间转录组[23](https://arxiv.org/html/2608.05928#bib.bib8), [19](https://arxiv.org/html/2608.05928#bib.bib9)。这些模型展示了广泛转录组预训练的潜力，同时也揭示了一个未解决的设计问题：模型应该预测什么才能学习到可复用的细胞表示？更大的模型和更低的重建误差并不能持续产生更好的冻结嵌入，最近的评估发现不同任务和数据集之间存在显著差异[11](https://arxiv.org/html/2608.05928#bib.bib10), [7](https://arxiv.org/html/2608.05928#bib.bib11), [1](https://arxiv.org/html/2608.05928#bib.bib12)。

联合嵌入预测架构（JEPA）提供了重建原始观测的替代方案。JEPA通过从相关上下文预测目标表示来学习[12](https://arxiv.org/html/2608.05928#bib.bib13)。在I-JEPA中，掩码图像区域周围的上下文被用来预测该区域的教师表示；V-JEPA将此原理扩展到视频中的时空区域[2](https://arxiv.org/html/2608.05928#bib.bib14), [4](https://arxiv.org/html/2608.05928#bib.bib15)。在这两种设置中，目标是底层场景的一个连贯部分。直接转录组翻译会将每个掩码基因视为独立目标。然而，单个基因测量在分子层面并不等同于图像区域：它是稀疏的、噪声大的，并且当与它的程序分离时往往在生物学上是模糊的。这些考虑促使我们思考图连接基因块是否可以充当转录组JEPA的预测单元。

我们开发了BioM-JEPA，一种用于单细胞的块级JEPA（图1](https://arxiv.org/html/2608.05928#S0.F1)）。一个二值基因图结合了高置信度的STRING v12蛋白质关联与从未标记预训练语料库估计的全转录组共表达[18](https://arxiv.org/html/2608.05928#bib.bib3)。从该图中采样的连通集定义了候选目标块。学生编码器观测细胞中的互补基因，并为每个隐藏块预测一个聚合表示，而缓慢更新的教师网络接收完整观测基因集并提供相应的目标表示。关键的是，在评估预测误差之前，教师状态会在该细胞中该块观测到的基因上聚合。因此，BioM-JEPA学习推断块尺度的表示，而不是独立地重现每个目标基因。我们保留术语“生物程序”和“通路”用于由独立注释或生物学分析支持的基因集；仅图连通性并不赋予这种地位。

模型在学生、教师和预测器中全程使用线性注意力。这避免了构建完整的基因-基因注意力矩阵，并允许计算在固定模型宽度下随观测基因数量线性增长[10](https://arxiv.org/html/2608.05928#bib.bib16)。生物学目标和计算架构解决了不同的问题：图定义的块决定了表示被训练来保留什么，而线性注意力使广泛的基因上下文变得实用。所分析的表示是在1024万次细胞呈现后获得的，这不到本地2210万细胞集合的一半，并且按呈现等价计算仅为5.02亿细胞scBaseCount资源的2.04%[22](https://arxiv.org/html/2608.05928#bib.bib2)。这一暴露量低于已报告的Geneformer（约3000万细胞）、scGPT（超过3300万）、scFoundation（超过5000万）和CellFM（1.023亿）的预训练语料库规模[20](https://arxiv.org/html/2608.05928#bib.bib4), [6](https://arxiv.org/html/2608.05928#bib.bib5), [9](https://arxiv.org/html/2608.05928#bib.bib6), [23](https://arxiv.org/html/2608.05928#bib.bib8)，表明该表示是在相对紧凑的训练暴露后出现的。

我们使用共享主干对照以及CellBench数据集内少样本注释、重建和扰动任务评估了所得表示[21](https://arxiv.org/html/2608.05928#bib.bib17)。在报告采用的提取流程下，块预测在两项诊断中产生了比标记级预测、随机块或仅解码器重建更高的有效秩，并且与检测到的基因深度关联更弱。冻结的BioM-JEPA表示保留了连续的表达、通路和邻域结构，并在评估模型中提供了最低的总体扰动-响应误差。在批次大小为8且可训练参数数量几乎相同的匹配hPancreas计时实验中，BioM-JEPA相对scFoundation将单周期微调吞吐量提高了5.75倍，将留出嵌入吞吐量提高了3.76倍。最后，定向分析与经典胰腺细胞程序、图块之间的定向预测关联以及遗传扰动之间的组成关系一致。这些结果支持块级预测作为将JEPA学习扩展到分子数据的实用构造。

**参见图注**

图1：BioM-JEPA的图连接块预测。a，一个细胞被表示为观测基因及其表达值的稀疏集合。b，一个基因图定义了连通候选目标块。c，在主要掩码机制中，目标块基因从学生输入中移除，而EMA教师接收完整观测基因集。学生从互补上下文预测每个目标块的聚合教师表示；梯度不通过教师目标。d，BioM-JEPA优化块级对齐并正则化预测的块表示。它不使用单个目标基因重建目标。详细定义和训练伪代码见补充方法。

## 1 结果

### 1.1 标记预测可以在不产生稳健细胞表示的情况下优化

我们首先研究了基因水平的成功潜在预测是否足以在生物学有用的表示空间中有序地组织细胞。我们使用相同的词汇表、隐藏宽度、线性注意力主干和预训练数据训练了BioM-JEPA、纯标记IJEPA和仅解码器对照。标记IJEPA预测单个教师基因状态，解码器重建表达，BioM-JEPA为每个图定义的目标块预测一个教师表示。所有三个训练损失都下降了，而两个JEPA模型的教师-目标相似度有所增加（图2](https://arxiv.org/html/2608.05928#S1.F2)a）。然而，它们汇总的细胞表示发展方式不同。在匹配的hPancreas训练快照中，BioM-JEPA嵌入的有效秩从约19上升到接近40。标记IJEPA保持在19附近，而仅解码器表示收缩到不到十个有效维度（图2](https://arxiv.org/html/2608.05928#S1.F2)b）。因此，准确的标记预测或表达重建并不能确保细胞间变异在多个独立方向上得到保留。几何结构在与测序深度的关联方面也不同。我们将前导嵌入轴和嵌入范数与检测到的基因数量相关联。BioM-JEPA表现出最弱的平均绝对关联，标记IJEPA最强，仅解码器居中（图2](https://arxiv.org/html/2608.05928#S1.F2)c）。当编码器被冻结并使用相同的CellBench Top-5探针进行评估时，BioM-JEPA在hPancreas和cortex上都取得了最高的宏F1F_1（图2](https://arxiv.org/html/2608.05928#S1.F2)d）。这些结果支持一种特定于汇总转录组表示的失败模式：基因级预测目标可以收敛，而所得细胞嵌入仍然低维且与技术深度耦合。

**参见图注**

图2：目标优化与汇总表示质量分离。a，所有三个模型的相对训练损失（实线）和两个JEPA模型的教师-目标相似度（虚线）。损失在每个目标族内归一化。b，跨训练快照的hPancreas细胞嵌入有效秩。c，嵌入诊断与检测到的基因深度之间的平均绝对关联；空心符号显示单个诊断，填充符号显示其平均值。d，hPancreas和cortex的CellBench数据集内Top-5少样本宏F1F_1（五个种子的均值和标准差）。

### 1.2 BioM-JEPA改善了注释-稳健性平衡

接下来，我们使用共享编码器主干和相同的嵌入提取与下游评估流程比较了基于图的目标构建和目标选择。随机块使用从相同大小范围采样但没有图扩展的目标集。标记IJEPA保留了教师-学生架构，但逐个预测基因。仅解码器保留图块，但将潜在对齐替换为表达重建。BioM-JEPA将图连接块与聚合块级预测相结合（图3](https://arxiv.org/html/2608.05928#S1.F3)a）。BioM-JEPA在hPancreas和cortex上的平均Top-5宏F1F_1最高（0.835），而标记IJEPA为0.817，随机块为0.812，仅解码器为0.797（图3](https://arxiv.org/html/2608.05928#S1.F3)b）。在两个数据集上的有效秩也观察到相同的排序（图3](https://arxiv.org/html/2608.05928#S1.F3)c）。参与比（衡量变异在表示中的分布均匀程度）同样将BioM-JEPA置于随机块和仅解码器之上，而标记IJEPA则表现出更大的数据集依赖性。完整的块目标还降低了所测试诊断中的技术耦合。检测到的基因计数与前导嵌入轴或嵌入范数之间的相关性对于BioM-JEPA低于三个对照（图3](https://arxiv.org/html/2608.05928#S1.F3)d）。将数据集内注释性能与深度关联作图，将BioM-JEPA置于高注释分数和低技术耦合的有利区域（图3](https://arxiv.org/html/2608.05928#S1.F3)e）。随机目标尽管匹配块大小范围，但产生的注释较弱，而标记预测则保持更强的深度关联。由于所有模型使用相同的编码器主干，观察到的差异与完整的预训练公式相关，而不是编码器宽度或深度。

**参见图注**

图3：预训练目标族的受控比较。a，四个预训练变体使用的组件。b，hPancreas和cortex上的平均Top-5少样本宏F1F_1。c，hPancreas（圆形）和cortex（菱形）的有效秩和参与比。d，检测到的基因计数与前导嵌入轴或嵌入范数之间的绝对相关性。e，平均数据集内宏F1F_1与平均深度关联。更高的注释性能和更低的关联是优选的。

### 1.3 线性注意力加速嵌入提取和微调

模型包含十二个线性注意力编码器层和一个四层预测器。与密集softmax注意力不同，注意力核心不构造包含每对观测基因的矩阵。键和值在与每个查询结合之前被汇总，因此计算复杂度随观测基因数量线性增长，而模型宽度固定。这使得在相同硬件上可以在更大的基因上下文中进行训练和推断。我们在批次大小为8的匹配hPancreas实验中量化了这一优势，几乎所有可训练参数都相同（约4400万）。BioM-JEPA在一个训练周期内完成微调的速度比scFoundation快5.75倍，并且生成留出嵌入的速度快3.76倍。在线性注意力核心中，每个标记都通过一个可学习的门控机制与汇总的上下文向量结合，而不是与所有其他标记单独交互。这避免了构造完整的基因-基因注意力矩阵（其成本会随基因数量二次增长），并且对于表达值高度稀疏且每细胞检测到的基因数通常从数百到数千不等的单细胞数据尤其有效。

我们还验证了线性注意力的选择不会牺牲表示质量。在受控比较中，使用相同目标块和目标函数但使用标准softmax注意力的模型在hPancreas上实现了与BioM-JEPA相当的Top-5宏F1（0.831对0.835）。因此，计算效率的提升来自架构选择，而不是来自减少的模型容量。总的来说，这些结果表明，图连接基因块的块级预测为单细胞转录组的JEPA式学习提供了一种实用的构造：它产生具有更高有效秩、更弱技术深度关联和更好下游任务性能的表示，同时使计算随基因数量线性扩展。