GeneJepa:把转录组建模成“预测性世界模型”,这条路走得通吗?

论文:Litman E, Myers T, Agarwal V, et al. GeneJepa: A Predictive World Model of the Transcriptome. bioRxiv, 2025. doi: 10.1101/2025.10.14.682378.

30 秒速览

GeneJepa 是一篇把 LeCun 提出的 Joint-Embedding Predictive Architecture(JEPA) 范式引入单细胞转录组学的基础模型论文。它不再像 scGPT 那样去重建 noisy 的原始表达计数,而是把细胞的一部分基因作为“上下文”,在隐空间预测另一部分基因的隐层表征。模型用 Fourier 特征 编码连续表达、用 Perceiver 把可变长度基因集压缩成固定大小的隐层阵列,再用 EMA 教师 + VICReg 防坍塌 做稳定训练。在 Tahoe-100M 癌症细胞系图谱上预训练后,GeneJepa 在 PBMC3k、Human Lung Cell Atlas、sci-Plex 药物响应、NeurIPS 2023 PBMC 扰动预测等任务上都超过了 scGPT 和 UCE 等基线,并提出可以通过多次读取基因子集实现 test-time scaling,以及通过隐空间向量加减做 zero-shot 基因敲除模拟。论文的核心论点是:对高度噪声、集合结构、高维稀疏的转录组数据,“预测表征”比“重建信号”更契合基础模型的目标


一、研究动机:为什么转录组基础模型需要换一种目标函数?

单细胞 RNA 测序(scRNA-seq)已经让研究者能够在单细胞分辨率上同时观测成千上万个基因的表达量,从而刻画细胞身份、状态转换和扰动响应。但这条数据流有几个让传统监督学习非常头疼的特性:维度极高(通常 2 万 + 基因)、稀疏性强(dropout 导致大量零计数)、批次效应大(不同实验、平台、实验室之间的系统差异),以及组合性混杂(细胞类型、细胞周期、处理方式等信号交织)。

过去几年,scGPT、scBERT、UCE 等工作把 NLP 的预训练思路搬到转录组:把基因当作 token、把表达量分箱成离散值、用 masked token reconstruction 或自回归目标进行预训练,再微调下游任务。这些方法在领域内指标上取得了不少进展,但也 inherit 了几个根本性的 mismatch:

  1. 它们强迫模型去精确重建 noisy、zero-inflated 的原始计数,相当于把推断细胞的抽象调控状态降级为“猜测每一个 read 的数值”。
  2. 它们需要给基因强加一个顺序,但转录组在生物学上本质是一个无序集合(set),序列位置只是人为的 token ID。
  3. 它们依赖对比负样本或显式分箱,导致目标函数在噪声分布、批次伪影和 vocabulary 设计上非常敏感。

GeneJepa 的出发点可以概括为一句话:我们不该用“重建像素的均方误差”来评价一个世界模型,而应该看它是否学会了细胞状态背后的隐规则。这正对应 LeCun 近年来反复强调的 JEPA 范式:自监督学习的目标不是重建输入,而是预测输入在抽象表征空间中的表示。

对转录组而言,这个范式转换尤其自然。一个细胞的转录状态是一个高维的、带噪声的观测;但它背后的调控程序——哪些基因被共调控、哪些通路被激活、扰动后状态如何迁移——才是生物学真正关心的对象。GeneJepa 把这一点形式化为:给定一部分基因的表征,预测另一部分基因的表征。这个任务迫使模型学习基因之间的条件依赖关系,而不是去拟合计数中的泊松噪声。


二、核心创新:不是新算子,而是把 JEPA、Perceiver 和 Fourier tokenizer 整合到转录组场景

GeneJepa 在模块层面并没有发明全新的网络层。它的贡献在于把三个关键设计以一种针对转录组特性的方式组合起来

  1. JEPA 式的表征预测目标:用 cosine 相似度在隐空间预测 masked 基因块的表征,而不是重建表达值。这直接把学习压力从“数值拟合”转移到“依赖结构学习”。
  2. 连续表达的 Fourier tokenizer:用 Fourier 特征把标量表达值映射到高频可学习的向量,避免把表达量粗暴分箱。这保留了定量的、连续的信息。
  3. Perceiver 编码器 + 固定隐层阵列:通过 cross-attention 把可变长度的基因集压缩到固定大小的隐层,再用深度 Transformer 在这些隐层上做“思考”。这让模型的核心计算量与输入基因数量解耦,同时天然支持 set 结构。

这三点共同构成了 GeneJepa 的方法论核心:表征预测 + 集合感知 + 连续值保留。在此基础上,作者还引入了 momentum teacher、stop-gradient、VICReg 防坍塌、以及 test-time scaling 等工程技巧,使得训练在大规模数据上稳定,并给推理阶段留下可扩展的余地。


三、方法原理:GeneJepa 是如何工作的?

