PaperlayerView exact revision on arXiv

Source-linked AI summary

GraphAF: a Flow-based Autoregressive Model for Molecular Graph Generation

Chence Shi, Minkai Xu, Zhaocheng Zhu, Weinan Zhang, Ming Zhang, Jian Tang

arXiv:2001.09382v2cs.LGstat.ML

TL;DR

Molecular graph generation must produce chemically valid structures while optimizing their properties in a vast discrete search space. GraphAF addresses this with a flow-based autoregressive model that combines flexible density estimation, parallel training, and iterative sampling with valency checking. It generates 68% valid molecules without chemical rules, 100% with them, trains twice as fast as GCPN, and achieves state-of-the-art optimization results after reinforcement-learning fine-tuning.

  • Problem

    Molecular graph generation requires producing chemically valid molecules while optimizing their properties across a vast discrete chemical space.

  • Method

    GraphAF is a flow-based autoregressive model that uses parallel likelihood computation and iterative graph sampling with optional valency constraints.

  • Results

    GraphAF generates 68% valid molecules without chemical rules and 100% with them, trains two times faster than GCPN, and achieves state-of-the-art property and constrained-property optimization after reinforcement learning.

  • Takeaways & Limitations

    GraphAF combines flexible molecular density modeling with chemically constrained generation and efficient training for molecular property optimization.

Abstract

from arXiv · show

Molecular graph generation is a fundamental problem for drug discovery and has been attracting growing attention. The problem is challenging since it requires not only generating chemically valid molecular structures but also optimizing their chemical properties in the meantime. Inspired by the recent progress in deep generative models, in this paper we propose a flow-based autoregressive model for graph generation called GraphAF. GraphAF combines the advantages of both autoregressive and flow-based approaches and enjoys: (1) high model flexibility for data density estimation; (2) efficient parallel computation for training; (3) an iterative sampling process, which allows leveraging chemical domain knowledge for valency checking. Experimental results show that GraphAF is able to generate 68% chemically valid molecules even without chemical knowledge rules and 100% valid molecules with chemical rules. The training process of GraphAF is two times faster than the existing state-of-the-art approach GCPN. After fine-tuning the model for goal-directed property optimization with reinforcement learning, GraphAF achieves state-of-the-art performance on both chemical property optimization and constrained property optimization.

1 INTRODUCTION

Molecular graph generation seeks chemically valid molecules with desirable properties across a vast discrete chemical space. GraphAF combines autoregressive generation with normalizing flows to improve density modeling, parallel training, validity, and property optimization.

  • Molecular design spans drug discovery and materials science but is difficult because chemical space is discrete and estimated at 10^33 structures.
  • Deep generative models increasingly represent molecules as graphs and generate them with VAEs, GANs, and autoregressive models.
  • GraphAF combines autoregressive sampling with flow-based density modeling, enabling flexible modeling, parallel likelihood computation, and stepwise valency checking.
  • 2 times faster training than GCPN was achieved on ZINC, while GraphAF generated 68% valid molecules without chemical rules and 100% with them.
  • Reinforcement-learning fine-tuning significantly outperformed GCPN on property optimization and constrained property optimization tasks.

2 RELATED WORK

Prior molecular graph generators include VAEs, GANs, and autoregressive models, while newer graph-flow methods provide an alternative generation framework. GraphAF differs by defining an end-to-end autoregressive flow directly over molecular graph structures.

  • VAE-based methods commonly optimize molecular properties by searching their latent spaces with Bayesian Optimization.
  • Autoregressive approaches such as GCPN and MolecularRNN model molecular graph generation as sequential decision processes and show competitive performance.
  • GNF generates hidden node representations before constructing graphs, whereas GraphAF maps a base distribution directly to molecular graphs end-to-end.
  • GraphNVP uses one-shot sampling, while GraphAF uses sequential generation to capture subgraphs and support molecular validity constraints.

3 PRELIMINARIES

