T OP -N: E QUIVARIANT SET AND GRAPH GENERATION WITHOUT EXCHANGEABILITY
A BSTRACT
This work addresses one-shot set and graph generation, and, more speciﬁcally, the parametrization of probabilistic decoders that map a vector-shaped prior to a distribution over sets or graphs.
Sets and graphs are most commonly generated by ﬁrst sampling points i.i.d.
from a normal distribution, and then processing these points along with the prior vector using Transformer layers or Graph Neural Networks.
This architecture is designed to generate exchangeable distributions, i.e., all permutations of the generated outputs are equally likely.
We however show that it only optimizes a proxy to the evidence lower bound, which makes it hard to train.
We then study equivariance in generative settings and show that nonexchangeable methods can still achieve permutation equivariance.
Using this result, we introduce Top-n creation , a differentiable generation mechanism that uses the latent vector to select the most relevant points from a trainable reference set.
Top-n can replace i.i.d.
generation in any Variational Autoencoder or Generative
Adversarial Network.
Experimentally, our method outperforms i.i.d.
generation by 15% at SetMNIST reconstruction, by 33% at object detection on CLEVR, generates sets that are 74% closer to the true distribution on a synthetic molecule-like dataset, and generates more valid molecules on QM9.
I NTRODUCTION
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 learn to sample stable drug-like molecules conditioned on some molecule-level properties such as solubility or synthesizability.
Set and graph generation are very similar problems, and one can navigate between set and graph representations using Set2Graph functions (Serviansky et al., 2020) or graph neural networks.
We therefore generally treat them together in this paper, but only use the terminology of sets to avoid overly abstract notations.
The distinctive properties of graph generation (e.g., graph matching problems) are discussed when needed.
There are two main classes of probabilistic decoders for sets: recursive and one-shot.
Recursive generators are conceptually simpler, since they add points one by one (Liu et al., 2018;
Liao et al., 2019; Sun et al., 2020;
Nash et al., 2020).
They are however slow and introduce an order in the points that does not exist in the data.
In this work, we instead focus on one-shot generation, which allows for more principled designs.
By deﬁnition, one-shot models 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 equivariant 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 by ﬁrst sampling points independently from a normal distribution, and then appending the latent vector to each point.
This design allows the generation of any number of points and decouples the set cardinality from the number of trainable parameters.
Furthermore, it generates exchangeable distributions (all permutations of a set are equally likely), a property which is commonly held as the equivalent of equivariance for generative models.
However, it was empirically observed that VAEs based on independent sampling are hard to train, which results in limited performance (Krawczuk et al., 2021).
In this work, we ﬁrst propose a theoretical argument to this empirical observation, by showing that
VAEs that use independent sampling only optimize a proxy to the evidence lower bound (ELBO).
As the standard deﬁnition of equivariance cannot be used in generative settings, we then propose a generalization of this notion called ( F, l ) -equivariance: informally, an architecture is ( F, l ) -equivariant if the parameter updates do not depend on the group elements used to represent training data.
We derive sufﬁcient conditions for equivariance in this setting.
They reveal that ( F, l ) -equivariance explains the loss functions commonly used both in generative and discriminative tasks, and suggest that exchangeability may not be useful in GANs and VAEs.
Based on these results, we ﬁnally propose a non-exchangeable set creation method called Top-n creation .
Our method relies on a trainable reference set where each point i has a representation r i ∈ R c and an angle φ i ∈ R a .
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 initially proposed for graph coarsening (Gao & Ji, 2019).
Top-n can eventually be integrated in any VAE or GAN to form a complete generative model for sets or graphs.
Our method is easier to train than stochastic generators, and has better generalization performance than other existing methods.
We benchmark Top-n on both set and graph generation tasks: it is able to reconstruct the data more accurately on a set version of MNIST, generalize better on the CLEVR object detection dataset, ﬁt more closely the true distribution on a dataset of synthetic molecule-like structures in 3D, and generate realistic molecular graphs on QM9.
T HE ONE - SHOT SET GENERATION PROBLEM
We consider the problem of learning a probabilistic decoder f that maps latent vectors z ∈ R l to multi-sets 1 X
= { x 1 , ..., x n } that contain a varying number n of points x i in R d .
Given sample sets from an unknown distribution D , f should be such that, if z is drawn from a prior distribution p Z ( z ) , then the push-forward measure f # ( p Z ) (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 ∈ X
= (cid:83) n ∈ N R n × d where each row represents a point x i ∈ R d .
Algorithms that return a set implicitly assume the use of a function mat-to-set that maps X to the corresponding set X .
Existing architectures for one-shot generation use one of the graphical models described in Figure 1.
First, a number of points for the set has to be sampled.
Most works assume that the set cardinalities are known during training.
At generation time, they sample n from the distribution of set cardinalities in the training data.
This method assumes that the latent vector z is independent of the number of points n , so that the generative mechanism writes p ( X| n, z ) p ( n ) p ( z ) .
Kosiorek et al.
(2020) instead propose 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.
Once n is sampled, one-shot generation models can formally be decomposed into several components: a ﬁrst function g (that we call creation function ) maps the latent vector to an initial set
X 0 ∈ R n × c .
This function is usually simple, and is therefore not able to model complex dependencies within each set.
For this reason, X 0 may then be reﬁned 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. M ETHODS FOR SET CREATION
MLP based Many existing methods (Achlioptas et al.
; Zhang et al.
(2020; 2021b) for sets,
Guarino et al.
(2017); De Cao & Kipf (2018); Simonovsky & Komodakis (2018) for graphs) learn a
MLP from R l to R n max c , where n max is the largest set size in the training data.
The output vector is then reshaped as a n max × c matrix, and masked to keep only the ﬁrst n rows (Figure 2).
MLPs ignore the symmetries of the problem and are only trained to generate up to n max points, with no ability to extrapolate to larger sets.
Despite these limitations, they often perform on par with more complex methods for small graphs (Madhawa et al., 2020; Mitton et al., 2021).
We explain in Section 3 this surprising phenomenon.
Independent sampling 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 (K ¨ ohler et al., 2020; Yang et al., 2019b; Kosiorek et al., 2020; Stelzner et al., 2020; Satorras et al., 2021; Zhang et al., 2021a;
Liu et al., 2021).
The main advantage of independent sampling 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 X 0 are equally likely.
This property is widely considered as the equivalent of equivariance for generative models (Yang et al., 2019a; Biloˇs & G¨unnemann, 2021;
Kim et al., 2021; K¨ohler et al., 2020; Li et al., 2020).
However, it was empirically observed that VAEs built with a i.i.d.
creation mechanism fail to ﬁt the training data correctly (Krawczuk et al., 2021), which reﬂects in the poor quality of the sampled sets.
To understand why, consider a VAE made of an encoder q φ ( z | X ) and a decoder f θ ( X | z ) parametrized by φ and θ .
Variational autoencoders maximize the evidence lower bound (ELBO) L , which is a proxy for the data likelihood under the model:
[Equation]
Probabilistic decoders f θ ( z ) based on independent sampling are stochastic, so that log p θ ( X , z ) cannot be computed in close form.
By conditioning on the initial set X 0 and using Jensen’s inequality, we have:
[Equation]
which gives in expectation
[Equation]
Methods based on independent sampling use a Monte-Carlo estimate for ∇ θ,φ L (cid:48)
( X ) that leverages the reparametrization trick.
They therefore only optimize a lower bound of the ELBO, which could explain why they are difﬁcult to train.
First-n Instead of sampling points, Zhang et al.
(2019) and Krawczuk et al.
(2021) propose to always start from the same learnable set X ref ∈ R n max × c , and mask this matrix to keep only the ﬁrst n rows: we therefore call this method First-n creation.
Similarly to the independent sampling method, the latent vector is then concatenated to each point of this set.
Empirically, First-n converges much faster than sampling-based methods (Krawczuk et al., 2021), but the network is only trained to generate up to n max points and has no ability to extrapolate to larger sets.
Furthermore, selecting the ﬁrst n rows of the reference set X ref introduces a bias because the ﬁrst 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 normal entries i.i.d.
An alternative is to ﬁrst generate a set, and then use a Set2Graph update function in order to learn the graph adjacency matrix as in (Bresson & Laurent, 2019; Krawczuk et al., 2021).
2. M ETHODS FOR SET UPDATE
Since all creation methods except MLPs can only generate very simple sets and adjacency matrices, additional layers are usually used to reﬁne 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 (Morris et al., 2019).
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.
A PERMUTATION EQUIVARIANCE VIEW ON SET GENERATION
Whereas exchangeability is usually considered as a key feature of independent sampling, we have seen that this method empirically does not outperform other strategies.
To understand why, we need to study equivariance in generative models and propose a relevant deﬁnition in this setting.
As set and graphs are unordered, the symmetric group S =
(cid:83) n ∈ N ∗ S n containing all permutations is a symmetry of these tasks.
A permutation π ∈ S n 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 ∈ R h unchanged ( π.
z = z ).
This symmetry constitutes a useless factor of variation in the data that should be factored out in the latent space (i.e., π.
z = z ).
In discriminative models, symmetries are accounted for when a neural network f is equivariant to the action of a group, which writes π.f ( X ) =
f ( π. X )
(Kondor, 2008).
When the input of f is a vector, imposing π.f ( z ) =
f ( π. z )
=
f ( z ) however only allows for solutions where all rows are equal, which is too restrictive.
To solve this issue, we propose a deﬁnition called ( F, l ) -equivariance which generalizes the common one, but provides more relaxed conditions in generative settings.
Our proposition is based on the assumption that the main role of equivariance is to make data augmentation useless. In discriminative settings, this is normally done by combining an equivariant model with an invariant loss function.
For example, the l 2 loss is commonly used to learn the future state of a n -body system, but not the l 1 loss, as it is not rotation invariant.
Formally, if F Θ = { f θ : X → Y ; θ ∈ Θ } is an hypothesis class of G -equivariant functions from X to Y (for example a neural architecture parametrized by θ ), then the loss functions l should satisfy
[Equation]
Furthermore, we observe that when l satisﬁes Eq.
, the gradients with respect to the parameters satisfy ∇ θ l ( f ( g. X ) , g. Y )
= ∇ θ 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.
We propose to use this property to deﬁne equivariance:
Deﬁnition 1 ( ( F, l ) -equivariance) .
Consider an hypothesis class F Θ ⊂ Y X , a group G that acts on X and Y and a loss function l deﬁned on Y .
We say that the pair ( F Θ , l ) is equivariant to the action of G if the dynamics of θ ∈ Θ trained with gradient descent on l do not depend on the group elements that are used to represent the training data.
By construction, using an equivariant architecture and an invariant loss is sufﬁcient for ( F, l ) - equivariance in discriminative settings.
For standard generative architectures for sets and graphs, we derive the following sufﬁcient conditions (proofs are given in Appendix A):
Lemma 1.
Sufﬁcient conditions for ( F, l ) -equivariance:
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.
VAEs: if F is an encoder-decoder architecture with a permutation invariant encoder, and the reconstruction loss l satisﬁes ∀ π ∈ S , l ( π.
X , ˆ X )
=
l ( X , ˆ X ) , then ( F, l ) is permutation equivariant.
No constraint is imposed on the decoder function.
Normalizing ﬂows: 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 K¨ohler et al.
As desired, ( F, l ) -equivariance does not impose π.f ( z ) =
f ( π. z ) in generative architectures, but still makes data augmentation unnecessary.
Furthermore, the constraints of Lemma 1 are satisﬁed by most existing architectures, including early ones (Simonovsky & Komodakis, 2018; De Cao & Kipf, 2018; K ¨ ohler et al., 2020).
In particular, the constraint on the loss for VAEs is satisﬁed by the two loss functions commonly used for sets, namely Chamfer loss and the Wasserstein-2 distance (deﬁned in Appendix B).
Our deﬁnition, introduced by observing common practice in discriminative settings, is therefore able to explain common practice for generative tasks as well.
We ﬁnally observe that exchangeability does not appear in the sufﬁcient conditions for GANs and
VAEs.
To understand why, 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.
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 proved 2 .
This observation explains why independent sampling creation does not outperform nonexchangeable set creation methods such as MLPs and First-n.
In the following section, we therefore design a new creation mechanism without worrying about the model exchangeability.
T HE T OP - N CREATION MECHANISM
We have seen in Section 2 that existing set creations methods suffer from important limitations: independent sampling makes it hard to train the model, while MLPs and First-n use a ﬁxed 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 3.
Similarly to First-n, Top-n also uses a reference set, but in Top-n this set can have an arbitrary size n 0 .
Each point in this set is a pair ( φ , r ) : the angle φ ∈ R a is used to decide when to select the point, and r ∈ R c contains the representation of the point.
Given a latent vector z ∈ R l × 1 , a reference set made of angles Φ ∈ R n 0 × a
and representations R ∈ R n 0 × c , as well as learnable matrices W 1 to W 4 (respectively of sizes 1 × c, 1 × c, l × c, l × c ), Top-n creation computes:
[Equation]
The crux of the Top-n creation module is to select the points that will be used to generate a set based on the value of the latent vector (Eq. 5, 6 ,7).
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.
In Top-n, we modulate the representation of the selected points R [ s ] with the cosines c (Eq.
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:
[Equation]
Equations (5) to (9) 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 tends 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 ( X
=
X ref [ s ] (cid:12) ˜ c ) in Eq.
(9) 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 exploit the equivalence between summation and concatenation when a linear layer is applied, which writes
[Equation]
for W = cat( W 1 , W 2 ) .
Contrary to the concatenation (left-hand side), the sum (right-hand side) does not compute z T W 2 several times, which reduces the complexity of this layer from O ( n ( c + l ) c ) to O ( nc 2 + cl + nl ) .
Again, we combine the sum and multiplicative modulation in a FiLM layer to build a more expressive model in Eq.
Our algorithm retains the advantages of First-n creation, but replaces the arbitrary selection of the ﬁrst n points by a mechanism that learns to select the most relevant points for each 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
[Table]
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 overﬁtting and generalize better.
Top-n creation can be used in any GAN or VAE architecture as a replacement for other set creation methods.
It is however not suited to normalizing ﬂows, because it is based on a hard selection process which is not invertible (the value of the reference points that are not selected is not used).
Since First-n and Top-n use a ﬁxed 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 (cid:15) , there is a 2-layer row-wise MLP f and reference points { x 1 , ..., x n } such that, for any set { y 1 , ..., y n } of points in R d with ∀ i,
||
y i
|| ≤ M there is a latent vector z of size nd × 1 that satisﬁes: || f (cat( X , 1 n z T )) − Y
[Equation]
) where W 2 denotes the Wasserstein-2 distance.
E XPERIMENTS
We compare Top-n to other set and graph creation methods on several tasks: autoencoding a set version of MNIST, detecting objects on CLEVR, generating realistic 3D structures on a synthetic molecule-like dataset, and generating varied valid molecules on the QM9 chemical dataset 3 .
All training curves are available in Appendix E.
5. S ET MNIST
We ﬁrst perform experiments on the SetMNIST benchmark, introduced in Zhang et al.
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., 2019)
TSPN (Kosiorek et al., 2020) 4 , which are respectively a sort of diffusion model and a transformer-based autoencoder.
Experiment details can be found in Appendix D.1.
The results for both methods are very similar (Figure 1): the model based on independent sampling has poor performance and needs more epochs to be trained than First-n and Top-n. Top-n performs consistently better for both DSPN and TSPN, which shows that it is able to select the most relevant reference points for each set.
5. O BJECT DETECTION ON CLEVR
We further benchmark Top-n on object detection with the CLEVR dataset, made of 70k training images and 15k validation images representing simple objects.
Again, we use the implementation of DSPN and the setting proposed in (Zhang et al., 2019).
Two tasks are evaluated.
In the ﬁrst one, the goal is to predict bounding boxes in each image.
In the second one, the full scene should be predicted (with the shape, color, size and material of the objects).
Images are encoded using a pretrained ResNet34 architecture – the resulting vector is used as input to the set generation model.
[Table]
[Table] [Table]
The model is trained on 10 DSPN iterations and evaluated on
[Table]
which is the setting that gave the best results in (Zhang et al., 2019).
Results are presented in Tables 2 and 3.
For bounding box prediction (which is a simpler task), independent sampling and Top-n outperform MLP and First-n creation.
As the metrics are computed for test data (which is not the case in SetMNIST), these results suggest that MLP and First-n creation may overﬁt the training images.
On the full scene prediction, Top-n outperforms all other methods.
5. S YNTHETIC DATASET
As previous datasets do not measure generation quality
[Table]
we further benchmark the different set generators on a synthetic dataset for which the quality of the generated sets can be assessed.
This dataset is a simpliﬁed 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 predeﬁned valency.
The generation procedure is described in Appendix D.2.
The goal is to reconstruct the atom positions and generate new realistic sets.
During training, we measure the Wasserstein ( W 2 ) distance between the input and reconstructed sets.
At generation time, we compute the W 2 distance between the distribution of valencies in the dataset and in the generated set, the proportion of generated atoms with valency 0 or more than 4, as well as a diversity score to ensure that a method does not always generate the same set.
We also measure the same metrics in an extrapolation setting where sets have on average 10 more points.
We train a VAE with different creation methods on this dataset.
Details about the model and the loss function can be found in Appendix D.2, as well as an ablation study on the number of points in the reference set.
The results are shown in Table 4.
We observe that the independent sampling generator generalizes well but reconstructs the training sets very poorly, which reﬂects the fact that the model is hard to train due to the stochastic i.i.d.
generation.
They tend to generate points that are too far apart, an issue which disappears when generating more points.
On the contrary, MLP and First-n
[Table]
overﬁt the training data at the expense of generation quality.
Finally, Top-n is able to generate points that have the right valency and generates the best new samples.
5. M OLECULAR GRAPH GENERATION
Finally
[Table]
we evaluate Top-n on a graph generation task.
We train a graph VAE (detailed in Appendix
D.3) on QM9 molecules and check its ability to generate a wide range of valid molecules.
As a generation metric
[Table]
we simply report the validity (i
[Table]
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 predeﬁned 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.
We also note that recursive methods such as Li et al.
(2018) can often generate higher rates of valid molecules because they 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.
While MolGAN and GraphVAE+RL both suffer from mode collapse, our method is able to generate a higher rate of valid and unique molecules.
We observe that the independent sampling method is not able to obtain a good train loss (Appendix E), which reﬂects 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 in Top-n.
C ONCLUSION
In this work, we strengthened the theoretical foundations of one-shot set and graph generation.
We showed 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 which 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 ﬁnally note that most of our experiments feature only small sets, which is in line with existing literature (Meldgaard et al., 2021;
Satorras et al., 2021; Mitton et al., 2021).
This observation leads to a simple question: 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 n max d , standard universal approximation results are likely to yield a sample complexity that is exponential in n max .
Whether 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 predeﬁned target.
These molecules can then be reﬁned by computational chemists in order to optimize 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 Code is included in the supplementary materials.
A CKNOWLEDGMENTS
Cl´ement Vignac would like to thank the Swiss Data Science Center for supporting him through the
PhD fellowship program (grant P18-11).
A P ROOF OF LEMMA
Generative adversarial networks Given a set generator f , a discriminator function d , and
X 1 , ..., X m a training dataset, the standard loss function for GANs is formulated as
[Equation]
In order to obtain l ( f, d, X 1 , ..., X m ) =
l ( f, d, π 1 .
X 1 , ..., π m .
X m ) for every choice of π i , it is therefore sufﬁcient to choose a permutation invariant discriminator.
Autoencoder based models We assume that the autoencoder is made of a permutation invariant encoder enc and an arbitrary decoder f .
For any set size n , set X ∈ R n × d and permutation π ∈ S n we have
[Equation]
B L OSS FUNCTIONS FOR SETS
Chamfer’s loss and the Wasserstein 2 distance are deﬁned as
[Equation]
where Γ is the set of couplings (i.e., bistochastic matrices) between X and X (cid:48) .
Both loss functions solve a matching problem over the space of permutations, which makes them invariant to permutations of one argument, as required by Lemma 1.
Chamfer’s loss runs in quadratic time, while the standard implementation of Wasserstein distance runs in O ( n 3 ) if both sets have the same size.
However, efﬁcient algorithms exist for approximating the Wasserstein distance, and the computations can be parallelized on GPUs (Cuturi, 2013; Feydy et al., 2019).
Note however that the equation l ( π.
X , ˆ X )
=
l ( X , ˆ X ) is may be difﬁcult to satisfy in other settings: matching graphs up to permutations, or sets up to the SE (3) group, leads to difﬁcult problems for which no polynomial time algorithm is known (M´emoli, 2007).
In these settings, the design of VAE is harder and other architectures may be better suited.
C P ROOF OF PROPOSITION
Proof.
We ﬁrst give the proof for First-n creation:
First-n Given a set Y =
{ y 1 , ..., y n } of points in R d , we propose to sort the points
y i by alphanumeric sort (sort the values using the ﬁrst feature, then sort points that have the same ﬁrst features along the second feature, etc.).
We denote the resulting matrix by Y (cid:48) .
We choose as a latent vector z
= ﬂatten ( Y (cid:48) )
∈ R nd .
It is the vector that contains the representation of y (cid:48) 1 in the ﬁrst d features, y (cid:48) 2 is the next d features...
By construction, z is a permutation invariant representation of the set
Y (this reﬂects the fact that there are canonical ordering for sets, which is not the case for general graphs).
We choose as a reference set the canonical basis in R n ( R = I n ).
After the latent vector is appended to the representation of each point, we have r i =
( e i , z ) .
We are now looking for a function f that allows to approximate the set Y , i.e., that satisﬁes ∀ i ≤ n, f ( e i , z ) =
z [ i d : ( i + 1) d ] .
We can choose f ( e i , z ) =
e Ti W 1 W 2 z , where
W 1 =
[Equation] [Equation]
[Equation]
1 1(cid:35) ∈ R n × nd and W 2 =           ...

 ∈ R nd × d .
This function is continuous, and its output is Y (cid:48)
= reshape n × d ( z ) .
Y (cid:48) is equal to Y up to a permutation of the rows, so that || Y − Y (cid:48)
||
W 2 = 0 .
If the entries of z are bounded, we can use standard approximation results for continuous functions over a compact space (Cybenko, 1989) to conclude that it be uniformly approximated by a 2-layer MLP.
Top-n Consider a Top-n network with n reference points such that:
• The angles of the reference points are 2d vectors such that φ i =
(cos( in π 4 ) , sin( in π 4 )) .
• The representations are r i =
e i / cos( in 0 π 4 ) , where ( e i ) is the canonical basis in R n .
• The MLP of equation 1 (that predicts an angle from the latent vector) always outputs (1 , 0) .
Then this Top-n creation module is equivalent to the First-n module build previously: for any set, it selects the ﬁrst n points and returns X
[ i ] =
e i .
The same MLP that is built for First-n can therefore be used for this network.
D D ETAILS ABOUT THE EXPERIMENAL SETTING D. S ET MNIST AND CLEVR
Since we use existing code for this task, we refer to the respective papers (Zhang et al., 2019) and (Kosiorek et al., 2020) for details on the model used and the loss function.
The code used for TSPN is not the original code (which is not available) but a reimplementation (not by one of the authors of the present paper).
The reader will notice that our results for DSPN are approximately 3 times worse than in the original paper.
The reason is that we ﬁxed what we believe to be a bug in the implementation of Chamfer’s loss in DSPN: a mean over channels was used instead of a sum, which explains this difference of a factor 3.
We also note that the results for TSPN are around 3 times worse than the original paper.
One possible reason could be that the authors of TSPN used the code of DSPN and had a similar bug.
For Top-n generation, we set the number of points in the reference set to twice the cardinality n of the generated sets.
We also experimented with n 0 = n which resulted in better performance for
DSPN (with a Chamfer loss of 6 .
14 ± 0 .
56 e-5), but not for TSPN ( 16 .
07 ± 0 . 47 ).
We observed that reducing the learning rate improves results for all methods: TSPN was therefore trained forepochs with a learning rate of 5e-4, and DSPN with a learning rate of 1e-4 for 200 epochs.
No other hyper-parameter was tuned.
D. S YNTHETIC SET GENERATION
Dataset generation procedure Each set is created via rejection sampling: points are drawn iteratively from a uniform distribution within a bounding box in R 3 .
The ﬁrst point is always accepted, and the next ones are accepted only if they satisfy a predeﬁned 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 simpliﬁcation of real molecules in several aspects: there are only single
[Table]
bonds, angles between bonds are not constrained and the atom types are only deﬁned by the valency, which reﬂects the fact that atoms with the same valency tend to play a similar role and be more interchangeable.
Model The set encoder is made of a 2-layer MLP, 3 transformer layers followed by a PNA global pooling layer (that computes the sum, mean, max and standard deviation over each channel)
(Corso et al., 2020) and a 2-layer MLP.
The decoder is made of a set generator followed by a linear layer, 3 transformer layers and a 2-layer MLP.
We use residual connections when possible and batch normalization between each layer.
The reference set contains 35 points.
We experimented with the two ways to sample the number of points presented in Section 2, but we found the results to be quite similar.
We therefore opted sampling the number of points from the data distribution, which is the simplest method.
Loss function We use a standard variational autoencoder loss with a Wasserstein reconstruction term and two additional regularizers.
Given an input set X and its reconstruction ˆ X , the loss can be written:
[Equation]
where
[Equation]
prevents atoms from being generated too close to each other.
reg 3 ( X ) penalizes atoms that have either no neighbor, or a too large valency.
It is computed in the following way:
• for each point i , compute s i = sort
(( d ij ) j ≤ n, j (cid:54) = i ) .
This vector contains the sorted distances between i and all other points.
Points that are at distance less than d 1 = neighbor-distance from i are considered as its neighbours.
• Compute l 1 ( i ) =
( s i 0 − d 1 ) + .
This term penalizes atoms that have no neighbour.
• Compute l 2 ( i ) =
(cid:80) n − 1 j = max-valency ( d 1 − s i j )
+ .
This term penalizes atoms that have too many neighbors.
• reg 3 ( X ) is deﬁned as (cid:80) 1 ≤ i ≤ n l 1 ( i ) + l 2 ( i ) .
Training details In order to train the model efﬁciently, mini-batches have to be used.
This may not be easy when dealing with sets and graphs, since they do not have all the same shape.
To circumvent this issue, we reorganise the training data in order to ensure that all sets inside a mini-batch have the same size.
At generation time, this method cannot be applied, so we simply generate sets one by one.
The optimizer is Adam with its default parameters.
We use a learning rate of 2 e − 4 and a scheduler that halves it when reconstruction performance does not improve signiﬁcantly after 750 epochs.
Experimentally, we found the learning rate decay to be important to achieve good reconstruction.
We also run a study with different reference set sizes.
For this purpose, we slightly modify our dataset so that each set has only up to 11 points (still with 9 points on average).
The reason is that there is more ﬂexibility in the choice of the reference size if the maximal size is not too large.
By training a Top-n network with several reference set sizes, we obtain the results of Table 6.
D. M OLECULAR GRAPH GENERATION ON QM
Model Our encoder is a graph neural network is made of 3 message-passing layers followed by a
PNA global pooling layer and a ﬁnal MLP.
For the decoder, we use a set creation method followed by transformer layers.
The resulting representation is then processed by i) a Set2Graph layer (Serviansky et al., 2020) followed by two MLPs to generate edge probabilities and edge features a MLP which generates node features ii) a MLP that takes as input the set representation and the valencies predicted for each atom, and returns an atom type.
Graph matching and loss function As explained in Appendix B, the loss function of the variational autoencoder should solve a graph matching problem, which is hard in general.
Instead of using a proper graph matching method, we propose to use the atom types to perform an imperfect but much cheaper alignment between the target and the predicted molecules.
For both molecules, we compute a score for each atom i deﬁned as:
[Equation]
This score cannot differentiate between all atoms in each molecule, but it reduces drastically the number of permutations that can represent the same input.
It is motivated by the fact that empirically, we observe that our method quickly learns to reconstruct the molecular formula very well.
Once they are all computed, we use these scores to sort the atoms reorder the adjacency matrix and the atom and edge types.
Our model learns to predict a probabilistic model for the atom types, edge presence and edge types.
For this purpose, we use standard cross entropy loss between the predicted probabilities and the ground truth.
However, these metrics can be hard to optimize because of the imperfect graph matching algorithm.
We therefore regularize these metrics with several other measures at the graph level, that do not depend on matching:
• The mean squared error between the real atomic formula and the predicted one.
The mean squared error between the average number of edges per atom in the input and predicted molecule.
• The mean squared error between the distribution of edge types in the real and predicted molecule.
Finally, we add a matching dependent term, which is the mean squared error between the valencies of the input and the target molecule.
Training details
The model is trained over 600 epochs with a batch size of 512 and a learning rate of 2 e − 3 .
It is halved after 100 epochs when the loss does not improve anymore.
The optimizer is Adam with default parameters.
The reference set has 12 points.
When using more points, we obtained overall similar results, but with a larger variance.
E T RAINING CURVES
[Figure]