# Graph-level Representation Learning with Joint-Embedding Predictive
Architectures

Geri Skenderi1 Hang Li 2 Jiliang Tang2 Marco Cristani1
1 University of Verona 2 Michigan State University
<EMAIL_ADDRESS>
{lihang4<EMAIL_ADDRESS>
###### 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 a context signal $x$. JEPAs bypass the need for data
augmentation and negative samples, which are typically required by contrastive
learning, while avoiding the overfitting issues associated with generative-
based pretraining. In this paper, we show that graph-level representations can
be effectively modeled using this paradigm and propose Graph-JEPA, the first
JEPA for the graph domain. In particular, we employ masked modeling to learn
embeddings for different subgraphs of the input graph. To endow the
representations with the implicit hierarchy that is often present in graph-
level concepts, we devise an alternative training objective that consists of
predicting the coordinates of the encoded subgraphs on the unit hyperbola in
the 2D plane. Extensive validation shows that Graph-JEPA can learn
representations that are expressive and competitive in both graph
classification and regression problems.

##  1 Introduction

Graph data is ubiquitous in the real world due to its ability to universally
abstract various concepts and entities (Ma & Tang, 2021; Veličković, 2023). To
deal with this peculiar data structure, Graph Neural Networks (GNNs)
(Scarselli et al., 2008; Kipf & Welling, 2016a; Gilmer et al., 2017;
Veličković et al., 2017) have established themselves as a staple solution.
Nevertheless, most applications of GNNs usually rely on ground-truth labels
for training. The growing amount of graph data in fields such as
bioinformatics, chemoinformatics, and social networks has made manual labeling
laborious, sparking significant interest in unsupervised graph representation
learning.

A particularly emergent area in this line of research is self-supervised
learning (SSL). In SSL, alternative forms of supervision are created stemming
from the input signal. This process is then typically followed by invariance-
based or generative-based pretraining (Liu et al., 2023; Assran et al., 2023).
Invariance-based approaches optimize the model to produce comparable
embeddings for different views of the input signal. A common term associated
with this procedure is contrastive learning (Tian et al., 2020). Typically,
these alternative views are created by a data augmentation procedure. The
views are then passed through their respective encoder networks (which may
share weights), as shown in Figure 1a. Finally, an energy function, usually
framed as a distance, acts on the latent embeddings. In the graph domain,
several works have applied contrastive learning by designing graph-specific
augmentations (You et al., 2020), using multi-view learning (Hassani &
Khasahmadi, 2020) and even adversarial learning (Suresh et al., 2021). This
pretraining strategy is effective but comes with several drawbacks i.e., the
necessity to augment the data and process negative samples. In order to learn
embeddings that are useful for downstream tasks, the augmentations must also
be non-trivial.

Generative methods on the other hand typically remove or corrupt portions of
the input and predict them using an autoencoding procedure (Vincent et al.,
2010; He et al., 2022), or rely on autoregressive modeling (Brown et al.,
2020; Hu et al., 2020). Figure 1b depicts the typical instantiation of these
methods: The input signal $x$ is fed into an encoder network that constructs
the latent representation and a decoder generates $\hat{y}$, the data
corresponding to the target signal $y$. The energy function is then applied in
data space, often as a reconstruction term (Bengio et al., 2013). Generative
models generally display strong overfitting tendencies since they have to
estimate the data distribution (implicitly or explicitly), so the latent
representations must be directly descriptive of the whole data space. This can
be very problematic for graphs given that they live in a non-Euclidean and
inhomogenous data space. Nevertheless, masked autoencoding has recently shown
promising results in the graph domain with appropriately designed models (Hou
et al., 2022; Tan et al., 2023).

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
different views of the input data. (b) Generative Architectures reconstruct a
signal $y$ from a compatible signal $x$ by conditioning the decoder network on
additional (potentially latent) variables $z$. (c) Joint-Embedding Predictive
Architectures act as a bridge: They utilize a predictor network that processes
the context $x$ and is conditioned on additional (potentially latent) factors,
to predict the embedding of the target $y$ in latent space.

Inspired by the innovative Joint-Embedding Predictive Architecture (JEPA)
(LeCun, 2022; Assran et al., 2023), we propose Graph-JEPA to learn graph-level
representations by bridging contrastive and generative models. As illustrated
in Figure 1c, a JEPA has two encoder networks that receive the input signals
and produce the corresponding representations. Notably, the two encoders can
be different models and don’t need to share weights. A predictor module
outputs a prediction of the latent representation of one signal based on the
other, possibly conditioned on another variable. Graph-JEPA does not require
any negative samples or data augmentation and by operating in the latent space
it avoids the pitfalls associated with learning high-level details needed to
fit the data distribution. However, the graph domain presents us with
additional challenges, namely: context and target extraction; designing a
latent prediction task that is optimal for graph-level concepts; and learning
expressive representations. In response to these questions, we design Graph-
JEPA with a specific masked modeling objective. The input graph is first
divided into several subgraphs, and then the latent representation of randomly
chosen target subgraphs is predicted given a context subgraph. The subgraph
representations are consequently pooled to create a graph-level
representation.

The nature of graph-level concepts is often assumed to be hierarchical (Ying
et al., 2018), thus we conjecture that the typical latent reconstruction
objective used in current JEPA formulations is not enough to provide optimal
downstream performance. To this end, we design a prediction objective that
starts by expressing the target subgraph encoding as a high-dimensional
description of the hyperbolic angle. The predictor module is then tasked with
predicting the location of the target in the 2D unit hyperbola. This
prediction is compared with the target coordinates, obtained by using the
aforementioned hyperbolic angle. Graph-JEPA outperforms popular contrastive
and generative graph-level SSL methods on different datasets, while
maintaining efficiency and ease of training. Notably, we observe from our
experiments that Graph-JEPA can run up to 1.45x faster than Graph-MAE (Hou et
al., 2022) and 8x faster than MVGRL (Hassani & Khasahmadi, 2020). Finally, we
empirically demonstrate Graph-JEPA’s ability to learn highly expressive graph
representations, showing it almost perfectly distinguishes pairs of non-
isomorphic graphs that the 1-WL test cannot differentiate.

##  2 Related work

###  2.1 Self-Supervised Graph Representation Learning

Graph Neural Networks (GNNs) (Scarselli et al., 2008; Kipf & Welling, 2016a;
Veličković et al., 2017; Hamilton et al., 2017; Xu et al., 2018; Wu et al.,
2019) are now established solutions to different graph machine learning
problems such as node classification, link prediction, and graph
classification. Nevertheless, the cost of labeling graph data is quite high
given the immense variability of graph types and the information they can
represent. To alleviate this problem, SSL on graphs has become a research
frontier. SSL methods on graphs can be divided into two major groups (Liu et
al., 2023):

Contrastive Methods. Contrastive learning methods usually consist of
minimizing an energy function (Hinton, 2002; Gutmann & Hyvärinen, 2010)
between different views of the same data. InfoGraph (Sun et al., 2019)
maximizes the mutual information between the graph-level representation and
the representations of substructures at different scales. GraphCL (You et al.,
2020) works similarly to distance-based contrastive methods in the imaging
domain. The authors first propose four types of graph augmentations and then
perform contrastive learning based on them. The work of Hassani & Khasahmadi
(2020) goes one step further by contrasting structural views of graphs. They
also show that a large number of views or multiscale training does not seem to
be beneficial, contrary to the image domain. Another popular research
direction for contrastive methods is learning graph augmentations (Suresh et
al., 2021) and AutoSSL (Jin et al., 2021), where the goal is to learn how to
automatically leverage multiple pretext tasks effectively. Contrastive
learning methods typically require a lot of memory due to data augmentation
and negative samples. Graph-JEPA is much more efficient than typical
formulations of these architectures given that it does not require any
augmentations or negative samples. Another major difference is that the
prediction in latent space in JEPAs is done through a separate predictor
network, rather than using the common Siamese structure (Bromley et al.,
1993)(Figure 1a vs c).

Generative Methods. The goal of generative models is to recover the data
distribution, an objective that is typically implemented through a
reconstruction process. In the graph domain, most generative architectures
that are also used for SSL are extensions of Auto-Encoder (AE) (Hinton &
Zemel, 1993) and Variational Auto-Encoder (VAE) (Kingma & Welling, 2013)
architectures. These models learn an embedding from the input data and then
use a reconstruction objective with (optional) regularization in order to
maximize the data evidence. Kipf & Welling (2016b) extended the framework of
AEs and VAEs to graphs by using a GNN as an encoder and the reconstruction of
the adjacency matrix as a training objective. However, the results on node and
graph classification benchmarks with these embeddings are often unsatisfactory
compared with contrastive learning methods, a tendency also observed in other
domains (Liu et al., 2023). A recent and promising direction is masked
autoencoding (MAE) (He et al., 2022), which has proved to be a very successful
framework for the image and text domains. GraphMAE (Hou et al., 2022) is an
instantiation of MAEs in the graph domain, where the node attributes are
perturbed and then reconstructed, providing a paradigm shift from the
structure learning objective of GAEs. S2GAE (Tan et al., 2023) is one of the
latest GAEs, which focuses on reconstructing the topological structure but
adds several auxiliary objectives and additional designs. Our architecture
differs from generative models in that it learns to predict directly in the
latent space, thereby bypassing the necessity of remembering and overfitting
high-level details that help maximize the data evidence (Figure 1b vs c).

###  2.2 Joint-Embedding Predictive Architectures

Joint-Embedding Predictive Architectures (LeCun, 2022) are a recently proposed
design for SSL. The idea is similar to both generative and contrastive
approaches, yet JEPAs are non-generative since they cannot directly predict
$y$ from $x$, as shown in Figure 1c. The energy of a JEPA is given by the
prediction error in the embedding space, not the input space. These models can
therefore be understood as a way to capture abstract dependencies between $x$
and $y$, potentially given another latent variable $z$. It is worth noting
that the different models comprising the architecture may differ in terms of
structure and parameters. An in-depth explanation of Joint-Embedding
Predictive Architectures and their connections to human representation
learning is provided by LeCun (2022). The most well-known works to use
particular instantiations of JEPAs before the term was officially coined are
BYOL (Grill et al., 2020) and SimSiam (Chen & He, 2021). There have recently
been two works that implement JEPAs in practice, both in the vision domain:
I-JEPA (Assran et al., 2023) and MC-JEPA (Bardes et al., 2023). I-JEPA
provides an instantiation in the context of images, where the goal is to
predict the latent embeddings of different image patches given a context
patch. MC-JEPA instantiates the architecture in the context of learning
optical flow and content features within a shared encoder. In this work, we
propose the first JEPA for the graph domain and use it to learn graph-level
representations in a self-supervised manner.

##  3 Method

Preliminary. Let a graph $G$ be defined as $G=(V,E)$ where $V=\\{v_{1}\ldots
v_{N}\\}$ is the set of nodes, with a cardinality $|V|=N$, and
$E=\\{e_{1}\ldots e_{M}\\}$ is the set of edges, with a cardinality $|E|=M$.
We consider symmetric, unweighted graphs, although our method can be
generalized to directed or weighted graphs. In this setting, $G$ can be
represented by an adjacency matrix $A\in\\{0,1\\}^{N\times N}$, with
$A_{ij}=1$ if nodes $v_{i}$ and $v_{j}$ are connected and $A_{ij}=0$
otherwise.

A general overview. Figure 2 provides an overview of the proposed
architecture. The high-level idea of Graph-JEPA is to first divide the input
graph into subgraphs (patches) (He et al., 2023) and then predict the
representation of a randomly chosen target subgraph from a context subgraph
(Assran et al., 2023). We would like to stress that this masked modeling
objective is realized in latent space, without the need for augmentations or
negative samples. The subgraph representations are then pooled to create a
vectorized graph-level representation. Therefore Graph-JEPA can be described
through a sequence of operations, namely: 1. Spatial Partitioning; 2. Subgraph
Embedding; 3. Context and Target Encoding; 4. Latent Target Prediction.

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
fed into their respective encoders (e.). The embeddings generated from the
target encoder are used to produce the target subgraphs’ coordinates
$\psi_{y}$. Finally, the predictor network is tasked with directly predicting
the coordinates $\hat{\psi}_{y}$ for each target subgraph based on the context
embedding and the positional embedding of each target subgraph (f.). A
regression loss acts as the energy function $D$ between the predicted and
target coordinates.

###  3.1 Spatial Partitioning

We base the initial part of the Graph-JEPA architecture on the recent work of
He et al. (2023), but similar ideas consisting of graph partitioning have been
proposed before for Graph SSL (Jin et al., 2020). This step consists of
creating different subgraphs (patches) of a graph, similar to how Vision
Transformers (ViT) (Dosovitskiy et al., 2020) operate on images. We rely on
the METIS (Karypis & Kumar, 1998) graph clustering algorithm, which partitions
a graph into a pre-defined, non-overlapping number of clusters $p$ (Figure
2a), such that the number of within-cluster links is much higher than between-
cluster links. Note that having non-overlapping subgraphs can be problematic
since edges can be lost in this procedure and it is possible to end up with
empty ”subgraphs”. We would also like to preserve the notion of locality in
each subgraph while relating it to others that are close in terms of shared
nodes. To create this overlap and avoid completely empty subgraphs, a one-hop
neighborhood expansion of all the nodes belonging to a subgraph is performed
(Figure 2b).

###  3.2 Subgraph Embedding

After partitioning the graph, we learn a representation for each subgraph
through a GNN (Figure 2c.). The specific choice of GNN architecture is
arbitrary and depends on what properties we wish to induce in the
representation. The learned node embeddings are mean pooled to create a vector
representation for each subgraph: $\\{h_{1}...h_{p}\\},h\in\mathbb{R}^{d}$.
Given that these embeddings will be used as context or target variables, it is
important to provide additional information regarding the subgraphs to help
guide the predictor network. Thus, we propose to use a positional embedding
for each subgraph, which is implemented as the maximum Random Walk Structural
Embedding (RWSE) of all the nodes in that subgraph. In this way, the position
is characterized in a global and consistent manner for each patch. Formally, a
RWSE (Dwivedi et al., 2021) for a node $v$ can be defined as:

$\displaystyle
P_{v}=(M_{ii},M^{2}_{ii},\ldots,M^{k}_{ii}),P_{v}\in\mathbb{R}^{k},~{}~{}~{}~{}~{}M^{k}=(D^{-1}A)^{k}$
(1)

where $M^{k}$ is the random-walk diffusion matrix of order $k$ and $i$ is the
index of node $v$ in the adjacency matrix. Therefore, $M^{k}_{ii}$ encodes the
probability of node $v$ landing to itself after a $k$-step random walk. Given
a subgraph $l$, let $V_{l}$ denote the set of nodes in $l$. The subgraph RWSE
is then defined as:

$\displaystyle P_{l}=\max\limits_{v\in V_{l}}P_{v}$ (2)

###  3.3 Context and Target Encoding

Given the subgraph representations and their respective positional embeddings,
we frame the Graph-JEPA prediction task in a similar manner to I-JEPA (Assran
et al., 2023). The goal of the network is to predict the latent embeddings of
randomly chosen target subgraphs, given one random context subgraph. The
prediction is conditioned on positional information regarding each subgraph.
At each training step, we choose one random subgraph as the context $x$ and
$m$ others as targets $Y=\\{y_{1},\ldots,y_{m}\\}$ (Figure 2d). These
subgraphs are processed by the context and target encoders (Figure 2e) which
are parametrized by Transformer encoder blocks (without self-attention for the
context) where normalization is applied at first (Xiong et al., 2020). The
target encoder uses Hadamard self-attention (He et al., 2023), but other
choices, such as the standard self-attention mechanism (Vaswani et al., 2017)
are perfectly viable. We can summarize this step as:

$z^{x}=E_{context}(x),z^{x}\in\mathbb{R}^{d},~{}~{}~{}Z^{y}=E_{target}(Y),Z^{y}\in\mathbb{R}^{m\times
d}$ (3)

At this stage, we could use the predictor network to directly predict $Z^{y}$
from $z^{x}$. This is the typical formulation of JEPAs, also followed by
Assran et al. (2023). We argue that learning how to organize concepts for
abstract objects such as graphs or networks directly in Euclidean space is
suboptimal. In the following section, we propose a simple trick to bypass this
problem using the encoding and prediction mechanisms in Graph-JEPA. The
following subsections and the ablations studies in Section 4.3 will provide
additional insights.

###  3.4 Latent Target Prediction

Different works in Neuroscience have shown the importance of learning
hierarchically consistent concepts (Deco et al., 2021), especially during
infancy and young age (Rosenberg & Feigenson, 2013). Networks in the real
world often conform to some concept of hierarchy (Moutsinas et al., 2021) and
this assumption is frequently used when learning graph-level representations
(Ying et al., 2018). Thus, we conjecture that Graph-JEPA should operate in a
hyperbolic space, where learned embeddings implicitly organize hierarchical
concepts(Nickel & Kiela, 2017; Zhao et al., 2023). This gives rise to another
issue: Hyperbolic (Poincarè) embeddings are known to have several tradeoffs
related to dimensionality (Sala et al., 2018; Guo et al., 2022), which
severely limits the expressive ability of the model. Given that graphs can
describe very abstract concepts, high expressivity in terms of model
parameters is preferred. In summary, we would ideally like to have a high-
dimensional latent code that has a concept of hyperbolicity built into it.

To achieve this, we think of the target embedding as a high-dimensional
representation of the hyperbolic angle, which allows us to describe each
target patch through its position in the 2D unit hyperbola. Formally, given a
target patch $l$, its embedding $Z^{y}_{l}$ and positional encoding $P_{l}$,
we first express the latent target as:

$\alpha^{y}_{l}=\frac{1}{N}\sum_{n=1}^{d}{Z^{y}_{l}}_{n},~{}~{}~{}\psi^{y}_{l}=\begin{pmatrix}\;cosh(\alpha^{y}_{l})\;\\\
\\\ \;sinh(\alpha^{y}_{l})\;\end{pmatrix}$ (4)

where $cosh(\cdot)$ and $sinh(\cdot)$ are the hyperbolic cosine and sine
functions respectively. The predictor module is then tasked with directly
locating the target in the unit hyperbola given the context embedding and the
target patch’s positional encoding.

$\hat{\psi}^{y}_{l}=MLP(LayerNorm(z^{x}+P_{l})),\hat{\psi}^{y}_{l}\in\mathbb{R}^{2}$
(5)

This allows us to frame the learning procedure as a simple regression problem
and the whole network can be trained end-to-end (Figure 2f). In practice, we
use the smooth L1 loss as the distance function, as it is less sensitive to
outliers compared to the L2 loss (Girshick, 2015):

$\displaystyle
L(y,\hat{y})=\frac{1}{N}\sum_{n=1}^{N}S_{n},~{}~{}~{}S_{n}=\begin{cases}0.5(y_{n}-\hat{y}_{n})^{2}/\beta,&\text{if
}|y-\hat{y}|<\beta\\\ |y-\hat{y}|-0.5\beta,&\text{otherwise}\end{cases}$ (6)

Thus, we are effectively measuring how far away the context and target patches
are in the unit hyperbola of the plane, but the targets are actually described
through a high dimensional latent code (Eq. 4). We explicitly show the
differences between this choice and using the Euclidean or Hyperbolic
distances as energy functions (in the latent space) in Section 4.3. Our
proposed pretraining objective forces the context encoder to understand the
differences in the hyperbolic angle between the target patches, which can be
thought of as establishing a hierarchy.

An asymmetric design of the predictor and target networks (in terms of
parameters) is used since it is reported to prevent representation collapse in
self-supervised contrastive techniques (Chen et al., 2020; Baevski et al.,
2022). We also utilize stop-gradient for the target encoder and update its
weights using an Exponential Moving Average of the context encoder, as done in
other works that utilize JEPAs (Assran et al., 2023; Grill et al., 2020; Chen
& He, 2021). In the next section, we provide insights into these choices and
why they are important. Finally, to produce the unique graph-level
representation when fine-tuning, we simply feed all the subgraphs through the
trained target encoder and then use mean pooling, obtaining a single feature
vector $z_{G}\in\mathbb{R}^{d}$ that represents the whole graph.

###  3.5 Why does Graph-JEPA work?

In this subsection, we aim to provide a discussion regarding the design
choices for Graph-JEPA and also why our formulation makes sense in the graph
domain. We make the following assumptions for our analysis: i) The predictor
network is linear; ii) We consider the encoded context features $z^{x}$ and
target coordinates $\psi^{y}_{l}$. Without loss of generality, we assume that
there is a one-to-one correspondence between context and target patches. (This
holds in practice due to Eq. 5; iii) The problem is treated as a least-squares
problem in finite-dimensional vector spaces over the field of reals
$\mathbb{R}$. Based on our assumptions, we can rewrite the context features as
$X\in\mathbb{R}^{n\times d}$, the target coordinates as
$Y\in\mathbb{R}^{n\times 2}$, and the weights of the linear model as
$W\in\mathbb{R}^{d\times 2}$. The objective of the predictor is:

