[Paper Review] Improving Small Molecule Generation using Mutual Information Machine

Danny Reidenbach, Micha Livne+3 more|arXiv (Cornell University)|Aug 18, 2022

Computational Drug Discovery Methods11 citations

TL;DR

MolMIM is a probabilistic auto-encoder for small molecule generation that uses Mutual Information Machine (MIM) learning to learn an informative, clustered, fixed-size latent space from SMILES strings. It enables high-quality, valid, unique, and novel molecule generation through simple latent space optimization with CMA-ES, achieving state-of-the-art results in single- and multi-property optimization without reinforcement learning or complex architectures.

ABSTRACT

We address the task of controlled generation of small molecules, which entails finding novel molecules with desired properties under certain constraints (e.g., similarity to a reference molecule). Here we introduce MolMIM, a probabilistic auto-encoder for small molecule drug discovery that learns an informative and clustered latent space. MolMIM is trained with Mutual Information Machine (MIM) learning, and provides a fixed length representation of variable length SMILES strings. Since encoder-decoder models can learn representations with ``holes'' of invalid samples, here we propose a novel extension to the training procedure which promotes a dense latent space, and allows the model to sample valid molecules from random perturbations of latent codes. We provide a thorough comparison of MolMIM to several variable-size and fixed-size encoder-decoder models, demonstrating MolMIM's superior generation as measured in terms of validity, uniqueness, and novelty. We then utilize CMA-ES, a naive black-box and gradient free search algorithm, over MolMIM's latent space for the task of property guided molecule optimization. We achieve state-of-the-art results in several constrained single property optimization tasks as well as in the challenging task of multi-objective optimization, improving over previous success rate SOTA by more than 5\% . We attribute the strong results to MolMIM's latent representation which clusters similar molecules in the latent space, whereas CMA-ES is often used as a baseline optimization method. We also demonstrate MolMIM to be favourable in a compute limited regime, making it an attractive model for such cases.

Motivation & Objective

• To address the challenge of controlled small molecule generation under property constraints, such as drug-likeness and target affinity.
• To learn a dense, informative, and chemically meaningful latent space that clusters similar molecules without explicit property supervision.
• To enable efficient and effective molecule optimization via simple, gradient-free search in the latent space.
• To achieve superior performance in both single- and multi-objective molecular optimization tasks under compute-limited conditions.
• To demonstrate that a simple optimization strategy like CMA-ES can outperform complex methods when combined with a well-structured latent space.

Proposed method

• MolMIM employs a Perceiver encoder and Transformer decoder to map variable-length SMILES strings into a fixed-size latent representation.
• The model is trained using Mutual Information Machine (MIM) learning, which maximizes mutual information between input molecules and their latent codes.
• A novel extension to MIM training is introduced to promote a dense latent space and reduce invalid samples from random latent perturbations.
• The latent space is optimized using CMA-ES, a gradient-free, black-box optimization algorithm, for property-guided molecule generation.
• The method does not require reinforcement learning or explicit chemical property labels during training, relying instead on unsupervised pretraining.
• Latent code perturbations are used to generate new molecules, with validity ensured through the model’s learned density and clustering.

Experimental results

Research questions

• RQ1 Can a self-supervised, MIM-trained auto-encoder learn a dense and informative latent space that clusters chemically similar molecules without explicit property supervision?
• RQ2 Does a simple, gradient-free optimization method like CMA-ES outperform more complex optimization strategies when applied to a well-structured latent space?
• RQ3 To what extent does MolMIM improve validity, uniqueness, and novelty in generated molecules compared to existing fixed- and variable-size models?
• RQ4 How effective is MolMIM in multi-objective optimization tasks with conflicting constraints, such as balancing QED, SA score, and target inhibition?
• RQ5 Can MolMIM maintain high performance under compute-limited regimes, making it practical for real-world drug discovery pipelines?

Key findings

• MolMIM achieved a 99.2% success rate in multi-objective optimization with exemplar-based initialization, outperforming JANUS and FaST in novelty and diversity.
• In single-property optimization, MolMIM achieved state-of-the-art results on logP, QED, and SA score, with 97.5% validity and 71.1% novelty in the best variant.
• The model demonstrated superior performance in multi-objective optimization with a 49-restart variant (MolMIM E†), achieving 49% success rate while maintaining high novelty and diversity.
• MolMIM’s latent space enables high-quality generation from random perturbations, with 98.3% validity and 55.1% novelty in the E variant, indicating a dense and robust representation.
• The model achieved competitive results without reinforcement learning, showing that the quality of the latent space is more critical than complex optimization algorithms.
• The latent space learned by MolMIM implicitly clusters similar molecules, enabling efficient optimization even with simple search methods like CMA-ES.