The quest for the GRAph Level autoEncoder (GRALE)

TL;DR

GRALE is a novel graph autoencoder leveraging optimal transport and Evoformer architecture, enabling versatile graph representation learning for diverse applications in chemistry and biology.

Contribution

Introduces GRALE, a graph autoencoder with a differentiable node matching module and an Evoformer-based architecture, supporting variable-sized graph encoding and decoding.

Findings

GRALE achieves high-quality graph reconstructions on simulated and molecular data.

Pre-training with GRALE improves performance across various downstream tasks.

GRALE supports complex graph manipulations like interpolation and editing.

Abstract

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.