全球关注何时能发挥作用：一项关于原子级图学习的统一实证研究

《Journal of Cheminformatics》：When does global attention help: a unified empirical study on atomistic graph learning

　　摘要：图神经网络（GNNs）被广泛用作昂贵实验和第一性原理模拟的替代品，用于在原子尺度上研究化合物的行为，其架构复杂性不断提高，以实现对复杂物理现象的建模。虽然最近的GNNs结合了更传统的消息传递神经网络（MPNNs）层来模拟短程相互作用，以及具有全局注意力机制的更先进的图变换器（GTs）来模拟长程相互作用，但由于实现不一致、特征差异或超参数调整，目前尚不清楚全局注意力机制何时能真正优于调整良好的MPNN层。我们引入了第一个统一、可复制的基准测试框架——基于HydraGNN——该框架能够无缝切换四种受控模型类别：MPNN、带有化学/拓扑编码器的MPNN、带有全局注意力的GPS风格混合模型，以及完全融合的局部-全局模型与编码器。使用七个不同的开源数据集进行回归和分类任务的基准测试，我们系统地分离了消息传递、全局注意力和基于编码器的特征增强的贡献。我们的研究表明，带有编码器的MPNN形成了一个稳健的基线，而融合的局部-全局模型在受长程相互作用效应影响的性质上提供了最明显的好处。我们进一步量化了注意力的准确性与计算之间的权衡，并报告了其在内存中的开销。这些结果共同建立了对原子图学习中全局注意力的第一个受控评估，并为未来的模型开发提供了一个可复制的测试平台。

引言：从原子结构预测原子和材料性质是计算化学、催化和材料发现中的一个核心挑战[1, 2]。传统上，性质估计依赖于（i）经验和半经验相关性——群加性和Hammett-Taft取代基分析，随后使用手工制作的描述符进行QSAR/QSPR[3,4,5,6]；（ii）基于物理的电子结构方法，如Hartree-Fock（HF）密度泛函理论（DFT）[107]，后HF方法如二阶M?ller-Plesset（MP2）微扰理论和耦合簇单双（三重）CCSD(T)[7,8,9,10,11]；（iii）使用经验或从头算势的分子动力学[12]；以及（iv）直接实验和高通量筛选[13]。尽管这些方法是基础性的，但它们存在局限性：QSAR和群加性依赖于适用范围狭窄的固定描述符，常常忽略非局部、构象或环境效应；DFT/后HF方法提供系统性的准确性，但扩展性极差（大约\(\mathcal {O}(N^{3{-}4})\)对于KohnSham DFT和\(\mathcal {O}(N^7)\)对于CCSD(T)），并且在没有特殊校正的情况下难以处理色散或强相关性[7, 8, 10]；MD的准确性取决于力场的保真度和长轨迹以采样罕见事件[106, 12]；而实验成本高昂、通量低且特定于条件[13]。相比之下，基于现代机器学习（ML）的方法将昂贵的计算/测量成本分摊到训练中，然后从2D/3D原子表示中快速、可扩展地推断出复杂的非局部结构-性质关系[1, 2, 14,15,16]。基于图的学习提供了一个自然的范式：分子可以被建模为一个图，其节点（原子）和边（键或基于距离阈值的连接）携带特定领域的属性，可以选择性地用3D坐标在几何图中增强[17,18,19,20,21]。这种结构支持归纳偏见——排列不变性和在需要时的欧几里得等变性——这对于建模标量、向量和张量目标至关重要[22,23,24,25,26]。图神经网络（GNNs）[27,28,29]包括消息传递神经网络（MPNNs）[18]及其几何扩展（例如，SchNet、DimeNet、EGNN）[9, 21, 24]，可以利用连接结构来预测原子性质。尽管取得了成功，但当目标性质依赖于长程相互作用（LRI）——如静电学、诱导和色散——这些作用超出了小k-hop邻域或键的局部连接范围时，MPNNs可能会遇到困难[30,31,32,33]。特别是，由于LRI引起的非局部效应发生在有影响力的原子相距较远的情况下，并概括了它们的物理缩放（例如，\(-C_6/R^6\)色散）[32, 33]。深化MPNNs以扩大接收场会引入两个众所周知的病理现象：过度平滑和过度压缩[34,35,36,37]。这些局限性表明需要一种能够在不牺牲局部结构细节的情况下全局传播信息的机制。变换器风格的全局注意力提供了这样的机制[38]。图变换器（GTs）用可学习的全局信息流替换了固定半径的聚合，通过多头自注意力在单层中潜在地捕获局部和全局依赖性[39,40,41]。然而，天真地在N个节点上应用密集注意力会导致每层\(\mathcal {O}(N^2)\)的时间和内存开销，这对于大型原子图来说可能是禁止性的[38]。已经提出了混合设计（例如，GPS[42]），将强大的局部消息传递能力与全局注意力机制结合起来，以稀疏图的连接性来控制计算成本。然而，这些方法仅在少数数据集上进行了评估，并且仅限于作为局部信息聚合器的MPNNs类别。在这项工作中，我们提出了一个基于HydraGNN[43]的统一、可复制的框架——一个高度并行的图多任务学习管道——用于系统评估基于变换器的图模型和一大类SOTA MPNNs，包括等变MPNNs。该框架实例化了四种受控配置：（i）MPNN，（ii）带有化学/拓扑编码器的MPNN，（iii）带有全局注意力的GPS风格MPNN，以及（iv）带有特定领域编码器的融合局部-全局模型。在这里，“受控”并不意味着在所有架构中固定超参数，因为这会导致结果偏向次优配置。相反，控制是在（i）相同的数据集和分割，（ii）相同的训练管道和预算，（iii）基于架构约束的相同HPO搜索空间，以及（iv）基于验证性能的相同模型选择标准层面实施的。这确保了性能差异反映了架构能力而非调整 artifacts。紧凑的编码器套件提供了化学和结构上丰富的输入：每个原子的物理化学描述符、节点/边结构特征（例如，度数、中心性、聚类、k-核、边指数）和拉普拉斯位置编码；边还携带键信息，如果有的话，还包括距离信息。轻量级线性嵌入整合了这些通道，训练使用分布式数据并行性（DDP）和自动超参数优化来实现公平的消融。在涵盖回归和分类的十二个基准测试中，我们评估了（1）全局注意力何时能带来超出调整良好MPNNs的收益，（2）领域/拓扑编码器的贡献，以及（3）与注意力二次成本相关的准确性与计算之间的权衡。实证表明，编码器系统地提高了性能，而单独的注意力并不总是更优越。融合的局部-全局模型在具有显著非局部或几何效应的任务上提供了最大的好处，而带有编码器增强的MPNNs在其他任务上仍然具有竞争力。结果是一个标准化的实验设置，并为未来方法的发展奠定了基础。

相关工作：之前已经直接将注意力机制整合到消息传递神经网络中，以调节相邻消息的聚合。值得注意的例子包括图注意力网络（GAT）Veli?kovi?等人[44]，它用学习到的边上的注意力权重替换了均匀的邻域聚合，使得在消息传递过程中能够自适应地加权邻居。后续工作通过将注意力整合到更通用的MPNN公式中扩展了这一想法，包括基于边的条件或方向性注意力，以及与学习到的边表示相结合的注意力机制Hussain等人[45]；Brody等人[46]。这些方法主要使用局部注意力，其中注意力权重是在直接邻居上计算的，并在消息传递操作符本身内集成。分子的GTs最近取得了几项进展。Graphormer注入了结构偏见（最短路径、中心性、边编码）并取得了强大的OGB结果，但依赖于密集的二次注意力以及特定于任务的启发式方法，而没有与同样调整好的MPNNs进行受控比较[47]。MAT通过RDKit特征和成对距离偏见增强了注意力，但同样继承了二次成本和繁重的特征工程[48]。等变变换器（例如，EquiformerV2）提高了原子能量的3D保真度，但计算要求高，并且针对连续的3D标签而不是这里考虑的更广泛的图级和节点/边目标[49]。大规模框架如Uni-Mol[50]利用大规模的SE(3)-等变预训练来获得强大的下游准确性，但将架构优势与数据规模混淆在一起，在典型的计算预算下难以复制[51]。混合局部-全局设计提供了一种将消息传递与全局注意力结合的通用方法，尽管评估强调了非原子级的基准测试，并没有隔离化学感知的编码器[52]；后续工作（GPS++）在PCQM4Mv2上表明，调整良好的消息传递可以在几乎没有或没有注意力的情况下保留大部分好处，但只关注一个大型数据集[53]。最后，长程基准测试（LRGB）激发了全局接收场的概念，但不是特定于原子的，并且省略了常见的化学设置，如节点级电荷或多标签评估[54]。因此，该领域缺乏在一个共同框架下比较这些架构选择的受控、可复制的研究。