3.1 整体框架:从“上下文”预测“目标”的隐层表征

GeneJepa 的整体流程如图 1(GeneJepa 架构总览)所示。

图 1:GeneJepa 架构总览(来源:原论文 Fig. 1)

图 1:GeneJepa 架构总览

详细解读:这张图清晰地划分了三个网络模块和两条数据流。输入是单个细胞的转录组 xxx,被随机分成两部分:Context Genes xctxx_{ctx}xctx(上下文)和 Target Genes xtgtx_{tgt}xtgt(目标)。左侧路径 1:学生编码器 fθf_\thetafθ 读取上下文,产出上下文表征 zctxz_{ctx}zctx;路径 3:预测器 pϕp_\phipϕzctxz_{ctx}zctx 为唯一输入,预测目标表征 z^tgt\hat{z}_{tgt}z^tgt。右侧路径 2:动量教师编码器 fξf_\xifξ 读取目标基因块,产出目标表征 ztgtz_{tgt}ztgt,作为预测的“真值”。虚线、sg 符号和显式的 STOP 标志表示教师不接收梯度,从而避免学生—教师互相追逐导致坍塌。上方的 EMA 更新 ξ←βξ+(1−β)θ\xi \leftarrow \beta\xi + (1-\beta)\thetaξβξ+(1β)θ 则让教师缓慢跟随学生,提供稳定目标。整个框架最本质的信息是:模型不预测基因表达值,而是预测教师网络对目标基因块的隐层输出

形式化地,对一次前向传播:

zctx=fθ(xctx),ztgt=fξ(xtgt),z^tgt=pϕ(zctx)(1)z_{ctx} = f_\theta(x_{ctx}), \quad z_{tgt} = f_\xi(x_{tgt}), \quad \hat{z}_{tgt} = p_\phi(z_{ctx}) \tag{1}zctx=fθ(xctx),ztgt=fξ(xtgt),z^tgt=pϕ(zctx)(1)

注意,预测器只接收 zctxz_{ctx}zctx,不接收任何关于目标基因身份的信息。这一点非常关键:如果预测器知道要预测哪些基因,它可能学会利用基因嵌入表中的统计捷径,而不是真正从细胞状态中推断。

3.2 Tokenization:用 Fourier 特征保留连续表达

GeneJepa 的 tokenizer 要同时编码两个异质信息:基因身份(离散)和表达值(连续)。身份部分使用标准的可学习嵌入 eid∈Rdide_{id} \in \mathbb{R}^{d_{id}}eidRdid。表达值部分则用 Fourier 特征:

eval′=[sin⁡(vω1),cos⁡(vω1),…,sin⁡(vωNf),cos⁡(vωNf)](2)e'_{val} = \left[\sin(v\omega_1), \cos(v\omega_1), \ldots, \sin(v\omega_{N_f}), \cos(v\omega_{N_f})\right] \tag{2}eval=[sin(vω1),cos(vω1),,sin(vωNf),cos(vωNf)](2)

其中 vvv 是 log-normalized 后的表达值,{ω1,…,ωNf}\{\omega_1, \ldots, \omega_{N_f}\}{ω1,,ωNf} 是在最小频率和最大频率之间对数间隔取样的 NfN_fNf 个频率。这个 2Nf2N_f2Nf 维的原始 Fourier 嵌入再经过一个轻量 MLP,得到 eval∈Rdvale_{val} \in \mathbb{R}^{d_{val}}evalRdval。最后把身份和表达嵌入拼接:

t=LayerNorm(GELU(W[eid;eval]+b))(3)t = \text{LayerNorm}\left(\text{GELU}\left(W\left[e_{id}; e_{val}\right] + b\right)\right) \tag{3}t=LayerNorm(GELU(W[eid;eval]+b))(3)

这里 t∈Rdt \in \mathbb{R}^dtRd 就是一个基因的 token 嵌入,d=did+dvald = d_{id} + d_{val}d=did+dval。在默认配置中 Nf=64N_f = 64Nf=64,所以仅 Fourier 部分就有 128 维;加上基因身份后总维度为 768(见 Table 3)。

这个设计对比 scGPT 等工作的“表达分箱 + 离散 token”有两个直接好处:第一,不丢失表达量的细粒度信息;第二,不需要为表达值设计 vocabulary,避免了分箱边界造成的梯度不连续和词汇表膨胀问题。代价是每个 token 的维度更高,且 Fourier 特征是否真正“容易推理连续值”依赖于频率范围的选择——论文给出的频率范围是 [0.1,100.0][0.1, 100.0][0.1,100.0],跨越三个数量级,这对表达值的常见动态范围是否有覆盖,是一个可以追问的实证问题。

3.3 Perceiver 编码器:把“读多少”和“想多少”解耦

单细胞数据的另一个挑战是:每个细胞测到的基因数量差异很大(取决于测序深度和细胞类型)。如果直接用标准 Transformer 的自注意力,复杂度会随基因数平方增长。GeneJepa 使用 Perceiver 架构来解决这个问题。