$\operatorname*{arg\,min}\limits_{W}\left\|XW-Y\right\|^{2}$ (7)

where $\left\|.\right\|$ indicates the Frobenius norm. The solution to this
system can be given by the (multivariate) least squares estimator:

$W=(X^{T}X)^{-1}X^{T}Y$ (8)

By plugging Eq. 8 into Eq. 7 and rewriting the difference by factorizing $Y$
we get:

$\operatorname*{arg\,min}\limits_{W}\left\|(X(X^{T}X)^{-1}X^{T}-I_{n})Y\right\|^{2}$
(9)

Thus, to find an optimal (linear) predictor we must find the orthogonal
projection of $Y$ onto the orthogonal complement of a subspace of $Col(X)$.
This simply translates into finding the linear combination that is closest (in
terms of $\left\|\cdot\right\|^{2}$) to $Y$ in a particular subspace.
Similarly to what was shown in Richemond et al. (2023), we argue that this
provides a key intuition behind the JEPA design. The target encoder which
produces $Y$ must not share weights or be optimized with the context encoder.
If that were the case, the easiest solution to Eq. 9 would be using a vector
that is orthogonal to itself, i.e., the $\mathbf{0}$ vector, leading to
representation collapse. Note that in practice, even though the predictor
network doesn’t have to be linear, it would still fail to perfectly learn a
linear relationship (as shown above). This would render the architecture
suboptimal.

