summaries: This work proposes a new deterministic set sampling mechanism, Top-n. Top-n learns to select the best 'n' points from a trainable reference set. Unlike the previous set sampling mechanisms such i.i.d. sampling, First-n and MLP projection, Top-n do not suffer from collision problem and can generate sets of various sizes (unseen during training). Top-n can be incorporated for one-shot sampling in VAE and GANs like generative models. Experimental results on standard benchmark for set and molecular graph generation, suggest improved performance in comparison to prior sampling mechanism.

summaries: This submission discusses probabilistic models that generate sets and graphs conditioned on latent vector representations. The paper discusses the following different approaches to generate sets:   * i.i.d sampling of set element representations, concatenated with the set's vector representation followed by (possibly equivariant) networks. * First-n generation, using a learnable reference set of a maximum number of nodes/set elements, and concatenating the latent set representation to each element of this reference set, and then picking the first n element of this reference set as node representations.  * MLP-based generation of node representations that don't take into account invariance of the generator under node permutation.   	 The authors claim that i.i.d sampling has two problems: the additional stochasticity makes it harder to train, and if two points are sampled too close together the resulting node representations will be similar (dubbed the collision problem). Furthermore, MLP-based generators don’t explicitly take into account permutation symmetries, can't generalize to arbitrary number of nodes and first-n generation prioritizes learning of the first elements of the reference list as they will be used more often than the last elements of this reference list.   The authors propose a new definition of equivariance with respect to permutations: instead of talking about equivariance of functions, they propose to define equivariance for a learning algorithm, by stating that  "a learning algorithm is equivariant to the action of a group if the training dynamics do not depend on the group elements that are used to represent the training data".  They furthermore adapt first-n to the method top-n with differentiable sorting of the representations of the elements in the reference set based on cosine similarity with a vector that depends on the latent vector representation.   The following claims are made:  *  top-n is easier to train than i.i.d generation because it doesn't involve an extra sampling step * top-n  captures complex dependencies in data better.  The proposed method is benchmarked on set and graph generation tasks : SetMNIST, synthetic molecule-like 3D structures, the QM9 dataset.

summaries: This paper considers the problem of “one-shot” set/graph generation, which involves learning a probabilistic decoder that maps latent vectors to sets.  First, the authors extend the usual definition of equivariance for a function to a learning algorithm.  This definition is used to show that exchangeability is not useful in GANs and VAEs when used for set generation.  Next, the authors propose Top-n, which is a new set creation mechanism which learns to select the most relevant points from a trainable reference set, in a deterministic and non-exchangeable fashion.  Top-n can replace iid generation in a VAE or GAN.  Experimental results are provided for SetMNIST reconstruction and generative tasks for a synthetic molecule dataset for sets, and the QM9 chemical dataset for graphs, demonstrating that Top-n is competitive with or outperforms a number of existing generative approaches.

summaries: In this paper, the author proposes that exchangeability is unnecessary for the generative model in the domain of set and graph generation. The definition of equivariance is generalized to learning algorithms, which is appropriate for generative modeling. Then, a method called Top-N which can be used in classical generative models, such as VAE, GANS, is proposed. In the author's argument, the proposed method has the ability to extrapolate to larger sets than those seen during training as well as the ability to train easily, which are not satisfied by the previous generative models.

abstractText: We consider one-shot probabilistic decoders that map a vector-shaped prior to a distribution over sets or graphs. These functions can be integrated into variational autoencoders (VAE), generative adversarial networks (GAN) or normalizing flows, and have important applications in drug discovery. Set and graph generation is most commonly performed by generating points (and sometimes edge weights) i.i.d. from a normal distribution, and processing them along with the prior vector using Transformer layers or graph neural networks. This architecture is designed to generate exchangeable distributions (all permutations of an output are equally likely) but it is hard to train due to the stochasticity of i.i.d. generation. In this work, we propose a new definition of equivariance and show that exchangeability is in fact unnecessary in VAEs and GANs. We then introduce Top-n, a deterministic, non-exchangeable set creation mechanism which learns to select the most relevant points from a trainable reference set. Top-n can replace i.i.d. generation in any VAE or GAN – it is easier to train and better captures complex dependencies in the data. Top-n outperforms i.i.d. generation by 15% at SetMNIST reconstruction, generates sets that are 64% closer to the true distribution on a synthetic molecule-like dataset, and is able to generate more realistic molecules when trained on the classical QM9 dataset. With improved foundations in one-shot generation, our algorithm contributes to the design of more effective molecule generation methods.

authors:
Panos Achlioptas
Olga Diamanti
Ioannis Mitliagkas
Leonidas Guibas

title: Learning representations and generative models for 3d point clouds

venue: In International conference on machine learning,

authors:
Marin Biloš
Stephan Günnemann

title: Equivariant normalizing flows for point processes and sets, 2021

authors:
Xavier Bresson
Thomas Laurent

title: A two-step graph convolutional decoder for molecule generation

venue: arXiv preprint arXiv:1906.03412,

authors:
Michael M Bronstein
Joan Bruna
Taco Cohen
Petar Veličković

title: Geometric deep learning: Grids, groups, graphs, geodesics, and gauges