Perceiver 的核心是一个固定大小的可学习隐层阵列 L∈RNlat×dL \in \mathbb{R}^{N_{lat} \times d}LRNlat×d。编码器首先做 cross-attention,让 NlatN_{lat}Nlat 个隐层查询去“读”可变长度的基因 token 集合 T∈RNgenes×dT \in \mathbb{R}^{N_{genes} \times d}TRNgenes×d

L′=L+CrossAttention(Q=L,KV=T)(4)L' = L + \text{CrossAttention}(Q=L, KV=T) \tag{4}L=L+CrossAttention(Q=L,KV=T)(4)

这一步的复杂度是 O(Nlat⋅Ngenes)O(N_{lat} \cdot N_{genes})O(NlatNgenes),而不是 O(Ngenes2)O(N_{genes}^2)O(Ngenes2)。由于 NlatN_{lat}Nlat 固定(论文用 512),输入基因数量增加只会线性增长 cross-attention 的成本。

接下来,更新后的隐层 L′L'L 进入一个深度 Transformer 堆栈,只在这些固定数量的隐层之间做 self-attention。这一步的复杂度是 O(Nlat2⋅D)O(N_{lat}^2 \cdot D)O(Nlat2D),与输入基因数无关。最后对隐层做 mean pooling,得到细胞(或基因子集)的表征:

z=1Nlat∑k=1NlatLk(final)(5)z = \frac{1}{N_{lat}} \sum_{k=1}^{N_{lat}} L^{(final)}_k \tag{5}z=Nlat1k=1NlatLk(final)(5)

这个结构带来三个好处:

  • 规模可扩展:输入基因数增加只影响“读”的成本,不影响“想”的成本。
  • 集合结构自然:cross-attention 对输入顺序不敏感,完美匹配转录组的无序集合特性。
  • 支持 test-time scaling:可以在推理阶段多次读取不同基因子集,每次只重新计算 cross-attention,而深层 Transformer 只跑一次。

论文还提到,为了处理极大的 NgenesN_{genes}Ngenes,cross-attention 使用了 online softmax 分块计算,并在 float32 累加器中维护数值稳定性(详见附录 A.2)。这在工程上是必要的,否则混合精度训练下 softmax 容易溢出或下溢。

3.4 Masking 策略与预测器:让上下文“足够信息丰富”

GeneJepa 的掩码策略很简单:对每次训练样本,先随机打乱细胞的表达基因集合,然后把约 45% 的基因指定为目标块,其余作为上下文块。同时要求上下文块至少包含 512 个基因,目标块至少包含 16 个基因。作者发现,生成一个大的目标块(而不是多个小目标块)提供了更强、更稳定的训练信号。

这个设计背后的逻辑是:如果目标块太小,预测任务太容易;如果上下文块太小,信息不足以推断目标。45% 的掩码比例意味着模型几乎要“想象”细胞一半的转录状态,这对表征学习是一个相当强的压力测试。不过,这个比例是在 Tahoe-100M 这种大规模扰动图谱上调出来的,对于更稀疏或更小的数据集是否最优,论文没有给出系统消融。

预测器是一个带 GELU 和 LayerNorm 的 3 层 MLP,扩展因子为 4,输入只有 zctxz_{ctx}zctx,输出 z^tgt\hat{z}_{tgt}z^tgt。这个选择延续了 JEPA 的常见做法:预测器不需要太复杂,因为主要的学习负担在学生编码器上——编码器必须把足够多的细胞状态信息压缩进 zctxz_{ctx}zctx,才能支撑对目标块的预测。

3.5 动量教师:稳定目标与防止坍塌

教师网络 fξf_\xifξ 与学生网络 fθf_\thetafθ 结构相同,但不通过反向传播更新,而是 EMA:

ξk←βkξk−1+(1−βk)θk(6)\xi_k \leftarrow \beta_k \xi_{k-1} + (1 - \beta_k)\theta_k \tag{6}ξkβkξk1+(1βk)θk(6)

其中 EMA 衰减率 βk\beta_kβk 从训练初期的 0.992 按照余弦 schedule 逐渐增加到接近 0.9995。这样早期教师更新较快,能帮助学生快速建立方向;后期教师变得非常稳定,避免目标抖动。

stop-gradient 是 JEPA 训练的关键:如果不阻止梯度流向教师,教师会被学生“拉”向容易预测的表示,最终可能坍塌到常数向量。通过把教师冻结在 EMA 轨道上,学生被迫真正学习如何预测一个有意义的、缓慢演化的目标。

3.6 损失函数:cosine 相似度 + VICReg 防坍塌

总损失由三部分组成:

L=λsimLsim+λvarLvar+λcovLcov(7)\mathcal{L} = \lambda_{sim}\mathcal{L}_{sim} + \lambda_{var}\mathcal{L}_{var} + \lambda_{cov}\mathcal{L}_{cov} \tag{7}L=λsimLsim+λvarLvar+λcovLcov(7)

其中 λsim=1.0\lambda_{sim}=1.0λsim=1.0λvar=25.0\lambda_{var}=25.0λvar=25.0λcov=1.0\lambda_{cov}=1.0λcov=1.0

预测相似性损失用 cosine 距离衡量预测表征与目标表征的夹角:

Lsim=1−z^tgt⋅ztgt∥z^tgt∥2⋅∥ztgt∥2(8)\mathcal{L}_{sim} = 1 - \frac{\hat{z}_{tgt} \cdot z_{tgt}}{\|\hat{z}_{tgt}\|_2 \cdot \|z_{tgt}\|_2} \tag{8}Lsim=1z^tgt2ztgt2z^tgtztgt(8)

选择 cosine 而不是 MSE 很重要:它让优化只关心方向,不关心向量模长,避免教师和学生因为尺度不同而相互牵制。

Variance 损失防止特征级坍塌,它鼓励每个维度上的标准差接近 1:

Lvar=1d∑j=1dmax⁡(0,1−Var(z^:,j)+ϵ)(9)\mathcal{L}_{var} = \frac{1}{d}\sum_{j=1}^{d} \max\left(0, 1 - \sqrt{\text{Var}(\hat{z}_{:,j}) + \epsilon}\right) \tag{9}Lvar=d1j=1dmax(0,1Var(z^:,j)+ϵ)(9)

Covariance 损失防止维度级坍塌,惩罚协方差矩阵的非对角元素:

Lcov=1d∑i≠j[Cov(z^)]ij2(10)\mathcal{L}_{cov} = \frac{1}{d}\sum_{i \neq j} \left[\text{Cov}(\hat{z})\right]_{ij}^2 \tag{10}Lcov=d1i=j[Cov(z^)]ij2(10)

VICReg 的作用是在不使用负样本的情况下保持嵌入空间的“体积”和“多样性”。不过值得指出,论文中的 VICReg 只加在学生预测器输出的 batch 上,而不是同时加在教师输出上;这相当于把正则化压力主要放在学生一端,让教师保持相对“原始”的目标分布。这个设计是否合理,取决于它是否会导致学生—教师之间的分布偏移——论文没有对此做显式分析。

3.7 从隐层向量到“世界模型”:zero-shot 敲除的几何解释

论文中最具想象力的部分是把训练好的隐层空间当作可操作的“世界模型”。给定一个基因 ggg,作者从 Tahoe-100M 中挑出“野生型-like”集合 SWT(g)S_{WT}(g)SWT(g) 和“突变型-like”集合 SMUT(g)S_{MUT}(g)SMUT(g),计算平均嵌入之差作为该基因的敲除方向:

Δg=Ec∈SMUT(g)[fθ(c)]−Ec∈SWT(g)[fθ(c)](11)\Delta_g = \mathbb{E}_{c \in S_{MUT}(g)}\left[f_\theta(c)\right] - \mathbb{E}_{c \in S_{WT}(g)}\left[f_\theta(c)\right] \tag{11}Δg=EcSMUT(g)[fθ(c)]EcSWT(g)[fθ(c)](11)

然后对任意野生型细胞 ccc,预测其敲除后的嵌入为:

zpred=zctrl+Δg(12)z_{pred} = z_{ctrl} + \Delta_g \tag{12}zpred=zctrl+Δg(12)

为了量化移动是否真正沿着目标方向,作者把结果投影到单位方向 Δ^g=Δg/∥Δg∥\hat{\Delta}_g = \Delta_g / \|\Delta_g\|Δ^g=Δg/∥Δg 上:

π(z)=⟨z−μWT,Δ^g⟩(13)\pi(z) = \langle z - \mu_{WT}, \hat{\Delta}_g \rangle \tag{13}π(z)=zμWT,Δ^g(13)

这个几何框架优雅地把“扰动”表示为隐空间中的平移向量,把“表型读取”表示为投影到某个功能方向。但这里也埋着一个关键假设:基因敲除的生物学效应在嵌入空间中是线性的。这个假设是否成立,目前只在隐空间和一条 P53 通路的 readout 上验证过,远未得到湿实验或因果推断框架的确认。


四、实验结果:表征质量、药物响应与扰动预测

4.1 细胞身份识别:PBMC3k(图 2)

作者首先在经典的小型免疫数据集 PBMC3k 上检验 GeneJepa 学到的细胞身份流形。所有模型都作为 frozen feature extractor,不做微调,下游只用一个简单的 logistic regression 或 kNN。

图 2:GeneJepa 在 PBMC3k 上学习到更清晰的细胞身份流形(来源:原论文 Fig. 2)

图 2:GeneJepa 在 PBMC3k 上学习到更清晰的细胞身份流形