We now argue about the existence of hyperbolicity due to our latent prediction
task and the necessity of the asymmetric predictor. The predictor network is
linearly approximating the behavior of the unit hyperbola such that it best
matches with the generated target coordinates (Eq 6). Thus, the network is
actually trying to estimate a space that can be considered a particular
section of the hyperboloid model (Reynolds, 1993). In the hyperboloid model,
hyperbolas appear as hyperbolic geodesics. We are therefore evaluating our
energy in a (very) restricted part of hyperbolic space. As mentioned before,
we find this task to offer great flexibility as it is straightforward to
implement, works well in practice (Section 4.2), and is computationally
efficient compared to the hyperbolic distance used to typically learn
hyperbolic embeddings (Nickel & Kiela, 2017). In our current setting, this
approximation cannot perfectly capture the dynamics of the unit hyperbola,
given the linearity of the predictor. In practice, with a non-linear predictor
network, we might immediately run into a degenerate solution where the target
encoder and predictor output collapse representations by immediately
minimizing the training objective. It is therefore important to parametrize
the predictor using a simpler network that is less expressive but can capture
the correct dynamics over the training process.

##  4 Experiments

The experimental section introduces the empirical evaluation of the Graph-JEPA
model in terms of downstream performance on different graph datasets and
tasks, as well as various ablation studies that provide more insights into our
design choices.