venue: arXiv preprint arXiv:2104.13478,

authors:
Gabriele Corso
Luca Cavalleri
Dominique Beaini
Pietro Liò
Petar Veličković

title: Principal neighbourhood aggregation for graph nets

venue: In Advances in Neural Information Processing Systems,

title: Sinkhorn distances: Lightspeed computation of optimal transport

venue: Advances in neural information processing systems,

title: Approximation by superpositions of a sigmoidal function

venue: Mathematics of control, signals and systems,

authors:
Nicola De Cao
Thomas Kipf

title: Molgan: An implicit generative model for small molecular graphs

venue: ICML Workshop on Theoretical Foundations and Applications of Deep Generative Models,

authors:
Jean Feydy
Thibault Séjourné
François-Xavier Vialard
Shun-ichi Amari
Alain Trouve
Gabriel Peyré

title: Interpolating between optimal transport and mmd using sinkhorn divergences

venue: In The 22nd International Conference on Artificial Intelligence and Statistics,

authors:
Hongyang Gao
Shuiwang Ji

title: Graph u-nets

venue: In international conference on machine learning,

authors:
Michael Guarino
Alexander Shah
Pablo Rivas

title: Dipol-gan: Generating molecular graphs adversarially with relational differentiable pooling

authors:
Jinwoo Kim
Jaehoon Yoo
Juho Lee
Seunghoon Hong

title: Setvae: Learning hierarchical composition for generative modeling of set-structured data

venue: arXiv preprint arXiv:2103.15619,

authors:
Diederik P Kingma
Max Welling

title: Auto-encoding variational bayes

venue: arXiv preprint arXiv:1312.6114,

authors:
Jonas Köhler
Leon Klein
Frank Noé

title: Equivariant flows: exact likelihood generative learning for symmetric densities

title: Group theoretical methods in machine learning

venue: Columbia University,

authors:
Adam R Kosiorek
Hyunjik Kim
Danilo J Rezende

title: Conditional set generation with transformers

venue: Workshop on Object-Oriented Learning at ICML 2020,

authors:
Igor Krawczuk
Pedro Abranches
Andreas Loukas
Volkan Cevher

title: Gg-gan: A geometric graph generative adversarial network, 2021

authors:
Youngchun Kwon
Jiho Yoo
Youn-Suk Choi
Won-Joon Son
Dongseon Lee
Seokho Kang

title: Efficient learning of non-autoregressive graph variational autoencoders for molecular graph generation

venue: Journal of Cheminformatics,

authors:
Juho Lee
Yoonho Lee
Jungtaek Kim
Adam Kosiorek
Seungjin Choi
Yee Whye Teh

title: Set transformer: A framework for attention-based permutation-invariant neural networks

authors:
Yang Li
Haidong Yi
Christopher Bender
Siyuan Shan
Junier B Oliva

title: Exchangeable neural ode for set modeling

venue: Advances in Neural Information Processing Systems,

authors:
Yibo Li
Liangren Zhang
Zhenming Liu

title: Multi-objective de novo drug design with conditional graph generative model

venue: Journal of cheminformatics,

authors:
Renjie Liao
Yujia Li
Yang Song
Shenlong Wang
Charlie Nash
William L Hamilton
David Duvenaud
Raquel Urtasun
Richard S Zemel

title: Efficient graph generation with graph recurrent attention

authors:
Meng Liu
Keqiang Yan
Bora Oztekin
Shuiwang Ji

title: GraphEBM: Molecular graph generation with energy-based models

venue: In Energy Based Models Workshop - ICLR 2021,

authors:
Qi Liu
Miltiadis Allamanis
Marc Brockschmidt
Alexander L Gaunt

title: Constrained graph variational autoencoders for molecule design

venue: arXiv preprint arXiv:1805.09076,

authors:
Kaushalya Madhawa
Katsuhiko Ishiguro
Kosuke Nakago
Motoki Abe

title: Graph{nvp}: an invertible flow-based model for generating molecular graphs, 2020

authors:
Haggai Maron
Heli Ben-Hamu
Nadav Shamir
Yaron Lipman

title: Invariant and equivariant graph networks

venue: arXiv preprint arXiv:1812.09902,

authors:
Søren Ager Meldgaard
Jonas Köhler
Henrik Lund Mortensen
Mads-Peter V Christiansen
Frank Noé
Bjørk Hammer

title: Generating stable molecules using imitation and reinforcement learning

venue: arXiv preprint arXiv:2107.05007,

title: On the use of gromov-hausdorff distances for shape comparison

authors:
Joshua Mitton
Hans M Senn
Klaas Wynne
Roderick Murray-Smith

title: A graph vae and graph transformer approach to generating molecular graphs

venue: arXiv preprint arXiv:2104.04345,

authors:
Ethan Perez
Florian Strub
Harm De Vries
Vincent Dumoulin
Aaron Courville

title: Film: Visual reasoning with a general conditioning layer

venue: In Proceedings of the AAAI Conference on Artificial Intelligence,

authors:
Raghunathan Ramakrishnan
Pavlo O Dral
Matthias Rupp
O Anatole von Lilienfeld

