The quest for the GRAph Level autoEncoder (GRALE)

Although graph-based learning has attracted a lot of attention, graph representation learning is still a challenging task whose resolution may impact key application fields such as chemistry or biology. To this end, we introduce GRALE, a novel graph autoencoder that encodes and decodes graphs of varying sizes into a shared embedding space. GRALE is trained using an Optimal Transport-inspired loss that compares the original and reconstructed graphs and leverages a differentiable node matching module, which is trained jointly with the encoder and decoder. The proposed attention-based architecture relies on Evoformer, the core component of AlphaFold, which we extend to support both graph encoding and decoding. We show, in numerical experiments on simulated and molecular data, that GRALE enables a highly general form of pre-training, applicable to a wide range of downstream tasks, from classification and regression to more complex tasks such as graph interpolation, editing, matching, and prediction. 1 f g NODE-LEVEL ! GRALE ! f g GRAPH-LEVEL (NAIVE) Graph Matching NP (exact) or P (approx.) Graph Reconstruction Loss Node-level tasks No Graph Matching Graph-level tasks Graph Matching Edge Reconstruction Loss Graph Reconstruction Loss ! f g m Graph-level tasks No Graph Matching Figure 1: The different classes of graph AutoEncoders.

(Left) Node-level AutoEncoders such as [36] provide node level embeddings. (Center) Naive graph-level AutoEncoders such as [53] directly provide graph-level embeddings but rely on a graph matching solver to compute the training loss. (Right) Matching free approaches, such as proposed in this work and in [68] use a learnable module to provide the matching. convolutional network (GCN) that returns node embeddings z i , and the decoder reconstructs the adjacency matrix from a dot product A i,j = σ(⟨z i , z j ⟩). Many extensions have been proposed for this model, such as adversarial regularization [48] or masking [56], and it has been shown to be efficient for many node-level tasks, such as node clustering [81, 82, 60] or node outlier detection [14]. Nodelevel models can also be used for graph generation, given that one knows the size (number of nodes) of the graph to generate; this includes GVAE [36] but also GraphGAN [61] and graph normalizing flows [44]. When a graph-level representation is needed, one can apply some pooling operation on the node embeddings, for instance z = z i . Yet, this strategy is not entirely satisfying, as it inevitably leads to information loss, and it becomes difficult to decode a graph back from this representation [59, 68]. ...to graph-level AutoEncoders. In contrast, very few works attempt to build graph-level AutoEncoders where the encoder embeds the full graph into a single vector and the decoder reconstructs the whole graph (including the number of nodes). The Graph Deconvolutional AutoEncoder [41] and Graph U-net [21] are close to this goal, but in both cases, the decoder takes advantage of the ground-truth information, including the graph size, in the decoding phase. Ultimately, we identify only two works that share a similar goal with this paper: GraphVAE [53] (not to be confused with GVAE [36]) and PIGVAE [68]. GraphVAE is a pioneering work, but it is heavily based on an expensive graph matching algorithm. In that regard, the main contribution of PIGVAE is the addition of a learnable module that predicts the matching instead. This was a major step forward, yet we argue that PIGVAE is still missing some key components: for instance, the ground truth size of the graph needs to be provided to the decoder. We detail our positioning with respect to PIGVAE in Section 4.

Contributions. We introduce a GRAph Level autoEncoder (GRALE) that encodes and decodes graphs of varying sizes into a shared Euclidean space. GRALE is trained with an Optimal Transportinspired loss, without relying on expensive solvers, as the matching is provided by a learnable module trained end-to-end with the rest of the model. The proposed architecture leverages the Evoformer module [30] for encoding and introduces a novel "Evoformer Decoder" for graph reconstruction. GRALE enables general pretraining, applicable to a wide range of downstream tasks, including classification, regression, graph interpolation, editing, matching, and prediction. We demonstrate these capabilities through experiments on both synthetic benchmarks and large-scale molecular datasets.