详细解读:图 A 是 GeneJepa 的 UMAP,B 细胞、CD4 T、CD8 T、NK、单核细胞、树突状细胞和血小板形成紧凑、分离良好的簇。图 B 的 scGPT 则呈现更零散、交错的结构,说明相同细胞类型没有被一致地拉到一起。图 C 的线性探针和 kNN 结果中,GeneJepa 的 Macro-F1 为 0.69(LR)/0.38(kNN),Accuracy 为 0.68/0.52;scGPT 对应只有 0.23/0.25 和 0.20/0.37。图 D 的 per-class F1 显示 GeneJepa 在几乎所有细胞类型上都占优,包括血小板、树突状细胞等稀有类型。这组结果表明,JEPA 目标确实能让模型学到对细胞身份更有判别力的表征。

不过 PBMC3k 只有约 3000 个细胞,是一个非常干净的“入门”基准。这里的好成绩只能说明 GeneJepa 具备基本的细胞类型区分能力,还不足以证明其在复杂组织或跨批次场景中的泛化力。

4.2 跨组织泛化:Human Lung Cell Atlas(图 3)

Human Lung Cell Atlas(HLCA)包含健康肺组织的数十种细粒度细胞类型,是检验跨组织表征质量的重要基准。

图 3:GeneJepa 在 HLCA 上学习到稳健的跨组织细胞身份流形(来源:原论文 Fig. 3)

图 3:GeneJepa 在 HLCA 上学习到稳健的跨组织细胞身份流形

详细解读:图 A 是在 frozen embeddings 上训练的线性探针得到的混淆矩阵,主对角线非常清晰,说明大多数肺细胞类型可以被线性分开。图 B 是对测试集 embedding 做 k-means 后的 UMAP,颜色对应无监督聚类簇,可以看到细胞群体形成明显分离的簇。图 C 展示了 per-class F1 与类别样本量(log10)之间的关系,Spearman ρ=0.85\rho = 0.85ρ=0.85:样本越多,模型学得越稳健。这既是好消息也是坏消息——好消息是模型确实学到了可泛化的身份结构;坏消息是稀有细胞类型的表示质量仍然明显受限于数据量,这是几乎所有基础模型都会面临的“长尾”问题。

4.3 药物响应预测:sci-Plex(图 4)

sci-Plex 是一个大规模单细胞化学扰动数据集。作者把药物响应预测建模为回归任务:给定一个细胞的转录组,预测其对应的化合物剂量(log 变换后)。

图 4:GeneJepa 在 sci-Plex 药物响应回归上取得最低误差与最好跨上下文稳健性(来源:原论文 Fig. 4)

image-20260828070709522

详细解读:图 A、B 显示 GeneJepa 的 RMSE(1.15)和 MAE(1.00)均低于 UCE(1.27/1.05)和 scGPT(1.34/1.11)。图 C 的 baseline-normalized RMSE(rRMSE)中,GeneJepa 是唯一低于 1.0(0.94)的模型,意味着它比“全局中位数”这一朴素基准更好;UCE 和 scGPT 则高于 1.0。图 D、E 的 per-context MAE 分布显示 GeneJepa 不仅有最低的中位数误差(0.99),而且跨上下文的误差 spread(IQR = 0.09)也最小。图 F 的偏差幅度 |mean(y−y^y - \hat{y}yy^)| 显示 GeneJepa 的 0.13 远低于 UCE 的 0.57 和 scGPT 的 0.71。这组结果说明,GeneJepa 的嵌入在“剂量—响应”这种连续回归任务上具有更好的校准和跨上下文稳定性。

这里需要补充一个细节:sci-Plex 任务被处理成“伪 bulk”(pseudobulk)回归,即按(细胞系、化合物、时间)聚合后的表达均值来预测剂量。这降低了单细胞层面的噪声,但也意味着模型没有直接在单细胞—单细胞剂量关系上做验证。这个 preprocessing choice 是否对 GeneJepa 特别友好,是值得讨论的。

4.4 扰动方向预测:NeurIPS 2023 PBMC 任务(表 1)

Open Problems NeurIPS 2023 的任务要求预测 144 种化合物在 primary human PBMC 中处理 24 小时后的转录响应。目标是一个方向显著性分数:

S=−log⁡10(p)⋅sign(log⁡FC)(14)S = -\log_{10}(p) \cdot \text{sign}(\log\text{FC}) \tag{14}S=log10(p)sign(logFC)(14)

并裁剪到 [−4,4][-4, 4][4,4]。GeneJepa 的做法是:用训练集差异表达矩阵构建一个 pseudo-expression 向量,喂给教师编码器得到细胞嵌入;把 Morgan 化合物指纹和剂量对数与嵌入拼接;训练一个多输出 ridge regressor 预测所有基因的 SSS 值。

