CellWorld：从基因级重建到空间转录组基础模型中的潜细胞预测

# CellWorld：从基因水平重建到空间转录组学基础模型中的潜在细胞预测
来源：https://arxiv.org/html/2608.06659

###### 摘要

本文表明，潜在空间预测式预训练可以为空间转录组学基础模型提供一条可扩展的路径。现有的空间转录组学基础模型主要重建被掩蔽的基因身份或表达值，这可能鼓励模型复现检测平台特有的技术变异，从而限制了表示的可迁移性。为了避免直接重建此类变异，我们将预测目标从可观测的基因测量值转变为潜在细胞表示，并引入CellWorld，该方法从可见的空间上下文和有限的基因部分表达提示中预测掩蔽细胞的潜在表示。我们在包含4600万个人类细胞的语料库上预训练了四个CellWorld变体，可训练参数从574万到9456万不等。我们的受控缩放实验表明，性能随模型容量提高而改善，尤其是在空间任务上；而空间迁移更多取决于充分的优化和广泛的生物来源多样性，而非单纯的细胞数量。在四个留出数据集中，即使只有574万可训练参数的CellWorld-Small，也在所有11个线性探针基准和全部7个微调空间基准上优于所有基线。最值得注意的是，仅在覆盖广泛生物来源的语料库5%子集上预训练的冻结CellWorld-Large，在全部7个空间基准上优于所有完全微调的基线。代码可在https://github.com/UoM-HealthAI/CellWorld获取。

## 引言

空间转录组学（ST）在保留组织空间结构的同时测量基因表达，从而能够在局部组织背景下表征细胞（Ståhl等人2016（https://arxiv.org/html/2608.06659#bib.bib17）；Chen等人2015（https://arxiv.org/html/2608.06659#bib.bib7））。ST数据集的规模和多样性快速增长，这为通过大规模预训练学习跨数据集和下游任务的可迁移表示创造了机会。

第一代ST基础模型（Tejada-Lapuerta等人2025（https://arxiv.org/html/2608.06659#bib.bib18）；Wen等人2024（https://arxiv.org/html/2608.06659#bib.bib20）；Zhang等人2025（https://arxiv.org/html/2608.06659#bib.bib21）；Madhu等人2025（https://arxiv.org/html/2608.06659#bib.bib13）；Wang等人2025（https://arxiv.org/html/2608.06659#bib.bib19））主要通过基因水平观测的掩蔽重建来追求这一机会。然而，由于靶向基因panel、测量噪声和批次效应的差异，ST测量在不同平台上是稀疏且异质的。因此，直接重建可能鼓励模型在复现生物学信号的同时复现检测平台特有的变异。

潜在预测则遵循以下原则：*通过预测未观测输入的抽象状态而非重建原始观测，可以学习到可迁移的表示*（Baevski等人2022（https://arxiv.org/html/2608.06659#bib.bib3））。联合嵌入预测架构（JEPAs）已在视觉领域展示了这一原则的潜力，其中模型从可见上下文预测被掩蔽目标区域的潜在表示（Assran等人2023（https://arxiv.org/html/2608.06659#bib.bib2），2025（https://arxiv.org/html/2608.06659#bib.bib1））。然而，直接将这一公式应用于ST并非易事。图像区域通常与其周围环境表现出强烈的局部连续性，而ST切片则由离散的细胞组成，这些细胞的身份和表达状态即使在看似相似的空间邻域中也可能不同。因此，仅凭可见的空间上下文可能无法唯一确定被掩蔽细胞的身份和分子状态。我们将这种欠定问题称为细胞级目标*歧义性*。迄今为止，尚不清楚潜在预测是否能够以及在何种公式下成为ST基础模型可靠且可扩展的预训练目标。

