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 will be made available upon acceptance.
Submission Length: Regular submission (no more than 12 pages of main content)
Changes Since Last Submission: The *changes* in this revised manuscript compared to the submitted version are *highlighted in blue*. Specifically, we have changed the following:
- Changed Fig. 1 such that it would be less crowded and contain more details in the caption, as entailed by the feedback from reviewer pezm. Along with the figure, we fixed any parts of the text referencing it, so that they would fit with the new one.
- Added a paragraph explaining possible reasons for the variance in the results based on the feedback and response from reviewer XWJw.
- Fixed the typo in Table 7 (now Table 8), as suggested by reviewer XWJw.
- Added a new table (Table 1) containing information about the datasets used for the experiments, as suggested by reviewer XWJw.
- Added a brief mention of D being the diagonal degree matrix when talking about the embedding method, as suggested by reviewer dx84
Assigned Action Editor: ~Serguei_Barannikov1
Submission Number: 2906
Loading