###  4.1 Experimental setting

We use the TUD datasets (Morris et al., 2020) as commonly done for graph-level
SSL (Suresh et al., 2021; Tan et al., 2023). We utilize seven different graph-
classification datasets: PROTEINS, MUTAG, DD, REDDIT-BINARY, REDDIT-MULTI-5K,
IMDB-BINARY and IMDB-MULTI. For all classification experiments, we report the
accuracy over five runs (with different seeds) of ten-fold cross validation.
It is worth noting that we retrain the Graph-JEPA model for each fold, without
ever having access to the testing partition both in the pretraining and fine-
tuning stage. For graph regression, we use the ZINC dataset and report the
Mean Squared Error over ten runs (with different seeds), given that the
testing partition is already provided. More details regarding the training
procedure are available in Appendix A.1.

###  4.2 Downstream performance

For the experiments on downstream performance, we follow Suresh et al. (2021)
and also report the results of a fully supervised Graph Isomorphism Network
(GIN) (Xu et al., 2018), denoted F-GIN. As can be seen in Table 1, Graph-JEPA
achieves competitive results on all datasets, setting the state-of-the-art as
a pretrained backbone on five different datasets and coming second on one.
Overall, our proposed framework learns semantic embeddings that work well on
different types of graphs, showing that Graph-JEPA can indeed be utilized as a
general pretraining method for graph-level SSL. Notably, Graph-JEPA works well
for both classification and regression and performs better than F-GIN on all
classification datasets.

We further explore the performance of our model on the synthetic EXP dataset
(Abboud et al., 2020). The goal behind this experiment is to empirically
verify if Graph-JEPA can learn highly expressive graph representations, in
terms of graph isomorphisms. The results in Table 2 show that our model is
able to perform much better than commonly used GNNs. Given its local and
global exchange of information, this result is to be expected. Most
importantly, Graph-JEPA closely rivals the flawless performance of Graph-MLP-
Mixer (He et al., 2023), which is trained in a supervised manner.

Table 1: Performance of different graph SSL techniques ordered by pretraining
type. The results of the competitors are taken as the best values from
(Hassani & Khasahmadi, 2020; Suresh et al., 2021; Tan et al., 2023). ”-”
indicates missing values from the literature. The best results are reported in
boldface, the second best in italics.

Model | PROTEINS $\uparrow$ | MUTAG $\uparrow$ | DD $\uparrow$ | REDDIT-B $\uparrow$ | REDDIT-M5 $\uparrow$ | IMDB-B $\uparrow$ | IMDB-M $\uparrow$ | ZINC $\downarrow$
---|---|---|---|---|---|---|---|---
F-GIN (Xu et al., 2018) | 72.39 $\pm$ 2.76 | 90.41 $\pm$ 4.61 | 74.87 $\pm$ 3.56 | 86.79 $\pm$ 2.04 | 53.28 $\pm$ 3.17 | 71.83 $\pm$ 1.93 | 48.46 $\pm$ 2.31 | 0.254 $\pm$ 0.005
RU-GIN (Xu et al., 2018) | 69.03 $\pm$ 0.33 | 87.61 $\pm$ 0.39 | 74.22 $\pm$ 0.30 | 58.97 $\pm$ 0.13 | 27.52 $\pm$ 0.61 | 51.86 $\pm$ 0.33 | 32.81 $\pm$ 0.57 | 0.809 $\pm$ 0.022
InfoGraph (Sun et al., 2019) | 72.57 $\pm$ 0.65 | 87.71 $\pm$ 1.77 | 75.23 $\pm$ 0.39 | 78.79 $\pm$ 2.14 | 51.11 $\pm$ 0.55 | 71.11 $\pm$ 0.88 | 48.66 $\pm$ 0.67 | 0.890 $\pm$ 0.017
GraphCL (You et al., 2020) | 72.86 $\pm$ 1.01 | 88.29 $\pm$ 1.31 | 74.70 $\pm$ 0.70 | 82.63 $\pm$ 0.99 | 53.05 $\pm$ 0.40 | 70.80 $\pm$ 0.77 | 48.49 $\pm$ 0.63 | 0.627 $\pm$ 0.013
AD-GCL-FIX (Suresh et al., 2021) | 73.59 $\pm$ 0.65 | 89.25 $\pm$ 1.45 | 74.49 $\pm$ 0.52 | 85.52 $\pm$ 0.79 | 53.00 $\pm$ 0.82 | 71.57 $\pm$ 1.01 | 49.04 $\pm$ 0.53 | 0.578 $\pm$ 0.012
AD-GCL-OPT (Suresh et al., 2021) | 73.81 $\pm$ 0.46 |  89.70 $\pm$ 1.03 |  75.10 $\pm$ 0.39 | 85.52 $\pm$ 0.79 |  54.93 $\pm$ 0.43 | 72.33 $\pm$ 0.56 | 49.89 $\pm$ 0.66 |  0.544 $\pm$ 0.004
MVGRL (Hassani & Khasahmadi, 2020) | - | - | - | 84.5 $\pm$ 0.6 | - | 74.2 $\pm$ 0.7 | 51.2 $\pm$ 0.5 | -
GraphMAE (Hou et al., 2022) | 75.30 $\pm$ 0.39 | 88.19 $\pm$ 1.26 | - |  88.01 $\pm$ 0.19 | - |  75.52 $\pm$ 0.66 |  51.63 $\pm$ 0.52 | -
S2GAE (Tan et al., 2023) |  76.37 $\pm$ 0.43 | 88.26 $\pm$ 0.76 | - | 87.83 $\pm$ 0.27 | - |  75.76 $\pm$ 0.62 |  51.79 $\pm$ 0.36 | -
Graph-JEPA |  75.67 $\pm$ 3.78 |  91.25 $\pm$ 5.75 |  78.64 $\pm$ 2.35 |  91.99 $\pm$ 1.59 |  56.73 $\pm$ 1.96 | 73.68 $\pm$ 3.24 | 50.69 $\pm$ 2.91 |  0.434 $\pm$ 0.014

Table 2: Classification accuracy on the synthetic EXP dataset (Abboud et al.,
2020), which contains 600 pairs of non-isomorphic graphs that are
indistinguishable by the 1-WL test. The competitor models are proposed in
(Kipf & Welling, 2016a; Bresson & Laurent, 2017; Xu et al., 2018; Dwivedi &
Bresson, 2020; He et al., 2023)

.  Model GCN GatedGCN GINE GraphTransformer Graph-MLP-Mixer Graph-JEPA
Accuracy 51.90 $\pm$ 1.96 51.73 $\pm$ 1.65 50.69 $\pm$ 1.39 52.35 $\pm$ 2.32
100.00 $\pm$ 0.00 98.77 $\pm$ 0.99

