GraViti: Graph-Level Variational Autoencoders with Relaxed Permutation Invariance

TL;DR

GraViti is a transformer-based graph autoencoder that learns a true graph-level latent space, enabling smooth interpolation, property-guided search, and high-quality graph generation, especially in domains with canonical node orderings.

Contribution

It introduces a novel graph-level variational autoencoder that supports permutation relaxation and achieves state-of-the-art reconstruction and generative performance.

Findings

Successfully decodes valid molecular graphs respecting chemical constraints.

Enforcing permutation invariance can hinder reconstruction in certain domains.

Achieves state-of-the-art accuracy and practical sample quality with single-step decoding.

Abstract

We introduce GraViti, a transformer-based graph-level variational autoencoder that maps entire graphs to compact latent vectors. This design produces a true graph-level latent space that supports smooth interpolation, property-guided search, and other downstream tasks beyond the constraints of node-level embeddings. On molecular benchmarks, GraViti learns to decode valid samples that follow the chemical constraints present in the training data, showing that the model recovers domain rules directly from graph-level representations. We also show that, in domains where a reliable canonical node ordering exists such as molecules or bayesian networks, enforcing permutation invariance can prove detrimental for consistent reconstruction. GraViti achieves state-of-the-art reconstruction accuracy on large datasets, and provides solid generative performance. Its single-step decoding offers a lightweight alternative to more complex generation pipelines while maintaining practical sample quality.