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