多任务学习在分子机器学习中被广泛探索，作为一种提高数据效率、泛化和跨相关化学性质表示共享的手段Ramsundar等人[55]；Martin等人[56]。早期研究表明，使用共享的基于图的表示共同学习多个分子端点可以胜过单任务模型，特别是当任务相关或个别数据集较小时Capela等人[57]。后续工作将多任务学习扩展到图神经网络和消息传递架构Li等人[58]，Chen等人[59]，表明共享的编码器结合特定于任务的预测头可以有效地捕获共同的化学结构-性质关系，同时保持任务特定的表达能力。然而，之前的研究在模型架构、训练协议和数据集选择上往往存在显著差异，使得难以区分多任务学习的效果与其他混淆因素。

贡献：与最先进的方法不同，我们的工作并没有为MPNNs或基于图的变换器架构提出新的注意力机制。相反，我们在统一的HydraGNN框架下提供了对多种注意力整合策略的受控和系统评估，这些策略包括局部消息传递、基于编码器的特征增强和全局注意力——在相同的训练和超参数优化协议下。此外，我们将HydraGNN的统一框架用于全局注意力与多任务学习策略相结合，并评估了这两种技术的整合。与之前在狭窄基准测试上测试孤立架构的研究不同，我们在单一实验管道下系统地评估了四种模型类别，从而能够清楚地将准确性增益归因于架构、编码器或它们的组合。为此，我们做出了以下关键贡献：开发了一个统一的框架来结合MPNN和GT。我们扩展了现有的基于HydraGNN的管道，提供了MPNN架构与GT架构，并提供了无缝的模型路径选择，即仅MPNN、MPNN + 编码器、GPS风格以及完全融合的（MPNN + GPS + 编码器）配置，并使用了DDP和自动化的HPO（图3）。将领域/拓扑编码器纳入图嵌入。化学信息丰富的每个原子的描述符、节点/边结构编码和拉普拉斯位置编码，通过轻量级嵌入集成（章节Encoders-Embeddings）。通过消融研究进行了全面评估。在7个数据集（表1）上进行了跨领域研究，包括回归、多类和多标签任务；定量比较（章节ZINC–OGB-molPCBA）。论文组织：第3节回顾了图、几何学习和不变性/等变性的初步知识，并通过LRI和MPNN的局限性激发了全局注意力的需求，然后介绍了GPS风格的混合模型。第4节描述了所提出的框架，包括编码器和嵌入模块。第5节报告了实验结果，第6节讨论了影响，第7节总结了工作。

背景和问题设置：我们首先介绍了一些将在本工作中广泛使用的基本定义和概念。\(\mathcal {X}\)表示集合。\([\textbf{X}]_{uv}\)和\([\textbf{x}]_{i}\)表示多维矩阵\(\textbf{X}\)的条目和向量\(\textbf{x}\)。通用下标: 表示整个维度，例如，矩阵\(\textbf{X}\)的第i行表示为\([\textbf{X}]_{i :}\)。标量由x或X表示。

