🤖 machine learning

GraViti: Graph-Level Variational Autoencoders with Relaxed Permutation Invariance

GraViti is a transformer-based graph-level variational autoencoder that maps entire graphs to compact latent vectors, achieving state-of-the-art reconstruction and generative performance by relaxing strict permutation invariance to enable smooth interpolation and property-guided search while directly learning domain constraints for valid sample decoding.

Original authors: Roman Bresson, Konstantinos Divriotis, Johannes F. Lutzeyer, Iakovos Evdaimon, Michalis Vazirgiannis

Published 2026-05-19

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Roman Bresson, Konstantinos Divriotis, Johannes F. Lutzeyer, Iakovos Evdaimon, Michalis Vazirgiannis

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a massive library of complex structures, like intricate LEGO models or molecular diagrams. In the world of artificial intelligence, these are called graphs. For a long time, computers struggled to understand these graphs as a whole. Instead, they looked at them piece by piece, like trying to understand a whole city by only studying individual street signs. This made it hard to compare two different cities or smoothly morph one into another.

The paper introduces GraViti, a new AI tool designed to solve this problem. Think of GraViti as a master "compression artist" that can look at an entire complex graph, understand its essence, and shrink it down into a single, tiny ID card (called a latent vector).

Here is how it works, using simple analogies:

1. The Problem: The "Seat Number" Confusion

Imagine you have a photo of a group of friends. If you swap the order of the people in the photo (Person A sits on the left, then Person B; or Person B sits on the left, then Person A), it's still the same group of friends. However, older AI models were like strict librarians who cared about who sat where. If you swapped the seats, the AI thought it was a completely different group. This made it impossible to create a smooth "map" of all possible groups.

GraViti solves this by realizing that the group matters, not the seating order.

2. The Solution: The "Magic ID Card"

GraViti uses a special type of brain called a Transformer (the same technology behind modern chatbots) to look at the whole graph at once.

The Encoder (The Compressor): It takes a messy, complex graph and squashes it into a single, compact ID card. This card holds all the essential information about the graph's shape and rules.
The Decoder (The Expander): When you give it an ID card, it can rebuild the original graph from scratch, perfectly following the rules of the universe it was trained on (like the rules of chemistry for molecules).

3. The Secret Sauce: "Relaxed Permutation Invariance"

This is the paper's biggest "aha!" moment.
Usually, AI tries to be so flexible that it ignores the order of nodes entirely. But the authors found that in specific worlds—like molecules (where atoms have a natural order) or Bayesian networks (where cause-and-effect has a flow)—being too flexible actually hurts performance.

The Analogy: Imagine trying to describe a sentence. If you ignore the order of words entirely, you get gibberish. But if you know the sentence follows a standard grammar (Subject-Verb-Object), you can describe it much more accurately.
GraViti says: "If the data has a natural order (like a molecule), let's use it!" By trusting this natural order, GraViti becomes much faster and more accurate at rebuilding the graph than its competitors. It doesn't need to waste time guessing which atom goes where; it just follows the map.

4. What Can GraViti Do?

The paper demonstrates three main superpowers:

The "Morphing" Machine (Interpolation): Because every graph has a unique ID card, you can take the ID card of a "Water" molecule and the ID card of an "Alcohol" molecule, and slide the dial between them. The AI will generate a smooth journey of molecules changing from one to the other, step-by-step, without breaking the laws of chemistry.
The "Noise-Canceling" Headphones (Denoising): If you take a valid molecule and randomly smash some atoms or break some bonds (adding noise), GraViti can look at the broken mess and say, "Ah, I know what this should be," and fix it back to a valid chemical structure.
The "Property Tweaker" (Optimization): You can tell GraViti, "I want a molecule that is more oily (hydrophobic)." The AI can navigate its ID card space to find a new molecule that is slightly more oily than the original, following real chemical rules to do so.

5. The Results

The paper tested GraViti on huge datasets of real-world molecules (millions of them).