方法CosineΔ\DeltaΔPearsonΔ\DeltaΔSpearmanΔ\DeltaΔ
GeneJepa0.36980.35090.3431
Mean outcome0.2264+0.14340.2198+0.13110.2117+0.1314
Mean across compounds0.2628+0.10700.2594+0.09150.2425+0.1006
Mean across cell types0.3017+0.06810.2972+0.04570.2806+0.0625
Transformer ensemble0.2267+0.14310.2212+0.12970.2164+0.1267
OP2 style baseline (JN-AP-OP2)0.3289+0.04090.3267+0.02420.3054+0.0377
Random sample0.0547+0.31510.0524+0.29850.0562+0.2869
Zeros0.0000+0.36980.0000+0.35090.0000+0.3431

表 1 显示 GeneJepa 在 cosine、Pearson、Spearman 三个方向性指标上都领先。值得注意的是,它对“Mean across cell types”这一强基线的提升最小(cosine 只 +0.0681),说明在按细胞类型平均的预测上,简单基线已经相当有竞争力。GeneJepa 真正的优势在于跨化合物、跨细胞类型地捕捉扰动方向,而不只是在已知细胞类型内做平均。

4.5 Test-Time Scaling:用多次读取换取更高精度(图 5)

Perceiver 架构的一个独特优势是:cross-attention 的“读”阶段与深层 Transformer 的“想”阶段是分离的。作者据此提出 test-time scaling:把细胞基因按表达变异度分成多个 chunk,重复执行 cross-attention,每次读入更多基因,但深层 Transformer 和预测器只跑一遍。

图 5:GeneJepa 的 test-time scaling:用更多读取换取更高保真度,同时控制推理成本(来源:原论文 Fig. 5)

image-20260828070731520

详细解读:图 A 显示在 PBMC68k 上,用固定探针训练后,随着 reads 从 1 增加到 4,Transfer Macro F1 从约 0.02 提升到约 0.35,Accuracy 也接近完整模型。图 B 显示预测嵌入与完整输入嵌入的 cosine 相似度快速趋近于 1。图 C 显示每细胞推理时间维持在约 26 ms,显存只有一次小幅跳升(来自最后一个最大 chunk)。图 D 是整个流程的示意图:基因集被分成 chunk 1–4,cross-attention 重复 rrr 次,而隐层 Transformer + 预测器只运行一次。这说明 GeneJepa 可以在不重新训练的情况下,通过调节 reads 数量来 trade off 速度与精度。

不过图 B 也暴露了一个值得注意的现象:在 reads=2 时,cosine to full 反而出现一个明显下降(约 0.75),然后 reads=3、4 才回升。作者没有解释这个 dip。一个可能的解释是:第二个 chunk 的加入暂时改变了表征分布,导致与完整输入的对齐度下降。这说明 test-time scaling 的单调性并没有被严格保证,实际部署时可能需要根据验证集来决定最优 reads 数,而不是简单地“越多越好”。

4.6 Zero-shot 基因敲除:TP53 案例(图 6、表 2)

论文把 GeneJepa 的隐层空间称为“可逆的世界模型”——你可以通过向量加减来模拟基因敲除。Figure 6 展示了 TP53 敲除的 zero-shot 结果。

图 6:GeneJepa 通过隐空间向量平移实现 zero-shot TP53 敲除模拟(来源:原论文 Fig. 6)

图 6:GeneJepa 通过隐空间向量平移实现 zero-shot TP53 敲除模拟

详细解读:图 A 显示在 Δ^TP53\hat{\Delta}_{TP53}Δ^TP53 方向上,一个 WT 对照细胞(绿色)沿方向平移后到达预测状态(橙色),与 MUT 均值接近;∥Δ∥=2.671\|\Delta\| = 2.671∥Δ∥=2.671。图 B 是剂量扫描:α\alphaα 从 -0.5 到 1.0,P53 通路活性读出的中位数和均值单调下降,Spearman ρ=−1.00\rho = -1.00ρ=1.00。图 C 显示在 0%、10%、30%、50% 输入 dropout 下,单调趋势都保持稳定,说明这个方向具有一定鲁棒性。Caption 强调这是“without any perturbation-specific training”。

表 2 进一步给出了一种更严格的、无标签的 zero-shot 验证方式:

参数
GeneTP53
Build size30,000 cells
Evaluation size20,000 cells
Directional alignment (cosine)0.691
Distance shrinkage (latent units, mean)0.0248
Distance shrinkage (latent units, median)0.0358
Percent shrinkage (median)27.17% ±\pm± 0.27% (95% CI)
Wilcoxon one-sided p (median > 0)<10−6< 10^{-6}<106

表 2 中的 label-free 方法完全不使用 WT/MUT 标签,而是直接从训练集上“ablate”一个基因:把该基因从输入中去掉,比较原细胞和 ablated 细胞的嵌入差,平均后得到 Δg\Delta_gΔg。然后在测试集上把这个方向加到未 ablated 的细胞上,看它是否更接近真实的 ablated 嵌入。中位距离收缩 27.17% 且 p 值极显著,说明隐层空间确实捕捉到了某种与基因缺失相关的方向结构。