原子图：图是一种数据结构，用于表示实体（节点）及其关系（边）。正式地，\(\mathcal {G} = \{\mathcal {V},\mathcal {E}\}\) 表示一个无向图，其中 \(\mathcal {V}\) 代表节点集，\(\mathcal {E}\) 代表边集。\(N = |\mathcal {V}|\) 是图中节点的数量。每个节点 \(v \in \mathcal {V}\) 都支持一个 \(p-\) 维的特征向量，表示为 \(\textbf{x}_v \in \mathbb {R}^p\)，使得 \([\textbf{X}]_{v:} = \textbf{x}_v\)，而 \(\textbf{e}_{uv} \in \mathbb {R}^f\) 代表每条边 \((u,v) \in \mathcal {E}\) 所支持的边属性，其中 \([\textbf{E}]_{uv:} = \textbf{e}_{uv}\)。此外，如果每个节点 \(v\) 代表 3D 欧几里得空间中的一个点，并且它们的连接具有相关的距离概念，那么图 \(\mathcal {G}\) 可以被定义为一个几何图。图 1。这张图片的替代文本可能是使用 AI 生成的。全尺寸图片。具有截断半径 \(r_c\) 的几何原子图。在这种情况下，增强后的图表示为 \(\mathcal {G} = \{\mathcal {V},\mathcal {E},\textbf{X},\textbf{E},\textbf{R}\}\)，其中 \(\textbf{R} \in \mathbb {R}^{N\times 3}\) 存储节点的 3D 坐标。通常，\(\mathcal {E}\) 是通过使用特定于给定下游应用的半径截断 (\(r_c\)) 对相邻节点之间的距离进行阈值处理来构建的。特别是，$$\begin{aligned} \mathcal {E} = \{ (u,v) \mid u \ne v, \ \Vert \textbf{r}_u - \textbf{r}_v \Vert < r_c \} \end{aligned}$$ 其中 \([\textbf{R}]_{u:} = \textbf{r}_u\)。例如，对于如图 1 所示的原子图，半径截断被仔细选择以捕捉化学和物理属性，包括特定的键长（12 ?）、非键合相互作用（如范德华力（35 ?）和金属-配体相互作用（56 ?）。利用图的连接结构来生成有用的数学抽象是发展图机器学习理论，特别是 GNN 的根本动机。长程相互作用（LRI）效应在分子中源于基本物理力——静电（多极相互作用）、诱导（极化）和色散（伦敦力）——这些力仅随原子间距离的倒数衰减，因此远远超出了局部模型中使用的典型邻居截断范围 [31,32,33]。为了形式化 LRI，我们考虑了原子 u 和 v 之间的两种互补的距离概念：原子键图上的最短路径（图）距离 \(d_G(u,v)\) 和 3D 空间中的欧几里得距离 \(r_{uv} = \Vert \textbf{r}_u - \textbf{r}_v \Vert _2\) [60]。如果存在至少一个原子 v，其对目标属性的贡献无法通过小 k 跳邻域内的信息来捕捉，则认为该相互作用是长程的：$$\begin{aligned} \exists v:\ d_G(u,v) \gg k \quad \text {或}\quad r_{uv} \gg r_c \quad \text {但}\quad \text {influence}(v\!\rightarrow \!u) \not \approx 0. \end{aligned}$$ 这两种距离捕捉了非局部性的不同表现：在折叠或空间紧凑的构象中，即使 \(d_G\) 很大，\(d_L\) 也可能很小，从而产生通过空间的非键合接触；在扩展或刚性的系统中，\(d_G\) 可能很小而 \(d_L\) 很大。相关概念出现在蛋白质接触顺序中，它衡量了尽管序列（图）分离较大但空间上的接近性 [61]。从物理上讲，这样的原子-原子关系可以归类为包含 u 和 v 的两个分子片段 A 和 B 之间的相互作用。在较大的分离距离下，A 和 B 之间的相互作用能量可以分解为多极-极化-色散：$$\begin{aligned} E_{\text {int}}(R) \;\approx \; E_{\text {elst}}(R) \;+\; E_{\text {ind}}(R) \;+\; E_{\text {disp}}(R) \;+\; \cdots , \end{aligned}$$ 其中 R 是片段中心的欧几里得分离距离（或代表性的原子间距离）。静电由多极展开描述，其中单极的势能随 \(R^{-1}\) 衰减，偶极的势能随 \(R^{-2}\) 衰减，一般来说，多极的势能随 \(R^{-(\ell +1)}\) 衰减，相应的场衰减得更快（例如，偶极场随 \(R^{-3}\) 衰减）[33]。当一个片段的场使另一个片段极化时，就会产生诱导作用；对于各向同性的片段，方向平均的诱导能量随 \(R^{-6}\) 规模变化 [31, 33]。由于电子波动的相关性，色散具有主要的吸引项 \(-C_6/R^6\) 以及更高阶的修正项 \(-C_8/R^8\), \(-C_{10}/R^{10}\) [32, 33]，在延迟的 Casimir–Polder 情况下，其规模为 \(R^{-7}\) [30]。相比之下，短程交换（Pauli）排斥力呈指数衰减，并在 SAPT 和基于重叠的模型中有所描述 [62, 63]。MPNN 架构 MPNN 架构 [27,28,29] 是一种深度学习架构，它通过利用图结构数据中固有的不规则连接模式对节点特征执行类似卷积的操作。一个 \(K-\) 层的 MPNN 表示为，$$\begin{aligned} \bar{\textbf{H}}^{(k)}&= \bigoplus (\phi (\textbf{H}^{(k-1)}, \omega _k), \mathcal {E}, \textbf{E}) \,\,\forall k \in \{1,\dots ,K\} \end{aligned}$$ （1） $$\begin{aligned} \textbf{H}^{(k)}&= \Psi (\bar{\textbf{H}}^{(k)},\textbf{H}^{(k-1)}) \end{aligned}$$ （2） 其中 \(\textbf{H}^{(k)} \in \mathbb {R}^{N\times d_{k}}\) 是经过 k 个 MPNN 层后的节点表示，\(\textbf{H}^{(0)} = \textbf{X}\)。\(\phi (\cdot ,\omega )\) 表示节点特征的线性变换，可训练的权重 \(\omega _k\) 由给定层 k 中的所有节点共享。图 2。详细展示了通用的 MPNN 层 k。聚合器 \(\bigoplus\) 通过结合来自相应局部邻域 \(\mathcal {N}_u = \{v \mid v\in \mathcal {V}, v \ne u, (u,v) \in \mathcal {E}\}\) 的信息来生成节点表示，同时利用可用的边属性 \(\textbf{E}\)。图 2 中显示了一个任意的 MPNN 层。根据 \(\bigoplus\) 的具体功能形式，已经提出了广泛的 MPNN 模型 [18, 27,28,29, 44, 64,65,66]。 \(\Psi\) 是一个节点更新函数，例如多层感知器（MLP）或门控循环单元（GRU）。MPNN 输出 \(\textbf{H}^{(K)}\) 存储用于节点级下游任务的完整节点表示集。对于图级任务，已经采用了多种池化操作（即最小值、最大值、求和、平均值）将 \(\textbf{H}^{(K)}\) 转换为密集的图级表示。在原子建模应用中，将几何信息（如原子坐标、原子间距离和角度）纳入消息传递过程至关重要。为此，像 SchNet [67]、DimeNet [9] 和 EGNN [68] 这样的化学或物理驱动的 MPNN 架构通过将连续的几何特征编码为边属性、应用基于距离的过滤器或在欧几里得变换下强制等变性来扩展标准 MPNN。这些受化学启发的 MPNN 架构通过以下方式构建几何嵌入：提取坐标为 \(\textbf{r}_u,\textbf{r}_v \in \mathbb {R}^3\) 的节点 u 和 v 之间的成对距离和三重键角，$$\begin{aligned} r_{uv} = \Vert \textbf{r}_u - \textbf{r}_v\Vert , \quad \theta _{uvw} = \angle (\textbf{r}_v-\textbf{r}_u,\ \textbf{r}_w-\textbf{r}_u), \end{aligned}$$ （3） 将这些几何量嵌入到潜在的边表示中：$$\begin{aligned} \bar{\textbf{e}}_{uv} = \phi _l(r_{uv}) \;+\; \sum _{w\in \mathcal {N}_u\setminus \{v\}} \phi _a(\theta _{uvw}), \end{aligned}$$ （4） 其中 \(\phi _r\) 和 \(\phi _a\) 分别是可学习的径向和角度嵌入函数。在任何 \(K-\) 层模型的第 k 层，从邻居 v 到节点 u 计算消息如下：$$\begin{aligned} \textbf{m}^{(k)}_{u\leftarrow v} = \xi \!\big (\textbf{h}^{(k-1)}_u, \textbf{h}^{(k-1)}_v,\ \bar{\textbf{e}}_{uv}\big ), \end{aligned}$$ （5） 其中 \(\xi\) 是一个可训练的消息函数，它同时考虑了节点特征和几何感知的边特征，以及 \([\textbf{H}]_v = \textbf{h}_v\)。使用传入的消息聚合节点特征：$$\begin{aligned} \textbf{h}^{(k)}_u = \Psi \!\Big (\textbf{h}^{(k-1)}_u,\ \bigoplus _{v\in \mathcal {N}_(u)} \textbf{m}^{(k)}_{u\leftarrow v}\Big ), \end{aligned}$$ （6） 对于图级预测任务，最终节点状态 \(\textbf{H}^{(K)}\) 通过池化层进行聚合。这些方法允许网络在尊重物理对称性的同时保持对空间关系的敏感性，从而提高其从数据中学习的能力，尤其是在几何直接控制目标属性的情况下。虽然 MPNN 在建模复杂的、相互连接的数据方面具有足够的表达能力，并且比非图 ML 模型提供更好的泛化能力（有例外；见第 3 节），但在大规模数据处理、模型可扩展性和不同基于图的 DL 架构之间的系统比较方面仍存在关键挑战。在 MPNN 架构中的不变性和等变性在几何 GNN 中，不变性意味着 GNN 的预测在输入的平移、旋转或反射下保持不变，而等变性意味着预测在对称群（如 E(n) 或 \(\textrm{SE}(3)\) 下以与输入相同的方式变换。对于作用在 \(\textbf{R}\) 上的刚性运动 \((R,\textbf{t})\in \textrm{SE}(3)\)，不变性意味着 \(f(R\textbf{R}+\textbf{t})=f(\textbf{R})\)，等变性意味着 \(g(R\textbf{R}+\textbf{t})=R\,g(\textbf{R})\)。通常，当目标是方向性的或张量性的（例如，原子力、偶极矩和矢量场）时，需要 E(n)/\(\textrm{SE}(3)\)–等变性——例如，EGNN、SE(3)-Transformer、PaiNN 和 NequIP 等架构被设计来尊重这些对称性 [23, 68,69,70]。对于不变的标量属性（例如，总能量、带隙、类别标签），强制不变性和注入不变的几何偏差——成对距离和角度/三重特征——通常就足够了；SchNet 使用连续的过滤器处理原子间距离，而 DimeNet/DimeNet++ 通过球形贝塞尔和球形谐波基来增强角度 [67, 71, 72]。请注意，即使不变的能量模型也可以通过对原子核坐标的预测能量取解析梯度来产生等变的力，尽管显式等变的模型通常可以提高精度和样本效率，适用于更一般的矢量/张量目标属性 [67,68,69]。在 MPNN 中捕获 LRI 的直接技术是深化 DL 架构，以确保有效的感受场覆盖正在处理的节点周围的整个图。然而，这种方法存在一些主要缺点：过度平滑。堆叠许多邻域平均层会反复对图应用（标准化的）平滑操作。随着深度的增加，来自同一连接组件的节点特征可能会变得几乎无法区分，从而损害模型的区分能力。本质上，聚合操作的低通特性随着模型深度的增加过滤掉了大部分高频率的区分信息。这种现象——过度平滑——已经通过拉普拉斯平滑解释和节点嵌入的渐近崩溃在实证和理论上进行了分析 [73, 74]。过度挤压（瓶颈和曲率）。即使 K 足够大以覆盖长图距离，来自指数级多个 k 跳节点的信息也可能被压缩成必须穿过狭窄图切割（小边分隔符）的固定宽度消息。这种过度挤压瓶颈限制了 MPNN 中的长程依赖性建模 [75]。最近的理论将过度挤压与离散曲率联系起来：负曲率（树状）区域会指数级扩展邻域，同时允许小切割，加剧了信息拥堵。因此，局部消息传递通常不足以捕获 LRI，特别是对于需要更大感受场来覆盖每个节点周围整个图的大型原子图。近年来，基于图的 DL 成功地超越了 MPNN 模型——后者本质上是局部信息聚合器，将图视为局部部分的组合——并利用了具有全局范围并且可以整体处理图的 DL 模型。具体来说，GT [39, 40, 76] 是一种使用自注意力机制 [38] 的 DL 架构。请注意，“图变换器”一词也可以指特定的全局注意力架构 [39]。然而，在这项工作中，我们将其用作一般基于注意力的模型类的描述符。在这个框架中，每个节点通过学习的注意力机制从图中的所有其他节点聚合信息，从而能够学习超出固定跳邻域的局部和全局依赖性。多头注意力（MHA）机制允许模型共同关注来自不同表示子空间的信息，从而增强了表达能力和鲁棒性。在 \(M-\) 层的任意层 m，节点嵌入矩阵 \(\textbf{H}^{m-1} \in \mathbb {R}^{N \times d_{m-1}}\) 通过注意力头投影到查询、键和值空间：$$\begin{aligned} \textbf{Q}^m = \textbf{H}^{m-1} \textbf{W}_Q^m, \quad \textbf{K}^m = \textbf{H}^{m-1} \textbf{W}_K^m, \quad \textbf{V}^m = \textbf{H}^{m-1} \textbf{W}_V^m, \end{aligned}$$ （7） 其中 \(\textbf{W}_Q^m, \textbf{W}_K^m, \textbf{W}_V^m \in \mathbb {R}^{d_{m-1} \times d_{m}}\) 是可学习的权重矩阵。一个头的缩放点积注意力计算如下：$$\begin{aligned} \textrm{AttHead}(\textbf{Q}^m, \textbf{K}^m, \textbf{V}^m) = \textrm{softmax} \left( \frac{\textbf{Q}^m {\textbf{K}^m}^\top }{\sqrt{d_{m}}} \right) \textbf{V}^m. \end{aligned}$$ （8） 其中对每一行独立应用 softmax。对于 B 个头，输出被连接并线性投影：$$\begin{aligned} & \textrm{MultiHead}(\textbf{Q}^m, \textbf{K}^m, \textbf{V}^m) = \textrm{Concat}(\textrm{AttHead}_1, \nonumber \\ & \quad \dots , \textrm{AttHead}_B) \textbf{W}_O^m, \end{aligned}$$ （9） 其中 \(\textbf{W}_O^m \in \mathbb {R}^{Bd_{m} \times d_m}\) 是一个可学习的投影矩阵。通过整合结构和位置编码，GT 有效地结合了关系归纳偏差，使其适用于多种任务，如节点分类、链接预测和图级预测。我们将在第 3 节中更详细地讨论这一点。HydraGNN HydraGNN [43] 是一种多任务学习 MPNN 架构，旨在从原子结构输入同时预测全局（图级）和原子（节点级）材料属性。HydraGNN 高效地整合了 DDP，用于大规模原子数据集（大小可以从几十万到几百万不等），并使用广泛的 MPNN 层作为基础学习器，这些层可以分布在分布式计算资源上，并具有模型并行性。HydraGNN 使用共享的 MPNN 层提取所有目标属性共有的特征，然后通过 MP每个任务都与其自身的损失函数相关联，全局目标被定义为各个任务损失的总和。HydraGNN支持在SOTA消息传递、几何和等变GNN架构之间无缝切换，从而通过集成的HPO（超参数优化）实现高效的架构搜索。该框架在预测准确性和可扩展性方面的有效性已经通过在多个材料科学应用中的严格评估得到了充分证明（Baker等人[77]；Lupo-Pasini等人[43]；Choi等人[78]；Baker等人[79]；Lupo-Pasini等人[80]；Lupo-Pasini等人[81]；Lupo-Pasini等人[82]）。混合GNN架构（GPS）是一个通用、强大且可扩展的框架，它结合了MPNN和GT的优势。在每一层，GPS通过利用位置和结构编码，将局部邻域聚合（通过MPNN捕获边级和结构信息）与全局多头注意力（实现长距离依赖性）结合起来，以克服标准GNN的表达能力限制[42]。实际上，这种方法使得在架构的每一层中，除了局部消息之外，还可以聚合全局信息。因此，它避免了更深层GNN架构的必要性（以及相应的陷阱）。形式上，第\(\ell ^{th}\)层GPS表示为：$$\begin{aligned} \textbf{X}^{\ell +1},\,\textbf{E}^{\ell +1}&= \textrm{GPS}^{\ell }(\textbf{X}^{\ell },\textbf{E}^{\ell },\mathcal {E}) \end{aligned}$$ （10）其中，$$\begin{aligned} {\hat{\textbf{X}}}^{\ell +1}_{\!M},\,\textbf{E}^{\ell +1}&= \textrm{GNN}^{\ell }(\textbf{X}^{\ell },\textbf{E}^{\ell },\mathcal {E}), \end{aligned}$$ （11）$$\begin{aligned} {\hat{\textbf{X}}}^{\ell +1}_{\!T}&= \textrm{MHA}^{\ell }(\textbf{X}^{\ell }), \end{aligned}$$ （12）$$\begin{aligned} \textbf{X}^{\ell +1}&= \textrm{MLP}^{\ell }\!\big ({\hat{\textbf{X}}}^{\ell +1}_{\!M} + {\hat{\textbf{X}}}^{\ell +1}_{\!T}\big ). \end{aligned}$$ （13）基于这些能力，GPS++ [83]改进了局部处理和全局处理之间的平衡，用于原子属性预测。GPS++强调了一个强大且调整良好的消息传递模块，并通过结构和几何先验（如最短路径和3D距离嵌入）来补充有偏自注意力。通过广泛的消融实验，GPS++表明即使没有全局自注意力，其大部分性能仍然可以保持，这突显了在有3D位置数据缺失的情况下，表达力强的消息传递仍然具有很高的竞争力。尽管局部和全局信息结合使用了协同效应，但GPS及其增强版本仅在有限的科学材料类别[42, 83]中得到了成功应用，并且应用的数量也较少。此外，基线结果是使用有限的MPNN架构获得的，没有比较等变模型。另外，研究表明，通过适当的超参数调整，MPNN模型仍然可以胜过GT模型和GPS [84]。因此，构建一个端到端框架来处理来自不同应用领域的大量科学材料，使用广泛的MPNN模型与全局注意力相结合——并通过自动超参数优化（HPO）选择最佳模型——仍然是研究社区的一个未解决问题。