Normalizing flows use invertible transformations for exact density estimation and sampling, while autoregressive flows make Jacobian computation efficient. Molecular graphs are represented with adjacency tensors and node features, then encoded using relational graph convolutions.

  • Normalizing flows transform a base distribution into observational data through an invertible parameterized mapping and compute data density with change of variables.
  • Density estimation in a flow inverts the transformation, while sampling draws from the base distribution and applies the feedforward transformation.
  • Autoregressive flows use triangular Jacobians whose determinants can be computed linearly, supporting efficient likelihood calculations.
  • Autoregressive flow density transformations can be computed in parallel, whereas sampling computes variables sequentially from previously generated variables.
  • A molecule is represented as G = (A, X), with adjacency tensor A encoding bond types and node-feature matrix X encoding atom features.
  • Relational graph convolutional networks learn atom representations from categorical edge types through message passing and aggregate node representations into graph representations.

4 PROPOSED METHOD

GraphAF models molecular graph generation as an autoregressive flow that sequentially generates nodes and edges, while supporting efficient likelihood training and validity-constrained sampling. It also fine-tunes generation with reinforcement learning to optimize molecular properties.

  • GraphAF framework: GraphAF sequentially generates each node from the current sub-graph, then generates its edges to existing nodes until graph construction terminates.The generation starts from an empty graph and uses p(X_i|G_i) and p(A_ij|G_i, X_i, A_i,1:j−1).
  • GraphAF framework: GraphAF converts discrete node and edge types into continuous data with dequantization, then defines conditional Gaussian distributions parameterized by neural networks.Relational GCN embeddings and Node-MLPs or Edge-MLPs provide the distribution parameters for node and edge generation.
  • Training: GraphAF uses an autoregressive invertible mapping whose triangular Jacobian enables efficient exact molecule-density computation and likelihood maximization.Masked feedforward computation permits parallel training, while BFS ordering further accelerates dependency processing.
  • Validity-constrained sampling: Valency checking rejects bonds that exceed allowed valency, resamples their latent variables, and terminates generation when size or connectivity conditions are met.Hydrogens are added to atoms whose valencies remain unfilled.
  • Reinforcement learning: GraphAF fine-tunes its autoregressive policy with reinforcement learning, combining intermediate valency penalties with final property and chemical-validity rewards.The policy uses the GraphAF autoregressive model and is trained with Proximal Policy Optimization and a clipped surrogate objective.

5 EXPERIMENTS

Experiments evaluate GraphAF on molecular density modeling, generation, property optimization, constrained optimization, and generic graph generation. Across these tasks, GraphAF demonstrates strong validity, efficiency, transfer across datasets and graph types, and property-optimization performance.

  • Evaluation setup: Experiments compare GraphAF with JT-VAE, GCPN, MolecularRNN, and GraphNVP across density modeling, generation, and optimization tasks.The evaluation uses standard molecular-generation tasks and includes generic graph benchmarks.
  • Density modeling and generation: 100% validity is achieved with valency checking, while validity without checking reaches 68%, compared with 20% for GCPN and 42.60% for GraphNVP.GraphAF also has perfect reconstruction among flow-based or latent-variable comparisons, and its training takes 4 hours versus approximately 8 hours for GCPN and 24 hours for JT-VAE.
  • Density modeling and generation: GraphAF generates valid and novel molecules on ZINC250k, QM9, and MOSES, including the larger and more challenging MOSES dataset.The datasets contain 250,000, 134,000, and 1.9 million molecules, respectively.
  • Generic graph generation: GraphAF produces comparable or better results than GraphRNN and GNF on generic graph datasets evaluated with MMD metrics.The model supports non-molecular graphs by modifying its node- and edge-generating functions.
  • Property optimization: GraphAF outperforms all baselines by a large margin on penalized logP and achieves comparable QED results after reinforcement-learning fine-tuning.Penalized logP is adjusted for ring size and synthetic accessibility, whereas QED measures molecular drug-likeness.
  • Constrained property optimization: GraphAF significantly outperforms previous approaches in constrained penalized-logP optimization while maintaining high similarity to the original molecules.The task modifies 800 low-scoring ZINC250k molecules under a Tanimoto-similarity constraint; Figure 2 shows examples with similarities 0.71 and 0.64.

6 CONCLUSION