###  4.3 Exploring the Graph-JEPA latent space

As discussed in Section 3.4, the choice of energy function has a big impact on
the learned representations. We argued against the use of simple Euclidean or
Poincarè embeddings previously and now provide empirical evidence for this
conjecture in Table 3. Our results reveal that learning the distance between
patches in the 2D unit hyperbola provides a simple way to get the advantages
of both embedding types. Hyperbolic embeddings must be learned in lower
dimensions due to stability issues (Yu & De Sa, 2021), while Euclidean ones do
not properly reflect the dependencies between subgraphs. Our results suggest
that the hyperbolic distance seems to be a better choice than the Euclidean
distance in lower dimensions, but it is very unstable and expensive
computationally in high dimensions. The proposed approach provides the best
overall results, with the Euclidean distance performing negligibly better in a
single case. We provide a qualitative example that describes how the embedding
space is altered when using our proposed latent objective in Appendix A.3.

Table 3: Comparison of Graph-JEPA performance for different distance
functions. We used a latent space of dimension 128 for Poincarè embeddings due
to optimization issues, which was still problematic for the IMDB-B dataset. LD
stands for Lower Dimension, where we use a much smaller embedding size (32).

Distance function | Ours | Euclidean | Hyperbolic | Euclidean (LD) | Hyperbolic (LD)
---|---|---|---|---|---
MUTAG | 91.25 $\pm$ 5.75 | 87.04 $\pm$ 6.01 | 89.43 +- 5.67 | 86.63 $\pm$ 5.9 | 86.32 $\pm$ 5.52
REDDIT-M | 56.73 $\pm$ 1.96 | 56.55 $\pm$ 1.94 | 56.19 +- 1.95 | 54.84 $\pm$ 1.6 | 55.07 $\pm$ 1.83
IMDB-B | 73.68 $\pm$ 3.24 | 73.76 $\pm$ 3.46 | NaN | 72.5 $\pm$ 3.97 | 73.4 $\pm$ 4.07
ZINC | 0.434 $\pm$ 0.01 | 0.471 $\pm$ 0.01 | 0.605 +- 0.01 | 0.952 $\pm$ 0.05 | 0.912 $\pm$ 0.04

###  4.4 Ablation studies

In the following, we present various ablation studies to get a better
understanding of the different design choices of Graph-JEPA. For all
ablations, we consider 4 out of the 8 datasets initially presented in Table 1,
making sure to choose different graph types for a fair comparison. All results
are summarized in Figure 3, but we also provide the results in tabular format
in Appendix A.2 for additional clarity.

Positional embedding. Following He et al. (2023), it is possible to use the
RWSE of the patches as conditioning information. Formally, let
$B\in\\{0,1\\}^{p\times N}$ be the patch assignment matrix, such that
$B_{ij}=1$ if $v_{j}\in p_{i}$. We can calculate a coarse patch adjacency
matrix $A^{\prime}=BB^{T}\in\mathbb{R}^{p\times p}$, where each
$A_{ij}^{\prime}$ contains the node overlap between $p_{i}$ and $p_{j}$. The
RWSE can be calculated as described in Eq. 1 (considering $A^{\prime}$). We
test Graph-JEPA with these relative positional embeddings and find that they
still provide good performance, but consistently fall behind the node-level
(global) RWSE that we employ in our formulation (Figure 3a). Indeed, a great
issue for the relative embeddings is that the number of shared neighbors
obscures the local peculiarities of each patch and this adds more variance to
the results.

Using different Self-Attention mechanisms. As stated in Section 3.3, Graph-
JEPA uses Hadamard self-attention (He et al., 2023), which provides a strong
inductive bias for graphs. It is possible to make no assumptions about the
data and render the prediction task more challenging, by using the standard
self-attention mechanism (Vaswani et al., 2017). We show the results of this
change in Figure 3b. The results reveal that the performance is slightly
better but more unstable, i.e., additional variance in the results. This is to
be expected since, as mentioned previously, using the standard self-attention
mechanism emphasizes a lack of inductive bias regarding the input data.

Random subgraphs. A natural question that arises in our framework is how to
design the spatial partitioning procedure. Using a structured approach like
METIS is intuitive and leads to favorable results. Another option would be to
simply extract random, non-empty subgraphs as context and target. As can be
seen in Figure 3c, the random patches provide strong performance as well but
do not show any improvement. Even though our results show that it might not be
necessary to use a structured way to extract the patches, not all graphs are
equally easy to sample randomly from. Thus, we advocate extracting subgraphs
with METIS as it is a safer option in terms of generalizability across
different graphs and what inductive biases it provides.

Figure 3: Ablation studies: a. Performance when using absolute (node-level) vs
relative (patch-level) RWSEs. b. Results when using the standard self-
attention mechanism vs a graph-specific one. c. Performance when using METIS
subgraphs vs using random subgraphs. We remind the reader that the MAE is
reported on ZINC and Accuracy on the other datasets. The error bars represent
the average standard deviation values over the runs in order to report results
similarly to the tables in previous sections. Best viewed in color.

##  5 Conclusion

In this work, we introduce the first Joint Embedding Predictive Architecture
(JEPA) for graph-level Self-Supervised Learning (SSL). A proper design of the
model both in terms of data preparation and pretraining objective reveals that
it is possible for a neural network model to self-organize the semantic
knowledge embedded in a graph, demonstrating competitive performance in both
graph classification and regression. Future research directions include
extending the method to node and edge-level learning, exploring the
expressiveness of Graph-JEPA, and gaining insights into the optimal geometry
of the embedding space for graph SSL.

