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

2024-07-12

arXiv.org

AI Summary

To address the challenge of multi-attribute controllable molecular generation in drug discovery, this work extends the Spanning Tree graph generation framework and proposes the first method enabling conditional modeling on arbitrary subsets of physicochemical properties (e.g., solubility, target activity). Methodologically, it introduces a novel multi-attribute random masking training strategy and a self-critical property prediction loss for end-to-end quality filtering, along with a classifier-free guidance mechanism to substantially improve out-of-distribution generalization. The model adopts a Transformer backbone and jointly optimizes generation via auxiliary property regression losses. Empirically, it achieves state-of-the-art performance on both multi-attribute conditional generation and reward maximization tasks: increasing valid molecule rate by 12.7%, improving target property satisfaction rate by 23.4% on average, and demonstrating strong generalization to unseen property combinations.

Technology Category

• Natural Language Processing: Generation
• Machine Learning: Deep Generative Models & Autoencoders
• Search and Optimization: Mixed Discrete/Continuous Search

Application Category

• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

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.

Problem

• Generate valid molecules with desired properties
• Improve conditional molecule generation using spanning trees
• Enhance performance with self-criticism and Transformer architecture

Innovation

• Transformer architecture for molecule generation
• Random masking enables multi-property conditioning
• Auxiliary loss for self-criticism and selection