Graph-level Representation Learning with Joint-Embedding Predictive Architectures
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.
Submission Length: Regular submission (no more than 12 pages of main content)
Changes Since Last Submission: Dear AE, Reviewers, and EiCs,
We are uploading, alongside this comment, the revised manuscript to be sent for publication. Firstly, we greatly thank the reviewers and the AE for the enjoyable and informative review process. We truly believe it helped elevate the quality of the manuscript. In this final version, we have made several changes, listed in the following:
- An appendix was added. The tables containing the dataset statistics and model hyperparameters, which were previously situated in the main paper, have been placed in the appendix. In said appendix, we provide new experimental results consisting of Graph-JEPA ablations on long-range tasks from the Long Range Graph Benchmark [1], both in multivariate graph regression and multilabel classification. Finally, we provide a figure showing the context and target patches from a small molecular graph in Fig. 4 in the Appendix.
- We used the free space remaining in the main manuscript to expand significantly the details on the spatial partitioning and subgraph positional embedding, as requested in the first three bullet items requested by the AE in point 1. of their decision.
- Minor details on the latent codes and the stochasticity regarding the subgraph partitioning were added.
- Uniformized the notation and typography of the paper, trying to pay close attention to the language used and the grammatical structure.
- Last but not least, we added the official code repository (Github) and deanonymized the paper to be ready for publication.
[1] Dwivedi, Vijay Prakash, et al. "Long range graph benchmark." Advances in Neural Information Processing Systems 35 (2022): 22326-22340.
Code: https://github.com/geriskenderi/graph-jepa
Assigned Action Editor: ~Serguei_Barannikov1
Submission Number: 2906
Loading