## References

  * Abboud et al. (2020) Ralph Abboud, Ismail Ilkan Ceylan, Martin Grohe, and Thomas Lukasiewicz.  The surprising power of graph neural networks with random node initialization.  _arXiv preprint arXiv:2010.01179_ , 2020.
  * Assran et al. (2023) Mahmoud Assran, Quentin Duval, Ishan Misra, Piotr Bojanowski, Pascal Vincent, Michael Rabbat, Yann LeCun, and Nicolas Ballas.  Self-supervised learning from images with a joint-embedding predictive architecture.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 15619–15629, 2023.
  * Baevski et al. (2022) Alexei Baevski, Wei-Ning Hsu, Qiantong Xu, Arun Babu, Jiatao Gu, and Michael Auli.  Data2vec: A general framework for self-supervised learning in speech, vision and language.  In _International Conference on Machine Learning_ , pp. 1298–1312. PMLR, 2022.
  * Bardes et al. (2023) Adrien Bardes, Jean Ponce, and Yann LeCun.  Mc-jepa: A joint-embedding predictive architecture for self-supervised learning of motion and content features.  _arXiv preprint arXiv:2307.12698_ , 2023.
  * Bengio et al. (2013) Yoshua Bengio, Aaron Courville, and Pascal Vincent.  Representation learning: A review and new perspectives.  _IEEE transactions on pattern analysis and machine intelligence_ , 35(8):1798–1828, 2013.
  * Bresson & Laurent (2017) Xavier Bresson and Thomas Laurent.  Residual gated graph convnets.  _arXiv preprint arXiv:1711.07553_ , 2017.
  * Bromley et al. (1993) Jane Bromley, Isabelle Guyon, Yann LeCun, Eduard Säckinger, and Roopak Shah.  Signature verification using a” siamese” time delay neural network.  _Advances in neural information processing systems_ , 6, 1993.
  * Brown et al. (2020) Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al.  Language models are few-shot learners.  _Advances in neural information processing systems_ , 33:1877–1901, 2020.
  * Chen et al. (2020) Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton.  A simple framework for contrastive learning of visual representations.  In _International conference on machine learning_ , pp. 1597–1607. PMLR, 2020.
  * Chen & He (2021) Xinlei Chen and Kaiming He.  Exploring simple siamese representation learning.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , pp. 15750–15758, 2021.
  * Deco et al. (2021) G. Deco, D. Vidaurre, and M. Kringelbach.  Revisiting the global workspace orchestrating the hierarchical organization of the human brain.  _Nature Human Behaviour_ , 5:497 – 511, 2021.  doi: 10.1038/s41562-020-01003-6.
  * Dosovitskiy et al. (2020) Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al.  An image is worth 16x16 words: Transformers for image recognition at scale.  _arXiv preprint arXiv:2010.11929_ , 2020.
  * Dwivedi & Bresson (2020) Vijay Prakash Dwivedi and Xavier Bresson.  A generalization of transformer networks to graphs.  _arXiv preprint arXiv:2012.09699_ , 2020.
  * Dwivedi et al. (2021) Vijay Prakash Dwivedi, Anh Tuan Luu, Thomas Laurent, Yoshua Bengio, and Xavier Bresson.  Graph neural networks with learnable structural and positional representations.  _arXiv preprint arXiv:2110.07875_ , 2021.
  * Gilmer et al. (2017) Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl.  Neural message passing for quantum chemistry.  In _International conference on machine learning_ , pp. 1263–1272. PMLR, 2017.
  * Girshick (2015) Ross Girshick.  Fast r-cnn.  In _Proceedings of the IEEE international conference on computer vision_ , pp. 1440–1448, 2015.
  * Grill et al. (2020) Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al.  Bootstrap your own latent-a new approach to self-supervised learning.  _Advances in neural information processing systems_ , 33:21271–21284, 2020.
  * Guo et al. (2022) Yunhui Guo, Haoran Guo, and Stella X Yu.  Co-sne: Dimensionality reduction and visualization for hyperbolic data.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 21–30, 2022.
  * Gutmann & Hyvärinen (2010) Michael Gutmann and Aapo Hyvärinen.  Noise-contrastive estimation: A new estimation principle for unnormalized statistical models.  In _Proceedings of the thirteenth international conference on artificial intelligence and statistics_ , pp. 297–304. JMLR Workshop and Conference Proceedings, 2010.
  * Hamilton et al. (2017) Will Hamilton, Zhitao Ying, and Jure Leskovec.  Inductive representation learning on large graphs.  _Advances in neural information processing systems_ , 30, 2017.
  * Hassani & Khasahmadi (2020) Kaveh Hassani and Amir Hosein Khasahmadi.  Contrastive multi-view representation learning on graphs.  In _International conference on machine learning_ , pp. 4116–4126. PMLR, 2020.
  * He et al. (2022) Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick.  Masked autoencoders are scalable vision learners.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , pp. 16000–16009, 2022.
  * He et al. (2023) Xiaoxin He, Bryan Hooi, Thomas Laurent, Adam Perold, Yann LeCun, and Xavier Bresson.  A generalization of vit/mlp-mixer to graphs.  In _International Conference on Machine Learning_ , pp. 12724–12745. PMLR, 2023.
  * Hinton (2002) Geoffrey E Hinton.  Training products of experts by minimizing contrastive divergence.  _Neural computation_ , 14(8):1771–1800, 2002.
  * Hinton & Zemel (1993) Geoffrey E Hinton and Richard Zemel.  Autoencoders, minimum description length and helmholtz free energy.  _Advances in neural information processing systems_ , 6, 1993.
  * Hou et al. (2022) Zhenyu Hou, Xiao Liu, Yukuo Cen, Yuxiao Dong, Hongxia Yang, Chunjie Wang, and Jie Tang.  Graphmae: Self-supervised masked graph autoencoders.  In _Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining_ , pp. 594–604, 2022.
  * Hu et al. (2019) Weihua Hu, Bowen Liu, Joseph Gomes, Marinka Zitnik, Percy Liang, Vijay Pande, and Jure Leskovec.  Strategies for pre-training graph neural networks.  In _International Conference on Learning Representations_ , 2019.
  * Hu et al. (2020) Ziniu Hu, Yuxiao Dong, Kuansan Wang, Kai-Wei Chang, and Yizhou Sun.  Gpt-gnn: Generative pre-training of graph neural networks.  In _Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining_, pp. 1857–1867, 2020.
  * Jin et al. (2020) Wei Jin, Tyler Derr, Haochen Liu, Yiqi Wang, Suhang Wang, Zitao Liu, and Jiliang Tang.  Self-supervised learning on graphs: Deep insights and new direction.  _arXiv preprint arXiv:2006.10141_ , 2020.
  * Jin et al. (2021) Wei Jin, Xiaorui Liu, Xiangyu Zhao, Yao Ma, Neil Shah, and Jiliang Tang.  Automated self-supervised learning for graphs.  In _International Conference on Learning Representations_ , 2021.
  * Karypis & Kumar (1998) George Karypis and Vipin Kumar.  A fast and high quality multilevel scheme for partitioning irregular graphs.  _SIAM Journal on scientific Computing_ , 20(1):359–392, 1998.
  * Kingma & Ba (2014) Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Kingma & Welling (2013) Diederik P Kingma and Max Welling.  Auto-encoding variational bayes.  _arXiv preprint arXiv:1312.6114_ , 2013.
  * Kipf & Welling (2016a) Thomas N Kipf and Max Welling.  Semi-supervised classification with graph convolutional networks.  _arXiv preprint arXiv:1609.02907_ , 2016a.
  * Kipf & Welling (2016b) Thomas N Kipf and Max Welling.  Variational graph auto-encoders.  _arXiv preprint arXiv:1611.07308_ , 2016b.
  * LeCun (2022) Yann LeCun.  A path towards autonomous machine intelligence version 0.9. 2, 2022-06-27.  _Open Review_ , 62, 2022.
  * Liu et al. (2023) Xiao Liu, Fanjin Zhang, Zhenyu Hou, Li Mian, Zhaoyu Wang, Jing Zhang, and Jie Tang.  Self-supervised learning: Generative or contrastive.  _IEEE Transactions on Knowledge and Data Engineering_ , 35(1):857–876, 2023.  doi: 10.1109/TKDE.2021.3090866.
  * Ma & Tang (2021) Yao Ma and Jiliang Tang.  _Deep Learning on Graphs_.  Cambridge University Press, 2021.  doi: 10.1017/9781108924184.
  * Morris et al. (2020) Christopher Morris, Nils M. Kriege, Franka Bause, Kristian Kersting, Petra Mutzel, and Marion Neumann.  Tudataset: A collection of benchmark datasets for learning with graphs.  In _ICML 2020 Workshop on Graph Representation Learning and Beyond (GRL+ 2020)_ , 2020.  URL www.graphlearning.io.
  * Moutsinas et al. (2021) Giannis Moutsinas, Choudhry Shuaib, Weisi Guo, and Stephen Jarvis.  Graph hierarchy: a novel framework to analyse hierarchical structures in complex networks.  _Scientific Reports_ , 11(1):13943, 2021.
  * Nickel & Kiela (2017) Maximillian Nickel and Douwe Kiela.  Poincaré embeddings for learning hierarchical representations.  _Advances in neural information processing systems_ , 30, 2017.
  * Reynolds (1993) William F. Reynolds.  Hyperbolic geometry on a hyperboloid.  _The American Mathematical Monthly_ , 100(5):442–455, 1993.  ISSN 00029890, 19300972.  URL http://www.jstor.org/stable/2324297.
  * Richemond et al. (2023) Pierre Harvey Richemond, Allison Tam, Yunhao Tang, Florian Strub, Bilal Piot, and Felix Hill.  The edge of orthogonality: A simple view of what makes byol tick.  In _International Conference on Machine Learning_ , pp. 29063–29081. PMLR, 2023.
  * Rosenberg & Feigenson (2013) R. Rosenberg and L. Feigenson.  Infants hierarchically organize memory representations.  _Developmental science_ , 16 4:610–21, 2013.  doi: 10.1111/desc.12055.
  * Sala et al. (2018) Frederic Sala, Chris De Sa, Albert Gu, and Christopher Re.  Representation tradeoffs for hyperbolic embeddings.  In Jennifer Dy and Andreas Krause (eds.), _Proceedings of the 35th International Conference on Machine Learning_ , volume 80 of _Proceedings of Machine Learning Research_ , pp. 4460–4469. PMLR, 10–15 Jul 2018.
  * Scarselli et al. (2008) Franco Scarselli, Marco Gori, Ah Chung Tsoi, Markus Hagenbuchner, and Gabriele Monfardini.  The graph neural network model.  _IEEE transactions on neural networks_ , 20(1):61–80, 2008.
  * Sun et al. (2019) Fan-Yun Sun, Jordan Hoffman, Vikas Verma, and Jian Tang.  Infograph: Unsupervised and semi-supervised graph-level representation learning via mutual information maximization.  In _International Conference on Learning Representations_ , 2019.
  * Suresh et al. (2021) Susheel Suresh, Pan Li, Cong Hao, and Jennifer Neville.  Adversarial graph augmentation to improve graph contrastive learning.  _Advances in Neural Information Processing Systems_ , 34:15920–15933, 2021.
  * Tan et al. (2023) Qiaoyu Tan, Ninghao Liu, Xiao Huang, Soo-Hyun Choi, Li Li, Rui Chen, and Xia Hu.  S2gae: Self-supervised graph autoencoders are generalizable learners with graph masking.  In _Proceedings of the Sixteenth ACM International Conference on Web Search and Data Mining_ , pp. 787–795, 2023.
  * Tian et al. (2020) Yonglong Tian, Chen Sun, Ben Poole, Dilip Krishnan, Cordelia Schmid, and Phillip Isola.  What makes for good views for contrastive learning?  _Advances in neural information processing systems_ , 33:6827–6839, 2020.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  _Advances in neural information processing systems_ , 30, 2017.
  * Veličković (2023) Petar Veličković.  Everything is connected: Graph neural networks.  _Current Opinion in Structural Biology_ , 79:102538, 2023.
  * Veličković et al. (2017) Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, and Yoshua Bengio.  Graph attention networks.  _arXiv preprint arXiv:1710.10903_ , 2017.
  * Vincent et al. (2010) Pascal Vincent, Hugo Larochelle, Isabelle Lajoie, Yoshua Bengio, Pierre-Antoine Manzagol, and Léon Bottou.  Stacked denoising autoencoders: Learning useful representations in a deep network with a local denoising criterion.  _Journal of machine learning research_ , 11(12), 2010.
  * Wu et al. (2019) Felix Wu, Amauri Souza, Tianyi Zhang, Christopher Fifty, Tao Yu, and Kilian Weinberger.  Simplifying graph convolutional networks.  In _International conference on machine learning_ , pp. 6861–6871. PMLR, 2019.
  * Xiong et al. (2020) Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu.  On layer normalization in the transformer architecture.  In _International Conference on Machine Learning_ , pp. 10524–10533. PMLR, 2020.
  * Xu et al. (2018) Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka.  How powerful are graph neural networks?  In _International Conference on Learning Representations_ , 2018.
  * Ying et al. (2018) Zhitao Ying, Jiaxuan You, Christopher Morris, Xiang Ren, Will Hamilton, and Jure Leskovec.  Hierarchical graph representation learning with differentiable pooling.  _Advances in neural information processing systems_ , 31, 2018.
  * You et al. (2020) Yuning You, Tianlong Chen, Yongduo Sui, Ting Chen, Zhangyang Wang, and Yang Shen.  Graph contrastive learning with augmentations.  _Advances in neural information processing systems_ , 33:5812–5823, 2020.
  * Yu & De Sa (2021) Tao Yu and Christopher M De Sa.  Representing hyperbolic space accurately using multi-component floats.  In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), _Advances in Neural Information Processing Systems_ , volume 34, pp. 15570–15581. Curran Associates, Inc., 2021.
  * Zhao et al. (2023) Wei Zhao, Federico Lopez, Maxwell J Riestenberg, Michael Strube, Diaaeldin Taha, and Steve Trettel.  Modeling graphs beyond hyperbolic: Graph neural networks in symmetric positive definite matrices.  In _Joint European Conference on Machine Learning and Knowledge Discovery in Databases_ , pp. 122–139. Springer, 2023.

