Vector-Quantised Variational Autoencoders (VQ-VAE) To learn useful representations of fragments, we employ the VQ-VAE architecture (van den Oord et al., 2017), which maps molecule fragment graphs to a discrete latent space through using categorical distributions for the prior and posterior. The VQ-VAE defines a dictionary of k embedding elements, [s 1 , s 2 , ...s k ] ∈ R k×l . Given an input x (here the graph for a molecular fragment), let z e (x) ∈ R d×l be the output of the encoder (a MPNN in our case). We define l to be the same dimension for both encoder output embeddings z e (x) and dictionary embeddings s i , because input z q (x) is computed by finding the l 2 nearest neighbor dictionary elements for each row of z e (x):