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

TMLR Paper5939 Authors

Abstract: Generating novel molecules is challenging, with most representations of molecules leading to generative models producing many invalid molecules. Spanning Tree-based Graph Generation (STGG) \citep{ahn2021spanning} is a promising approach to ensure the generation of valid molecules, outperforming state-of-the-art generative models for unconditional generation. In practice, it is desirable to generate molecules conditional on one or multiple target properties rather than unconditionally. Thus, we extend STGG to multi-property conditional generation. Our approach, \highlight{\textbf{STGG+}}, incorporates a modern Transformer architecture, random masking of properties during training (enabling conditioning on \emph{any} subset of properties and classifier-free guidance), an auxiliary property-prediction loss (allowing the model to \emph{self-criticize} molecules and select the best ones), and other improvements. We show that \highlight{\textbf{STGG+}} achieves state-of-the-art performance on in-distribution and out-of-distribution conditional generation, as well as reward maximization.

Submission Length: Regular submission (no more than 12 pages of main content)

Assigned Action Editor: ~Emmanuel_Bengio1

Submission Number: 5939