##  Appendix A Appendix

###  A.1 Experimental setting

For the datasets that do not natively have node and edge features, we use a
simple constant (0) initialization. The subgraph embedding GNN (Figure 2c)
consists of the GIN operator with support for edge features (Hu et al., 2019)
a.k.a GINE. For fine-tuning, we feed the high-dimensional representation of
the hyperbolic angle produced by the target encoder to a linear model with L2
regularization. For the classification datasets, that corresponds to a
Logistic Regression with L2 regularization, while for ZINC it corresponds to a
Ridge Regression. We also provide details regarding the hyperparameters of the
JEPA design for Graph-JEPA in Table 4. The neural network is trained using the
Adam optimizer (Kingma & Ba, 2014). All experiments were performed on a single
Nvidia RTX 3090 GPU.

Table 4: Hyperparameters regarding the Graph-JEPA design for the TUD datasets.

Hyperparameter | PROTEINS | MUTAG | DD | REDDIT-B | REDDIT-M5 | IMDB-B | IMDB-M | ZINC
---|---|---|---|---|---|---|---|---
# Subgraphs | 32 | 32 | 32 | 128 | 128 | 32 | 32 | 32
# GNN Layers | 2 | 2 | 3 | 2 | 2 | 2 | 2 | 2
# Encoder Blocks | 4 | 4 | 4 | 4 | 4 | 4 | 4 | 4
Embedding size | 512 | 512 | 512 | 512 | 512 | 512 | 512 | 512
RWSE size | 20 | 15 | 30 | 40 | 40 | 15 | 15 | 20
# context - # target | 1 - 2 | 1 - 3 | 1 - 4 | 1 - 4 | 1 - 4 | 1- 4 | 1- 4 | 1- 4

Table 5: Performance when using absolute (node-level) vs relative (patch-level) RWSEs. Patch PE | Node-level RWSE | Patch-level RWSE
---|---|---
MUTAG | 91.25 $\pm$ 5.75 | 91.23 $\pm$ 5.86
REDDIT-M | 56.73 $\pm$ 1.96 | 56.01 $\pm$ 2.1
IMDB-B | 73.68 $\pm$ 3.24 | 73.58 $\pm$ 4.47
ZINC | 0.434 $\pm$ 0.01 | 0.505 $\pm$ 0.005
Table 6: Performance when using METIS subgraphs vs. using random subgraphs. Partitioning method | Metis | Random
---|---|---
MUTAG | 91.25 $\pm$ 5.75 | 91.58 $\pm$ 5.82
REDDIT-M | 56.73 $\pm$ 1.96 | 56.08 $\pm$ 1.69
IMDB-B | 73.68 $\pm$ 3.24 | 73.52 $\pm$ 3.08
ZINC | 0.434 $\pm$ 0.01 | 0.43 $\pm$ 0.01
Table 7: Results when using the standard self-attention mechanism vs. a graph-specific one. Attention mechanism | Standard | Hadamard
---|---|---
MUTAG | 91.46 $\pm$ 6.1 | 91.25 $\pm$ 5.75
REDDIT-M | 56.92 $\pm$ 2.09 | 56.73 $\pm$ 1.96
IMDB-B | 73.92 $\pm$ 3.85 | 73.68 $\pm$ 3.24
ZINC | 0.465 $\pm$ 0.01 | 0.434 $\pm$ 0.01

