Graph-level Representation Learning with Joint-Embedding Predictive Architectures

We would like to sincerely thank the reviewer for the detailed comments and valuable questions, and also for positively evaluating our loss function design and analysis. We provide details to clarify the major concerns.

Q1: Missing literature in related work

A1: We thank the reviewer for this attentive comment, we have indeed missed these important works talking about latent graph prediction. All three suggested papers have been added in the Related Works section (Section 2) of the rebuttal revision, along with a small discussion. As it is the most recent and competitive work, we have added results from LaGraph in Table 1 of the rebuttal revision as well.

Q2: The performance improvement is marginal based on the main results in Table 1. Could you add a quantitative comparison for the efficiency?

A2: We provide a discussion regarding our performance in the response to Q2 of reviewer 15Jn. We provide the same explanation here in order for the reviewer to have a self-contained response. Given that another big advantage of Graph-JEPA is efficiency, and the reviewer had asked about experiments regarding this, we provide in this answer the results of the comparison for the efficiency: Graph-JEPA provides a new way of thinking about SSL on graphs, by essentially bridging the advantages of both contrastive and generative approaches. Contrastive approaches present several computational bottlenecks, while generative models such as GraphMAE tend to strongly overfit because of the generative objective. As mentioned in our paper, we can bypass the need for data augmentation and negative samples, while not needing to fit the data distribution. Not only does this bring about a new paradigm and research direction for graph SSL, but we also show that it can be beneficial in terms of performance and expressiveness of representations. Additionally, we provide new results in this response and in the appendix of the rebuttal revision showcasing the efficiency of Graph-JEPA compared to GraphMAE and MVGRL. The table below presents the training time needed for the three methods in order to produce the representations that give rise to optimal downstream performance. Our model clearly shows superior efficiency and convergence, especially as the size of the dataset grows ( REDDIT-MULTI-5K).

Q3: Most existing SSL methods can handle both graph-level and node-level tasks. However, the proposed Graph-JEPA only supports graph-level downstream tasks.

A3: While we agree with the statement of the reviewer, our method is framed with the graph-level representation learning problem in mind, therefore performance on node-level tasks was not taken into account. Graph-level learning is crucial in different domains and most importantly, has different dynamics, particularly in terms of hierarchy, compared to node-level concepts. Our self-predictive task and objective function are both implemented considering these aspects for an improved graph-level representation. Nevertheless, this is definitely a direction of future work for future applications of JEPAs in the graph domain, for which we hope our work can provide inspiration.

Q4: The proposed Graph-JEPA uses Transformer encoder blocks. However, most baseline models are based on simpler models like GIN and GCN. Is it an unfair comparison? Can you use GIN/GCN encoder?

A4: We utilize the Transformer to exchange information between the subgraph embeddings because of the modeling flexibility, without necessarily requiring an adjacency matrix between the various subgraphs. Firstly, we argue that the comparison is not unfair since we utilize a MP-GNN to encode the node features before moving on to the subgraph encoding. To demonstrate the benefits of our architecture and SSL objective, we provide the downstream results of a Graph-JEPA that only utilizes MLP modules after the MP-GNN subgraph encoding. The results in the Table below, show that the model remains competitive, maintaining state-of-the-art performance on MUTAG, REDDIT-BINARY, and ZINC, albeit marginally in the first two cases. Nevertheless, even when parametrized through simple MLPs, Graph-JEPA is still able to outperform most contrastive approaches presented in Table 1 of the submission.

We would like to express our gratitude to the reviewer for acknowledging our writing and analysis. We proceed by answering the listed weaknesses of our paper.

Q1: The method is not novel. The proposed Graph-JEPA is very similar to MLM in BERT, which utilizes the context to predict the masked word type.

A1: While the idea of masking a part of the input can be considered similar to BERT, Graph-JEPA is quite different as a pretraining strategy. MLM with BERT is an autoencoding procedure, thus it is actually part of the generative pretraining family of SSL, like most Language Modeling approaches. Our approach works in latent space only and it’s predictive of the latent embeddings, thereby it does not fit the data distribution. Naturally, this completely alters learning dynamics. A general overview, both schematic and written was made available in the Related Work (Section 2) of our submission.

Q2: The proposed method is too simple and the motivation is not clear. We have graph MAE and contrastive learning. Why do we need JEPA for the graph domain?

A2: Graph-JEPA provides a new way of thinking about SSL on graphs, by essentially bridging the advantages of both contrastive and generative approaches. Contrastive approaches present several computational bottlenecks, while generative models such as GraphMAE tend to strongly overfit because of the generative objective. As mentioned in our paper, we can bypass the need for data augmentation and negative samples, while not needing to fit the data distribution. Not only does this bring about a new paradigm and research direction for graph SSL, but we also show that it can be beneficial in terms of performance and expressiveness of representations. Additionally, we provide new results in this response and in the appendix of the rebuttal revision showcasing the efficiency of Graph-JEPA compared to GraphMAE and MVGRL. The table below presents the training time needed for the three methods in order to produce the representations that give rise to optimal downstream performance. Our model clearly shows superior efficiency and convergence, especially as the size of the dataset grows (REDDIT-MULTI-5K).

Q3: Compared with graph MAE and S2GAE, the performance is not good enough to show it can inspire future research.