为填补这一空白，我们引入了CellWorld，一种通过潜在细胞预测进行预训练的大规模ST基础模型。给定一个包含空间邻近细胞的局部patch，CellWorld将每个细胞的基因表达谱和元数据映射为细胞token，并使用空间Transformer建模细胞-细胞相互作用。它随机将一部分细胞指定为掩蔽目标，其余细胞作为可见上下文。上下文编码器仅处理可见细胞token，而EMA更新的目标编码器处理完整patch以为每个掩蔽细胞生成潜在目标。空间预测器被训练为从编码的可见上下文和对应的目标坐标中恢复每个潜在目标。为解决细胞级目标歧义性，预测器还额外以掩蔽细胞的有限部分表达提示为条件。预训练后，目标编码器和预测器被丢弃，而细胞tokenizer和上下文编码器被保留用于下游任务。

我们在涵盖三个平台和11个器官的4600万个人类细胞上预训练了四个CellWorld变体，即Small、Base、Large和Huge。如图1（https://arxiv.org/html/2608.06659#Sx1.F1）所示，即使仅有574万可训练参数的CellWorld-Small，也在线性探针下的全部11个基准和微调下的全部7个空间基准上优于所有现有方法，更大的变体进一步扩大了这一领先优势。我们的受控缩放实验表明，性能随模型容量增加而持续提升，尤其是在空间任务上。数据缩放实验表明，关于细胞身份的可迁移信息可以从相对较少的预训练数据中学习到，而空间迁移更多取决于*广泛的生物来源多样性和充分的优化*，而非原始细胞数量。最值得注意的是，仅在覆盖广泛生物来源的语料库5%子集上预训练的冻结CellWorld-Large，在全部7个空间基准上优于所有完全微调的基线。

参见说明图1：在(a)线性探针和(b)微调两种协议下，与11个留出任务-数据集对上的现有方法进行比较。每条轴均归一化到相应协议下全数据CellWorld-Large的得分（100%）；数字标注给出其绝对得分。(b)中的虚线轮廓显示了在广泛采样的5%语料库子集上预训练的冻结CellWorld-Large，归一化到微调后全数据模型。越高越好。

## 相关工作

#### ST基础模型。

Nicheformer（Tejada-Lapuerta等人2025（https://arxiv.org/html/2608.06659#bib.bib18））从表达排序的基因序列中预测被掩蔽的基因身份，而scGPT-spatial（Wang等人2025（https://arxiv.org/html/2608.06659#bib.bib19））、CellPLM（Wen等人2024（https://arxiv.org/html/2608.06659#bib.bib20））和BrainBeacon（Zhang等人2025（https://arxiv.org/html/2608.06659#bib.bib21））重建被掩蔽的表达值。HEIST（Madhu等人2025（https://arxiv.org/html/2608.06659#bib.bib13））还额外引入了空间和对比目标。SToFM（Zhao等人2025（https://arxiv.org/html/2608.06659#bib.bib22））重建被掩蔽的、由表达派生的细胞嵌入（由领域自适应的单细胞编码器产生），同时引入空间重建目标，因此尽管在嵌入空间中运行，仍属于重建式方法。相比之下，CellWorld通过预测掩蔽细胞的上下文化潜在表示进行预训练。

#### 潜在预测表示学习。

潜在预测已成为视觉模态中重建式自监督学习的替代方案。BYOL（Grill等人2020（https://arxiv.org/html/2608.06659#bib.bib9））训练在线网络预测指数移动平均目标网络在不同增强下生成的表示。I-JEPA（Assran等人2023（https://arxiv.org/html/2608.06659#bib.bib2））从可见上下文预测被掩蔽图像区域的潜在表示，后续工作则在多样化的架构、模态和应用领域中对这一范式进行了探索和扩展（Assran等人2025（https://arxiv.org/html/2608.06659#bib.bib1）；Mur-Labadia等人2026（https://arxiv.org/html/2608.06659#bib.bib14）；Chen等人2025（https://arxiv.org/html/2608.06659#bib.bib6）；Saito, Kudeshia, and Poovvancheri2025（https://arxiv.org/html/2608.06659#bib.bib16）；Balestriero and LeCun2025（https://arxiv.org/html/2608.06659#bib.bib4）；Klindt, LeCun, and Balestriero2026（https://arxiv.org/html/2608.06659#bib.bib11）；Litman等人2025（https://arxiv.org/html/2608.06659#bib.bib12）；ElSheikh等人2026（https://arxiv.org/html/2608.06659#bib.bib8））。同期工作ST-JEPA（Birk等人2026（https://arxiv.org/html/2608.06659#bib.bib5））将潜在预测应用于由局部细胞邻域构建的序列中的被掩蔽基因token，其表示学习主要通过聚类而非大规模跨数据集预训练和迁移进行评估。CellWorld则将每个细胞表示为单个token，使用空间编码器建模细胞-细胞相互作用，并通过预测掩蔽细胞的上下文化潜在表示进行大规模预训练，以支持跨不同数据集和下游任务的迁移。

