GFlowNet Pretraining with Inexpensive Rewards

TL;DR

This paper introduces Atomic GFlowNets, a new generative model that uses individual atoms and inexpensive molecular descriptors for pretraining, enabling more comprehensive exploration of drug-like chemical space and targeted property optimization.

Contribution

The work presents a novel A-GFN model that leverages atomic building blocks and proxy rewards for unsupervised pretraining and goal-conditioned fine-tuning in drug discovery.

Findings

A-GFNs outperform baseline methods in chemical space exploration.

Pretraining with proxy rewards improves property optimization.

Effective adaptation to target-specific drug design tasks.

Abstract

Generative Flow Networks (GFlowNets), a class of generative models have recently emerged as a suitable framework for generating diverse and high-quality molecular structures by learning from unnormalized reward distributions. Previous works in this direction often restrict exploration by using predefined molecular fragments as building blocks, limiting the chemical space that can be accessed. In this work, we introduce Atomic GFlowNets (A-GFNs), a foundational generative model leveraging individual atoms as building blocks to explore drug-like chemical space more comprehensively. We propose an unsupervised pre-training approach using offline drug-like molecule datasets, which conditions A-GFNs on inexpensive yet informative molecular descriptors such as drug-likeliness, topological polar surface area, and synthetic accessibility scores. These properties serve as proxy rewards, guiding A-GFNs towards regions of chemical space that exhibit desirable pharmacological properties. We further our method by implementing a goal-conditioned fine-tuning process, which adapts A-GFNs to optimize for specific target properties. In this work, we pretrain A-GFN on the ZINC15 offline dataset and employ robust evaluation metrics to show the effectiveness of our approach when compared to other relevant baseline methods in drug design.