title: Quantum chemistry structures and properties of 134 kilo molecules

venue: Scientific Data,

authors:
Lars Ruddigkeit
Ruud Van Deursen
Lorenz C Blum
Jean-Louis Reymond

title: Enumeration of 166 billion organic small molecules in the chemical universe database gdb-17

venue: Journal of chemical information and modeling,

authors:
Victor Garcia Satorras
Emiel Hoogeboom
Fabian B Fuchs
Ingmar Posner
Max Welling

title: E (n) equivariant normalizing flows for molecule generation in 3d

venue: arXiv preprint arXiv:2105.09016,

authors:
Hadar Serviansky
Nimrod Segol
Jonathan Shlomi
Kyle Cranmer
Eilam Gross
Haggai Maron
Yaron Lipman

title: Set2graph: Learning graphs from sets

authors:
Martin Simonovsky
Nikos Komodakis

title: Graphvae: Towards generation of small graphs using variational autoencoders

venue: In International conference on artificial neural networks,

authors:
Karl Stelzner
Kristian Kersting
Adam R Kosiorek

title: Generative adversarial set transformers

venue: In Workshop on Object-Oriented Learning at ICML,

authors:
Yongbin Sun
Yue Wang
Ziwei Liu
Joshua Siegel
Sanjay Sarma

title: Pointgrow: Autoregressively learned point cloud generation with self-attention

venue: In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision,

authors:
Ashish Vaswani
Noam Shazeer
Niki Parmar
Jakob Uszkoreit
Llion Jones
Aidan N Gomez
Łukasz Kaiser
Illia Polosukhin

title: Attention is all you need

venue: In Advances in neural information processing systems,

authors:
Carl Yang
Peiye Zhuang
Wenhan Shi
Alan Luu
Pan Li

title: Conditional structure generation through graph variational generative adversarial nets

venue: In NeurIPS,

authors:
Guandao Yang
Xun Huang
Zekun Hao
Ming-Yu Liu
Serge Belongie
Bharath Hariharan

title: Pointflow: 3d point cloud generation with continuous normalizing flows

venue: In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 4541–4550,

authors:
Manzil Zaheer
Satwik Kottur
Siamak Ravanbakhsh
Barnabas Poczos
Russ R Salakhutdinov
Alexander J Smola

title: URL https://proceedings.neurips.cc/ paper/2017/file/f22e4747da1aa27e363d86d40ff442fe-Paper.pdf

venue: Deep sets

authors:
David W Zhang
Gertjan J. Burghouts
Cees G.M. Snoek

title: Set prediction without imposing structure as conditional density estimation

venue: In International Conference on Learning Representations,

authors:
Yan Zhang
Jonathon Hare
Adam Prugel-Bennett

title: Deep set prediction networks

authors:
Yan Zhang
Jonathon Hare
Adam Prügel-Bennett

title: Fspool: Learning set representations with featurewise sort pooling

venue: arXiv preprint arXiv:1906.02795,

heading: 1 INTRODUCTION

text:
In recent years, many architectures have been proposed for learning vector representations of sets and graphs. Sets and graphs are unordered and therefore have the same symmetry group, which offers a fruitful theoretical framework for the development of new models (Bronstein et al., 2021). The opposite problem, i.e., learning to generate sets and graphs from a low-dimensional vector prior, has been less explored. Yet, it has important applications in drug discovery, for example, where neural networks can be used to sample from the space of stable drug-like molecules.
As set and graph generation are very similar problems, we generally treat them together in this paper1. In particular, one can navigate between set and graph representations using Set2Graph functions (Serviansky et al., 2020) or graph neural networks.
There are two main classes of probabilistic decoders for sets: recursive and one-shot. Recursive generators are conceptually simpler, since they add nodes, edges or substructures one by one (Liu et al., 2018; Liao et al., 2019; Sun et al., 2020). They are however slow and introduce an order in the points that does not exist in the data. In this paper, we instead focus on one-shot generation, which allows for more principled designs. These probabilistic decoders can be incorporated into any architecture (generative adversarial networks, variational autoencoders, normalizing flows or diffusion models) and trained to generate sets from a latent vector.
By definition, one-shot generative models for sets feature a layer that maps a vector to an initial set. We group all layers until this one into a creation module, and the subsequent layers into an update module. The design of the update is well understood: as it maps a set to another set, any permutation
1In the following, we only use the terminology of sets in order to avoid overly abstract notations. We discuss the distinctive properties of graph generation (e.g., graph matching problems) when needed.
equivariant neural network suits this task. In contrast, the creation step poses two major challenges, which are i) the need to generate sets of various sizes, and ii) the generation of different points from a single prior vector, which cannot be done with a permutation equivariant function.
Set creation is typically performed in two steps: first, points are drawn i.i.d. from a low-dimensional normal distribution, and then the latent vector is appended to each point. Architectures that use this design can generate any number of points, and decouple the number of generated points from the number of trainable parameters in the model. Furthermore, when combined with a permutation equivariant update module, they produce exchangeable distributions, a property which is commonly held as the equivalent of permutation equivariance for generative models. However, the stochasticity of i.i.d. creation makes learning difficult, as the network has to learn to reconstruct a set for any combination of values drawn from the base distribution (Krawczuk et al., 2021).
In this work, we first extend the usual definition of equivariance. Given a neural network architecture F⇥ parametrized by ⇥ and a loss function l, we say that the pair (F, l) is equivariant to the action of a group if the dynamics of ⇥ trained on gradient descent on l are independent of the group elements used to represent the data. Using this definition, we provide sufficient conditions for equivariance and show that exchangeability is actually not needed for equivariance in GANs and VAEs, thus explaining why i.i.d. generation is empirically not superior to other approaches.
We then propose a new set creation method, called Top-n creation, which relies on a trainable reference set where each of the point has a representation ri 2 Rc and an angle i 2 Ra. To generate a set with n elements, we select the n points whose angles have the largest cosine with the latent vector. In order to make this process differentiable, we build upon the Top-K pooling mechanism (Gao & Ji, 2019) initially proposed for graph coarsening. Top-n can eventually be integrated in any VAE or GAN to form a complete generative model for sets or graphs. As it is deterministic, it is easier to train than i.i.d. generators; in contrary to MLPs, it can extrapolate to larger sets than those seen during training.
We benchmark our creation method on both set and graph generation tasks: it reveals that Top-n is able to reconstruct the data more accurately on a set version of MNIST, to extrapolate to larger sets and fit more closely the true distribution (on a dataset of synthetic molecule-like structures in 3D), and to generate realistic graphs (on the QM9 chemical dataset). Overall, our work provides new theoretical foundations for one-shot set and graph generation and a more efficient method that contributes to the development of effective learning-based methods for drug discovery.