## 方法

### 空间采样策略

我们通过三阶段空间采样过程构建包含4096个细胞的局部patch。首先，使用二维细胞坐标上的kk均值将每个切片划分为空间上连贯的子切片，对于一个包含nn个细胞的切片，子切片数量为K_sub=max⁡(1,⌊n/50,000⌋)K_{\mathrm{sub}}=\max(1,\lfloor n/50{,}000\rfloor)。在每个子切片ss内，最远点采样（FPS）选择Ks=2⌈ns/4096⌉K_s=2\lceil n_s/4096\rceil个patch中心，其中nsn_s是子切片中的细胞数量。FPS在每个epoch重新运行，距每个中心最近的4096个细胞构成一个训练样本。

### CellWorld预训练框架

图2（https://arxiv.org/html/2608.06659#Sx3.F2）展示了整体架构。在每个采样的空间patch内，CellWorld将基因表达谱转换为细胞token，使用空间注意力建模它们之间的相互作用，并从编码的可见上下文和有限的部分表达提示中预测被掩蔽目标细胞的潜在表示。下面我们详细介绍每个组件。

参见说明图2：CellWorld预训练架构概览。

#### 细胞tokenization。

我们首先将每个细胞的基因符号映射到一个由预训练语料库构建的统一词汇表（包含19227个基因），每个基因在平台上共享单一嵌入。对于每个细胞ii，Gi\mathcal{G}_{i}包含表达值最高的最多K=512K=512个基因；这保留了95%细胞所检测到的完整panel。CellWorld随后通过按表达值聚合这些嵌入并添加器官和平台嵌入来构建细胞token：

ui=∑g∈Gixigeg+eiorgan+eiplatform.\mathbf{u}_{i}=\sum_{g\in\mathcal{G}_{i}}x_{ig}\mathbf{e}_{g}+\mathbf{e}^{\mathrm{organ}}_{i}+\mathbf{e}^{\mathrm{platform}}_{i}。(1) 其中xigx_{ig}是基因gg的表达值，eg∈Rd\mathbf{e}_{g}\in\mathbb{R}^{d}是其基因嵌入，eiorgan\mathbf{e}^{\mathrm{organ}}_{i}和eiplatform\mathbf{e}^{\mathrm{platform}}_{i}是元数据嵌入。上下文编码器、目标编码器和提示条件空间预测器共享同一细胞tokenizer。

#### 空间注意力。

由于不同切片的坐标系和空间尺度不同，我们避免使用绝对位置嵌入，而是通过注意力偏置编码相对空间关系。具体而言，在掩蔽之前，坐标在每个空间patch内被归一化为s~i=(si−s ̄)/δ\tilde{\mathbf{s}}_{i}=(\mathbf{s}_{i}-\bar{\mathbf{s}})/\delta，其中s ̄\bar{\mathbf{s}}是patch中心，δ=maxj⁡‖sj−s ̄‖∞\delta=\max_{j}\|\mathbf{s}_{j}-\bar{\mathbf{s}}\|_{\infty}是最大的绝对坐标位移。CellWorld随后采用ALiBi（Press, Smith, and Lewis2021（https://arxiv.org/html/2608.06659#bib.bib15)）的二维变体，该变体使用固定的、按头特定的斜率将归一化坐标之间的成对距离转换为加性自注意力偏置。对于注意力头hh，偏置为Bij(h)=−mh‖s~i−s~j‖2B_{ij}^{(h)}=-m_{h}\|\tilde{\mathbf{s}}_{i}-\tilde{\mathbf{s}}_{j}\|_{2}，并加到标准注意力logits上：

