Mol-JEPA: A multimodal Joint Embedding Predictive Architecture for Molecules
Florian Rottach, Sebastian Schieferdecker, William Rudman, Randall Balestriero, Carsten Eickhoff
cs.LG, cs.AI
2026-08-24
Boehringer 等提出分子世界模型 Mol-JEPA,不改结构,把 14 种生化模态整条掩掉再在潜空间预测,23 项 ADMET 与活性任务平均 MAE 0.402,优于 Chemprop 的 0.519。
分子预训练经常用原子或键的掩码、图扰动来造「同一分子的多个视角」,再逼表征彼此靠近。这个假设在图像里说得通,在分子上经常是错的。药化里有一类现象叫 property cliff,中文常译活性悬崖:结构只改一个基团,结合活性或毒性就能差一个数量级。把这种扰动当正样本,等于教模型把化学上不该一样的东西看成一样。
更硬的问题是,分子的药学行为并不写在结构式里。吸收、代谢、细胞表型、脱靶结合都发生在生化环境中。只在 SMILES 或分子图上做自监督,学到的是拓扑,学不到环境。已有多篇工作发现,预训练表征经常打不过 ECFP 加随机森林这类老方法。对比学习能接上细胞图像等模态,但负样本选不好就会把相近分子推开,模态一多还容易 collapse。
Mol-JEPA 把 JEPA(Joint Embedding Predictive Architecture,联合嵌入预测架构)搬到分子上。关键改动是不扰动分子图,直接把整条模态掩掉,让 Transformer 预测被掩模态的潜表示。
预训练集约 469 万个去重化合物,来自 ChEMBL、TDC、PCBA-1328、∇²DFT 和 Enamine REAL。每条分子最多对应 14 条模态,分三类:
三种可学习编码器(原子、图、向量)把不同形状的输入投到同一维度。随机按掩码率关掉若干已有模态,至少留一条可见;缺失模态用 token 填上。预测器用 PyTorch TransformerEncoder,带 CLS token 做下游读出。防 collapse 用 LeJEPA 的 SIGReg(Sketched Isotropic Gaussian Regularization,随机投影后逼近各向同性高斯),加在 target 嵌入上;加在预测嵌入上会训崩。最终约 5000 万参数,掩码率 0.3,正则权重 λ=0.7。12 张 RTX-6000ADA、有效 batch 1024,大约 3 天、150 epoch 收敛,再往后在线探针开始过拟合。
下游三种用法:CLS 线性或非线性探针,以及 TabICLv2;各模态嵌入送进两层 Transformer 聚合;LoRA(rank 16)微调预测器。代码和权重已随论文公开。
评测走 OpenADMET ExpansionRx、ASAP-Polaris、PXR 盲测和 Biogen ADME,共 23 个终点。划分用 Taylor-Butina 聚类(ECFP4,Tanimoto 0.65),整簇进训练或测试,做三套 split。基线包括随机森林、LightGBM、TabICLv2、CheMeleon、Chemprop、CLAMP。
跨任务平均误差如下:
| 方法 | 平均 MAE | 平均 R² |
| Mol-JEPA Best | 0.402 | 0.42 |
| Mol-JEPA Transformer | 0.427 | 0.38 |
| TabICLv2 + AlvaDesc | 0.453 | 0.37 |
| CheMeleon | 0.471 | 0.34 |
| LightGBM + AlvaDesc | 0.468 | 0.33 |
| Chemprop | 0.519 | 0.28 |
| CLAMP | 0.576 | 0.27 |
| 随机森林 + ECFP4 | 0.559 | 0.13 |
优势集中在小数据(ASAP、Biogen)和 Tanimoto 相似度更低的分布外分子。Wilcoxon 成对比较里,Mol-JEPA 赢面约 38%,TabICLv2 加 AlvaDesc 约 47%,后者主要吃 ExpansionRx 这类更大的集。公开竞赛的时间划分上,MAE 和 Wilcoxon 仍是 TabICLv2 更常赢,Mol-JEPA 的 R² 最好。
全量数据上,只保留 Graph 加 ECFP4 对比全 14 模态,线性探针 MAE 降 14%,非线性降 13%。10 万子集消融:半监督直接预测实验标签差过 JEPA(线性 MAE 0.467 对 0.436);去掉实验模态也会变差(0.449);嵌入维度 256 升到 512 收益大,到 1024 基本饱和;对比学习常用的投影头在这里反而有害。leave-one-out 里图模态贡献最大,它是唯一可训的 GNN 骨干,MOE 和 ECFP 次之。
对做 ADMET 或活性模型的人,这是一条可复用的多模态预训练配方:结构扰动在分子上不可靠,就把「别的测量和别的模型怎么看这个分子」当成视角。JEPA 不需要负样本,模态稀疏也能训,这对药化数据很对症。小数据和分布外更好,也更接近真实项目:高质量标签贵,新骨架经常落在训练集外面。
这仍是渐进工作。TabICLv2 加传统描述符在大数据集和竞赛划分上并没有被拉开,Wilcoxon 甚至略占上风。Mol-JEPA 的卖点是把细胞表型、结合预测、量子化学和实验向量装进同一套潜空间预测,不是在每个终点上都刷新榜。想跟的人可以从探针开始,不必一上来微调。
作者自己写了四条。重建目标没建模「有的模态根本还原不了另一些」;分布外仍难;多模态训练动态(稀疏、维度、冗余)还没想清楚;消融大多在 10 万子集上跑,和 469 万全量可能对不上。
读下来还有几处要打折。Mol-JEPA Best 是所有探针和嵌入变体里挑最好的,和「一个固定头」的基线不完全同台;LoRA 微调反而较差,说明表征并不即插即用。Boltz-2 只覆盖约 10 万分子、9 个靶点。CLOOME 这类模态有时会拖下游。实验向量极稀疏。竞赛时间划分上传统表格模型仍很硬,离「分子世界模型」能外推新化学空间还有距离。