A3: Graph-JEPA outperforms Graph-MAE and S2GAE by a significant margin on MUTAG, DD, and REDDIT-BINARY. We have added the results of using GraphMAE with the officially provided code on REDDIT-MULTI-5K and ZINC in the revised version of the paper, where Graph-JEPA still shows superior performance. Our model remains competitive in all datasets, setting the state of the art in 5 datasets and coming second in 1, out of 8 total datasets. We obtain these results while also showing great efficiency, as demonstrated in the table above. We believe all of these factors make Graph-JEPA a valid candidate for the exciting new frontier of JEPAs in graph SSL.

Dear reviewer 15Jn,

As the discussion deadline approaches, we kindly seek clarification on whether the three weaknesses highlighted in the initial review have been adequately addressed in our rebuttal. We have endeavored to furnish comprehensive details in each response to facilitate your assessment. If further clarification is needed or if there are additional concerns regarding our paper, we would greatly appreciate the opportunity to engage in further discussion. Thank you once again for your valuable comments and feedback.

We would like to thank the reviewer for their positive comments regarding the performance and soundness of our method. In the following, we aim to provide detailed responses to the raised weaknesses and questions.

Q1: The overall method seems to be a direct application of JEPA to graphs. What makes applying JEPA to graphs special and non-trivial?

A1: The Joint-Embedding Predictive Architecture is a recently proposed, general design for SSL [1]. As such, applying this concept to the graph domain brings about several challenges. This is very similar to previous challenges faced when initially applying contrastive learning or autoencoding in the graph domain, tackled by related work. We state the main challenges that Graph-JEPA tries to solve below:

1. How do we design the self-supervised task in order for it to be self-predictive and expressive? This includes creating supervision from the input graphs such that an appropriate design can capture feature similarities via latent reconstruction. There are multiple possibilities: Is it appropriate to work only at node level, or at structure level? If we decide to operate at subgraph level, how do we extract subgraphs? Should there be an overlap between them? How should the SSL pretraining task be concretely defined? Given all these possibilities, developing an intuitive and simple to implement JEPA is not straightforward.
2. How can we make sure that the SSL learning procedure defined in step 1 is optimal for graph-level concepts? First of all, we want our representations to not collapse in the latent space, and we also wish to impose a particular geometry to the latent space that reflects the hierarchical nature of graph-level concepts. Both of these problems require an appropriately designed solution
3. Can we learn expressive graph representations through a JEPA, such that we bypass the 1-WL expressiveness limit of common Message Passing Graph Neural Networks (MP-GNNs)? Given that we wish to lean graph-level representations, it is imperative that these are expressive in terms of graph isomorphisms and can go beyond simple MP-GNNs.

Our proposed architecture tackles all these questions and reveals that it is possible to build a JEPA architecture for graph-level SSL, as shown by the competitive performance both in downstream classification and expressiveness.

Q2: The discussion of "why does graph-JPEA works" is useful, but not information. Any theoretical analysis here will be useful.

A2: The title of the section might have been misleading so we have removed it in the rebuttal revision. The content has been merged in two different parts: the additional theoretical analysis on the necessity of the stop-gradient operation and the asymmetric predictor network has been added in the end of Section 3.4, while the hyperbolic properties of the latent space are discussed in Section 4.3. Our aim with these discussions is to show that the design choices behind our proposed approach make sense theoretically and not just empirically. We agree with the reviewer that an additional theoretical analysis as to why a JEPA architecture can learn in the first place would be useful, but that is beyond the scope of this work. We believe that the insights from our paper provide information regarding key design aspects of the architecture and we would happily discuss further the conclusions drawn in the paper if the reviewer would kindly provide more details regarding what aspects they find not informative.

Q3: The experimental settings are confusing. It is unclear to me why "GCN", a GNN model, can be compared with "Graph-JEPA", which is a graph self-supervised training method.

A3: We apologize for the lack of clarity or confusion caused by the experimental setting. GCN only appears in Table 2, where we empirically evaluate the expressiveness of the learned representations of the different models. Given that most MP-GNNs are known to be as expressive as the 1-WL test, the experiments indeed show that they are incapable of recognizing 1-WL indistinguishable graphs, with an accuracy level (approximately 50%) that shows almost a tendency to randomly guess. In contrast, representations learned by Graph-JEPA are able to perform almost perfectly. Thus the use of GCN, similar to the use of GIN in Table 1 of the submission, is necessary simply to put the results into perspective. The same study is performed in [2].

Q4: All the figures and tables are not professional and could be improved to be more appealing. Font sizes and colors should be improved.

A4: We are sorry for the seemingly unprofessional presentation of our work. We have updated all of the tables in the rebuttal revision and tried to provide a better layout.

[1] LeCun, Yann. "A path towards autonomous machine intelligence version 0.9. 2, 2022-06-27." Open Review 62 (2022).

[2] He, Xiaoxin, et al. "A generalization of vit/mlp-mixer to graphs." International Conference on Machine Learning. PMLR, 2023.

Dear reviewer b3AV,

As the discussion deadline approaches, we kindly seek clarification on whether we have adequately answered your question and addressed the four weaknesses highlighted in the initial review. We have made diligent efforts, including modifications to the paper, adjustments to the paper's formatting, and answering every single weak point. If additional clarification is required or if there are further concerns with our paper, we would greatly appreciate the opportunity for additional discussion.

Thank you for your constructive feedback.