图3
此图像的替代文本可能是使用AI生成的。
全尺寸图像
流程图展示了我们提出的框架中的计算步骤。它描述了一个\(K-\)层模型。每个块并行处理单个图。四个独立的流水线由开关S1和S2的配置表示。例如，如果S1和S2都打开，则获得HydraGNN流水线。如果S1和S2都关闭，则将全局注意力模块与局部消息传递模块融合，同时编码器为两个模块提供特定领域和位置信息。输入特征在输入到可学习模块之前被嵌入到合适的子空间中。第K层的输出通过超参数优化用于下游任务。

我们扩展了现有的HydraGNN框架（该框架仅实现了MPNN机制，我们仍然将其作为本工作中结果的基线），以集成MPNN和全局注意力，从而允许用户采用升级后的HydraGNN框架，无缝切换不同的模型配置，这些配置提供了各种跨图数据聚合信息的技术。图3显示了新扩展的HydraGNN框架提供的不同模型配置选项，这些配置通过开关S1和S2进行切换。1. 通过同时打开S1和S2，模型类似于HydraGNN [43]架构，具有其核心功能，如用于大规模数据集的DDP和大量可选择的MPNN模型——包括不变型和等变型。2. 当S1打开而S2关闭时，我们获得了HydraGNN的增强版本，其中通过编码器模块为每个图计算额外的化学和拓扑属性作为预处理步骤，这些属性首先使用可学习的转换嵌入到用户定义的子空间中，然后由MPNN模块用于最终的下游任务。3. 进一步地，通过打开S2并关闭S1，模型类似于嵌入在HydraGNN流水线中的GPS [42]架构。适当嵌入的拉普拉斯特征向量被用作全局注意力模块的位置编码。这种方案本质上在一个端到端的设置中结合了两种模型的优势。4. 最后，通过同时关闭S1和S2，我们获得了一个完全集成的模型，该模型结合了编码器、嵌入、全局注意力和MPNN模块，以及用于大规模数据处理的DDP。我们注意到，这个模型在原理上类似于GPS++ [83]。然而，我们不使用graphormer风格的编码，并且我们使用比前述工作更广泛的MPNN类来聚合局部信息。

对于该框架的操作来说，编码器模块至关重要，它预先计算分子的某些基本化学属性（在组成原子和键的层面上）以及相应图形的拓扑属性。除了对全局注意力模块有用（该模块不使用原子图的基础局部连接结构）来区分特定化学类型的节点和边以及图中的拓扑结构和位置外，这些特征还允许MPNN模块利用否则将通过消息传递隐式计算的现有信息。

每个原子图都增加了组成原子的3D坐标（如果几何信息可用的话）。图的最终模型输入包括：（i）原始原子/键属性，（ii）基于化学信息的每个原子的描述符，（iii）节点和边上的结构和位置编码，以及（iv）简单的几何边属性。所有特征块都针对每个数据集分割进行了标准化（零均值，单位方差）。

对于提供丰富原子编码的数据集（例如，OGB [85] mol任务），我们采用标准的9维原子特征向量，包括原子序数、度数、形式电荷、连接的氢原子数、手性、杂化度、芳香性和环成员身份。