Attn(h)=softmax⁡(Q(h)K(h)⊤dh+B(h))V(h).\operatorname{Attn}^{(h)}=\operatorname{softmax}\!\left(\frac{\mathbf{Q}^{(h)}\mathbf{K}^{(h)\top}}{\sqrt{d_{h}}}+\mathbf{B}^{(h)}\right)\mathbf{V}^{(h)}。(2) 这里，mhm_{h}控制有效空间尺度：较大的值强调邻近细胞，而较小的值保留更广泛的的空间上下文。上下文编码器、目标编码器和预测器均使用该机制。

#### 随机掩蔽和上下文编码器。

对于每个空间patch，我们应用随机掩蔽，通过不放回均匀采样选择比例为r=0.6r=0.6的细胞作为预测目标M\mathcal{M}，其余细胞构成可见上下文C\mathcal{C}。在预训练期间，目标集在每次呈现patch时都会重新采样。目标细胞在上下文编码之前被移除，而不是用可学习的掩蔽token替换。因此，上下文编码器fθf_{\theta}仅处理可见细胞token及其归一化坐标：

zC=fθ(uC,s~C).\mathbf{z}_{\mathcal{C}}=f_{\theta}\left(\mathbf{u}_{\mathcal{C}},\tilde{\mathbf{s}}_{\mathcal{C}}\right)。(3)

#### 目标编码器。

目标编码器fθ ̄f_{\bar{\theta}}镜像上下文编码器，但处理完整patch，其在目标位置的输出定义潜在预测目标：

zM=sg⁡[fθ ̄(uC∪M,s~C∪M)M],\mathbf{z}_{\mathcal{M}}=\operatorname{sg}\left[f_{\bar{\theta}}\left(\mathbf{u}_{\mathcal{C}\cup\mathcal{M}},\tilde{\mathbf{s}}_{\mathcal{C}\cup\mathcal{M}}\right)_{\mathcal{M}}\right],(4) 其中sg\operatorname{sg}表示停止梯度。这些表示作为预测器被训练去恢复的潜在目标。

#### 提示条件空间预测器。

为减少细胞级目标歧义性，我们以有限的部分表达提示作为每个目标查询的条件。对于每个目标细胞ii，我们将其保留的非零基因集定义为Gi+={g∈Gi:xig>0}\mathcal{G}_{i}^{+}=\{g\in\mathcal{G}_{i}:x_{ig}>0\}，并以概率κ=0.10\kappa=0.10独立地包含每个基因，使用ξig∼Bernoulli⁡(κ)\xi_{ig}\sim\operatorname{Bernoulli}(\kappa)。由此产生的提示表示为

hi=∑g∈Gi+ξigxigeg+eiorgan+eiplatform.\mathbf{h}_{i}=\sum_{g\in\mathcal{G}_{i}^{+}}\xi_{ig}x_{ig}\mathbf{e}_{g}+\mathbf{e}^{\mathrm{organ}}_{i}+\mathbf{e}^{\mathrm{platform}}_{i}。(5)
提示仅由预测器使用，不传递给任一编码器。每个目标查询为qi=m+Wcphi\mathbf{q}_{i}=\mathbf{m}+W_{\mathrm{cp}}\mathbf{h}_{i}，其中m\mathbf{m}是可学习的掩蔽token，WcpW_{\mathrm{cp}}将提示投影到预测器空间。相同的投影将可见上下文表示映射到预测器空间。