NeVAE: A Deep Generative Model for Molecular Graphs

Motivation

• SMILES-based generation lacks structural similarity and permits multiple representations per molecule, while graph models also omit atom coordinates and rely on external optimization to seek high-property molecules.The introduction describes graph-based generation as an alternative to SMILES, but notes that prior property-search procedures have not found a sizeable set of high-value candidates.Source: §1 · pp. 1–2
• Existing graph models may require a fixed atom count, lack node-permutation invariance, have quadratic training complexity, and limit molecular diversity by combining a small set of graphlets.The paper notes these limitations constrain training on large molecules and generation of molecules with varied sizes.Source: §1 · p. 2

NeVAE

• NeVAE’s variational encoder aggregates information from different-hop neighborhoods, enabling variable-size, permutation-invariant graph encoding, while its joint multinomial edge decoder yields permutation-invariant inference with O(l) complexity.The paper contrasts this edge-decoding approach with prior Bernoulli modeling of potential edges, which has O(n^2) inference complexity.Source: §1 · p. 2
• NeVAE masks its edge distribution to prevent undesirable edges and models each generated atom’s spatial position with a Gaussian whose parameters depend on the atom and its neighbors.The masking mechanism is intended to guarantee local structural and functional properties in generated molecules.Source: §1 · p. 2

Property optimization

• A gradient-based algorithm optimizes NeVAE’s decoder to generate molecules with higher values of a target property and can optimize a molecule’s atom configuration for greater stability.The authors state that the algorithm benefits from the original decoder’s inductive bias and may be adapted to generative models for other data types.Source: §1 · pp. 2–3

Experiments

• NeVAE beats state-of-the-art models on validity, novelty, and uniqueness, while its optimized decoder identifies property values 121% higher than those found by several state-of-the-art Bayesian optimization and reinforcement learning methods.Experiments also report smooth interpolation between molecules, generation of valid molecules larger than those in the datasets, and optimization of atom configurations for lower potential energy.Source: §1 · p. 3

• A variational autoencoder uses a probabilistic decoder pθ(x|z), latent prior p(z), and approximate encoder qφ(z|x), with neural networks parameterizing the distributions.The decoder maps latent variables to observations, while the encoder maps observations to latent variables.Source: §2 · p. 3
• Because marginalizing latent variables to compute the observed-data log-likelihood is typically intractable, training maximizes a variational lower bound, or ELBO.Source: §2 · p. 3
• The lower-bound quality depends on the expressive ability of qφ(z|x), typically modeled as a normal distribution whose mean and variance are neural-network outputs conditioned on x.Source: §2 · p. 3

Prior

• A Poisson prior determines graph size, and each node has its own latent variable, so both the graph and its latent representation can vary in size.The latent prior for a graph with n nodes is N(0, I).Source: §3 · pp. 4, 6

Inference model

• NeVAE’s encoder symmetrically aggregates information from multiple hop distances, producing node embeddings invariant to node-label permutations and supporting graphs with varying numbers of nodes and edges.The encoder parameters do not depend on graph size.Source: §3 · pp. 4–5

Generative model

• The decoder samples atom types, edge counts, masked edges and bond types, and 3D coordinates whose distributions depend on each atom’s latent vector, neighbors, and chemical bonds.Masks can prevent edges or bond types that violate predefined local structural or functional properties.Source: §3 · p. 6

Training

• Randomized-tie-breaking BFS and label-independent source-node sampling yield a lower-bound training objective whose learned parameters are invariant to node-label permutations.Training maximizes the ELBO plus the log-likelihood of a Poisson distribution modeling graph size.Source: §3 · pp. 6–7

Scalability and implementation details

• Negative sampling reduces edge-likelihood computation from O(|V|^2) to O(l), giving complexity O(Sl) for S sampled source nodes.Here l is the number of true edges; the passage notes that real-world graphs are typically sparse.Source: §3 · p. 8

Property-oriented decoder design using variational inference

• The decoder objective balances higher target-property values against divergence from the original decoder, with ρ controlling the trade-off and encouraging plausible molecules.The objective penalizes low property values and gives an inductive bias toward molecules the original decoder would generate.Source: §4 · p. 9
• The optimal decoder can be sampled by accepting original-decoder molecules according to their exponentiated property value, though this may be impractical for high-dimensional molecules.The section motivates a practical approximation because original-decoder samples may have low probability under the optimal decoder.Source: §4 · pp. 8–9

A stochastic gradient-based algorithm