heading: 2 THE ONE-SHOT SET GENERATION PROBLEM

text:
We consider the problem of learning a probabilistic decoder f that maps latent vectors z 2 Rl to multi-sets2 X = {x1, ...,xn} where each xi 2 Rd for some varying number of points n. Given sample sets from an unknown distribution D, f should be such that, if z is drawn from a prior distribution pZ(z), then the push-forward measure f#(pZ) (i.e., the law of f(z)) is close to D. In practice, representing sets is not convenient on standard hardware, and sets are internally represented by matrices X 2 X = S n2N Rn⇥d where each row represents a point xi 2 Rd. Algorithms that return a set implicitly assume the use of a function mat-to-set that maps a matrix X to the corresponding set X . Existing architectures for one-shot set or graph generation use one of the two generation mechanisms described in Figure 1. First of all, a set cardinality n has to be sampled. Two strategies have been proposed:
• Most of the literature assumes that n is known during training. The distribution of the number of points in the training data is computed, and at generation time, a variable is sampled from this multinomial distribution. This method assumes that the latent vector z is independent of n, so that the generative mechanism writes p(X|n, z) p(n) p(z). • Kosiorek et al. (2020) instead proposed to learn the value of n from the latent vector using a MLP. This layer is trained via an auxiliary loss, but the predicted value is used only at generation time (the ground truth cardinality is used during training). The generative model is in this case p(X|z, n) p(n|z) p(z), i.e., z and n are not independent anymore.
2For the sake of simplicity, we will refer to sets instead of multi-sets in the rest of the paper.
Once n is sampled, one-shot generation models can formally be decomposed into several components: a first function g (that we call creation function) maps the latent vector to an initial set X0 2 Rn⇥c. This function is usually simple, and is therefore not able to model complex dependencies within each set. For this reason, X0 may then be refined by a second function h that we call update, so that the whole model can be written f = mat-to-set h g. We now review the parametrizations that have been proposed for the creation and update modules.

heading: 2.1 METHODS FOR SET CREATION

text:
MLP based generation Early works (Achlioptas et al. (2018); Zhang et al. (2019b) for sets, Guarino et al. (2017); De Cao & Kipf (2018); Simonovsky & Komodakis (2018) for graphs) proposed to learn a MLP from Rl to Rnmaxc, where nmax is the largest set size in the training data. The output vector is then reshaped as a nmax ⇥ c matrix, and masked to keep only the first n rows. MLPs ignore the symmetries of the problem and are only trained to generate up to nmax points, with no ability to extrapolate to larger sets. Despite these limitations, they often perform on par with more complex methods, and still constitute the standard way to generate small graphs (Madhawa et al., 2020; Mitton et al., 2021). We explain in Section 3 this surprising phenomenon.
Random i.i.d. generation Along with MLP based generation, the most popular method for set and graph creation is to draw n points i.i.d. from a low dimensional normal distribution, and to concatenate the latent vector to each sample (Köhler et al., 2020; Yang et al., 2019b; Kosiorek et al., 2020; Stelzner et al., 2020; Satorras et al., 2021; Zhang et al., 2021; Liu et al., 2021). The main advantage of random i.i.d. creation is that it does not constrain the number of points that can be generated. Furthermore, it is exchangeable, i.e., all permutations of the rows of X0 are equally likely. This property is widely considered as the equivalent of equivariance for generative models (Yang et al., 2019a; Biloš & Günnemann, 2021; Kim et al., 2021; Köhler et al., 2020; Li et al., 2020).
Unfortunately, i.i.d. generation makes training very difficult. Generative models based on i.i.d. creation indeed have two stochastic components: the prior vector that is sampled through the reparametrization trick (Kingma & Welling, 2013) and the values for the initial set drawn from the normal distribution – both components make it harder to generate realistic sets or reconstruct the training data (Krawczuk et al., 2021), which we also observe experimentally (Appendix D).
First-n creation While MLP and random generation are the most common methods for set creation, Zhang et al. (2019a) and Krawczuk et al. (2021) use another design: they always start from the same learnable set Xref 2 Rnmax⇥c, and concatenate the latent vector to each point of this set. Similarly to MLP based creation, they then mask this matrix to keep only the first n rows: we therefore call this method First-n creation. First-n converges much faster than with i.i.d. generation (Krawczuk et al., 2021), but the network is only trained to generate up to nmax points and has no ability to extrapolate to larger sets. Furthermore, selecting the first n rows of the reference set Xref introduces a bias because the first rows are selected more often than the last ones.
Creation methods for graph generation For graph generation, edge weights (or edge features) also need to be learned. To the best of our knowledge, First-n creation has not been used for this purpose yet, and the weights are generated either by a MLP or by sampling i.i.d. normal entries. Alternatively, Bresson & Laurent (2019) and Krawczuk et al. (2021) proposed to generate a set at creation time, and use a Set2Graph update function in order to learn the graph adjacency matrix.