Accuracy: It rebuilt these complex molecules with higher accuracy than previous state-of-the-art models, especially when the molecules got larger.
Speed: Because it doesn't need to run expensive, slow math to figure out which atom matches which, it is much more efficient.
Generality: It worked not just on molecules, but also on other types of graphs, proving it's a versatile tool.

Summary

GraViti is a new way for computers to understand complex networks. Instead of getting lost in the details of "who is next to whom," it creates a compact summary of the whole picture. By trusting that some things (like molecules) have a natural order, it builds a better, smoother, and more accurate map of the world of graphs, allowing us to generate new designs and fix broken ones with ease.

Technical Summary: GraViti – Graph-Level Variational Autoencoders with Relaxed Permutation Invariance

Problem Statement

Graph-level autoencoding aims to map entire graphs into compact latent vectors to facilitate tasks such as interpolation, property-guided search, and generative modeling. While node-level Variational Autoencoders (VAEs) are well-established, they suffer from a fundamental limitation: the latent space is defined at the node level and depends on arbitrary node ordering. Consequently, two isomorphic graphs can map to entirely different sets of latent vectors, preventing the definition of meaningful distances between graphs or the capture of graph-level manifolds.

Existing graph-level approaches face a trade-off. Methods like PIGVAE and GRALE attempt to enforce strict permutation invariance, often relying on computationally expensive matching procedures (e.g., graph matching, Sinkhorn algorithms) to align input and reconstructed graphs. These alignment steps introduce significant computational overhead (typically cubic complexity) and can limit scalability on large datasets. Furthermore, many graph-level decoders rely on padding strategies that impose a fixed maximum graph size or introduce "dummy node" classes, which restrict generalization to larger graphs and bias classification tasks.

Methodology: The GraViti Architecture

GraViti is a transformer-based graph-level VAE designed to encode entire graphs into single latent vectors while relaxing the strict requirement for permutation-invariant reconstruction when a consistent canonical node ordering exists.

1. Graph Representation

The model operates on graphs with discrete node and edge classes (e.g., atom types, bond types).

Input: A node class matrix $X$ , an edge class tensor $E$ , and a global feature vector $y$ encoding graph-level statistics (counts of node/edge types and graph size).
Positional Encodings: To inject structural information, the model utilizes Laplacian eigenvectors and random-walk features as positional encodings ( $P$ ).
Domain Knowledge: For specific domains like molecules, additional node-specific features ( $K$ ) such as atomic mass or electronegativity are incorporated.
Batching: Graphs are padded to the size of the largest graph in a batch, but the model is designed to ignore these padded elements entirely, avoiding the need for dummy node classes.

2. Encoder

The encoder maps the input triplet $(X, E, y)$ to a latent distribution parameters $(z_\mu, z_\sigma)$ .

Embedding: Node classes, edge classes, and global features are projected into embeddings.
Graph Transformer Stack: The embeddings are processed through a stack of graph transformer layers (based on DiGress and CatFlow architectures). This allows simultaneous updates of node, edge, and graph-level features while ensuring permutation equivariance during the encoding phase.
Readout: Unlike traditional sum/mean pooling, GraViti employs a sequence transformer-based readout. Two parallel sequence-decoder transformers process node and edge representations respectively, using the global feature vector $y_{rep}$ as a target to inject global context. This ensures the output is independent of node ordering at the readout stage.

3. Graph Size Prediction

A crucial innovation is the explicit disentanglement of graph size from the rest of the latent representation.

The model forces the first $d_s$ dimensions of the latent space to encode the log-size of the graph.
An MLP predicts the log-size $\hat{s}$ from these dimensions.
During decoding, the ground-truth log-size $s$ is re-encoded and reinjected into the latent representation to guide the generation of the correct number of nodes, ensuring the decoder knows the target graph size without relying on masking.

4. Decoder

The decoder reconstructs the graph from the latent vector $z$ .