• The practical approximation iteratively updates a parameterized property-oriented decoder with stochastic gradient descent, initialized from the original decoder parameters.The updates aim to approximate the optimal decoder within the class of parameterized probabilistic decoders.Source: §4 · p. 10
• The log-derivative trick expresses the gradient as an expectation under the learned decoder, enabling an unbiased finite-sample Monte Carlo estimator.The estimator uses molecules sampled from the property-oriented decoder.Source: §4 · p. 10

5.1 Experimental setup

• The experiments sampled about 10,000 molecules from each of ZINC and QM9, then generated 10^6 graphs per trained model using prior or posterior latent samples.Bond masking enforced valid valences during training and sampling; masking during training increased training time by 5%.Source: §5.1 · p. 11

5.2 Quality of the generated molecules

• Across ZINC and QM9, NeVAE significantly outperformed several baselines in validity even without masking and had validity comparable to JTVAE, CGVAE, and GCPN.Novelty was nearly perfect for most methods, but NeVAE generated significantly more novel molecules than GraphVAE, GrammarVAE, CVAE, and SDVAE; most baselines also had poor uniqueness. Novelty and uniqueness are summarized in Table 1, and validity in Table 2.Source: Table 1 · p. 11, §5.2 · p. 12, Table 2 · p. 12
• Decoding sampled latent representations produced molecules topologically similar to the source, while increasing a single node’s latent perturbation made outputs differ more from it.The examples demonstrate smoothness across molecules from ZINC and QM9.Source: §5.2 · pp. 12–13, Figure 3 · p. 13, Figure 4 · p. 13

5.3 Property oriented molecule generation

• NeVAE’s property-oriented decoder found molecules with property values 121% higher than the best competitor, Bayesian optimization over JTVAE’s latent space.The decoder targets penalized logP and QED and, unlike the competitors, also provides plausible spatial configurations for generated molecules.Source: Table 3 · p. 14, §5.3 · pp. 14–15, Figure 5 · p. 14
• For a majority of 100 molecular graphs absent from training, optimizing spatial coordinates lowered potential energy and thus increased molecular stability.The optimization changes the decoder’s spatial-coordinate component rather than the entire decoder.Source: §5.3 · p. 15, Figure 7 · p. 16

• For several properties of interest, the optimized decoder identified molecules with property values 121% higher than those identified by several state-of-the-art methods.Source: §6 · p. 16
• The variational autoencoder supports permutation-invariant molecular graphs with varying numbers of nodes and edges, atom coordinates, and masked local structural and functional properties.Source: §6 · p. 16
• The authors propose adapting the design to dynamic graphs using recurrent neural networks or LSTM units, applying it to other real-world graphs, and exploring machine learning for retrosynthesis.Source: §6 · p. 16

Architecture details

• NeVAE’s learned parameters do not depend on graph size; its architecture uses linear forms and a symmetric sum aggregator.The authors chose the sum aggregator for simplicity.Source: A Implementation Details · p. 19

Hyperparameter tuning

• Cross-validation selected D = 7 (5), K = 3 (5), L = 10 (10), and lr = 0.005 (0.005) for synthetic (real) data, respectively.Training used stochastic gradient descent with Adam; the learning rate was searched logarithmically and the other hyperparameters arithmetically.Source: A Implementation Details · p. 19

Training with minibatch

• Minibatches grouped graphs with equal node counts, reducing computation-graph construction overhead from per sample to per batch.At each iteration, the method randomly selected a batch, built its computation graph using parameters from the previous iteration, and updated the parameters.Source: A Implementation Details · p. 19

B Additional details on Bayesian optimization

• The BO setup trains SGPs on 100 inducing points from 3,000 ZINC molecules, then runs five expected-improvement iterations with 50 new latent vectors per iteration.The 3,000 molecules are split 90%/10% into training and test sets; measures include test-set LL and RMSE, plus discovered-molecule property averages and valid and good fractions.Source: B Additional details on Bayesian… · p. 20
• JTVAE’s BO finds a few molecules with higher property values than NeVAE’s, but not a sizeable set of unique molecules with high values.Figure 8 plots uniqueness across BO iterations and sorts y1(m) and y2(m) values in decreasing order.Source: B Additional details on Bayesian… · p. 20, Figure 8 · p. 20
• NeVAE’s SGP outperforms all baselines in log-likelihood and RMSE; BO under NeVAE leads in E[y1(m)] and valid and good molecule fractions, while JTVAE leads in E[y2(m)].The good-molecule fraction counts molecules with y(m) > 0.Source: Table 5 · p. 20, B Additional details on Bayesian… · p. 20

C Additional Experiments on Synthetic Graphs

