Any-Property-Conditional Molecule Generation with Self-Criticism using Spanning Trees

TL;DR

This paper introduces STGG+, an advanced molecule generation method that uses spanning trees and Transformer architecture to generate valid molecules conditioned on multiple properties, outperforming previous models in quality and flexibility.

Contribution

STGG+ extends spanning tree-based graph generation to multi-property conditional generation with self-criticism, Transformer integration, and classifier-free guidance, achieving state-of-the-art results.

Findings

Outperforms existing models in conditional molecule generation

Effective in out-of-distribution scenarios

Enables flexible property conditioning

Abstract

Generating novel molecules is challenging, with most representations leading to generative models producing many invalid molecules. Spanning Tree-based Graph Generation (STGG) is a promising approach to ensure the generation of valid molecules, outperforming state-of-the-art SMILES and graph diffusion models for unconditional generation. In the real world, we want to be able to generate molecules conditional on one or multiple desired properties rather than unconditionally. Thus, in this work, we extend STGG to multi-property-conditional generation. Our approach, STGG+, incorporates a modern Transformer architecture, random masking of properties during training (enabling conditioning on any subset of properties and classifier-free guidance), an auxiliary property-prediction loss (allowing the model to self-criticize molecules and select the best ones), and other improvements. We show that STGG+ achieves state-of-the-art performance on in-distribution and out-of-distribution conditional generation, and reward maximization.