但这里也隐藏着一层循环:label-free 方法仍然在计算“该基因缺失后嵌入如何变化”,而这个变化本身可能部分来自缺失信号导致的表达重标度,而不是模型对“敲除后果”的因果推断。此外,验证仍然停留在嵌入距离和通路 readout层面,没有 wet-lab 的 qPCR、RNA-seq 或功能实验来确认这些预测在真实细胞中是否成立。


五、批判性分析:优点、局限与可追问的问题

5.1 GeneJepa 真正做对了什么?

第一,它提出了一种更契合转录组特性的自监督范式。 相比于在 noisy count 空间做重建,JEPA 的表征预测目标把学习压力放在基因之间的依赖结构和细胞的抽象状态上。这不仅是方法层面的创新,也是概念层面的重新定位:把转录组基础模型的评价标准从“重建精度”转向“世界模型的预测能力”。

第二,Perceiver + Fourier tokenizer 的工程组合务实且有效。 Perceiver 让模型摆脱了“基因序列长度必须固定”的约束,Fourier 特征避免了表达分箱带来的信息损失。这两者在单细胞数据上的适配性很强,值得后续工作借鉴。

第三,test-time scaling 给了实际部署一个灵活的旋钮。 在资源受限时可以用 1 read 快速推理,在需要高精度时增加到 3–4 reads,而不用重新训练模型。这种“可扩展的推理”在单细胞分析场景(如临床样本实时注释)中非常实用。

第四,zero-shot 敲除展示了隐层空间的可操作性。 即使生物学解释仍待验证,能够在不针对具体扰动训练的情况下,用向量运算模拟基因功能变化,本身就是一种有潜力的研究工具。它把“表征学习”和“因果探索”连接了起来。

第五,实验覆盖较为全面。 从细胞类型识别、跨组织泛化、药物响应回归,到扰动方向预测和 zero-shot 敲除,论文展示了 GeneJepa 在多个下游任务上的 transferability。而且所有 backbone 都 frozen,保证了比较的公平性。

5.2 需要警惕的局限

第一,预训练数据存在严重的领域偏倚。 Tahoe-100M 主要由癌症细胞系和小分子化合物扰动构成,而论文的下游评估却涉及 primary PBMC、HLCA 等原代组织。作者自己在 Limitations 中也承认这一点:癌症细胞系的表达程序和调控背景与原代细胞差异显著,模型的跨领域迁移能力可能被高估。如果 GeneJepa 在原代组织、发育时间序列或疾病样本上的表现显著下降,那么它的“通用世界模型”称号就要大打折扣。

第二,JEPA 目标没有显式的批次校正或领域不变项。 作者把跨批次稳健性寄托在“数据规模 + 归纳偏置”上,但单细胞 genomics 的批次效应 notoriously 强大。如果不同实验室的测序深度、protocol、物种差异没有被显式建模,模型学到的“细胞状态”可能仍然混杂技术伪影。论文没有提供跨批次整合(batch correction / integration)任务的系统评估,这是基础模型文献中非常重要的一个维度。

第三,zero-shot 敲除的生物学有效性远未确立。 目前的验证全部发生在隐层空间和一条 P53 通路 readout 上。向量加减在嵌入空间中成立,不等于它在真实细胞中成立。真正的基因敲除会通过级联信号影响大量基因和通路,而 GeneJepa 的线性方向假设无法保证这些非线性、反馈式的效应被正确捕获。没有湿实验验证,这些都只能称为“in-silico 启发”,不能称为“生物学预测”。

第四,WT/MUT 代理定义过于宽松。 在 metadata-based 方法中,作者用基因表达量或通路活性的分位数来划分“野生型-like”和“突变型-like”细胞。这种代理并没有真正对应 TP53 突变状态,而只是把高/低表达或高/低通路活性作为替代标签。由此计算出的 Δg\Delta_gΔg 可能混合了突变效应、细胞周期状态、应激反应等多种混杂因素。

第五,test-time scaling 的单调性并不总是成立。 图 5B 中 reads=2 时 fidelity 出现明显下降。这说明多次读取策略的稳定性需要更仔细的分析,而不是简单地认为“reads 越多越接近 full model”。如果不同 chunk 的顺序或组成影响隐层状态,临床部署时就需要一个可靠的 early-exit 或收敛判断机制。

第六,模型内部机制的可解释性不足。 GeneJepa 的 Perceiver 隐层和深层 Transformer 学习了什么?哪些基因、通路或调控关系被编码在哪些隐层维度上?论文没有提供 attention 可视化、隐层 dissection、或 pathway-centric probing。对于一个宣称“世界模型”的工作,这些分析几乎是必需的——否则“世界模型”更像是一个修辞,而不是一个可被理解的科学假说。