• NeVAE is demonstrated on generating graphs without triangles, a predefined local topological property.The section also examines whether a property-oriented decoder designed using variational inference can generate synthetic graphs with structural properties.Source: C Additional Experiments on Synthetic… · p. 21
• The experiments show that NeVAE can learn smooth latent representations of Kronecker graphs.Source: C Additional Experiments on Synthetic… · p. 21
• Quantitative experiments assess NeVAE’s ability to learn and mimic node- and edge-presence processes in Kronecker and Barabási-Albert graphs.Source: C Additional Experiments on Synthetic… · p. 21
• The section includes a scalability analysis and examines how node-label permutations affect decoder parameter estimation.Source: C Additional Experiments on Synthetic… · p. 21

C.1 Experimental setup

• Models were trained by maximizing the ELBO on synthetic graph datasets, then used to sample three sets of 1,000 graphs from their decoders.Each dataset contained 100 graphs with up to 1,000 nodes: triangle-free graphs or a 50%-50% mixture of Kronecker graphs with initiator matrices Θ1 = [0.9, 0.6; 0.3, 0.2] and Θ2 = [0.6, 0.6; 0.6, 0.6].Source: §C.1 · p. 21

C.2 Quality of the generated graphs

• With masking during training and testing, the model generates triangle-free graphs with 100% validity; training-only masking yields 68%, and no masking yields 57%.Test-time masking does not significantly increase generation time, while training-time masking increases training time by 18%. Figure 10 shows example generated graphs.Source: §C.2 · p. 21, Figure 10 · p. 21
• As the latent code interpolates between two Kronecker graphs, sampled graphs become similar to the corresponding endpoint graphs and inferred initiator matrices interpolate smoothly.The endpoint graphs use initiator matrices Θ0 = [0.9, 0.6; 0.5, 0.1] and Θ1 = [0.6, 0.6; 0.6, 0.6]; node labels are arbitrary.Source: §C.2 · p. 21, Figure 11 · p. 22
• Quantitative evaluation reports that the model’s generated graphs are indistinguishable from true Kronecker graphs.Table 6 reports rank correlation and top and bottom precision for models trained on Barabási-Albert or Kronecker graphs.Source: §C.2 · p. 22, Table 6 · p. 23
• The quality metrics compare model rankings with true-model likelihood rankings using Spearman rank correlation and top and bottom precision.For Barabási-Albert graphs, the evaluation uses 100 graphs with up to 1,000 nodes sampled with generation parameter m = 1; the ranking procedure focuses on the top and bottom 10% by true-model likelihood.Source: §C.2 · p. 22

C.3 Effect of permuting node labels on decoder parameter estimation

• The relabeling experiment uses 1,000-node Barabási-Albert and 1,024-node Kronecker graphs, comparing degree-proportional, maximum-degree, and uniform source-node sampling.It tracks mean differences between estimated decoder parameters over isomorphic graph pairs as training iterations increase.Source: §C.3 · p. 23
• Degree-based sampling performs best for Barabási-Albert graphs, while uniform sampling performs best for Kronecker graphs.The evaluation compares variation in estimated decoder parameters across isomorphic relabelings.Source: §C.3 · p. 23
• Barabási-Albert graphs have skewed degrees, so repeated sampling of a few sources yields similar parameter estimates; Kronecker graphs have more uniform degrees.This explains why degree-based sampling performs best in the former and worse in the latter.Source: §C.3 · p. 23

C.4 Effect of K (search-depth in encoder) on model performance

• Low K performs consistently well for Barabási-Albert graphs, whereas higher K improves performance for Kronecker graphs.The comparison uses rank correlation with respect to decoder search depth K.Source: Figure 13 · p. 24, §C.4 · p. 24

C.5 Scalability

• At 1,000 nodes, inference takes 67 + 20 seconds per SGD iteration (batch size 10), and graph sampling takes 5 seconds; both procedures scale to about 1,000 nodes.The graphs have average degree 3.Source: Figure 14 · p. 24, §C.5 · p. 24
• For 50-node graphs of average degree 3, inference takes less than 10 seconds per iteration and decoder sampling takes 0.5 seconds.Inference uses a batch of 10 graphs.Source: §C.5 · p. 24

C.6 Property oriented graph generation

• Lower ρ yields lower graph diameters and higher clustering coefficients in property-oriented generation.NeVAE is trained on 10,000 Barabási–Albert graphs, then optimized for either property; Figure 15 shows distributions across ρ values.Source: §C.6 · pp. 24–25, Figure 15 · p. 25
• For graphs constrained to diameter below the specified bound, lower ρ produces a higher success rate.The decoder is trained with a loss based on the excess over the maximum allowed diameter; success rate measures graph-generation quality.Source: Table 7 · p. 25, §C.6 · p. 25