heading: 2.2 METHODS FOR SET UPDATE

text: Since all creation methods except MLPs can only generate very simple sets and adjacency matrices, additional layers are usually used to refine these objects – we gather these layers into the update module. Because these layers map a set or a graph to another set/graph, the update module falls into the standard framework of permutation equivariant representation learning and all existing equivariant layers can be used: Deep sets (Zaheer et al., 2017), self-attention (Vaswani et al., 2017; Lee et al., 2019), Set2Graph (Serviansky et al., 2020), graph neural networks (Battaglia et al., 2018) or higher-order neural networks (Maron et al., 2018). Recently, Transformer layers have constituted the most popular method for set and graph update (Bresson & Laurent, 2019; Kosiorek et al., 2020; Stelzner et al., 2020; Krawczuk et al., 2021) – we also use such layers in our experiments.

heading: 3 AN EQUIVARIANCE PERSPECTIVE ON SET AND GRAPH GENERATION

text:
In this section, we analyse one-shot set and graph generation from the perspective of permutation equivariance. As set and graphs are unordered, the symmetric group S = S n2N⇤ Sn containing all permutations is a symmetry of these tasks. A permutation ⇡ 2 Sn acts on a n ⇥ n matrix A by permuting its rows and columns (which we write ⇡.A = ⇡A ⇡T ), on a n⇥c matrix by permuting its rows (⇡.X = ⇡X), and leaves a vector z 2 Rh unchanged (⇡.z = z). This symmetry constitutes a useless factor of variation in the data that should be factored out in the latent space. It is achieved when the latent representation z is invariant to the action of S (⇡.z = z). In discriminative models, symmetries are accounted for when a neural network f is equivariant to the action of the permutation group (Kondor, 2008), which writes ⇡.f(X) = f(⇡.X) . This definition can unfortunately not be used for set and graph generation from a latent vector. Indeed, imposing ⇡.f(z) = f(⇡.z) = f(z) would only allow for trivial solutions where all rows are equal, which is too restrictive. Another definition, that provides more relaxed conditions, is therefore needed for taking symmetries into account. For this purpose, we propose a definition called (F, l)-equivariance.
Our proposition is based on the observation that equivariant models for discriminative tasks are almost always used along with invariant loss functions. For example, the l2 loss ||Y Ŷ ||2 is commonly used to learn the future state of a n-body system, but not the l1 loss ||Y Ŷ ||1, as it is not rotation invariant. Formally, if F⇥ = {f✓ : X ! Y; ✓ 2 ⇥} is an hypothesis class of Gequivariant functions from X to Y (for example a neural architecture parametrized by ✓), then the loss functions l should satisfy
8f 2 F, 8g 2 G, 8(X,Y ) 2 X⇥ Y, l(g.f(X), g.Y ) = l(f(X),Y ) (1)
When l is chosen accordingly, the gradients with respect to the parameters satisfy r✓ l(f(g.X), g.Y ) = r✓ l(f(X),Y ), i.e., each parameter update is independent of the group elements that are used to represent X and Y . It follows that the training dynamics as a whole become independent of the group elements used to represent the data, which makes data augmentation unnecessary. We therefore propose the following definition: Definition 1 ((F, l)-equivariance). Consider an hypothesis class F⇥ ⇢ YX, a group G that acts on X and Y and a loss function l defined on Y. We say that the pair (F⇥, l) is equivariant to the action of G if the dynamics of ✓ 2 ⇥ trained with gradient descent on l do not depend on the group elements that are used to represent the training data.
With our definition, the use of an equivariant architecture and an invariant loss is a sufficient condition for (F, l)-equivariance in discriminative settings. As for standard generative architectures for sets and graphs, we can derive the following sufficient conditions (proofs are given in Appendix A): Lemma 1. Sufficient conditions for (F, l)-equivariance:
1. GANs: if F is a GAN architecture with a permutation invariant discriminator, and l the
standard GAN loss, then (F, l) is permutation equivariant. No constraint is imposed on the generator.
2. Autoencoders, VAEs, DSPN: if F is an encoder-decoder architecture with a permutation
invariant encoder, and the reconstruction loss l satisfies 8⇡ 2 S, l(⇡.X, X̂) = l(X, X̂), then (F, l) is permutation equivariant. No constraint is imposed on the decoder function.
3. Normalizing flows: if F is an architecture such that the set creation yields an exchangeable
distribution, the update is permutation equivariant and invertible, and p✓ denotes the model likelihood, then (F, log p✓) is permutation equivariant (proved in Köhler et al. (2020)).
As desired, (F, l)-equivariance is a notion of equivariance that does not impose the equivariance of the probabilistic decoder in generative architectures. Furthermore, for GANs and VAEs, we observe that the set generator does not need to yield exchangeable distributions either. To understand why exchangeability is not needed, recall that in GANs and VAEs a mat-to-set function is implicitly applied to the model output: a method that generates matrices that are always permuted in the same way is therefore equivalent to one that generates exchangeable matrices (as the corresponding sets are the same). In other words, the fact that the output of the model is not a matrix but a set is an assumption, not something that needs to be proved3. This observation explains why random i.i.d. generation empirically does not outperform non-exchangeable set creation methods such as MLPs and First-n (Krawczuk et al., 2021).
We also note that the constraints of Lemma 1 are satisfied by most existing architectures, including early ones (Simonovsky & Komodakis, 2018; De Cao & Kipf, 2018; Köhler et al., 2020). For example, the constraint on the loss for VAEs (8⇡ 2 S, l(⇡.X, X̂) = l(X, X̂)) is satisfied by the two loss functions commonly used for sets, namely Chamfer loss and the Wasserstein-2 distance, defined as
dChamfer(X,X 0) =
X
1in min
1jn0 ||xi x0j ||22 +
X
1jn0 min
1in ||xi x0j ||22 (2)
dW2(X,X 0) = inf
u2{ (X,X0)}
X
1in
X
1jn0 u(xi,x
0 j) kxi x0jk22, (3)
where is the set of couplings (i.e., bistochastic matrices) between X and X 0. While efficient algorithms exist for computing the Wasserstein distance (Cuturi, 2013; Feydy et al., 2019), the equation l(⇡.X, X̂) = l(X, X̂) may be difficult to satisfy in other settings: matching graphs up to permutations, or sets up to the SE(3) group, leads to difficult problems for which no polynomial time algorithm is known (Mémoli, 2007). In these settings, the design of VAE is harder and other architectures may be better suited.
Overall, the notion of (F, l)-equivariance provides justifications to the design of standard architectures for both generative and discriminative tasks. Contrary to common belief, it also appears that exchangeability not needed to build equivariant GANs and VAEs for sets and graphs. In the following section, we therefore design a new creation mechanism without worrying about the model exchangeability.

