[Paper Review] The general theory of permutation equivarant neural networks and higher order graph variational encoders

Erik H. Thiede, Truong Son Hy+1 more|arXiv (Cornell University)|Apr 8, 2020

Advanced Graph Neural Networks39 references4 citations

TL;DR

This paper develops a general theory of permutation-equivariant neural networks for higher-order interactions, introducing second-order equivariant layers that simultaneously permute rows and columns of matrices—critical for graph and relational learning. The authors propose a second-order graph variational autoencoder that achieves 98.4% reconstruction accuracy on ZINC molecular generation, outperforming methods with explicit chemical prior knowledge, and demonstrate that latent distributions in such models must be exchangeable.

ABSTRACT

Previous work on symmetric group equivariant neural networks generally only considered the case where the group acts by permuting the elements of a single vector. In this paper we derive formulae for general permutation equivariant layers, including the case where the layer acts on matrices by permuting their rows and columns simultaneously. This case arises naturally in graph learning and relation learning applications. As a specific case of higher order permutation equivariant networks, we present a second order graph variational encoder, and show that the latent distribution of equivariant generative models must be exchangeable. We demonstrate the efficacy of this architecture on the tasks of link prediction in citation graphs and molecular graph generation.

Motivation & Objective

• To extend the theory of permutation-equivariant neural networks beyond first-order actions (permuting vector elements) to higher-order actions, such as simultaneous row and column permutations of matrices.
• To formalize the structure and parameterization of higher-order equivariant layers, particularly for second-order actions arising in graph representation learning.
• To develop a second-order graph variational autoencoder (VAE) that respects permutation symmetry in adjacency matrices and generates valid molecular graphs.
• To investigate the connection between permutation-equivariant generative models and exchangeable distributions in Bayesian nonparametric statistics.
• To empirically validate the architecture on link prediction and molecular graph generation, showing performance gains despite no explicit chemical knowledge.

Proposed method

• Derives general formulae for permutation-equivariant linear layers acting on tensors under arbitrary symmetric group actions, including second-order actions on matrices.
• Applies representation theory of the symmetric group to classify equivariant layers and determine the number of learnable parameters in each layer.
• Constructs a second-order graph variational autoencoder where the encoder and decoder are equivariant to simultaneous row and column permutations of the adjacency matrix.
• Uses a variational inference framework with a latent space that is exchangeable (but not necessarily i.i.d.), ensuring symmetry in the generative process.
• Employs a continuous relaxation of discrete bond features to enable differentiable generation of molecules from the latent space.
• Trains the model using a VAE objective, with ablation showing improved performance when the KL divergence term is disabled.

Experimental results

Research questions

• RQ1 How can permutation-equivariant neural networks be generalized beyond first-order actions to higher-order actions, such as simultaneous row and column permutations of matrices?
• RQ2 What is the mathematical structure of equivariant layers under second-order symmetric group actions, and how many learnable parameters do they have?
• RQ3 Can a second-order permutation-equivariant graph variational autoencoder generate valid molecular graphs without explicit chemical priors or rule-based constraints?
• RQ4 What constraints does permutation equivariance impose on the latent distribution in generative models, and how does this relate to exchangeability?
• RQ5 Does higher-order equivariance lead to improved performance in graph representation learning tasks like link prediction and molecular generation?

Key findings

• The proposed second-order graph variational autoencoder achieves 98.4% reconstruction accuracy on the ZINC molecular generation benchmark, significantly outperforming methods like CVAE (44.6%) and GVAE (53.7%).
• Despite lacking explicit chemical knowledge (e.g., valency rules or functional group constraints), the model generates 33.4% valid molecules from the prior, outperforming SD-VAE (43.5%) and GraphVAE (13.5%) in validity.
• The model generates molecules with reasonable carbon backbones and bonding patterns, though it shows a bias toward small molecules and strained ring systems like 4-membered rings.
• Disabling the KL divergence term in the VAE loss improved reconstruction accuracy, suggesting that the inductive bias of the model may be better preserved without posterior collapse regularization.
• The latent space of the model is exchangeable, confirming a theoretical connection between permutation-equivariant generative models and exchangeable sequences in Bayesian nonparametrics.
• The results demonstrate that higher-order equivariance enables models to learn complex structural patterns—such as ring formation—directly from data, even without explicit chemical rules.