GraphAF models realistic and diverse molecular graphs with a flow-based autoregressive approach, supports reinforcement-learning property optimization, and outperforms prior state-of-the-art baselines on standard tasks.

  • 6 CONCLUSION: GraphAF is presented as the first flow-based autoregressive model for generating realistic and diverse molecular graphs.Its normalizing-flow component models complex molecular distributions, while its autoregressive process generates graphs.
  • 6 CONCLUSION: GraphAF generates novel and 100% valid molecules in empirical experiments.
  • 6 CONCLUSION: Reinforcement learning fine-tunes GraphAF to optimize generated-molecule properties.
  • 6 CONCLUSION: GraphAF outperforms all previous state-of-the-art baselines on standard property-optimization tasks.

A DISCCUSIONS ON DEQUANTIZATION TECHNIQUES

GraphAF dequantizes discrete molecular data by adding uniform noise, creating non-overlapping continuous ranges that can be mapped back through argmax; variational dequantization is identified as future work.

  • A DISCCUSIONS ON DEQUANTIZATION TECHNIQUES: Adding noise from U[0, 1) maps discrete categories into non-overlapping continuous intervals.A one-hot 1-entry lies in [1, 2), while a 0-entry lies in [0, 1).
  • A DISCCUSIONS ON DEQUANTIZATION TECHNIQUES: Argmax operation converts generated dequantized continuous data back to discrete one-hot data.
  • A DISCCUSIONS ON DEQUANTIZATION TECHNIQUES: Variational dequantization treats noise as a hidden variable and infers the optimum noise for each discrete datum.The authors identify this technique as a direction for future work.

B PARALLEL TRAINING ALGORITHM

The parallel training algorithm samples molecules, converts them into BFS-reordered graph representations, and iteratively computes latent variables across nodes and edges until convergence.

  • B PARALLEL TRAINING ALGORITHM: Training repeatedly samples molecules from the dataset and obtains each graph's size.
  • B PARALLEL TRAINING ALGORITHM: Each molecule is converted to G = (A, X) using BFS re-ordering before latent-variable computation.
  • B PARALLEL TRAINING ALGORITHM: The algorithm iterates over graph nodes and their preceding BFS dependencies when computing latent variables.The edge loop considers indices from max{1, i − P} through i − 1.

C EXPERIMENT DETAILS

Experiments use a fixed graph network architecture and specified density-modeling and reinforcement-learning procedures, including ZINC250K training and constrained-property subgraph sampling.

  • C EXPERIMENT DETAILS: The R-GCN uses 3 layers with 128-dimensional embeddings, batch normalization before pooling, and sum-pooling for graph readout.
  • C EXPERIMENT DETAILS: Node and edge MLPs each contain two fully connected layers with tanh nonlinearities.
  • C EXPERIMENT DETAILS: GraphAF density modeling trains on ZINC250K for 10 epochs with batch size 32, Adam, learning rate 0.001, one V100 GPU, and 32 CPU cores.
  • C EXPERIMENT DETAILS: Property optimization pretrains GraphAF for 300 epochs, then fine-tunes it toward a desired molecular distribution through reinforcement learning.Rewards combine step-wise validity rewards with discounted final property-targeted and chemical-validation rewards.
  • C EXPERIMENT DETAILS: Fine-tuning runs for 200 iterations with batch size 64, Adam, learning-rate warm-up, and grid-searched hyperparameters.
  • C EXPERIMENT DETAILS: Constrained property optimization samples connected subgraphs from 800 ZINC molecules by dropping 0–5 final BFS nodes and their induced edges.

D VISUALIZATION OF GENERATED GENERIC GRAPHS

GraphAF generates graphs that model different structures in generic graph datasets, as illustrated by comparisons with training-set graphs.

  • GraphAF visualizations compare training-set graphs with generated graphs for generic graph datasets.

E MORE MOLECULE SAMPLES

GraphAF generates diverse molecule samples and supports constrained penalized logP optimization through structural modifications to input molecules.

  • 50 molecules sampled from a multivariate Gaussian demonstrate GraphAF’s ability to generate novel, realistic, and unique molecules.
  • Penalized logP scores of generated molecules range from 5 to 10 and are diverse.
  • GraphAF improves constrained penalized logP scores by reducing ring size, removing large rings, or growing carbon chains.
Loading 2001.09382v2…