###  A.2 Ablation results

Tables 5, 6, and 7 present the results from Figure 3 in tabular format,
providing additional clarity.

###  A.3 Visualizations

In this section, we provide two qualitative examples that give a more visual
understanding of how and what Graph-JEPA learns. Firstly, we plot a low-
dimensional representation of the high-dimensional latent code obtained from
the target encoder when using the Euclidean objective vs using our objective
in Figure 4. The modified prediction task greatly alters the geometry and
distribution of the latent codes, avoiding the formation of a blob-like
structure like in the Euclidean distance task.

Figure 5 depicts how the Graph-JEPA prediction task is framed. In 5a we can
see a particular molecule from the ZINC dataset and the corresponding context
and target nodes. The target encoder learns a high-dimensional embedding which
then can represent the nodes in the 2D unit hyperbola, as shown in 5b. The
predictor network tries to guess these coordinates and this training procedure
gives rise to the representations learned in Figure 4b.

Figure 4: 3D representation (via t-SNE) of the latent codes used to finetune
the linear classifier for the DD dataset. Best viewed in color.

Figure 5: Visualization of a partitioned graph from the ZINC dataset and its
corresponding embedding in the 2D unit hyperbola, as detailed in Eq. 4


d^{3}\mbox{\boldmath$x$}\,\,e^{i\mbox{\boldmath$k$}\cdot\mbox{\boldmath$x$}}\langle
u_{z}^{2}\delta\delta^{\prime}\rangle-2\int
d^{3}\mbox{\boldmath$x$}\,\,e^{i\mbox{\boldmath$k$}\cdot\mbox{\boldmath$x$}}\langle
u_{z}\delta u_{z}^{\prime}\delta^{\prime}\rangle\right.$
$\displaystyle\left.+\int
d^{3}\mbox{\boldmath$x$}\,\,e^{i\mbox{\boldmath$k$}\cdot\mbox{\boldmath$x$}}\langle\delta
u_{z}^{\prime 2}\delta^{\prime}\rangle\right]$ $\displaystyle
M_{b}=\frac{1}{2}j^{2}_{1}\left[\int
d^{3}\mbox{\boldmath$x$}\,\,e^{i\mbox{\boldmath$k$}\cdot\mbox{\boldmath$x$}}\langle
u_{z}^{2}\delta\nabla_{z}u_{z}^{\prime}\rangle+\int
d^{3}\mbox{\boldmath$x$}\,\,e^{i\mbox{\boldmath$k$}\cdot\mbox{\boldmath$x$}}\langle
u_{z}^{2}\nabla_{z}u_{z}\delta^{\prime}\rangle\right.$ $\displaystyle-2\int
d^{3}\mbox{\boldmath$x$}\,\,e^{i\mbox{\boldmath$k$}\cdot\mbox{\boldmath$x$}}\langle
u_{z}\delta u_{z}^{\prime}\nabla_{z}u_{z}^{\prime}\rangle-2\int
d^{3}\mbox{\boldmath$x$}\,\,e^{i\mbox{\boldmath$k$}\cdot\mbox{\boldmath$x$}}\langle
u_{z}\nabla_{z}u_{z}u_{z}^{\prime}\delta^{\prime}\rangle$
$\displaystyle\left.+\int
d^{3}\mbox{\boldmath$x$}\,\,e^{i\mbox{\boldmath$k$}\cdot\mbox{\boldmath$x$}}\langle\delta
u_{z}^{\prime 2}\nabla_{z}u_{z}^{\prime}\rangle+\int
d^{3}\mbox{\boldmath$x$}\,\,e^{i\mbox{\boldmath$k$}\cdot\mbox{\boldmath$x$}}\langle\nabla_{z}u_{z}u_{z}^{\prime
2}\delta^{\prime}\rangle\right]$ $\displaystyle
M_{L}=\frac{1}{2}j^{2}_{1}\left[\int
d^{3}\mbox{\boldmath$x$}\,\,e^{i\mbox{\boldmath$k$}\cdot\mbox{\boldmath$x$}}\langle
u_{z}^{2}\nabla_{z}u_{z}\nabla_{z}u_{z}^{\prime}\rangle\right.$
$\displaystyle\left.-2\int
d^{3}\mbox{\boldmath$x$}\,\,e^{i\mbox{\boldmath$k$}\cdot\mbox{\boldmath$x$}}\langle
u_{z}\nabla_{z}u_{z}u_{z}^{\prime}\nabla_{z}u_{z}^{\prime}\rangle+\int
d^{3}\mbox{\boldmath$x$}\,\,e^{i\mbox{\boldmath$k$}\cdot\mbox{\boldmath$x$}}\langle\nabla_{z}u_{z}u_{z}^{\prime
2}\nabla_{z}u_{z}^{\prime}\rangle\right].$

In this work, we measure these terms from individual N-body simulations. The
measured statistics are then utilised to train individual emulators for
$A_{b^{2}}$, $A_{b}$, $A_{L}$, $M_{b^{2}}$, $M_{b}$, and $M_{L}$ respectively.
For each of these six terms, we build one emulator at all five of our
redshifts of interest. It is important to highlight that if we set the linear
bias $b_{1}$ to one, then $A=A_{g}$ and $M=M_{g}$. Therefore, our emulators
can be used to compute the higher-order corrections for the underlying matter
distribution: $A$ and $M$. This methodology was followed to generate the
results analysed in Section 3.1.3.

We now introduce the scaling relations we utilise to predict the values of
higher-order terms in an alternative cosmological models with different values
of the scale independent parameters. As have stated before, these
relationships allow us to calculate the higher-order corrections in a new
cosmology based on their values in a reference cosmology, as well as their
dependence on the variables $G_{\delta}$ and $G_{\Theta}$. The underlying
proposition is that the temporal evolution of each term is primarily
determined by the growth factor dependence of $u_{z}$ and $\delta$ at the
leading order. Using our standard notation where bared quantities correspond
to a fiducial cosmology, we can write the equations to compute higher order
corrections in a new cosmology as:

$\displaystyle
A_{g}(k,\mu)=j_{1}[b_{1}^{2}\left(G_{\delta}/\bar{G}_{\delta}\right)^{2}\left(G_{\Theta}/\bar{G}_{\Theta}\right)\bar{A}_{b^{2}}$
$\displaystyle+b_{1}\left(G_{\delta}/\bar{G}_{\delta}\right)\left(G_{\Theta}/\bar{G}_{\Theta}\right)^{2}\bar{A}_{b}$
$\displaystyle+\left(G_{\Theta}/\bar{G}_{\Theta}\right)^{3}\bar{A}_{L}],$
$\displaystyle
M_{g}(k,\mu)=\frac{1}{2}j_{1}^{2}[b_{1}^{2}\left(G_{\delta}/\bar{G}_{\delta}\right)^{2}\left(G_{\Theta}/\bar{G}_{\Theta}\right)^{2}\bar{M}_{b^{2}}$
$\displaystyle+b_{1}\left(G_{\delta}/\bar{G}_{\delta}\right)\left(G_{\Theta}/\bar{G}_{\Theta}\right)^{3}\bar{M}_{b}$
$\displaystyle+\left(G_{\Theta}/\bar{G}_{\Theta}\right)^{4}\bar{M}_{L}].$ (50)