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. 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.   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. 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.
references.authors: Panos Achlioptas Olga Diamanti Ioannis Mitliagkas Leonidas Guibas
references.title: Learning representations and generative models for 3d point clouds
references.venue: In International conference on machine learning,
references.authors: Marin Biloš Stephan Günnemann
references.title: Equivariant normalizing flows for point processes and sets, 2021
references.venue: URL https://openreview.net/forum?id=LIR3aVGIlln
references.authors: Xavier Bresson Thomas Laurent
references.title: A two-step graph convolutional decoder for molecule generation
references.venue: arXiv preprint arXiv:1906.03412,
references.authors: Michael M Bronstein Joan Bruna Taco Cohen Petar Veličković
references.title: Geometric deep learning: Grids, groups, graphs, geodesics, and gauges
references.venue: arXiv preprint arXiv:2104.13478,
references.authors: Gabriele Corso Luca Cavalleri Dominique Beaini Pietro Liò Petar Veličković
references.title: Principal neighbourhood aggregation for graph nets
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Marco Cuturi
references.title: Sinkhorn distances: Lightspeed computation of optimal transport
references.venue: Advances in neural information processing systems,
references.authors: George Cybenko
references.title: Approximation by superpositions of a sigmoidal function
references.venue: Mathematics of control, signals and systems,
references.authors: Nicola De Cao Thomas Kipf
references.title: Molgan: An implicit generative model for small molecular graphs
references.venue: ICML Workshop on Theoretical Foundations and Applications of Deep Generative Models,
references.authors: Jean Feydy Thibault Séjourné François-Xavier Vialard Shun-ichi Amari Alain Trouve Gabriel Peyré
references.title: Interpolating between optimal transport and mmd using sinkhorn divergences
references.venue: In The 22nd International Conference on Artificial Intelligence and Statistics,
references.authors: Hongyang Gao Shuiwang Ji
references.title: Graph u-nets
references.venue: In international conference on machine learning,
references.authors: Michael Guarino Alexander Shah Pablo Rivas
references.title: Dipol-gan: Generating molecular graphs adversarially with relational differentiable pooling
references.authors: Jinwoo Kim Jaehoon Yoo Juho Lee Seunghoon Hong
references.title: Setvae: Learning hierarchical composition for generative modeling of set-structured data
references.venue: arXiv preprint arXiv:2103.15619,
references.authors: Diederik P Kingma Max Welling
references.title: Auto-encoding variational bayes
references.venue: arXiv preprint arXiv:1312.6114,
references.authors: Jonas Köhler Leon Klein Frank Noé
references.title: Equivariant flows: exact likelihood generative learning for symmetric densities
references.venue: In International Conference on Machine Learning,
references.authors: Imre Risi Kondor
references.title: Group theoretical methods in machine learning
references.venue: Columbia University,
references.authors: Adam R Kosiorek Hyunjik Kim Danilo J Rezende
references.title: Conditional set generation with transformers
references.venue: Workshop on Object-Oriented Learning at ICML 2020,
references.authors: Igor Krawczuk Pedro Abranches Andreas Loukas Volkan Cevher
references.title: Gg-gan: A geometric graph generative adversarial network, 2021
references.venue: URL https://openreview.net/forum? id=qiAxL3Xqx1o
references.authors: Youngchun Kwon Jiho Yoo Youn-Suk Choi Won-Joon Son Dongseon Lee Seokho Kang
references.title: Efficient learning of non-autoregressive graph variational autoencoders for molecular graph generation
references.venue: Journal of Cheminformatics,
references.authors: Juho Lee Yoonho Lee Jungtaek Kim Adam Kosiorek Seungjin Choi Yee Whye Teh
references.title: Set transformer: A framework for attention-based permutation-invariant neural networks
references.venue: In International Conference on Machine Learning,
references.authors: Yang Li Haidong Yi Christopher Bender Siyuan Shan Junier B Oliva
references.title: Exchangeable neural ode for set modeling
references.venue: Advances in Neural Information Processing Systems,
references.authors: Yibo Li Liangren Zhang Zhenming Liu
references.title: Multi-objective de novo drug design with conditional graph generative model
references.venue: Journal of cheminformatics,
references.authors: Renjie Liao Yujia Li Yang Song Shenlong Wang Charlie Nash William L Hamilton David Duvenaud Raquel Urtasun Richard S Zemel
references.title: Efficient graph generation with graph recurrent attention
references.authors: Meng Liu Keqiang Yan Bora Oztekin Shuiwang Ji
references.title: GraphEBM: Molecular graph generation with energy-based models
references.venue: In Energy Based Models Workshop - ICLR 2021,
references.authors: Qi Liu Miltiadis Allamanis Marc Brockschmidt Alexander L Gaunt
references.title: Constrained graph variational autoencoders for molecule design
references.venue: arXiv preprint arXiv:1805.09076,
references.authors: Kaushalya Madhawa Katsuhiko Ishiguro Kosuke Nakago Motoki Abe
references.title: Graph{nvp}: an invertible flow-based model for generating molecular graphs, 2020
references.venue: URL https:// openreview.net/forum?id=ryxQ6T4YwB
references.authors: Haggai Maron Heli Ben-Hamu Nadav Shamir Yaron Lipman
references.title: Invariant and equivariant graph networks
references.venue: arXiv preprint arXiv:1812.09902,
references.authors: Søren Ager Meldgaard Jonas Köhler Henrik Lund Mortensen Mads-Peter V Christiansen Frank Noé Bjørk Hammer
references.title: Generating stable molecules using imitation and reinforcement learning
references.venue: arXiv preprint arXiv:2107.05007,
references.authors: Facundo Mémoli
references.title: On the use of gromov-hausdorff distances for shape comparison
references.authors: Joshua Mitton Hans M Senn Klaas Wynne Roderick Murray-Smith
references.title: A graph vae and graph transformer approach to generating molecular graphs
references.venue: arXiv preprint arXiv:2104.04345,
references.authors: Ethan Perez Florian Strub Harm De Vries Vincent Dumoulin Aaron Courville
references.title: Film: Visual reasoning with a general conditioning layer
references.venue: In Proceedings of the AAAI Conference on Artificial Intelligence,
references.authors: Raghunathan Ramakrishnan Pavlo O Dral Matthias Rupp O Anatole von Lilienfeld
references.title: Quantum chemistry structures and properties of 134 kilo molecules
references.venue: Scientific Data,
references.authors: Lars Ruddigkeit Ruud Van Deursen Lorenz C Blum Jean-Louis Reymond
references.title: Enumeration of 166 billion organic small molecules in the chemical universe database gdb-17
references.venue: Journal of chemical information and modeling,
references.authors: Victor Garcia Satorras Emiel Hoogeboom Fabian B Fuchs Ingmar Posner Max Welling
references.title: E (n) equivariant normalizing flows for molecule generation in 3d
references.venue: arXiv preprint arXiv:2105.09016,
references.authors: Hadar Serviansky Nimrod Segol Jonathan Shlomi Kyle Cranmer Eilam Gross Haggai Maron Yaron Lipman
references.title: Set2graph: Learning graphs from sets
references.venue: Advances in Neural Information Processing Systems,
references.authors: Martin Simonovsky Nikos Komodakis
references.title: Graphvae: Towards generation of small graphs using variational autoencoders
references.venue: In International conference on artificial neural networks,
references.authors: Karl Stelzner Kristian Kersting Adam R Kosiorek
references.title: Generative adversarial set transformers
references.venue: In Workshop on Object-Oriented Learning at ICML,
references.authors: Yongbin Sun Yue Wang Ziwei Liu Joshua Siegel Sanjay Sarma
references.title: Pointgrow: Autoregressively learned point cloud generation with self-attention
references.venue: In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision,
references.authors: Ashish Vaswani Noam Shazeer Niki Parmar Jakob Uszkoreit Llion Jones Aidan N Gomez Łukasz Kaiser Illia Polosukhin
references.title: Attention is all you need
references.venue: In Advances in neural information processing systems,
references.authors: Carl Yang Peiye Zhuang Wenhan Shi Alan Luu Pan Li
references.title: Conditional structure generation through graph variational generative adversarial nets
references.venue: In NeurIPS,
references.authors: Guandao Yang Xun Huang Zekun Hao Ming-Yu Liu Serge Belongie Bharath Hariharan
references.title: Pointflow: 3d point cloud generation with continuous normalizing flows
references.venue: In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 4541–4550,
references.authors: Manzil Zaheer Satwik Kottur Siamak Ravanbakhsh Barnabas Poczos Russ R Salakhutdinov Alexander J Smola
references.title: URL https://proceedings.neurips.cc/ paper/2017/file/f22e4747da1aa27e363d86d40ff442fe-Paper.pdf
references.venue: Deep sets
references.authors: David W Zhang Gertjan J. Burghouts Cees G.M. Snoek
references.title: Set prediction without imposing structure as conditional density estimation
references.venue: In International Conference on Learning Representations,
references.authors: Yan Zhang Jonathon Hare Adam Prugel-Bennett
references.title: Deep set prediction networks
references.venue: Advances in Neural Information Processing Systems,
references.authors: Yan Zhang Jonathon Hare Adam Prügel-Bennett
references.title: Fspool: Learning set representations with featurewise sort pooling
references.venue: arXiv preprint arXiv:1906.02795,
sections.heading: 1 INTRODUCTION
sections.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.
sections.heading: 2 THE ONE-SHOT SET GENERATION PROBLEM
sections.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.
sections.heading: 2.1 METHODS FOR SET CREATION
sections.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.
sections.heading: 2.2 METHODS FOR SET UPDATE
sections.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.
sections.heading: 3 AN EQUIVARIANCE PERSPECTIVE ON SET AND GRAPH GENERATION
sections.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.
sections.heading: 4 THE TOP-N CREATION MECHANISM
sections.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.
sections.heading: 5 EXPERIMENTS
sections.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.
sections.heading: 5.1 SET MNIST
sections.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.
sections.heading: 5.2 SYNTHETIC DATASET
sections.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.
sections.heading: 5.3 MOLECULAR GRAPH GENERATION
sections.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.
sections.heading: 6 CONCLUSION
sections.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