化学信息描述符（CE）。除了上述原始属性外，我们还使用以下门捷列夫原子属性计算一个15维的每个原子描述符向量：原子重量、族、周期、块（s/p/d/f）、价电子数、共价半径、范德华半径、保罗inging和艾伦电负性、第一电离能、熔点、沸点、密度和原子体积。这个描述符矩阵被标准化并作为一个单独的通道提供给编码器。

拓扑编码（TE）。我们添加了两种互补的、位置和结构感知的编码：（i）节点PE：一个包含经典中心性/结构度量的9个特征的堆栈——度数、接近度、介数、特征向量中心性、PageRank、局部聚类系数、k-核数、调和中心性和偏心率——在无向原子图上计算并进行了标准化。（ii）边PE：每个边包含4个特征的向量，包括边介数、Jaccard系数、AdamicAdar分数和优先连接性，都进行了标准化。此外，我们还包括拉普拉斯特征向量位置编码（LPE）：（对称化）图的top k个特征向量，我们每次为图计算一次（超参数k在配置中设置）并进行了标准化。

根据所使用的特定数据集，键和几何特征进行了定制。例如，OGB [85]提供键类型、立体化学、共轭信息，而ZINC [31]和QM9 [86, 87]提供离散键属性。对于3D结构，当坐标可用时，我们根据数据集的键列表构建图，或者对于晶体，通过笛卡尔空间中的半径截止来构建图，并将成对原子距离\(||\textbf{r}_i-\textbf{r}_j||\)作为连续的边属性。在当前的实验中，没有使用显式的角度或二面角特征；相反，角度信息是通过在距离增强的图上的消息传递隐式捕获的。

因此，给定一个有N个节点和\(|\mathcal {E}|\)条边的图，模型消耗：\(\textbf{X}\in \mathbb {R}^{N\times p}\)（原始原子特征），\(\textbf{C}\in \mathbb {R}^{N\times 15}\)（门捷列夫描述符），\(\textbf{P}\in \mathbb {R}^{N\times 9}\)（节点拓扑编码），\(\textbf{L}\in \mathbb {R}^{N\times d_l}\)（LPE），\(\textbf{E}\in \mathbb {R}^{|\mathcal {E}|\times f}\)（原始键特征和/或距离），以及\(\textbf{G}\in \mathbb {R}^{M\times 4}\)（边拓扑编码）。此外，如果模型是等变的，它还消耗3D坐标\(\textbf{R}\in \mathbb {R}^{M\times 3}\)。所有通道在训练前都进行了标准化，并且计算出的编码中包含无效值的样本被丢弃，这对每个数据集来说通常是可以忽略的量。

我们引入了一个轻量级的嵌入模块，它将原始节点/边属性与多种编码类统一起来，并将节点和边属性映射到用户选择的维度的适当子空间。特别是，这个模块用于所有方案，除了原始HydraGNN流水线（方案1），后者直接在X上操作，可能还在E上操作。如果只激活全局注意力（方案3），它会在节点上嵌入X和L的组合，在边上嵌入L；而方案2和4需要将所有可用的编码嵌入到公共子空间中。

