TriForces: Augmenting Atomistic GNNs for Transferable Representations

TriForces 是一种模型无关的三流框架，它将自监督学习与分离的组成和结构表示相结合，以显著提升用于机器学习原子间势的原子图神经网络的迁移能力和数据效率。

要点

想象一下，你正在教一位机器人厨师如何烹饪。

问题：“一刀切”的厨师
目前，科学家使用强大的 AI 模型（称为 MLIPs）来预测原子的行为，例如材料的能量有多少，或者推动原子有多难。这些模型是在超级计算机（DFT）产生的海量数据上训练而成的。

然而，这些模型存在一个缺陷：它们就像一位记住了某道菜确切味道，却忘记了为什么是这种味道的厨师。如果你让它们做一道略有不同的菜（一种新材料），它们就会束手无策。它们会将食材（成分）与锅的形状（结构）混淆。如果你改变食材，它们就会对形状感到困惑，反之亦然。这使得它们难以快速学习新任务，尤其是在缺乏大量数据用于教学时。

解决方案：TriForces（三流厨房）
作者提出了TriForces，这是一种构建这些 AI 厨师的新方法。与其让一个巨大的大脑试图一次性记住所有事情，不如将大脑拆分为三个专门的“流”或部门：

1. 食材流（成分）： 该部门只关注锅里有什么（例如，“我们有 2 个氢原子和 1 个氧原子”）。它完全忽略形状。它学习的是化学。
2. 形状流（结构）： 该部门只关注原子在空间中如何排列（例如，“它们呈三角形排列”）。它忽略原子实际是什么。它学习的是几何学。
3. 相互作用流： 这是主厨，它接收来自食材和形状部门的笔记，并将它们结合起来以预测最终结果（能量或力）。

秘诀：自监督学习
在模型被要求预测特定属性之前，作者通过一种名为“自监督学习”的游戏对其进行训练。可以将其想象为一次练习环节，AI 必须：

• 去噪： 观察一张略有破损或带有噪声的分子图片并将其修复。
• 掩蔽： 遮盖住一种食材，并根据邻居猜测它是什么。
• 匹配： 观察同一分子的两种略有不同的版本，并意识到它们是同一个东西。

这种训练迫使 AI 将其知识有条理地组织起来。它学会了将“食材”归入一个文件夹，将“形状”归入另一个文件夹，而不是将它们混杂在一起。

为何重要（结果）
论文表明，这种新的“三流”厨房比旧的“单脑”厨房表现好得多：

• 更快的学习速度： 当给定少量新数据（例如 2 万个示例而非数百万个）时，TriForces 学习速度快得多，且错误更少。这就像一位只需尝一次就能学会新食谱的厨师，而不是需要烹饪一千次。
• 更好的记忆： AI 不会忘记它所学到的东西。它可以将知识从一种材料转移到另一种材料，而不会感到困惑。
• 可检索的知识： 由于 AI 将“食材”和“形状”分开保存，你可以要求它寻找外观相同但成分不同的材料，或者成分相同但形状不同的材料。旧模型无法做到这一点，因为它们的知识过于混杂。

总结
TriForces 是一个框架，它将理解原子的复杂任务分解为三个更简单的任务：了解成分、了解形状以及了解它们如何协同工作。通过训练 AI 将这些任务分开，并通过“猜谜游戏”（自监督学习）进行练习，该模型成为了一个更灵活、更高效、更准确的工具，用于发现新材料。

作者已发布其代码和预训练模型，以便其他科学家可以利用这个“三流厨房”为材料科学构建更好的 AI。

技术摘要

技术摘要：TriForces

问题陈述

基于几何图神经网络（GNN）的机器学习原子间势（MLIPs）在大规模密度泛函理论（DFT）数据集上训练时，已在预测原子尺度性质方面取得了高精度。然而，其实际效用受到两个主要限制因素的阻碍：

1. 迁移性不一致：MLIPs 往往难以有效地迁移到新的化学体系或任务中，特别是在针对小型、昂贵且特定任务的微调数据集进行微调时。预训练模型在处理简单的诊断任务（例如识别晶体系统或主要元素）时经常表现不佳，并表现出“灾难性遗忘”或跨领域的迁移性能不稳定。
2. 表征重用性差：当前的 MLIPs 针对特定的预测目标（能量和力）优化表征，而非为了通用重用。因此，这些表征往往将成分（化学）与结构（几何）纠缠在一起，使其不适合探索性分析、最近邻检索或分解式相似性搜索。标准的监督目标鼓励足以进行回归的表征，但并未组织成保留可访问的成分和结构信息的形式。

方法论：TriForces 框架

作者提出了TriForces，这是一个与模型无关的框架，通过三流架构和多目标自监督学习（SSL）增强现有的几何 GNN。核心创新在于将原子表征显式分解为三个不同的组件：

