沙丘社区
论文

Atom-JEPA：原子系统的联合嵌入预测预训练

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
模型训练预训练

摘要

大规模自监督预训练重塑了现代机器学习，极大地提高了语言和视觉模型泛化下游任务的能力。尽管近年来深度学习在原子系统建模方面取得了长足进步，但该领域的自监督预训练尚未实现可比的下游泛化。为了解决这个问题，我们引入了 Atom-JEPA，这是一种自监督预训练框架，受联合嵌入预测架构的启发，通过互补的原子级和子结构级目标，从未标记的 3D 结构中学习潜在表示。我们在大规模分子和晶体数据集上预训练 Atom-JEPA，并通过微调在各种下游属性预测任务上评估其传输性能。 Atom-JEPA在分子ADMET和量子化学性质预测任务上实现了最先进的性能，并且在预测晶体材料的物理性质方面具有很强的竞争力。这些结果证明了潜在空间预测预训练仅支持结构数据的广泛下游泛化的潜力。代码和预训练模型检查点可在 https://github.com/khelverskovp/atom-jepa 公开获取

Atom-JEPA：原子系统的联合嵌入预测预训练：论文原图
图 1：Atom-JEPA 概述。通过对锚原子周围的 $k$-hop 邻域进行采样，分子和晶体结构被划分为互补上下文（红色）和目标（蓝色）子图。 3D 等变图神经网络对上下文进行编码，然后用于预测两个尺度的目标表示：单个原子和均值池子结构。使用更新为上下文编码器的指数移动平均值的目标编码器后，从完整结构中提取目标表示。预训练最大限度地减少了 潜空间 中的预测误差，两个子图交替作为上下文和目标。随后对预训练的上下文编码器进行微调以进行下游属性预测。