heading: 4 THE TOP-N CREATION MECHANISM

text:
We have seen in Section 2 that existing set creations methods suffer from important limitations: i.i.d. creation is slow to train due to the random initial points, while MLPs and First-n use a fixed mask to select the correct number of points and cannot extrapolate to larger sets. In order to solve these limitations, we propose a new method called Top-n creation, which is summarized in Figure 2.
Similarly to First-n, Top-n also uses a reference set, but in Top-n this set can have an arbitrary size n0. Each point in this set is a pair ( , r) : the angle 2 Ra is used to decide when to select the point, and r 2 Rc contains the representation of the point. Given a latent vector z 2 Rl⇥1, a
3On the contrary, normalizing flows typically compute probability distributions on the space of matrices and not on sets. It is therefore natural that exchangeability (which is a property defined for matrices) appears for normalizing flows and not for the other architectures.
reference set made of angles 2 Rn0⇥a and representations R 2 Rn0⇥c and learnable matrices W1 to W4 (respectively of size 1⇥ c, 1⇥ c, l ⇥ c, l ⇥ c), Top-n creation computes:
a = MLP1(z) 2 Ra angle of the latent vector (4) c = a / vec((|| i||2)1in0) 2 Rn0 cosine computation (5) s = argsort#(c)[: n] 2 Nn selection of the top n points (6) c̃ = softmax(c[s]) 2 Rn⇥1 modulation coefficients (7)
X0 = R[s] c̃ W1 + c̃ W2 2 Rn⇥c modulation of R with cosine values (8) X0 = X0 1n zT W3 + 1n zT W4 2 Rn⇥c condition on the latent vector (9)
The crux of the Top-n creation module consists in selecting the points that will be used to generate a set based on the value of the latent vector (Eqns (4, 5, 6)). Unfortunately, the gradient of the argsort operation is 0 almost everywhere (@R[s]/@ = 0), and a mechanism has to be used in order to train the angles and the MLP of Eq. (4). In Top-n, we modulate the representation of the selected points R[s] with the cosines c (Eq. 7). This operation provides a path in the computational graph that does not go through the argsort, so that the gradients of and a are not always 0. For example:
dX0
d =
@X0
@R[s]
dR[s]
d +
@X0
@c̃
dc̃ d =
@X0
@c̃
dc̃ d ,
Equations (4) to (8) build upon the Top-K pooling mechanism used by Gao & Ji (2019) for graph coarsening, but differ in several aspects. First, Gao & Ji (2019) compute cosines between the angle a and the representations R. This method may tend to select points that are similar. On the contrary, we parametrize the angles and the representations independently, so that two points can have similar angles (in which case they will usually be selected together) but very diverse representations at the same time. Second, Gao & Ji (2019) uses a multiplicative modulation (X0 = Xref[s] c̃) in Eq. (8) while we use a more expressive FiLM layer (Perez et al., 2018) that combines both additive and multiplicative modulation.
Finally, while previous works usually append the latent vector to each point, we propose to exploit the equivalence between summation and concatenation when a linear layer is applied. This property can be written cat(X0, 1nzT ) W = X0W1 + 1n(zTW2), for W = cat(W1,W2). Contrary to the concatenation (left-hand side), the sum (right-hand side) does not compute several times zTW2, which reduces the complexity of this layer from O(n(c+l)c) to O(nc2 + cl + nl). Again, we combine the sum and multiplicative modulation in a FiLM layer to build a more expressive model in Eq. (9).
Our algorithm retains the advantages of First-n creation, but replaces the arbitrary selection of the first n points by a mechanism that can select the most relevant points for a given set. Since it
also decouples the number of points in the reference set from the number of points in the training examples, the size of the reference set becomes a hyperparameter of the model. Empirically, we observe a tradeoff: when using more points in the reference set, each point is updated less often which makes training slower; however, the model tends to avoid overfitting and generalize better.
Since First-n and Top-n use a fixed set of reference points, one may wonder if they restrict the model expressivity. We however show that it is not the case: used with a two-layer MLP, the First-n and Top-n modules are universal approximators over sets. Proposition 1 (Expressivity). For any set size n, maximal norm M and precision parameter ✏, there is a 2-layer row-wise MLP f and reference points {x1, ...,xn} such that, for any set {y1, ...,yn} of points in Rd with 8i, ||yi||  M there is a latent vector z of size nd⇥ 1 that satisfies:
||f(cat(X,1nzT )) Y ||W2  ✏
Top-n creation can be used in any GAN or autoencoder-based architecture as a replacement for other set creation methods. It is however not suited to normalizing flows, because it is based on a hard selection process (the value of the reference points that are not selected is not used) and is therefore non invertible.