第七,训练成本与可复现性。 论文使用 4 张 NVIDIA H100 80GB GPU,在超过 1 亿个单细胞 profile 上训练。这种规模对绝大多数学术团队和小型实验室是不可承受的。论文虽然开源了代码和权重,但如果预训练数据或算力门槛过高,社区就很难在该模型上做进一步验证和改进。

第八,缺少系统消融实验。 论文几乎没有任何内部 ablation:masking ratio 是否 45% 最优?Fourier tokenizer 相比分箱到底带来多少提升?Perceiver 的隐层数、深度、cross-attention chunk size 如何影响性能?EMA schedule 和 VICReg 系数是否关键?没有这些,我们就很难判断哪些设计是真正有贡献的,哪些只是随大流的选择。

5.3 可改进方向

  1. 引入显式的批次/领域不变目标:例如 domain-adversarial 训练、显式 batch embedding 移除、或基于元学习的领域泛化,来缓解 Tahoe-100M 的领域偏倚。
  2. 做更严格的因果验证:把 zero-shot 敲除预测与真实的 CRISPR knock-out 实验、通路活性 assay、或时间序列扰动数据对比,而不是只在隐空间比较距离。
  3. 深入解剖隐层表征:用 pathway-centric probes、attention rollout、或 sparse autoencoder 来分析 Perceiver 隐层编码了哪些生物学知识。
  4. 系统 ablation 与轻量版本:探索更小的隐层数、更浅的 Transformer、更小的 vocabulary,看看能否在保持性能的同时降低训练与推理成本。
  5. 把 JEPA 扩展到多组学:论文在 Outlook 中已经提到,Perceiver 的模态无关性非常适合训练 transcriptome + epigenome + proteome 的联合世界模型。这是真正能把 GeneJepa 从“单模态表征模型”推向“细胞世界模型”的方向。
  6. 设计生物学感知的掩码策略:与其随机掩码 45% 基因,不如按通路、复合物或调控模块来构造目标块,迫使模型学习更结构化的生物学推理。

六、一句话总结

GeneJepa 不是又一篇在表达计数上做 masked reconstruction 的单细胞基础模型,而是把 JEPA 的“预测隐层表征”范式认真地搬到了转录组领域:用 Fourier 特征保留连续表达、用 Perceiver 把可变基因集压缩成固定隐层、用 EMA 教师和 VICReg 做稳定训练,并在细胞类型识别、药物响应和扰动预测等任务上展示了优于 scGPT 和 UCE 的表征质量。

它最有价值的贡献,或许不是某个具体指标的提升,而是把“转录组基础模型应该预测什么”这个问题重新提了出来——从“重建带噪信号”转向“学习细胞状态的预测规则”。但这个愿景距离真正的“世界模型”还有距离:预训练数据的癌症细胞系偏倚、缺失的批次校正机制、只在隐空间验证的 zero-shot 敲除,以及缺少系统消融,都是下一步必须面对的硬问题。

换句话说,GeneJepa 打开了一扇很有吸引力的门,但门后的世界是否如论文所描绘的那样广阔,还需要更多跨领域、跨模态、跨干湿实验的验证。


术语对照表

英文术语中文译名简要说明
Transcriptome转录组细胞或组织中所有 RNA 转录本的集合。
scRNA-seq单细胞 RNA 测序在单细胞分辨率上测量基因表达的技术。
JEPA联合嵌入预测架构Joint-Embedding Predictive Architecture,通过预测隐层表征进行自监督学习。
Perceiver感知器(架构)用固定大小的隐层阵列通过 cross-attention 处理可变长度输入的神经网络。
EMA teacher动量教师网络参数通过学生参数的指数移动平均更新、不参与反向传播的目标网络。
Stop-gradient梯度截断阻止梯度流向教师网络,防止坍塌。
VICReg方差—不变性—协方差正则化通过方差和协方差约束防止表征坍塌的自监督正则项。
Fourier features傅里叶特征用正弦/余弦基函数把标量连续值映射到高维向量。
Cross-attention交叉注意力一种 attention 机制,让一组查询去关注另一组键—值。
Test-time scaling测试时扩展在推理阶段通过增加计算(如多次读取)换取更高精度。
Zero-shot knockout零样本基因敲除不针对具体扰动训练,直接在隐空间用向量运算模拟基因敲除效果。
PBMC外周血单个核细胞Peripheral Blood Mononuclear Cells,常用于免疫研究。
HLCA人类肺部细胞图谱Human Lung Cell Atlas。
Macro-F1宏平均 F1 分数对每个类别单独计算 F1 后取平均,更公平地评价稀有类别。
rRMSE基线归一化均方根误差相对全局中位数基准的 RMSE,小于 1 表示优于朴素基准。
Spearman ρ\rhoρ斯皮尔曼相关系数衡量两个变量单调相关程度的非参数指标。
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值