本文介绍了一项由国际机器学习会议ICML 2024接受的原创性研究工作。以下是一份详细的学术报告,旨在向中国研究人员介绍此项研究。
研究报告:Graph Positional and Structural Encoder (GPSE) —— 一种用于增强图神经网络的通用位置与结构编码器
一、 作者、机构与发表信息
本研究的通讯作者为 Semih Cantürk (Mila – Quebec AI Institute, Université de Montréal) 和 Renming Liu (Michigan State University)。其他主要作者包括 Olivier LaPointe-Gagné, Vincent Létourneau, Guy Wolf (均来自 Université de Montréal 和 Mila – Quebec AI Institute), Dominique Beaini (Valence Labs, Université de Montréal) 以及 Ladislav Rampášek (Isomorphic Labs)。该论文发表于第41届国际机器学习会议 (International Conference on Machine Learning, ICML 2024) 的会议录中,收录于Proceedings of Machine Learning Research (PMLR) 第235卷。
二、 研究背景与目标
本研究属于图表示学习 (Graph Representation Learning) 领域,具体聚焦于图神经网络 (Graph Neural Networks, GNNs) 的增强技术。
背景知识: 图神经网络,特别是图变换器 (Graph Transformers, GTs),在处理图结构数据(如分子、社交网络、推荐系统)方面取得了巨大成功。然而,标准消息传递神经网络 (Message-Passing Neural Networks, MPNNs) 存在表达能力有限(受限于1-WL测试)、远距离依赖捕捉困难(欠触及问题)和信息瓶颈(过挤压问题)等固有限制。图变换器通过全局注意力机制部分解决了这些问题,但同时也丢弃了图结构本身带来的重要归纳偏置 (Inductive Bias)。为了弥补这一损失,位置与结构编码 (Positional and Structural Encodings, PSEs) 被引入,它们为图中的节点提供唯一可识别的、反映其位置(在图中相对何处)和结构(局部连接模式)的信息,这对于图变换器的成功至关重要。
研究动机与问题: 尽管已有多种手工设计的PSEs(如拉普拉斯特征向量、随机游走编码、静电势编码等),但尚无一种“万能”的PSE能在所有图预测任务上都表现最优。例如,随机游走编码在分子性质预测任务中更有效,而拉普拉斯特征向量则在涉及长距离依赖的任务中更有用。简单地拼接多种PSEs不仅不能带来预期增益,有时甚至会导致性能下降。因此,研究人员不得不依赖启发式方法和试错来选择最适合其任务的单一编码,这增加了研究负担并限制了模型性能。
研究目标: 本研究旨在解决这一核心挑战,提出首个通用的、可学习的图位置与结构编码器 (GPSE)。其核心目标是:1) 设计一个能够从图结构本身学习并融合多种PSEs丰富语义的统一编码器;2) 该编码器应具有高度可迁移性,能够有效应用于与训练数据分布和模态显著不同的下游图数据集;3) 作为特征提取器,GPSE应能无缝增强任何GNN(包括MPNNs和GTs),成为手工计算PSEs的高效且强大的替代方案。
三、 研究方法与工作流程
本研究的工作流程主要分为三个核心阶段:GPSE模型设计、自监督预训练以及下游任务评估。详细流程如下:
第一阶段:GPSE模型架构设计(图1a,第2.2节及附录B.1) 研究团队设计了一个基于MPNN的深度编码器来从图结构中提取丰富的PSE表示。其核心思想是:给定一个仅包含结构信息的查询图,模型能够输出一个蕴含了多种PSE信息的节点级潜在表示。 * 输入处理: 模型输入仅使用图的结构(邻接矩阵)。每个节点被赋予一个从20维标准正态分布中采样的随机特征向量。这种随机节点特征被证明有助于MPNN突破1-WL的表达能力限制。 * 编码器主干: 编码器采用带残差连接和门控机制的GatedGCN层堆叠而成,深度为20层。如此深的网络是为了捕获拉普拉斯特征向量等PSE所需的全局图视图。为了解决深度GNN中常见的过平滑和过挤压问题,架构中集成了关键技术: * 残差连接与门控机制: 减少过平滑,门控机制可以动态调整边的重要性,起到类似图稀疏化的作用,缓解过平滑。 * 虚拟节点 (Virtual Node): 引入一个与图中所有节点相连的虚拟节点,极大地增加了图的连通性,促进了全局信息交换,有效缓解了过挤压问题,并加速了信息传播。 * 输出: 经过L层(默认为20层)消息传递后,得到每个节点的最终隐藏表示 h_i^(L),这便是GPSE学习到的通用PSE表示。隐藏维度d设为512。
第二阶段:自监督预训练(图1b,第2.1,2.3节及附录A) 为了让GPSE学习到丰富的PSE语义,研究团队设计了一个多样化的PSE集合作为自监督训练目标,并通过重建这些PSE来训练编码器。 * 训练目标(PSE集合): 设计了六类共58维的PSE任务,包括: 1. 位置编码 (Positional Encodings): 编码节点在图中的相对位置。 * 拉普拉斯特征向量 (Laplacian Eigenvectors, LAPPE): 使用前4个非平凡特征向量的绝对值。 * 静电势编码 (Electrostatic Potential Encodings, ELSTATICPE): 基于拉普拉斯伪逆计算,生成7个统计量(如最小、平均、标准差等)。 2. 结构编码 (Structural Encodings): 编码节点周围的局部或全局连接模式。 * 随机游走编码 (Random Walk Structural Encodings, RWSE): 计算返回起始节点的k步随机游走概率,使用前20步。 * 热核对角编码 (Heat Kernel Diagonal Structural Encodings, HKDIGSE): 基于热核矩阵对角线的值,使用前20个。 * 环计数编码 (Cycle Counting Structural Encodings, CYCLESE): 计算图中不同长度环(如三角形、四元环等)的数量,作为图级回归任务,共7维。 * 拉普拉斯特征值 (Laplacian Eigenvalues, EIGVALSE): 作为图级回归任务。 * 训练流程: 1. 数据集: 使用分子图数据集MolPCBA中的323,555个独特图结构进行训练。仅使用图结构,忽略原子和键类型等域特征。 2. 前向传播: 将带随机特征的图输入GPSE编码器,得到节点表示 h_i。 3. 解码与损失计算: 为每一类PSE任务设置独立的2层MLP解码头,将 h_i 解码为对应的PSE预测值 y_hat。损失函数结合了L1损失和余弦相似度损失,确保模型既能捕捉PSE信号的方向(余弦损失),也能捕捉其幅值(L1损失)。 4. 优化: 通过最小化总重建损失来优化GPSE模型的所有参数。训练完成后,编码器权重被冻结,用于下游任务的特征提取。
第三阶段:下游任务评估与应用(图1c,第3节) 预训练好的GPSE模型作为一个冻结的特征提取器,用于增强各种下游GNN模型。 * 应用方式: 对于任何下游图数据集,仅需输入其图结构,GPSE即可生成对应的节点级PSE表示。这些表示随后被拼接(或添加)到下游模型原有的节点特征中,然后下游模型(如GPS、GIN、GCN等)从头开始训练进行特定任务预测。 * 评估策略: 研究团队在广泛的下游基准上进行了系统性评估,以验证GPSE的有效性、通用性和可迁移性。评估包括: 1. 分子图预测任务: 在ZINC、PCQM4Mv2、MolHIV、MolPCBA等数据集上,比较GPS模型使用GPSE编码与使用各种手工PSE(如LAPPE、RWSE)及它们组合的性能。 2. 通用性验证: 测试GPSE在非分子图(如图像超像素图CIFAR10、MNIST)和大规模节点分类图(OGB-Arxiv、OGB-Proteins)上的表现,这些数据在规模、连接模式上与训练数据存在分布偏移。 3. 模型通用性验证: 将GPSE编码用于增强多种不同的GNN架构,包括GCN、GIN、GatedGCN、Graph Transformer等。 4. 效率与可扩展性分析: 对比GPSE与手工计算PSEs在不同图数量和图规模下的计算时间。 5. 表达能力分析: 通过合成图基准(CSL, EXP)测试GPSE编码是否能为基础MPNN(如GIN)提供超越1-WL的表达能力。 6. 消融研究: 分析GPSE架构中各组件(深度、虚拟节点、门控、卷积类型)的重要性,以及不同PSE训练目标和不同预训练数据集对下游性能的影响。
四、 主要研究结果
研究结果全面支持了GPSE的有效性和优越性。
1. PSE重建能力(表1): GPSE在自监督预训练中表现出色。在仅使用5% MolPCBA数据训练时,对六类PSE的平均测试确定系数 (R²) 达到0.979。当使用90%数据训练时,平均R²进一步提升至0.9979,证明GPSE能够高度准确地从图结构中恢复多样化的PSE信息,且其性能随训练数据量增加而渐近提升。
2. 下游任务性能提升: * 分子图基准(表2): GPSE增强的GPS模型在多个分子基准上显著优于或至少持平于使用最佳手工PSE的模型。例如,在ZINC数据集上,GPS+GPSE取得了0.0648的MAE,不仅显著优于其他PSE,甚至挑战了当时的最先进结果。在PCQM4Mv2上,GPSE也取得了不使用3D信息和非集成方法中的最佳结果。 * 通用性与可迁移性(表5,表6): GPSE展现了卓越的跨域迁移能力。 * 图规模迁移: 在节点数远超训练分子图的肽类图数据集 (Peptides-struct/func) 上,GPSE仍能提升GPS和GCN等模型的性能,甚至在Peptides-struct上挑战了SOTA。 * 连接模式迁移: 在连接模式与分子图迥异的超像素图 (CIFAR10, MNIST) 上,GPS+GPSE取得了与使用计算PSEs相当的性能。 * 极端分布外迁移: 在包含数十万节点的大规模引文网络和蛋白质相互作用网络 (OGB-Arxiv, OGB-Proteins) 的节点分类任务中,GPSE在10个模型-数据集组合中的8个上取得了最佳或并列最佳结果,显著优于无编码或使用LAPPE编码的基线。 * 模型通用性(表3): GPSE编码能够普遍提升各类GNN模型。在ZINC数据集上,将GPSE用于GCN、GatedGCN、GIN、GINE和Transformer等模型,平均带来了56.24%的MAE降低,其改善幅度超过了使用任何单一或组合手工PSE。
3. 与自监督学习预训练方法的比较(表4): 与专门为分子图设计的自监督学习 (SSL) 预训练方法(如GraphCL, GraphMAE)相比,GPSE作为通用特征提取器,在八个MoleculeNet数据集中的三个上取得了最佳性能,并在所有八个数据集上均相对于无预训练的基线有提升,而某些SSL方法在某些数据集上出现了负迁移。这证明了富含PSE信息的特征本身即能提供强大增益。
4. 表达能力验证(表7,附录E): 在需要超越1-WL表达能力的合成图基准(CSL, EXP)上,为GIN模型添加GPSE编码能显著提升其区分图同构类的能力,而仅添加随机特征则提升有限。这证实了GPSE编码确实为基础MPNN注入了额外的表达能力。
5. 效率优势(附录G): 相对于手工计算PSEs(尤其是计算全部PSEs的组合),GPSE在推理时具有显著的计算效率优势。其计算时间不随PSE类型数量增加而增长,且在处理大规模图或大量图时,速度优势更加明显。
6. 消融研究关键发现: * 架构组件: 深度(20层)、虚拟节点、GatedGCN的门控机制和残差连接对于GPSE成功学习全局PSE(如拉普拉斯特征向量)并避免过平滑至关重要(图2)。 * PSE多样性: 排除任何一类PSE进行训练通常会导致下游任务性能下降,表明多样化的PSE语义对下游任务均有贡献。 * 预训练数据: 使用更大、更多样的数据集(如MolPCBA, ChEMBL)预训练GPSE通常能带来更好的下游性能,体现了数据缩放定律。
五、 研究结论与价值
本研究成功提出了GPSE,这是首个旨在从图结构学习并提取丰富、通用位置与结构编码的图编码器。
科学价值: 1. 统一框架: GPSE为解决“如何为不同任务选择最佳PSE”这一长期问题提供了一个统一的、数据驱动的解决方案,推动了图学习从手工特征工程向学习通用特征提取器的范式转变。 2. 理论验证: 研究通过精心设计的架构,实证了深度MPNN结合随机特征、虚拟节点和门控机制能够有效平衡过平滑与过挤压,并学习到超越1-WL限制的图表示。 3. 可迁移性证明: 研究系统性地证明了基于图结构学习的编码器具有强大的跨分布、跨模态迁移能力,为开发图基础模型提供了重要启示。
应用价值: 1. 性能提升: GPSE能够显著且稳定地提升各类GNN在广泛下游任务上的性能,且无需针对特定任务进行繁琐的PSE选择和调优。 2. 高效便捷: 作为一次预训练、多次使用的特征提取器,GPSE避免了每次运行都需要重新计算复杂PSE的开销,尤其在大规模场景下效率优势巨大。 3. 易于集成: GPSE已集成到PyG库中,方便研究者和开发者轻松应用于各种图学习任务。
六、 研究亮点
七、 其他有价值内容
研究在附录中进行了深入的讨论(附录I),探讨了GPSE为何在某些任务上优于手工PSE,以及为何简单的PSE拼接无效。分析指出,GPSE学习的是PSE的潜在语义表示,可能对噪声更鲁棒,并避免了手工PSE中可能存在的冗余或冲突信息。此外,研究也指出了当前GPSE模型的局限性,如其规模尚不足以高效训练于超大规模(十亿级)图数据集,这为未来的研究方向——开发更高效、更强大的PSE提取器——指明了道路。