Any-Property-Conditional Molecule Generation with Self-Criticism using Spanning Trees (2407.09357v2)

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.

• The paper introduces STGG+, extending spanning tree methods to enable flexible multi-property conditional molecule generation through self-criticism.
• It employs a two-layer MLP for mixed-data conditioning and enhanced Transformer architectures to improve efficiency and performance.
• Experiments on datasets like QM9 and Zinc250K demonstrate that STGG+ achieves near-perfect validity, novelty, and high fidelity in generated molecules.

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

The paper "Any-Property-Conditional Molecule Generation with Self-Criticism using Spanning Trees," authored by Alexia Jolicoeur-Martineau et al., presents a comprehensive framework for the generation of novel molecules conditioned on multiple desired properties. The primary innovation is the enhancement and extension of Spanning Tree-based Graph Generation (STGG) for property-conditional generation, which the authors denote as STGG+.

The challenges in generating valid and novel molecules using generative models are substantial due to the complexities in molecular representations. The use of SMILES strings often leads to many invalid molecules, while graph-based diffusion models face similar difficulties. STGG has shown promise for unconditional molecule generation by ensuring the production of valid molecular structures through the smart masking of invalid tokens. The current work builds upon STGG to support multi-property-conditional generation.

Key Contributions and Methodology

The authors highlight several contributions that address the limitations of previous methods:

1. Mixed-Data Property-Conditioning: The model uses a two-layer MLP to handle standardized continuous features and embeddings for categorical features. By masking some properties during training, the model can condition on varying subsets of properties at test time, enhancing flexibility.
2. Improved Transformer Architecture: Various architectural improvements from advancements in LLMs are incorporated. These include RMSNorm, rotary embeddings, Flash-Attention, SwiGLU activation, no bias terms, and RMSProp, leading to enhanced efficiency and performance.
3. Enhanced Spanning-Tree Approach: STGG+ includes several enhancements to the original STGG:
   • New tokens and masking conditions for unconnected compound structures.
   • Dynamic vocabulary and valency calculation based on the dataset.
   • Special masking to prevent incomplete samples and ring overflow.
   • Randomization of graph order during training for better generalization.
4. Self-Criticism Through Property Prediction: By predicting properties during training and sampling, the model can filter generated molecules, selecting those that best meet the desired properties without external predictors.
5. Random Guidance for Extreme Value Conditioning: The authors propose using random guidance for out-of-distribution properties to improve generative performance. By randomly sampling the guidance parameter and using best-out-of-$k$ filtering, the model can better handle extreme conditioning values.

Experiments and Results

The paper includes extensive experiments on various datasets (QM9, Zinc250K, HIV, BACE, BBBP, Chromophore DB):

Unconditional Generation

For unconditional generation, the proposed STGG+ achieves performance comparable to state-of-the-art methods such as STGG and GEEL on datasets like QM9 and Zinc250K, maintaining high validity, uniqueness, and novelty.

Conditional Generation

In the context of conditional generation on datasets like HIV, BBBP, and BACE, STGG+ demonstrates strong performance, achieving near-perfect validity and high fidelity on specified properties. It outperforms many recent methods, including Graph DiT, MOOD, and DiGress, particularly in Fréchet ChemNet Distance (FCD), a key metric indicating the distributional similarity between generated and real molecules.

Out-of-Distribution Conditional Generation

The experiments involving extreme out-of-distribution properties showcase STGG+'s ability to generate molecules with high fidelity to desired properties, outperforming other methods like ControlVAE and multi-decoder VAEs in terms of Minimum Mean Absolute Error (MinMAE).

Reward Maximization

For the task of reward maximization on the QM9 dataset, the proposed STGG+ achieves comparable or superior results to online learning methods (e.g., GFlowNet, MOA2C), producing high-reward molecules efficiently and with high diversity.

Practical and Theoretical Implications

The paper presents both practical and theoretical advancements:

• Practically, STGG+ offers a robust framework for generating valid, diverse molecules with specified properties, significantly enhancing the efficiency and applicability of generative models in materials science and drug discovery.
• Theoretically, it provides insights into the integration of self-criticism and guidance mechanisms in generative models, potentially paving the way for more sophisticated generative frameworks that can adapt to a wide range of conditional settings.

Future Developments

Future work could extend these methods to even larger datasets and more complex molecular structures, including stereoisomerism, which was not addressed in the current framework. Additionally, further refinement of the property-predictor could enhance the fidelity of generated molecules in extreme out-of-distribution scenarios.

In summary, the authors successfully extend the capabilities of Spanning Tree-based Graph Generation for multi-property conditions, presenting a versatile and effective approach for molecule generation in real-world applications.