Atom-JEPA: Joint-Embedding Predictive Architecture for 3D Atomistic Systems
作者: Kasper Helverskov Petersen, Rasmus Hannibal Tirsgaard, François R J Cornet, Mikkel Jordahn, Mikkel N. Schmidt
分类: cs.LG, cs.AI, physics.comp-ph
发布日期: 2026-10-06
🔗 代码/项目: GITHUB
💡 一句话要点
提出Atom-JEPA以解决3D原子系统自监督预训练问题
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 自监督学习 3D原子系统 潜在表示 分子预测 晶体材料 机器学习 深度学习
📋 核心要点
- 现有的自监督预训练方法在3D原子系统建模中未能实现良好的下游泛化能力，限制了其应用。
- Atom-JEPA框架通过互补的原子级和子结构级目标，从未标记的3D结构中学习潜在表示，提升了模型的学习效果。
- 在多样化的下游属性预测任务中，Atom-JEPA在分子ADMET和量子化学属性预测上实现了最先进的性能，展示了其有效性。
📝 摘要（中文）
大规模自监督预训练已重塑现代机器学习，显著提升了语言和视觉模型在下游任务中的泛化能力。然而，深度学习在建模原子系统方面的进展尚未在自监督预训练中实现类似的泛化效果。为此，本文提出了Atom-JEPA框架，通过互补的原子级和子结构级目标，从未标记的3D结构中学习潜在表示。我们在大规模分子和晶体数据集上进行预训练，并通过微调评估其在多样化下游属性预测任务中的迁移性能。Atom-JEPA在分子ADMET和量子化学属性预测任务上达到了最先进的性能，并在预测晶体材料的物理属性方面表现出强竞争力。这些结果展示了潜在空间预测预训练在仅依赖结构数据时支持广泛下游泛化的潜力。
🔬 方法详解
问题定义：本文旨在解决现有自监督预训练方法在3D原子系统中的泛化能力不足的问题。现有方法在处理未标记数据时，未能有效捕捉到结构信息，导致下游任务性能不佳。
核心思路：Atom-JEPA的核心思想是通过互补的原子级和子结构级目标，利用联合嵌入预测架构，从未标记的3D结构中学习潜在表示。这种设计旨在增强模型对复杂结构的理解能力。
技术框架：Atom-JEPA的整体架构包括数据预处理、潜在表示学习和下游任务微调三个主要模块。首先，对大规模分子和晶体数据集进行预处理，然后通过自监督学习获取潜在表示，最后在多样化的下游任务上进行微调以评估性能。
关键创新：Atom-JEPA的主要创新在于其结合了原子级和子结构级的学习目标，形成了一个新的自监督预训练框架。这一方法与传统的单一目标学习方法相比，能够更全面地捕捉结构信息，从而提升模型的泛化能力。
关键设计：在参数设置上，Atom-JEPA采用了多层神经网络结构，并设计了特定的损失函数以平衡原子级和子结构级的学习目标。此外，模型在预训练阶段使用了大规模的分子和晶体数据集，以确保学习到的潜在表示具有广泛的适用性。
🖼️ 关键图片
📊 实验亮点
在实验中，Atom-JEPA在分子ADMET和量子化学属性预测任务上达到了最先进的性能，显著优于现有基线方法，展示了其在预测晶体材料物理属性方面的强竞争力，进一步验证了潜在空间预测预训练的有效性。
🎯 应用场景
该研究的潜在应用领域包括药物发现、材料科学和化学工程等。通过提高对3D原子系统的理解，Atom-JEPA能够加速新材料和药物的开发，具有重要的实际价值和未来影响。
📄 摘要（原文）
Large-scale self-supervised pretraining has reshaped modern machine learning, substantially advancing the ability of language and vision models to generalize across downstream tasks. While deep learning has driven considerable progress in modeling atomistic systems in recent years, self-supervised pretraining in this domain has not yet achieved comparable downstream generalization. To address this, we introduce Atom-JEPA, a self-supervised pretraining framework that learns latent representations from unlabeled 3D structures through complementary atom-level and substructure-level objectives inspired by joint-embedding predictive architectures. We pretrain Atom-JEPA on large-scale molecular and crystalline datasets and evaluate its transfer performance by fine-tuning on a diverse set of downstream property prediction tasks. Atom-JEPA achieves state-of-the-art performance on molecular ADMET and quantum-chemical property prediction tasks, and is highly competitive in predicting the physical properties of crystalline materials. These results demonstrate the potential of latent-space predictive pretraining to support broad downstream generalization from structural data alone. Code and pretrained model checkpoints are publicly available at https://github.com/khelverskovp/atom-jepa