[Paper Review] GraphAF: a Flow-based Autoregressive Model for Molecular Graph Generation

Chence Shi, Minkai Xu+4 more|arXiv (Cornell University)|Jan 26, 2020

Computational Drug Discovery MethodsComputer Science45 references205 citations

TL;DR

GraphAF is a flow-based autoregressive model for molecular graph generation that achieves high validity and efficiency, enabling rapid density modeling and superior property optimization after reinforcement learning.

ABSTRACT

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.

Motivation & Objective

• Motivate molecular graph generation for drug discovery and material science under chemical validity and property optimization constraints.
• Develop a flow-based autoregressive framework that can model molecular data density and enable efficient training and sampling.
• Incorporate chemical domain knowledge via valency constraints during generation and enable goal-directed optimization through reinforcement learning.
• Evaluate GraphAF on standard molecular datasets to compare validity, uniqueness, novelty, and reconstruction against state-of-the-art methods.

Proposed method

• Formulate molecular graph generation as a sequential decision process with nodes and edges generated autoregressively.
• Define an invertible flow from a base Gaussian to discrete graph structures using dequantization and Gaussian conditionals.
• Use Relational GCN (R-GCN) to generate node and edge conditionals via neural networks for mean and scale parameters (g_mu, g_alpha).
• Enable efficient parallel training by masking and BFS-based graph ordering, allowing exact likelihood computation in a single forward pass.
• Guarantee chemical validity by integrating valency checks during sampling and adding hydrogens to satisfy valence.
• Optional reinforcement learning fine-tuning (PPO) to optimize penalized logP and QED with similarity constraints for constrained property optimization.

Experimental results

Research questions

• RQ1 Can GraphAF accurately model the data density of molecular graphs while enabling efficient parallel training?
• RQ2 Does the autoregressive flow approach deliver higher validity and reconstruction accuracy than existing graph generative models?
• RQ3 How does GraphAF perform on density modeling, generation, and property/ constrained property optimization compared to baselines?
• RQ4 What is the impact of valency-based constraints and reinforcement learning on the quality and validity of generated molecules?

Key findings

• GraphAF achieves 100% validity when valency rules are applied during generation.
• Without valency checks, GraphAF still achieves 68% validity, higher than prior autoregressive methods.
• Training GraphAF is significantly faster (about 4 hours) than GCPN (about 8–24 hours) on the same hardware.
• GraphAF matches or surpasses state-of-the-art in density modeling and generation across ZINC250k, QM9, and MOSES datasets.
• Fine-tuning with reinforcement learning yields state-of-the-art results for penalized logP and comparable QED scores, with strong performance on constrained property optimization.
• On generic graphs, GraphAF shows competitive MMD metrics against GraphRNN and GNF.