为了构建方案2、3和4的输入节点嵌入，我们通过连接可用的编码来形成节点输入张量$$\begin{aligned} \textbf{Z}_\text {node} \;=\; \big [\,\textbf{X}\;\Vert \; \textbf{L}\;\Vert \; \textbf{P}\;\Vert \;\textbf{C}\,\big ] \in \mathbb {R}^{N\times d_\text {node}^{\text {in}}, \end{aligned}$$其中根据所使用的特定方案，缺失的张量被简单地省略。一个没有偏置的单一线性投影映射到一个隐藏宽度\(d_h\)：$$\begin{aligned} \textbf{H} \;=\; \textbf{Z}_\text {node}\,\textbf{W}_\text {node}, \qquad \textbf{W}_\text {node}\in \mathbb {R}^{d_\text {node}^{\text {in}}\times d_h}. \end{aligned}$$这保持了节点之间的对齐，同时最小化了参数。对于方案1，该模块在节点编码上退化为恒等操作，并返回\(\textbf{X}\)不变，从而实现了即插即用的消融实验。

当使用设计用于处理边特征的MPNN时，我们通过两种模式之一构建边输入：（1）当使用编码器时直接使用\(\textbf{G}\)；或者（2）如果激活了全局注意力但没有使用编码器，对于每条边（u,?v）及其索引（i,?j），我们计算节点间的LPE的绝对差异，$$\begin{aligned} \textbf{r}_{(u,v)} \;=\; \big |\,\textbf{L}_{i:} - \textbf{L}_{j:}\,\big | \in \mathbb {R}^{d_\ell }, \end{aligned}$$这是排列不变的，并捕获了简单的光谱位移概念。如果有的话，原始边属性\(\textbf{E}\)与编码连接在一起。结果边张量$$\begin{aligned} \textbf{Z}_\text {edge} \in \mathbb {R}^{E\times d_\text {edge}^{\text {in}}} \end{aligned}$$被线性投影（无偏）到一个边嵌入\(\textbf{A}\in \mathbb {R}^{|\mathcal {E}|\times d_e'}\)，并用于替换模型的输入边属性\(\textbf{E} = \textbf{A}\)。对于方案1，原始边特征直接传递给模型。这使得下游卷积对编码源保持不可知，同时受益于统一的边表示。

在嵌入操作的复杂性方面，主要的成本是两个密集乘法：$$\begin{aligned} \mathcal {O}\!\left( N\,d_\text {node}^{\text {in}}\,d_h\right) \quad \text {和}\quad \mathcal {O}\!\left( |\mathcal {E}|\,d_\text {edge}^{\text {in}}\,d_e'\right) , \end{aligned}$$加上当禁用编码器时的\(\mathcal {O}(|\mathcal {E}|\,d_\ell )\) LPE差异。由于MPNN和全局注意力已经执行了\(\mathcal {O}(N^2)\)操作，模型的整体复杂性不受额外嵌入层的影响。此外，我们注意到我们方法的主要优势有：（i）单投影融合。对节点和边各使用一个无偏线性投影，减少了参数并避免了每个通道的冗余MLP，同时让主干在需要的地方分配深度。（ii）解耦的切换。四个标志允许进行清晰的消融实验：仅原始特征；原始+LPE；原始+TE(+CE)；以及考虑边与仅考虑节点的。 （iii）光谱相关性。LPE差异路径在显式RPE不可用时提供了一个廉价且符号不变的关系信号。

训练是在HydraGNN框架下进行的，该框架由DeepHyper [89, 90]驱动的超参数优化（HPO）控制。让\(\mathcal {H}\)表示超参数的搜索空间（例如，学习率、隐藏宽度\(d_h\)、边宽度\(d_e'\)、丢弃率）。对于每个候选\(h \in \mathcal {H}\)，我们通过最小化训练损失来训练模型\(f_{\theta ,h}\) T个周期$$\begin{aligned} \mathcal {L}(\theta ;h) \;=\; \frac{1}{|\mathcal {D}_\text {train}|}\sum _{(x,y)\in \mathcal {D}_\text {train}} \ell \!\left( f_{\theta ,h}(x), y\right) , \end{aligned}$$其中\(\ell\)通常选择为平方误差。验证性能 \(\mathcal {L}_\text {val}(h)\) 决定了最佳配置$$\begin{aligned} h^*\;=\; \arg \min _{h\in \mathcal {H}} \; \mathcal {L}_\text {val}(h). \end{aligned}$$然后我们使用扩展的预算 \(T' \gg T\) 重新训练 \(f_{\theta ,h^*}\)，并在连续 p 个周期内 \(\mathcal {L}_\text {val}\) 无改进时提前停止训练。模型检查点确保每个周期参数 \(\theta _t\) 的持久性，从而提供恢复性和可重复性。HPO 阶段的成本按以下方式缩放$$\begin{aligned} \mathcal {O}\!\big (|\mathcal {H}| \, T \, C_\text {epoch}\big ), \end{aligned}$$其中 \(C_\text {epoch}\) 表示在数据集上进行单个训练周期的成本。对于每个样本有 N 个节点和 \(|\mathcal {E}|\) 条边的图数据，我们近似计算$$\begin{aligned} C_\text {epoch} = \mathcal {O}\!\left( \sum _{(x,y)\in \mathcal {D}_\text {train}} (N d_h^2 + |\mathcal {E}| d_e'^2) \right) , \end{aligned}$$其中 \(d_h\) 是隐藏节点宽度，\(d_e'\) 是边嵌入宽度（参见第节嵌入）。HPO 之后，重新训练阶段需要 \(\mathcal {O}(T' \, C_\text {epoch})\)。实验我们的实验协议在一系列开源的原子级数据集上评估所提出的框架。对于每个数据集，我们采用提供的官方训练/验证/测试分割；否则，我们应用 \(80\%\!/\!10\%\!/\!10\%\) 的分割。对于每个数据集和学习方案，我们执行超参数优化（HPO）以最大化验证目标，选择最佳配置，然后在保留的测试集上报告该配置。在整个过程中使用适合任务的指标：回归目标的平均绝对误差（MAE），OGB-PPA 的分类准确率，以及 OGB-PCBA 的平均精度（mAP）。这个程序产生了跨数据集的标准化、可重复的比较，并提供了对未见图的泛化能力的无偏估计。数据集我们在 \(7\) 个数据集上进行评估，这些数据集涵盖了原子级和生化图谱（表 1）。我们有意使用广泛可用的、社区标准的数据集，以确保可重复性，并在相同的 HPO/训练设置下进行受控的跨架构比较，而不是专注于特定应用的专有数据集。该集合包括大约 \(4.21\)M 个图，平均节点数从大约 \(18\) 个到大约 \(243\) 个不等，涵盖了小分子（QM9、ZINC、OGB-PCQM）和更大、更复杂的结构（TMQM、NIAID、OGB-PPA）。监督包括四个图级回归目标——在 \(298.15\,\textrm{K}\) 下的自由能（QM9）、受限溶解度（ZINC）、色散能（TMQM）和 HOMO–LUMO 能隙（OGB-PCQM）——一个节点级回归目标（NIAID 中的部分电荷），以及两个图级分类基准：OGB-PPA 上的多类准确率和 OGB-PCBA 上的多标签 mAP。我们将 LPE 维度 \(d_\ell\) 设置为每个数据集中图一致可用的最小非零拉普拉斯特征向量的数量；具有少于 \(d_\ell\) 个非平凡特征向量的图被省略。表 1 列出了每个数据集的 \(d_\ell\)。表 1 数据集摘要和 LPE 维度完整表格QM9 提供了从 GDB 化学宇宙中策划的小有机分子的平衡几何结构和标准化的量子化学属性 Ramakrishnan 等 [91]；Ruddigkeit 等 [87]。节点编码原子序数；边对应于具有独热键类型（单/双/三/芳香）的共价键。QM9 提供了 \(19\) 个回归目标：偶极矩、各向同性极化率、最高占据分子轨道能量、最低未占据分子轨道能量、HOMO 和 LUMO 之间的能隙、电子空间范围、0?K 时的内能、298.15K 时的内能、298.15K 时的焓、298.15K 时的自由能、298.15K 时的原子化能、298.15K 时的原子化焓、298.15K 时的原子化自由能、旋转常数 A、旋转常数 B、旋转常数 C。在我们的实验中，我们预测 298.15K 时的自由能。ZINC 是一个免费提供的存储库，其中包含为虚拟筛选准备的三维、质子化和互变异构状态的小分子；我们参考了原始版本和 ZINC15 Irwin 和 Shoichet [31]；Sterling 和 Irwin [92]。节点代表重原子，其特征包括原子序数（原子类型）、手性、原子度数、形式电荷、连接的氢原子数量、自由基电子数量、杂化状态和环成员资格。边是具有键级（单/双/三/芳香）、立体构型（例如，顺/反）和二元共轭指示器的共价键。任务是对图级受限溶解度（logP）进行回归。TMQM 包括来自 Cambridge 结构数据库 Balcells 和 Skjelstad [93] 的大量结构多样的单核过渡金属复合物的量子优化几何结构和电子属性。我们使用 \(5.0\)? 半径截止值构建几何图。节点特征包括原子序数、形式电荷、原子价指数和 3D 坐标；边携带截止值内原子对之间的实数值欧几里得距离。该数据集提供了 \(11\) 个目标（例如，分子电荷、自旋、金属配位数、电子能量、色散能、偶极矩、金属中心的自然电荷、HOMO–LUMO 能隙、HOMO/LUMO 能量、极化率）；在这里我们专注于图级色散能回归。NIAID 包含实验和计算金属-有机框架，具有通过 NIAID 数据发现门户 Burner 等 [94]；Burner 等 [95] 导出的 REPEAT DFT 衍生的部分电荷和预计算的描述符。图是用 \(5.0\)? 半径截止值构建的。节点特征是原子序数；边编码欧几里得距离。任务是节点级部分电荷的回归。OGB-PCQM（PCQM4Mv2）定义了从 PubChemQC 策划的 2D 分子图中预测 DFT HOMO-LUMO 能隙，具有大规模图学习的标准分割和评估协议 Pcqm4mv2 [96]；Hu 等 [97]；Nakata 和 Shimazaki [98]。OGB-PCBA 是 OGB 对 MoleculeNet molPCBA 套件的改编，源自 PubChem BioAssay，提供了许多二元端点和评估分类器校准和迁移的强大基线 Hu 等 [99]；Wu 等 [100]。我们使用这个数据集进行图级多标签分类。OGB-PCQM 和 OGB-PCBA 的特征化。这两个数据集共享相同的分子特征化。每个节点都标注有其原子序数、手性类型（例如，R/S 立体中心）、原子度数、形式电荷、连接的氢原子数量、自由基电子数量、杂化状态、芳香性标志和环成员资格指示器。边对应于共价键，并携带键类型（单/双/三/芳香）、键的立体构型（例如，顺/反）和二元共轭指示器。OGB-PPA 由跨越许多物种和分类群的蛋白质-蛋白质关联（PPA）网络邻域组成；图是从 STRING 关联中派生的，标签反映了来源群体，使得可以评估生物图的泛化 Ogb [101]；Szlarczyk 等 [102]。我们使用这个数据集进行蛋白质-图分类（37 类别，通过准确率进行评估）。节点代表蛋白质，不携带内在的原子级特征。边编码关联；在我们的设置中，每条边包括一个 7 维的二元向量，指示证据通道（例如，基因共现、基因融合事件、共表达），其中 \(1\) 表示相应证据类型的存在。超参数搜索空间（HPO）我们指定了一个基于 MPNNs 和全局注意力的条件搜索空间。该空间根据以下条件分支：(i) 3D 坐标的可用性（has_pos），(ii) 是否启用了全局注意力块（global_attn_engine），以及 (iii) 是否使用了特征编码（use_encodings）。离散范围表示整数值网格；大括号 \(\{\cdot \}\) 列出了分类选项。表 2 按坐标可用性划分的模型家族完整表格非等变架构（例如，PNA、CGCNN）不消耗 3D 坐标（pos），因此在 2D 图上操作相同。尽管如此，我们仍将它们包含在 HPO 搜索空间中，与等变模型一起（如表 2 所示，has_pos = True），以实现数据驱动的选择，并量化显式几何信息——以及 SE(3) 等变性——相对于标准消息传递（表 3）的边际优势。表 3 当 global_attn_engine = False 时的搜索空间完整表格表 4 当 global_attn_engine = True 时的搜索空间完整表格多头注意力施加了约束 \(\texttt {hidden\_dim} \,\%\, \texttt {global\_attn\_heads}=0\)，即每个头的宽度 \(d_{\text {head}}=\texttt {hidden\_dim}/\texttt {global\_attn\_heads}\) 必须是整数。因此，在表 4 中，我们将 hidden_dim 限制为 \(\textrm{lcm}(2,4,8)=8\) 的倍数，确保每个列出的 hidden_dim 与所有允许的 global_attn_heads 值兼容，并防止无效的 HPO 配置。当 global_attn_engine 被启用时，我们扩大了 hidden_dim 的候选集，以允许更大的层宽度。这反映了全局注意力引入的增加的有效感受野和信息聚合，使模型更具表现力。等变性切换。设置 \(\texttt {edge\_embed\_dim}=0\) 使用原始的原子间距离作为唯一的边属性，保持等变架构（例如，SchNet、DimeNet、EGNN、PAINN）的 SE(3) 等变性。选择 \(\texttt {edge\_embed\_dim}>0\) 激活了完整的边特征流程（例如，当可用时包括键/类型/辅助编码），这通常会打破严格的等变性，但可以提高表现力。对于非等变模型（GAT、GINE、PNA 和 CGCNN），\(\texttt {edge\_embeb\_dim}=0\) 允许模型直接使用原始的边特征，而无需可学习的投影。硬件设置所有数值实验都在 OLCF Frontier 超级计算机上进行，这是美国第一个 exascale 系统。Frontier 基于 HPE Cray EX 架构，由 9,472 个计算节点组成，每个节点配备一个 AMD EPYC? 7A53 64 核 CPU 和四个 AMD Instinct? MI250X GPU。每个节点提供了一个高带宽内存子系统，CPU 上有 512 GB 的 DDR4 内存，GPU 上有 128 GB 的 HBM2e 内存，通过 AMD 的 Infinity Fabric 连接。节点通过 HPE Slingshot-11 互连连接，为大规模并行应用提供高带宽和低延迟。该系统由一个多拍字节的 Lustre 基础并行文件系统支持，以实现大规模模拟的高吞吐量 I/O。数值结果我们在所有数据集上评估了四种训练方案：(S1) GPS 禁用，编码器禁用，这完全对应于原始的 HydraGNN 架构；(S2) GPS 禁用，编码器启用；(S3) GPS 启用，编码器禁用；(S4) GPS 启用，编码器启用。对于每种方案，我们报告最佳 HPO 试验及其测试性能（MSE/MAE/Pearson r，如图标题所示）。表 5 和表 6 表明，在不使用 GPS 的情况下启用学习编码器（S2）为该数据集提供了最有利的偏差-方差权衡。相对于非编码器基线（S1），MSE 减少了约 \(62\%\)，MAE 减少了约 \(37\%\)，参数增加适中（185k 对 111k）。相比之下，激活 GPS 但不启用编码器（S3）将参数数量减少到 62k，但显著降低了准确性，表明额外的全局信息无法补偿较弱的标记特征。在编码器之上添加 GPS（S4，207k 参数）恢复了 S3 失去的许多性能，但仍低于 S2，表明全局信息的收益递减。对称性图（图 4）显示残差紧密聚集，只有少数异常值；S2 与身份的相似度最高，在高幅度范围内分散最小。表 5 每种方案的 ZINC 最佳 HPO 配置完整表格表 6 每种方案的 ZINC 测试性能完整表格图 4此图像的替代文本可能是使用 AI 生成的。完整图像ZINC 对称性图（预测值 vs. 真实值 logP）对于每种方案的最佳 HPO 试验QM9在这种情况下，所有方案都实现了 \(r\approx 1\)（图 5），然而误差幅度强烈依赖于架构选择（表 7 和表 8）。没有 GPS 的增强编码器 PAINN（S2，82.9k 参数）获得了最低的 MSE/MAE，比紧凑型 PAINN 基线（S1，15.1k）提高了约 \(20\%\) 的 MSE 和约 \(9\%\) 的 MAE；因此 S1 在参数效率上很高，但 S2 实现了最佳的绝对误差。启用 GPS 的变体（S3 DimeNet，S4 PAINN）尽管相关性接近完美，但显示出明显的较大误差，这与校准偏差而不是排名排序失败一致——在图 5 中表现为轻微的斜率偏移。结果表明，对于 QM9，局部信息的质量和隐藏层维度起着重要作用，而浅层全局注意力（单层）在相当或更小的模型大小下没有带来好处。表 7 每种方案的 QM9 最佳 HPO 配置完整表格表 8 每种方案的 QM9 测试性能完整表格图 5此图像的替代文本可能是使用 AI 生成的。完整图像QM9 对称性图（预测值 vs. 真实值自由能）对于每种方案的最佳 HPO 试验TMQM对于这个数据集，所有配置都导致了类似的高质量结果（图 6a，b），并且配置之间的性能变化非常小（表 9 和表 10）。值得注意的是，最佳模型（S2，PAINN，63.7k 参数，1 个卷积层）的性能优于大幅更大的 GPS 变体（S3，184.6k）甚至更深的 PAINN（S1，65.9k），表明一旦有足够表现力的局部编码器，增加深度和全局头是不必要的。最小的模型（S4，27.5k）接近 S2，强调了这个目标适合紧凑型网络；额外的容量主要带来边际收益。表 9 每种方案的 TMQM 最佳 HPO 配置完整表格表 10 每种方案的 TMQM 测试性能完整表格图 6此图像的替代文本可能是使用 AI 生成的。完整图像TMQM 对称性图（预测值 vs. 真实值色散能）对于每种方案的最佳 HPO 试验NIAID在这种情况下，当与适当的宽度配对时，GPS 提供了适度但一致的改进（表 11 和表 12）。S4 实现了最低的 MSE 和最高的相关性，而 S2（282.6k）获得了最低的 MAE；两者都明显优于较小的无 GPS 基线（S1，98.3k）。S4与S2之间的差距很小（例如，均方误差约为1.8%），这意味着每个参数的收益递减，然而等价性图表（图7）显示S4在整个动态范围内的范围略紧，这与GPS有助于长距离电荷重新分布的效果一致。不需要更深的堆栈（34个卷积层就足够了），但更大的隐藏层/边缘嵌入似乎有益。表11显示了每种方案的NIAID最佳HPO配置；表12显示了每种方案的NIAID测试性能。图7的替代文本可能是使用AI生成的。

表13和表14显示，没有GPS的编码器增强型PAINN（S2，71.1k参数）在均方误差/平均绝对误差（MSE/MAE）和准确率（r）方面表现最佳，同时其规模比DimeNet基线（S1，95.1k参数）小得多，也比依赖GPS的模型（S4，130.2k参数）小得多。与S1相比，S2在参数数量减少约27%的情况下，MSE降低了约18%，尽管其规模小了45%，但仍表现出色。等价性图表显示，在这种模式下，S2的残差变化最小，表明在没有GPS额外复杂性的情况下，其校准和长距离交互的平滑效果更好。表13显示了每种方案的OGB-PCQM4Mv2最佳HPO配置；表14显示了每种方案的OGB-PCQM4Mv2测试性能。

表15和表16显示，最佳准确率是由一个中等深度的GPS模型（S4，3个卷积层，2个头）在369.8k参数下获得的，比最强的无GPS基线（S1，86.2k参数）高出1.77个百分点。然而，这种改进的代价是参数数量增加了约4倍，而且一个更深/更宽的无GPS模型（S2，340.7k参数）的表现不如S1和S4，这表明仅靠容量无法替代全局信息混合。表15显示了每种方案的OGB-PPA最佳HPO配置；表16显示了每种方案的OGB-PPA测试准确率。

表17和表18显示，对于molPCBA，结合PNA的GPS提高了准确率和参数效率。最佳配置S3（3个卷积层，2个头）在215.8k参数下达到了0.184的mAP，优于更深的无GPS PNA（S2，252.1k参数；绝对值提高了0.017%，相对值提高了10.2%）和更大的GAT基线（S1，304.3k参数；绝对值提高了0.029%，相对值提高了18.7%），同时使用的参数更少。将头数从2增加到4（S4，217.0k参数）并没有带来额外的收益，这与一旦有适度的长距离混合，全局容量的收益递减是一致的。这些结果表明，对于molPCBA来说，适度的深度和有限的GPS就足以利用跨任务结构，而进一步的扩展主要只是增加了模型规模而没有相应的改进。

为了评估所提出扩展的实际影响，我们将所有架构与HydraGNN基线（方案S1）进行了比较，后者对应于没有全局注意力或基于编码器的特征增强的标准消息传递。在多个数据集和任务中，我们观察到增强注意力和增强编码器的方案（S2至S4）始终与HydraGNN基线相当或表现更好，特别是在图规模较大或结构依赖性更复杂的 数据集中收益最大。特别是基于编码器的增强（S2）和组合方案（S4）在几个基准测试中提高了预测准确率，同时保持了相当的训练稳定性，这证实了所提出的扩展可以提供超出原始HydraGNN架构的实际性能优势。这些结果表明，在统一框架内结合注意力机制和特征编码器可以增强模型的表达能力和准确率，而不会牺牲HydraGNN所追求的可扩展性和鲁棒性。

在多个数据集和架构中，HPO过程经常选择相对较小的隐藏维度（通常小于64），这表明在给定的训练和数据条件下，更宽的表示并没有一致地带来性能提升。这种行为与先前的观察结果一致，即分子图数据集在模型宽度适中时可能会在性能上达到饱和状态。

我们的工作目标并不是声称全局注意力机制对于原子图学习是普遍必要的，也不是说仅靠消息传递本身就根本不够。相反，我们的目的是确定在什么情况下全局注意力能够为调整良好的MPNN基线提供实际的好处，以及在什么数据和架构条件下这些好处能够证明其额外的计算成本是合理的。通过在单一的HydraGNN框架内进行受控比较，这项研究将通常被引用的MPNN的局限性重新定义为可以通过实验验证的假设，而不是假设。

广泛引用的MPNN的局限性——特别是它们在捕捉长距离交互方面的困难——往往是基于孤立案例研究或跨异构架构、数据集和训练协议的比较得出的。因此，这些局限性是在消息传递范式本身还是由于架构选择、特征表示或优化设置等混淆因素造成的，这一点并不清楚。我们工作的核心目标是通过受控实验系统地检验这些说法。

与以往的工作不同，我们的贡献不是引入另一种架构，而是提供了一个统一的框架，用于可复制的原子级基准测试，该框架结合了MPNN、全局注意力、混合GPS风格的模型以及基于编码器的特征增强，在相同的训练、实现和超参数优化设置下进行。通过在单一的HydraGNN框架内实现消息传递、全局注意力和基于编码器的特征增强的所有组合，我们能够在相同的训练和超参数优化程序下进行直接和可复制的比较。这种统一的设置使我们能够评估标准MPNN在何时足够，以及何时额外的机制提供了可测量的好处，从而将关于长距离建模的定性假设转化为可以通过实验验证的假设。这种设计使我们能够分离出每个建模组件的效果——局部消息传递、全局注意力和编码器——在不同的数据集上，为了解注意力机制何时真正有益提供了原则性的指导。据我们所知，没有现有的研究提供了这样的受控评估。

在各种数据集上，我们的实验明确了编码器增强和全局注意力（GPS）在什么计算成本下提供了可测量的好处。在小到中等规模的回归任务（ZINC、QM9、TMQM）中，没有GPS的编码器增强型PAINN始终具有竞争力或表现最佳：在ZINC和QM9上，仅使用编码器的设置（S2）产生了最低的错误，同时保持了紧凑性；在TMQM上，带有编码的浅层PAINN（S2）再次最小化了MSE/MAE，尽管其规模小于带有GPS的变体。等价性图表（图4、5和6）显示了方案之间的相关性接近于1，残差差异集中在少数高幅度的异常值上；在这种情况下，更丰富的节点/边缘编码带来了主要的收益，而在类似的预算下，额外的全局信息几乎没有好处。相比之下，具有较弱局部特征或明显长距离依赖性的任务（OGB-PPA、OGB-PCBA）在使用GPS时受益适中，当与足够的宽度和边缘容量结合时：在PPA上，一个中等深度的PNA+GPS模型（S4；三个卷积层，两个GPS头）在更高的参数成本下提高了top-1准确率（表15和16），而较浅的GPS模型（S3）表现较差；在PCBA上，PNA+GPS with two heads（S3）在参数数量上少于更深的无GPS PNA（S2）和更大的GAT基线（S1）的情况下达到了最高的mAP，而将头数增加到四个（S4）则带来了收益递减（表17和18）。总体而言，这些结果表明，适度的全球信息（两个头）和适度的深度（三个消息传递层）足以捕捉大型原子图中的非局部效应，而编码器质量是化学局部回归的主要驱动因素。经常观察到的仅使用编码器的模型（S2）表现优于仅使用GPS或融合变体（S3S4）并不是框架的局限性，而是该框架所支持的一个核心实证发现，表明架构融合并不保证在没有互补特征表示的情况下就能提高性能。

第二个方面是参数效率。在各种数据集中，最佳或接近最佳的方案通常在数万到几十万参数之间。例如，在OGB-PCQM4Mv2上，没有GPS的编码器增强型PAINN（S2；71k参数）实现了最低的MSE/MAE和最高的相关性，表现优于更深层次的DimeNet基线（S1；95k参数）和更大的GPS模型（S4；130k参数）（表13和14；图8）。在PPA上，S4相对于参数增加了4倍以上的情况下，准确率的提升是适度的；因此，当参数预算受限时，不使用GPS的基线仍然有用。在NIAID上，GPS在MSE和相关性上带来了小但一致的改进（S4），而最低的MAE是在不使用GPS的情况下实现的（S2），这强调了最佳配置可能取决于目标指标（表11和12；图7）。

GraphGPS [42]推广了将局部MPNN层与GPS块和丰富的位置/结构编码相结合的混合方案。为ZINC、PPA、PCBA和PCQM4Mv2报告的参考配置虽然强大，但它们的容量和深度远高于我们的研究；例如，在PCQM4Mv2上，gps-medium设置使用了10个GPS层，隐藏层大小在数百个左右，参数预算达到了数百万（例如，约950万参数），即使是“小型”变体也处于数百万级别。这些设计虽然取得了出色的绝对分数，但由于内存占用大以及对超参数和训练基础设施的敏感性，复制起来可能很困难。我们的实验有意采用了小型模型，并具有透明的切换选项（编码器开启/关闭；GPS开启/关闭；头数和深度），并报告了参数数量和指标。这种关注容量的协议揭示了GPS何时真正有帮助（PPA/PCBA），以及何时仅使用编码器就足够了（ZINC/QM9/PCQM），并且比与高度调整的、更大的系统进行比较时提供了更公平、更可复制的比较。

对于已有最佳实践基线的数据集（ZINC、PCQM4v2、molPCBA等），我们避免强调与特定排行榜的原始数字差距，而是提出了一个标准化的、受预算控制的评估方法：每个结果都附有参数数量和层预算（上面的表格），并且等价性图表可视化了超出单一数字摘要的错误结构。对于没有广泛使用的最佳实践排行榜的数据集（TMQM、NIAID等），我们提供了紧凑的、可复制的参考以及等价性图表，为未来的工作建立了统一的衡量标准。在各种任务中，有三个实际发现：（i）编码器模块在化学局部回归上占据了大部分收益；（ii）适度的GPS（两个头）在与适当的宽度和边缘嵌入结合时改善了大型原子图的分类；（iii）将GPS的深度/头数扩展到这一点之后，相对于参数的增长，收益递减。总体而言，同时展示架构切换和参数数量为评估进展提供了更清晰、更可复制的依据。

表19总结了我们在不同数据集特征上的发现。我们观察到以下模式：（1）图的大小本身并不决定GPS的效用——任务类型很重要；（2）对于局部分子属性的回归，无论规模如何增加，编码器都足够了；（3）分类任务从GPS中受益更多，尤其是在大型图上；（4）当GPS有帮助时，2个头加上3个卷积层就足够了；更多的容量会带来收益递减。这些指导原则来自七个数据集，阈值是近似的；对于新的化学系统，进行实证验证将是至关重要的。

结论和未来工作。我们提出了一个基于容量控制的、可复制的评估方法，用于在统一的HydraGNN流程中评估局部（MPNN）、混合和融合的局部/全局图架构，隔离了编码器和全局注意力（GPS）在回归和分类基准测试中的效果。有三个发现。首先，域编码器在化学局部回归中起着主导作用：在ZINC、QM9、TMQM和PCQM模型上，添加编码器的PAINN算法在没有全局结构（GPS）的情况下能够持续匹配或超越其他方法，同时保持模型的紧凑性。对比图表显示，剩余的误差主要集中在少数高幅度的尾部，而非系统性校准错误。其次，全局结构（GPS）是对消息传递机制的补充而非替代：在OGB-PPA和OGB-PCBA模型中，适度的全局信息混合（两个GPS节点与三个消息传递层相结合，并配备足够的隐藏层/边容量）能够带来最可靠的性能提升；而增加更多的节点或深度则会导致收益递减，尺寸过小的GPS模型反而表现不佳。第三，参数效率至关重要：使用数万到数十万个参数就能获得出色的结果，这些模型通常优于更深层次的基线模型，并且能够接近或超越那些参数数量更多的GPS变体；通过报告指标和参数数量，可以揭示真实的准确性与计算能力之间的权衡关系，从而提高模型间的可比性。在实际应用中，我们建议将添加编码器的MPNN算法作为中小型分子回归任务（平均节点数<50）的默认选择；当满足以下条件时，可以使用GPS算法：（i）图结构较大（平均节点数>150）；（ii）任务涉及分类；（iii）目标依赖于跨越多个化学键或较大空间距离的长期相互作用。在确实需要使用GPS算法的情况下，应从较小的配置开始（两个节点、三层结构），然后再逐步增加节点数或层数。对于中等规模的模型（50-150个节点，包含局部和非局部效应），建议对比S2和S4两种配置。除了现有的领先算法之外，我们还为那些缺乏广泛使用的研究数据集提供了紧凑且透明的基线模型和诊断工具，为未来的研究提供了可复制的评估标准，并为在固定计算预算下实现合理扩展提供了明确的方向。