Broadcasting: The latent vector is duplicated $N$ times (where $N$ is the predicted or ground-truth size) and augmented with sinusoidal positional encodings to break symmetry.
Sequence Transformer: A sequence transformer encoder processes this sequence, splitting it into node-oriented and edge-oriented sequences.
Refinement: A second graph transformer stack refines the node and edge representations.
Prediction Heads: Final MLPs classify node types and edge types. The model predicts the graph size directly, eliminating the need for a "no-edge" dummy class.

5. Training Objective

The model is trained using a standard VAE objective:

Reconstruction Loss: A combination of Focal Loss for node and edge classification (to handle class imbalance) and a smooth L1 loss for graph size prediction.
Variational Loss: A KL-divergence term regularizes the latent space toward a standard normal prior.
Relaxed Permutation Invariance: Crucially, the reconstruction loss is applied based on the input node order. The authors argue that in domains with a reliable canonical ordering (e.g., molecules, Bayesian networks), enforcing strict permutation invariance via matching is detrimental. By relying on consistent ordering, the model achieves quadratic complexity and higher reconstruction fidelity.

Key Contributions

Scalable Architecture: GraViti achieves state-of-the-art reconstruction accuracy on large graph datasets (PubChem16, PubChem32) by avoiding expensive permutation-invariant matching procedures.
Structured Latent Space: The learned latent space is shown to be well-structured, supporting smooth interpolation, controlled editing, and property-driven optimization. On molecular benchmarks, the model recovers established chemical principles (e.g., the relationship between alkyl chain length and hydrophobicity) directly from the latent space.
Relaxed Permutation Invariance: The paper demonstrates empirically that strict permutation invariance is not always necessary. When a consistent canonical node ordering is available, relaxing this constraint significantly improves reconstruction quality and scalability compared to methods that enforce invariance through complex matching.

Experimental Results

The authors evaluated GraViti on several benchmarks:

Reconstruction: On molecular datasets (QM9, PubChem16, PubChem32), GraViti outperforms baselines (GRALE, PIGVAE) by up to 25% in reconstruction accuracy, particularly on larger graphs where matching-based methods struggle.
Ablation Studies: Removing the transformer encoder or the transformer-based readout significantly degrades performance, confirming the importance of these components.
Denoising: The model effectively restores chemically valid structures from noisy inputs, with the VAE variant showing superior robustness compared to the AE variant.
Property Optimization: Using CMA-ES in the latent space, the model successfully maximizes logP (lipophilicity), generating molecules with progressively longer carbon chains or specific hydrophobic moieties, aligning with known chemical trends.
Interpolation: Linear interpolation in the latent space yields smooth transitions between molecular structures.
Downstream Tasks: Latent vectors trained on PubChem successfully predict molecular properties (dipole moment, HOMO-LUMO gap) and Bayesian Information Criterion (BIC) scores for Bayesian networks with high accuracy.
Generation: While generation validity is lower on PubChem (approx. 50-70% depending on aromatic representation) compared to QM9, the model produces diverse and novel molecules. The authors attribute lower validity on larger datasets to the global compactness of the latent space rather than local structural failures.

Significance and Claims

The paper claims that GraViti provides a practical foundation for graph-level representation learning. Its primary significance lies in demonstrating that strict permutation invariance is not a prerequisite for effective graph-level autoencoding when a stable node ordering exists. By relaxing this constraint, GraViti avoids the computational bottlenecks of graph matching, enabling scalable training on massive datasets while maintaining high reconstruction quality.

The authors position GraViti as a lightweight alternative to complex generation pipelines, offering a single-step decoding mechanism that balances generative performance with practical sample quality. The work suggests that for domains like molecular graphs and Bayesian networks, where canonical orderings are inherent or can be constructed, graph-level models can achieve superior performance without the overhead of alignment-based losses. Future work is proposed to integrate domain-aware constraints to further improve generation validity and explore transfer to other graph domains such as circuits and proteins.