Source-linked AI summary
GraphDF: A Discrete Flow Model for Molecular Graph Generation
Youzhi Luo, Keqiang Yan, Shuiwang Ji
TL;DR
Molecular graph generation must model a huge discrete space, whereas most existing approaches use continuous latent variables and dequantization. GraphDF introduces a discrete normalizing-flow model that maps latent variables to graph nodes and edges, and reports improved performance across random, property, and constrained optimization tasks.
Problem
Existing molecular graph generators use continuous latent variables for naturally discrete structures, while dequantization distorts the original distribution and complicates its modeling.
Method
GraphDF sequentially samples discrete latent variables and maps them to molecular graph nodes and edges using invertible modulo shift transforms.
Results
GraphDF achieves new state-of-the-art results across three molecule-generation tasks and outperforms prior methods over most metrics on three datasets.
Takeaways & Limitations
Discrete latent variables let GraphDF avoid Jacobian computation and dequantization effects while supporting accurate graph modeling and exploration of chemical space.
Abstract
from arXiv · showhide
We consider the problem of molecular graph generation using deep models. While graphs are discrete, most existing methods use continuous latent variables, resulting in inaccurate modeling of discrete graph structures. In this work, we propose GraphDF, a novel discrete latent variable model for molecular graph generation based on normalizing flow methods. GraphDF uses invertible modulo shift transforms to map discrete latent variables to graph nodes and edges. We show that the use of discrete latent variables reduces computational costs and eliminates the negative effect of dequantization. Comprehensive experimental results show that GraphDF outperforms prior methods on random generation, property optimization, and constrained optimization tasks.
1. Introduction
Molecular graph generation seeks novel molecules with specific properties, but the enormous discrete molecular space and continuous-latent modeling create challenges. GraphDF addresses these issues with discrete latent variables and invertible modulo shift transforms.
- Molecule design seeks novel structures with specific properties across a naturally discrete molecular space estimated at 10^33.Deep generative models increasingly represent molecules as graphs and generate molecular graphs.
- Existing graph-generation methods map discrete molecular structures into continuous latent spaces for generation.The supplied introduction describes this as a common conceptual and computational pattern.
- Dequantization adds real-valued noise to discrete graph data, preventing accurate capture of the original discrete distribution and complicating training.This limitation makes modeling graph-structure distributions and generating diverse molecules difficult.
- GraphDF uses discrete latent variables and invertible modulo shift transforms to map sampled variables sequentially to molecular graph nodes and edges.The discrete transform also eliminates Jacobian-matrix computation costs and addresses dequantization limitations.
- GraphDF outperforms prior state-of-the-art methods across multiple molecule-generation tasks.The introduction reports this as the main experimental conclusion.
2. Background and Related Work
Molecular generation research has progressed from SMILES sequences to graph and flow-based models, while the paper frames generation, property optimization, and constrained optimization as distinct objectives. Flow models provide invertible mappings and exact likelihood computation, but efficient training and sampling require tractable Jacobians.
- Molecule Generation: Molecule-generation objectives include learning a sampling distribution, maximizing property scores, and optimizing a molecule under a similarity constraint.The constrained objective requires SIM(M, M′) > δ.
- Related Work: SMILES-based sequence models struggled to learn complicated grammatical rules and generate syntactically valid sequences.Later methods instead represented molecules as graphs.
- Related Work: Graph-generation methods include VAE-based approaches, sequential node-or-edge addition, and one-shot graph construction.The related work surveys these families without reducing them to a single generation strategy.
- Flow Models: Normalizing flows define invertible transformations between latent variables and data, enabling exact likelihood computation and molecular-graph density modeling.GraphDF builds its discrete flow model on this normalizing-flow framework.
- Flow Models: Flow models train by maximizing exact data log-likelihoods and sample by transforming latent draws through x = fθ(z).The mapping must be invertible, and its Jacobian must be tractable for efficient training and sampling.
- Related Work: Prior flow-based graph models generate structures either one-shot from node and adjacency matrices or sequentially through autoregressive graph construction.The cited examples include GraphNVP, GRF, MoFlow, GNF, and GraphAF.
3. A Discrete Flow Model for Graphs
GraphDF models molecular graphs with discrete latent variables and sequentially maps them to nodes and edges through invertible modulo shift transforms. Its design avoids dequantization and Jacobian computation while incorporating graph-structure information into conditional generation.
- Motivation and novelty: GraphDF uses discrete latent variables to model molecular graph structures, addressing the mismatch between discrete graphs and continuous latent-variable methods.Prior discrete-flow work handled 1D sequences but not richly structured graphs.
- Sequential generation: The model generates an empty graph sequentially, adding a node and then edges to existing nodes until no further edge is added.Node and edge types are generated autoregressively from previously generated elements, with invalid bond-valency choices re-sampled.
- Discrete flow construction: GraphDF samples node and edge latents from multinomial distributions and reversibly maps them to graph elements using compositions of modulo shift modules.The modulo shifts use conditional factors derived from previously generated nodes and edges.
- Advantages of discreteness: Discrete invertible mappings remove the Jacobian term, allowing sampled node and edge probabilities to be obtained directly and reducing computational cost.The paper attributes this reduction to omitting Jacobian computation from the discrete change-of-variable formula.
- Advantages of discreteness: GraphDF avoids dequantization, which otherwise distorts discrete graph distributions and can map the same graph to distant latent points during training.The paper argues that this distortion can cause model oscillation and failure to converge.
- Invariant conditional generation: Graph convolutional networks extract conditional information from intermediate sub-graphs while guaranteeing invariance to the ordering of previously generated elements.This addresses the limitation of sequence models that lack rich graph structure and assume an order absent from intermediate sub-graphs.
4. Experiments
GraphDF is evaluated on random, property, and constrained molecular graph generation, plus generic graph generation. Across these experiments, it generally outperforms prior methods while maintaining validity, reconstruction, and optimization performance.
- Experimental scope: GraphDF is evaluated against strong baselines and flow-based methods on random, property, constrained, and generic graph-generation tasks.The random-generation evaluation uses ZINC250K, QM9, MOSES, COMMUNITY-SMALL, and EGO-SMALL datasets.
- Random generation: GraphDF achieves 100% validity and reconstruction rate, with reasonably high uniqueness and novelty across three molecular datasets.The five random-generation metrics are computed from 10,000 generated molecular graphs.
- Random generation: GraphDF outperforms GraphAF by a large margin in validity w/o check and performs well on generic graph datasets against GNF and GraphAF.The authors attribute the strong molecular performance to discrete flow and argue that discrete latent variables model chemical rules and graph distributions more accurately.
- Property optimization: GraphDF outperforms all baselines on penalized logP optimization and achieves comparable performance on QED optimization.The evaluation summarizes top-3 property scores on ZINC250K, using reinforcement learning to optimize the targets.
- Constrained optimization: GraphDF achieves higher average property improvement than JT-VAE, GCPN, and MoFlow under all four similarity constraints, and outperforms GraphAF under three.Its average similarities are comparable with baselines, while its success rate does not decrease dramatically as the similarity threshold increases.
5. Conclusion
GraphDF uses discrete latent variables and invertible modulo shift transforms to generate molecular graphs, avoiding Jacobian computation and dequantization effects. It achieves state-of-the-art results across three molecule-generation tasks, while relying on BFS ordering and generating more slowly than one-shot methods.
- GraphDF is a molecular graph generation framework that uses discrete latent variables mapped to nodes and edges through invertible modulo shift transforms.
- Discrete latent variables provide more accurate graph-structure modeling and stronger chemical-space exploration capacity.
- GraphDF avoids expensive Jacobian computation and eliminates the negative effects of dequantization found in many prior methods.
- GraphDF achieves new state-of-the-art results across random generation, property optimization, and constrained optimization tasks.
- GraphDF relies on BFS node ordering during training, and its generation speed is slower than one-shot methods.
A. Generation Algorithm
GraphDF generates graphs sequentially by sampling discrete latent variables, transforming them with graph-conditioned modulo shifts, and adding nodes and edges while enforcing connectivity and valency constraints.
- The generation algorithm initializes an empty graph and iteratively samples a discrete latent variable for each potential node.The node latent is transformed using graph representations computed by an R-GCN and summed node features.
- Each node latent is passed through repeated modulo shift transforms to produce a node type before the node is added.
- For each existing node, the algorithm samples an edge latent and transforms it into an edge type conditioned on the partially generated graph.
- Candidate edges are repeatedly sampled until the resulting bond satisfies the chemical valency check.
- If the newly added node remains disconnected after edge generation, it is deleted; otherwise the updated graph becomes the next generation state.
B. Training Algorithm
GraphDF training maximizes molecular-graph log-likelihood by converting BFS-ordered graph elements into discrete latent variables and optimizing their latent-distribution likelihoods.
- Training samples batches of molecular graphs and determines each graph's node count and BFS-ordered sequence representation.
- For every node and edge in the BFS sequence, the algorithm inverts the graph transformations to infer corresponding discrete latent variables.
- The training objective accumulates negative log-probabilities under the node and edge latent distributions.
- Model parameters are updated by gradient descent with learning rate α until the objective converges.
C. Data Information
The molecular experiments use ZINC250K, QM9, and MOSES, while generic graph experiments use COMMUNITY-SMALL and EGO-SMALL. The datasets span drug-like, small-molecule, and synthetic or network-derived graph collections.
- ZINC250K contains around 250K drug-like molecules, with at most 38 nodes and 9 heavy-atom types.
- QM9 contains around 130K molecules with up to 9 heavy atoms.
- The random-generation experiments use the ZINC250K, QM9, and MOSES molecular datasets.
- COMMUNITY-SMALL contains 100 synthetic two-community graphs, while EGO-SMALL contains 200 small subgraphs from the Citeseer network dataset.
- GraphDF is evaluated on the generic graph datasets using maximum mean discrepancy under two generated-graph comparison settings.
D. Experiment Details
GraphDF is evaluated for random generation, property optimization, and constrained optimization with dataset-specific training schedules and temperature settings. Optimization tasks use ZINC250K pretraining followed by reinforcement-learning fine-tuning or iterative optimization.
- Random generation: GraphDF samples from a discrete prior using tunable temperature parameters, following a strategy also used by GraphAF and MoFlow.The temperature parameters for each dataset are reported in Table 9; only t2 is needed for COMMUNITY-SMALL and EGO-SMALL because each has one node type.
- Property optimization: Property optimization pretrains on ZINC250K for 1000 epochs, then uses reinforcement learning for 200 iterations with learning rate 0.0001 and batch size 8.Generation uses t1 = 0.8 and t2 = 0.1.
- Constrained optimization: Constrained optimization pretrains on ZINC250K for 1000 epochs, fine-tunes for 200 iterations, and optimizes each molecule 200 times.Fine-tuning uses learning rate 0.0001, batch size 16, and temperature parameters t1 = 1.0 and t2 = 1.0.