Top-N: Equivariant Set and Graph Generation without Exchangeability

1 INTRODUCTION . 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 can not 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 . 2 THE ONE-SHOT SET GENERATION PROBLEM . 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 . 2.1 METHODS FOR SET CREATION . 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 .

Summary: 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.