1. 三流架构

TriForces 将节点级表征 $h_i$ 分解为三个串联的流，而不是单个潜在向量：

• 成分流（$h^{comp}$）：编码不含坐标的化学信息。它使用带有计数加权注意力的 Transformer 处理唯一原子元素集合及其化学计量计数。该流保留绝对元素计数以编码系统尺寸和能量尺度，独立于几何结构。
• 结构流（$h^{struct}$）：编码不含元素身份（类型无关）的几何信息。它构建受 SOAP（原子位置平滑重叠）启发的旋转不变局部描述符，利用径向基函数、球谐函数和多尺度截断。该流捕捉可重用的几何基元和拓扑模式，并通过不变消息传递得到增强。
• 相互作用流（$h^{int}$）：一个标准的基线几何 GNN（例如 MACE、eSEN、Orb-v3），用于捕捉成分与几何之间的耦合，保留原始架构的表达力。

2. 自监督预训练

为了组织潜在空间并提高迁移性，TriForces 采用多目标 SSL 预训练策略，使用随机增强（位置噪声、元素掩码、图变异和旋转）。该框架结合了三个互补的目标：

• 非重建（LeJEPA）：在节点和图级别对齐同一结构的两个增强视图的嵌入。这强制了对增强的不变性，并组织了全局潜在空间，而无需停止梯度或动量编码器。
• 去噪：训练模型从噪声输入中恢复干净的原子位置。这稳定了几何表征，并隐式提供了旋转增强。
• 掩码：基于周围几何和成分预测被掩码的原子类型。这鼓励模型学习成分模式和上下文。

最终的预训练损失是这三个目标的加权和。

主要贡献

1. 架构分解：一种三流设计，显式分离成分、结构和相互作用，确保这两个因素在设计上被保留而非纠缠。
2. 混合预训练策略：一种自监督方法，结合基于重建的目标（去噪、掩码）与潜在预测学习（LeJEPA），以构建嵌入空间，从而更好地服务于下游迁移。
3. 可解释的检索：能够在成分、结构或联合嵌入空间中执行有针对性的相似性搜索，使基于特定标准（例如仅化学或仅结构）的材料比较成为可能。
4. 实证验证：在多种架构（Orb-v3、eSEN、MACE）和基准测试（OMat24、MatBench、QM9）上进行了广泛的实验，证明了数据效率、迁移性能和表征质量的提升。

结果

• 迁移性能（OMat24）：在数据受限的情况下，TriForces 显著优于基线。在 20K 个样本时，与基线模型相比，它将能量平均绝对误差（MAE）降低了 57%。它在所有样本量下都提高了力的 MAE，并降低了应力误差。
• 数据效率：TriForces 在每个数据集规模（20K 到 2M 个样本）下都实现了更低的误差，其中在低数据设置中观察到的增益最为显著。
• 基准测试性能：
  • MatBench：TriForces 变体在 8 个任务中的 6 个上取得了最佳总体结果，优于自监督预训练和 DFT 标记预训练的基线。例如，声子 MAE 从 57.8 提升至 19.5 cm$^{-1}$。
  • MatBench Discovery：TriForces eSEN-sm 实现了与更大的 eSEN-30M-OAM 模型相当的能量 MAE，同时使用了少 60% 的参数，并且训练速度快达 5 倍。
  • QM9：在多样化的化学输入（块体 + 分子）上进行预训练，与仅块体或无 SSL 基线相比，一致降低了 MAE。
• 表征质量：在冻结嵌入上进行线性探测表明，TriForces 保留了标准 MLIPs 所丢失的基本信息（晶体系统、主要元素、配位数）。TriForces 在晶体系统和主要元素分类上达到了 96–100% 的准确率，而基线模型则表现挣扎（55–73%）。
• 检索：该框架实现了有效的 k-NN 检索，其中成分流在元素集召回方面表现出色，结构流在空间群召回方面表现出色，这是单流模型所不具备的能力。

意义与主张

该论文将 TriForces 定位为不仅仅是一种自监督方法，而是一个架构框架，其表征通过 SSL 得到进一步增强。

• 机制依赖性：作者声称，流分解在大规模监督设置中提供了主导性的增益，而 SSL 在低数据迁移、表征组织和检索任务中最为有价值。
• 解耦：通过分离成分和结构，TriForces 解决了当前基础模型“脆弱的迁移”和“难以重用”的问题。它允许模型学习不仅用于预测，而且用于分析（例如探测、检索）的有序表征。
• 实用性：该框架与模型无关且可即插即用，使其能够立即应用于现有或新的原子架构。作者发布了预训练检查点和代码，以促进在下游原子建模中的重用。

这项工作表明，未来的原子基础模型应超越单流预测架构，转向显式保留化学系统不同物理因素的分解表征。