Top-N: Equivariant set and graph generation without exchangeability

TL;DR

This paper introduces Top-n, a novel non-exchangeable generative method for sets and graphs that improves training stability and performance over traditional exchangeable models, with applications in molecule and object generation.

Contribution

We propose Top-n creation, a differentiable, non-exchangeable generation mechanism that replaces i.i.d. sampling in variational models, enabling better permutation equivariance and improved results.

Findings

Outperforms i.i.d. generation by 15% on SetMNIST

Achieves 33% improvement in object detection on CLEVR

Generates 74% closer to true distribution on synthetic datasets

Abstract

This work addresses one-shot set and graph generation, and, more specifically, the parametrization of probabilistic decoders that map a vector-shaped prior to a distribution over sets or graphs. Sets and graphs are most commonly generated by first sampling points i.i.d. from a normal distribution, and then processing these points along with the prior vector using Transformer layers or Graph Neural Networks. This architecture is designed to generate exchangeable distributions, i.e., all permutations of the generated outputs are equally likely. We however show that it only optimizes a proxy to the evidence lower bound, which makes it hard to train. We then study equivariance in generative settings and show that non-exchangeable methods can still achieve permutation equivariance. Using this result, we introduce Top-n creation, a differentiable generation mechanism that uses the latent vector to select the most relevant points from a trainable reference set. Top-n can replace i.i.d. generation in any Variational Autoencoder or Generative Adversarial Network. Experimentally, our method outperforms i.i.d. generation by 15% at SetMNIST reconstruction, by 33% at object detection on CLEVR, generates sets that are 74% closer to the true distribution on a synthetic molecule-like dataset, and generates more valid molecules on QM9.