[论文解读] Graph-level Representation Learning with Joint-Embedding Predictive Architectures

Geri Skenderi, Hang Li+2 more|arXiv (Cornell University)|Sep 27, 2023

Advanced Graph Neural NetworksComputer Science被引用 3

一句话总结

本文提出 Graph-JEPA，这是首个用于图级别自监督表征学习的联合嵌入预测架构（Joint-Embedding Predictive Architecture）。它通过上下文子图预测被掩码的子图嵌入，无需数据增强或负样本，在图分类和回归任务中取得具有竞争力的性能，其核心是提出了一种新颖的双曲坐标预测目标，能够捕捉图的层次结构概念。

ABSTRACT

Joint-Embedding Predictive Architectures (JEPAs) have recently emerged as a novel and powerful technique for self-supervised representation learning. They aim to learn an energy-based model by predicting the latent representation of a target signal y from the latent representation of a context signal x. JEPAs bypass the need for negative and positive samples, traditionally required by contrastive learning while avoiding the overfitting issues associated with generative pretraining. In this paper, we show that graph-level representations can be effectively modeled using this paradigm by proposing a Graph Joint-Embedding Predictive Architecture (Graph-JEPA). In particular, we employ masked modeling and focus on predicting the latent representations of masked subgraphs starting from the latent representation of a context subgraph. To endow the representations with the implicit hierarchy that is often present in graph-level concepts, we devise an alternative prediction objective that consists of predicting the coordinates of the encoded subgraphs on the unit hyperbola in the 2D plane. Through multiple experimental evaluations, we show that Graph-JEPA can learn highly semantic and expressive representations, as shown by the downstream performance in graph classification, regression, and distinguishing non-isomorphic graphs. The code is available at https://github.com/geriskenderi/graph-jepa.

研究动机与目标

• 为解决对比学习和生成式自监督学习在图结构中的局限性，如对数据增强和负样本的依赖，或因分布估计导致的过拟合问题。
• 开发一种基于 JEPA 的图级别表征学习框架，该框架在潜在空间中运行，避免显式重建或对比目标。
• 设计一种预测目标，通过在二维潜在空间中的单位双曲面上预测子图坐标，捕捉图级别的层次结构概念。
• 在标准图基准数据集上评估 Graph-JEPA 在图分类和回归任务中的有效性。
• 对关键设计选择进行消融分析，包括位置编码、注意力机制以及子图采样策略。

提出的方法

• Graph-JEPA 使用掩码建模方法，基于从输入图中提取的上下文子图，预测随机选择的目标子图的潜在表征。
• 通过池化操作将子图表征聚合为全局图级别嵌入，从而支持下游的分类和回归任务。
• 引入一种新颖的预测目标，将目标子图嵌入编码为二维空间中单位双曲面上的坐标，以建模图的层次结构概念。
• 模型采用预测网络，将上下文子图表征映射至目标子图的潜在嵌入，采用联合嵌入范式。
• 使用相对小波谱嵌入（Relative Wavelet Spectral Embeddings, RWSE）作为位置编码，并通过消融实验比较节点级与块级 RWSE 的性能差异。
• 架构采用哈达玛自注意力（Hadamard self-attention）以提供对图结构数据的归纳偏置，并与标准自注意力进行消融比较。

Figure 1: Illustration of the three main SSL approaches: (a) Joint-Embedding Architectures learn to create similar embeddings for inputs x and y that are compatible with each other and dissimilar embeddings for inputs that are not compatible. This compatibility is implemented in practice by creating

实验结果

研究问题

• RQ1 基于 JEPA 的方法是否能在无需数据增强或负样本的情况下，有效学习图级别表征？
• RQ2 在图自监督学习中，通过在单位双曲面上预测子图坐标是否优于标准潜在表征重建？
• RQ3 不同的子图采样策略（结构化 METIS 与随机采样）如何影响模型性能与稳定性？
• RQ4 位置编码类型（节点级 vs. 块级 RWSE）对 Graph-JEPA 性能有何影响？
• RQ5 距离函数的选择（欧氏距离 vs. 双曲距离）如何影响模型在不同图数据集上的泛化能力？

主要发现

• Graph-JEPA 在图分类任务中达到最先进性能，在 MUTAG 数据集上准确率为 91.25±5.75%，在 IMDB-B 上为 73.68±3.24%，优于对比学习和生成式基线方法。
• 在 ZINC 回归基准上，Graph-JEPA 实现了 0.434±0.01 的平均绝对误差（MAE），显著优于次优方法（0.471±0.01，使用欧氏距离）。
• 使用双曲距离函数在 MUTAG 和 REDDIT-M 上提升性能，准确率分别达到 89.43±5.67%，高于使用欧氏距离的 87.04±6.01%。
• 消融实验表明，节点级（全局）RWSE 始终优于块级（相对）RWSE，后者因邻近节点信息重叠导致方差更高。
• 将哈达玛自注意力替换为标准自注意力后，性能略有提升但更不稳定，表明归纳偏置有助于提升训练的一致性。
• 随机子图采样性能与 METIS 基于采样相当，但泛化能力较差；因此推荐使用 METIS 以获得更强的归纳偏置和在多样化图结构中的稳定性。

Figure 2: An overview of Graph-JEPA. We first extract non-overlapping subgraphs (patches) (a.), perform a 1-hop neighborhood expansion (b.), and encode the subgraphs with a GNN (c.). After the subgraph encoding, one is randomly picked as the context and $m$ others as the targets (d.) and they are fe