模型无关图提示学习用于晶体性质预测

# 模型无关的图提示学习用于晶体属性预测 来源: https://arxiv.org/html/2607.08996

Shrimon Mukherjee (https://arxiv.org/html/2607.08996v1/mailto:[email protected])  
数学与计算科学学院，印度科学培育协会，印度加尔各答  
*同等贡献*

Kishalay Das (https://arxiv.org/html/2607.08996v1/mailto:[email protected])  
计算机科学与工程系，印度理工学院卡哈拉格普尔，印度  
*同等贡献*

Partha Basuchowdhuri  
数学与计算科学学院，印度科学培育协会，印度加尔各答

Pawan Goyal, Niloy Ganguly  
计算机科学与工程系，印度理工学院卡哈拉格普尔，印度

###### 摘要

图神经网络已成为快速准确预测各种晶体属性的强大工具。这些模型通常将领域特定知识编码到其图编码模块中，这增加了参数量，并使其性能严重依赖于领域专业知识。此外，将所有可能影响特定晶体属性的化学和结构特征显式地纳入GNN编码器是一项具有挑战性的任务。在这项工作中，我们提出了一种软提示学习框架，用于捕获对属性预测至关重要的潜在特征，而这些特征并未显式提供给GNN。我们引入了一种新颖的多级图提示学习框架，包含节点级和图级软提示。在节点级，我们捕获不同原子类型的局部化学语义；在图级，我们编码晶体图的全局结构对称性。我们提出的提示学习框架轻量且无缝集成到任何现有的GNN编码器中。在流行基准数据集上的大量实验表明，加入提示学习显著提升了最先进GNN模型在晶体属性预测任务中的性能（提升3% - 15%）。此外，学习到的软提示支持跨属性知识迁移，增强了训练数据有限属性的预测性能。代码可在 https://github.com/shrimonmuke0202/Prompt.git 获取。

## 1 引言

预测晶体属性仍是材料科学中的一个重大挑战。与通常表示为规则图的分子不同，晶体材料被建模为原子在有序重复结构中的三维点云。准确高效地预测晶体属性有助于识别具有所需特性的新材料，减少了大量物理合成或测试的需求，显著加速了发现过程。

一种自然的晶体建模方法是将它们表示为图。因此，大量研究专注于开发用于晶体属性预测的图神经网络（GNN）模型。早期方法 xie2018crystal, chen2019graph, louis2020graph, Wolverton2020, schmidt2021crystal 从三维晶体结构构建多边图，并应用GNN模型编码原子周围的邻域结构信息。随后，许多研究提出了各种GNN架构变体，将特定领域知识集成到编码器中以增强晶体表示学习。ALIGNN choudhary2021atomistic 结合边间的键角信息以编码多体相互作用，Matformer yan2022periodic 被设计为对周期性不变以显式捕获重复模式，PotNet lin2023efficient 基于物理原理直接建模原子间势能。

这些模型的一个统一原则是假设晶体属性（如形成能、带隙、总能等）与构成原子的化学语义以及3D周期性结构固有地相关。因此，预测这些属性需要同时捕获节点级的局部化学语义和图级的全局3D结构对称性。多年来，最先进的GNN模型通过采用复杂架构逐渐改进了这些方面的编码，从而实现了更准确的属性预测，但同时也增加了参数数量，导致更长的训练和推理时间以及更高的计算开销。

然而，将所有可能影响特定晶体属性的化学和结构特征显式地纳入GNN模型是一项重大挑战。这个过程需要专门的领域专业知识，并且可能涉及尚未完全理解或发现的特征。因此，这些模型中使用的特征尽管可能很广泛，但不能被认为是穷尽的，这固有的限制了它们的有效性。一个潜在的解决方案是将可学习的潜在特征——在局部节点级和全局图级——纳入建模过程。这些学习到的潜在特征将解释这些缺失的属性，并进一步增强简单模型（如CGCNN）和高级模型（如ALIGNN、Matformer和PotNet）的表达能力。

在这项工作中，我们提出将**软提示或基于嵌入的提示学习**作为一种学习潜在特征的新方法。基于像素级视觉提示（VP）技术 bahng2022exploring, wu2023approximated, wu2022unleashing 以及在NLP任务中使用可学习软提示 lester2021power, liu2021p 的成功，我们引入了一种新的用于晶体属性预测的多级图提示学习框架。该框架结合了节点级（局部）和图级（全局）软提示。在节点级，我们学习一组 \(k\) 个独立向量，每个向量可能代表一个（潜在）化学特征。这些特征对晶体图中节点的影响通过从这 \(k\) 个向量及其线性投影计算出的注意力分数来确定。在图级，我们设计提示模板以捕获关键结构特征，如晶体几何和取向。我们提出的提示学习方法高效，仅为Matformer增加了约 *0.32%* 的额外参数，并且与任何现有最先进的GNN编码器兼容。

我们在两个广泛使用的基准数据集上进行了大量实验，以评估我们提出的提示学习框架相对于六种最先进GNN模型在晶体属性预测任务中的有效性。我们表明，我们提出的提示学习框架为现有基于GNN的架构在多种属性（跨越这些基准材料数据集）上带来了显著的性能提升。具体来说，我们观察到加入这种轻量软提示使简单模型（如CGCNN）的性能提升了15.11%，使其与高级模型（如原始ALIGNN，参数更多）具有竞争力。此外，即使对于像ALIGNN、Matformer和PotNet这样的复杂模型，我们也观察到使用这些轻量软提示后性能分别提升了5.52%、6.64%和3.15%。这些结果凸显了我们提出框架的重要性和创新性，它不仅提升了简单模型的性能，也提升了先进复杂模型的性能。此外，我们将提示学习与几种替代架构修改策略进行了比较，发现提示学习既更有效也更轻量。我们进一步观察到学习到的软提示具有跨属性知识迁移能力，增强了稀疏/有限训练数据属性的性能。此外，消融研究验证了节点级和图级提示学习在有效捕获晶体材料的化学和结构语义方面的重要性。我们在补充材料中提供了代码库。

参见图注

图1: (a) 一个周期性晶体结构，展示原子在重复排列中的点云。(b) 单个晶胞的放大视图，清晰地描绘了晶格向量 \(L=[l_1, l_2, l_3]\)，突出了晶体几何排列的基本构建块。(c) 晶胞的多图表示。

## 2 预备知识

### 2.1 晶体表示

晶体材料被视为具有有序重复模式的三维原子点云。它可以被建模为最小的**晶胞**，包含所有组成原子在特定坐标中，并在三维空间中沿着规则晶格无限重复，从而产生周期性结构，如图1(a) 和 1(b) 所示。给定一个晶胞中有 \(N\) 个原子的材料，我们可以用两个矩阵来描述晶胞：**原子类型矩阵 \(A\)** 和**坐标矩阵 \(C\)**。

原子类型矩阵 \(A = [a_1, a_2, ..., a_N]^T \in \mathbb{R}^{N \times k}\) 表示原子类型的独热表示集合，其中 k 是最大可能的原子类型数量。坐标矩阵 \(C = [c_1, c_2, ..., c_N]^T \in \mathbb{R}^{N \times 3}\) 表示原子坐标位置，其中 \(c_i \in \mathbb{R}^3\) 对应晶胞中第 i 个原子的坐标。

此外，还有一个额外的**晶格矩阵** \(L = [l_1, l_2, l_3]^T \in \mathbb{R}^{3 \times 3}\)，它描述了晶胞如何在三维空间中沿着 \(l_1, l_2\) 和 \(l_3\) 方向重复以形成材料的周期性三维结构。形式上，给定材料可以定义为 \(\mathbf{M} = (A, C, L)\)，我们可以将其无限周期性结构表示为：
\(\hat{C} = \{\hat{c}_i | \hat{c}_i = c_i + \sum_{j=1}^3 k_j l_j\}\) 和 \(\hat{A} = \{\hat{a}_i | \hat{a}_i = a_i\}\)，其中 \(k_1, k_2, k_3, i \in Z, 1 \leq i \leq N\)。

参见图注

图2: 我们通用提示学习方法的整体框架。对于节点级提示，我们使用一个根据化学语义区分不同原子类型的模板。在图级，我们的提示模板捕获图表示中固有的晶体结构对称性。

## 3 方法论

给定一个用于属性预测的带标签晶体数据集 \(\mathcal{D}_t = \{(\mathcal{G}_1, y_1), (\mathcal{G}_2, y_2), \dots, (\mathcal{G}_m, y_m)\}\)，其中晶体结构表示为多图 \(\mathcal{G}_i = (\mathcal{V}_i, \mathcal{E}_i, \mathcal{X}_i, \mathcal{F}_i)\)，属性预测任务的目标是基于晶体图的图表示 \(h_{\mathcal{G}_i}\) 预测 \(y_i \in \mathbb{R}\)。形式上，这涉及学习一个映射函数：\(f: \mathcal{G}_i \rightarrow y_i\)。建模上述函数 \(f\) 的传统方法是使用图神经网络（GNN）架构，因为它们具有出色的编码图结构数据的能力。为了进一步提高预测性能，我们将图提示学习集成到GNN模型中，以学习针对特定属性定制的潜在化学和结构语义。接下来在本节中，我们首先简要介绍晶体属性预测器的背景，然后详细解释所提出的提示学习框架（包括节点级和图级提示），最后是学习过程。

### 3.1 晶体属性预测器

传统的基于GNN的晶体属性预测器有两个构建模块：(a) 一个由多个GNN层组成的图嵌入模块 \(f_\theta\)，以及 (b) 一个预测头 \(z_\phi\)，其中 \(\theta\) 和 \(\phi\) 分别是各自模块的参数。图嵌入模块将多边图结构以及节点和边特征作为输入，并生成一组节点嵌入 \(\mathcal{H}\) 作为输出。

\[\mathcal{H} = \{h_1, ..., h_{|\mathcal{V}|}\} = f_\theta(\mathcal{A}, \mathcal{X}, \mathcal{F}) \quad (1)\]

其中 \(\mathcal{A}\) 是晶体图的邻接矩阵。进一步，执行 READOUT 操作（对称聚合，例如平均、求和或图池化）以获得整个晶体图 \(\mathcal{G}\) 的表示 \(h_\mathcal{G}\)。然后将该图表示输入到预测头 \(z_\phi\)，该预测头预测属性值 \(\hat{y}\)。

\[
\begin{split}
\hat{y} = z_\phi(h_\mathcal{G}); \quad h_\mathcal{G} = READOUT \biggl\{ h_u^{(L)} | u \in \mathcal{V} \biggr\}
\end{split}
\quad (2)
\]

整个晶体属性预测器架构，包括GNN层和预测头，以端到端的方式训练，优化其参数以完成下游属性预测任务。

### 3.2 晶体提示学习

基于GNN的晶体属性预测器的有效性在很大程度上取决于其将组成原子的化学语义和晶体结构拓扑编码到低维嵌入空间的能力，从而生成图表示 \(h_\mathcal{G}\)。然而，将影响特定晶体属性的所有化学和结构语义显式地提供给GNN模型本质上具有挑战性。因此，某些影响特定晶体属性的潜在特征，无论是在局部节点级还是全局图级，通常在现有模型中未被捕获。为了捕获这些潜在特征的影响，我们提出了一种 **多级图提示学习** 方法（图2），在节点级和图级引入可学习的软提示，以捕获用于特定晶体属性预测的基本化学和结构潜在特征。

最近，软提示已被有效地用于将额外的元信息集成到各种计算机视觉和NLP任务的输入数据中。在计算机视觉中，像素级视觉提示（VP）技术 bahng2022exploring, wu2022unleashing, wu2023approximated 通过添加可学习的软视觉提示来修改输入图像，这些提示作为视觉线索引导机器学习模型的决策。类似地，在NLP中，研究人员 liu2021p, lester2021power 展示了通过向句子中引入可学习的软提示（例如向输入标记添加位置编码以传达位置信息）来合并额外元信息的巨大潜力。基于这些最近的成功，在我们提出的方法中，我们通过将可学习的软提示引入晶体图的特征空间（输入节点特征）和潜在空间（图嵌入）中，将输入图 \(\mathcal{G}\) 转换为提示图 \(\psi(\mathcal{G})\)。我们将变换后的图 \(\psi(\mathcal{G)}\) 输入GNN模型，并以端到端监督的方式将可学习的软提示与GNN模型参数一起训练。在接下来的部分中