heading: 5 EXPERIMENTS

text: We compare Top-n to other set and graph creation methods on several tasks: autoencoding point clouds on a set version of MNIST, generating realistic 3D structures on a synthetic molecule-like dataset, and generating varied valid molecules on the QM9 chemical dataset4. All training curves are available in Appendix D.

heading: 5.1 SET MNIST

text:
We first perform experiments on the SetMNIST benchmark, introduced in Zhang et al. (2019a). The task consists in autoencoding point clouds that are built by thresholding the pixel values in MNIST images, adding noise on the locations and normalizing the coordinates. Our goal is to show that Topn can favorably replace other set creation methods without having to tune the rest of the architecture extensively. For this purpose, we use existing implementations of DSPN (Zhang et al., 2019a) TSPN (Kosiorek et al., 2020)5, which are respectively a sort of diffusion model and a transformer-based autoencoder. Experiment details can be found in Appendix C.
The results for both methods are very similar (Figure 1): the standard exchangeable i.i.d. creation has poor performance and needs more epochs to be trained than First-n and Top-n. Since all sets in this dataset have the same cardinality (360 points), the better performance of Top-n cannot come from its better ability to process varying sizes, but rather shows that it is able to select the most relevant reference points for each set.

heading: 5.2 SYNTHETIC DATASET

text:
Since SetMNIST does not allow to measure generation quality easily, we further benchmark the different set generators on a synthetic dataset for which the quality of the generated sets can be as-
4Source code is available at [linked masked for anonymity] 5github.com/LukeBolly/tf-tspn (reimplementation by someone else) and github.com/Cyanogenoid/dspn
sessed. This dataset is a simplified model of molecules in 3D, and retains some of its characteristics: i) atoms are never too close to each other ii) there is a bond between two atoms if and only if they are closer than a given distance (that depends on the atom types) iii) if formal charges are forbidden, each atom has a predefined valency.
Each set is created via rejection sampling: points are drawn iteratively from a uniform distribution within a bounding box in R3. The first point is always accepted, and the next ones are accepted only if they satisfy a predefined set of constraints: i) they are not closer to any other point than a given threshold min-distance. ii) they are connected to the rest of the set, i.e., have at least one neighbor than a distance neighbor-distance iii) they do not have too many neighbors. Our dataset is made of 2000 sets that have between 2 and 35 points, with 9 points per set on average. It is a simplification of real molecules in several aspects: there are only single bonds, angles between bonds are not constrained and the atom types are only defined by the valency, which reflects the fact that atoms with the same valency tend to play a similar role and be more interchangeable.
The goal is to reconstruct the atom positions and generate new realistic sets. During training, we measure the Wasserstein (W2) distance between the input and reconstructed sets. At generation time, we compute the W2 distance between the distribution of valencies in the dataset and in the generated set, as well as a diversity score to ensure that a method does not always generate the same set. For that purpose we use the average W2 distance between 1000 pairs of generated sets. We also measure the same metrics in an extrapolation setting where sets have on average 10 more points. Since the valency distribution may not be exactly the same on bigger sets, we also report the proportion of generated atoms whose valency is not present in the dataset (0, or more than 4). In practice, we find that the two measures are very correlated.
We train a VAE with different creation methods on this dataset. Details about the model and the loss function can be found in Appendix C, as well as an ablation study on the number of points in the reference set. The results are shown in Table 2. We observe that the random i.i.d. generator generalizes well but reconstructs the training sets very poorly, which reflects the fact that the model is hard to train due to the stochastic i.i.d. generation. The generated sets also lack diversity, which may reveal that the network tends to ignore the features that are drawn and rely on the layers biases. On the contrary, the MLP and First-n generators seem to overfit the training data at the expense of generation quality. Finally, Top-n is able to generate points that have the right valency and performance does not degrades when extrapolating to larger sets.

heading: 5.3 MOLECULAR GRAPH GENERATION

text:
Finally, we evaluate Top-n on a graph generation task. We train a graph VAE (detailed in Appendix C) on QM9 molecules and check its ability to generate a wide range of valid molecules. As a generation metric, we simply report the validity (i.e., the proportion of molecules that satisfy a set of basic chemical rules) and uniqueness (the proportion non-isomorphic graphs among valid ones) of the generated molecules. We do not report novelty, as QM9 is an enumeration of all possible molecules up to 9 heavy atoms that satisfy a predefined set of constraints (Ruddigkeit et al., 2012; Ramakrishnan et al., 2014): in this setting, generating novel molecules is therefore not an indicator of good performance, but rather a sign that the distribution of the training data has not been properly captured. Note that most recently proposed methods proposed on this task use a non-learned validity
correction code to generate almost 100% of valid molecules, which obfuscates the real performance of the learned model. We therefore only compare to works that do not correct validity: GraphVAE (Simonovsky & Komodakis, 2018), an extension of GraphVAE that incorporates reinforcement learning (Kwon et al., 2019), MolGAN (De Cao & Kipf, 2018) and GraphTransformerVAE (Mitton et al., 2021).
We observe that both MolGAN and the Graph VAE that incorporates reinforcement learning suffer from mode collapse: they can generate only a few molecules, but these molecules are often valid. On the contrary, our method is able to generate a higher rate of valid and unique molecules. We observe that the random i.i.d. method is not able to obtain a good train loss (Appendix D), which reflects in the poor generation performance as well. The three other set creation methods seem to perform similarly. Our interpretation is that almost all molecules in QM9 have the same size (9 heavy atoms, since hydrogens are not represented), which makes it less important to properly handle varying graph sizes as is proposed in Top-n.
We finally note that recursive methods such as Li et al. (2018) can often generate higher rates of valid molecules. The reason is that recursive methods can easily check at each step that the added edge will not break valency constraints. They have therefore an unfair advantage over one-shot models which do not incorporate these checks.

heading: 6 CONCLUSION

text:
In this work, we strengthened the theoretical foundations of one-shot set and graph generation by proposing a definition of equivariant learning algorithms. We showed in particular that, contrary to common belief, exchangeability is not a required property in GANs and VAE. We then proposed Top-n, a non-exchangeable model for one-shot set and graph creation. Our method consistently achieves better performance on set generation, and is able to generate a wide range of valid molecular graphs. Overall, Top-n is easy to train (as it is deterministic), it extrapolates well and is able to select the most relevant points for each set. Our method can be incorporated in any GAN or VAE architecture and replace favorably other set creation methods.
We finally note that our experiments on graph generation and molecule generation in 3D feature only small graphs with 9 nodes on average. This set size is in line with most literature on the topic (Meldgaard et al., 2021; Satorras et al., 2021; Mitton et al., 2021), and we observe on the synthetic dataset that good reconstruction is significantly harder to achieve on large sets. This observation leads to a simple question that has not yet been investigated: can set and graph generation methods avoid the curse of dimensionality? For MLP-based generators, the answer is probably no: as they learn vectors of size nmaxd that are then reshaped to matrices, standard universal approximation results will most likely yield a sample complexity that is exponential in nmax. Whether and under what circumstances other set generators can overcome this problem is an important question that conditions the success of one-shot models for larger sets and graphs.
Ethics statement Although our method can be used to generate any type of sets or graphs, it was primarily built and benchmarked with molecule generation applications in mind. While learning based methods will probably not replace expert knowledge in the short term, they can already be used to generate hit molecules, i.e., small and simple molecules that have promising activity on a predefined target. These molecules can then be refined by computational chemists in order to optimise activity and develop new effective drugs. We are therefore convinced that research on set and graph generation can have a direct positive impact on society.
Reproducibility statement All code is included in the supplementary materials. Further details can be found in Appendix C.

title: TOP-N: EQUIVARIANT SET AND GRAPH GENERATION WITHOUT EXCHANGEABILITY