[Paper Review] Text-Guided Molecule Generation with Diffusion Language Model

Haisong Gong, Qiang Liu+2 more|arXiv (Cornell University)|Feb 20, 2024

Chemical Synthesis and Analysis4 citations

TL;DR

This paper proposes TGM-DLM, a diffusion language model for text-guided molecule generation that uses a two-phase process: first optimizing SMILES token embeddings from noise using text guidance, then correcting invalid SMILES strings without text input. It outperforms autoregressive models like MolT5-Base in validity and molecular similarity metrics without additional data, demonstrating the superiority of diffusion-based generation for controlled, precise molecular design.

ABSTRACT

Text-guided molecule generation is a task where molecules are generated to match specific textual descriptions. Recently, most existing SMILES-based molecule generation methods rely on an autoregressive architecture. In this work, we propose the Text-Guided Molecule Generation with Diffusion Language Model (TGM-DLM), a novel approach that leverages diffusion models to address the limitations of autoregressive methods. TGM-DLM updates token embeddings within the SMILES string collectively and iteratively, using a two-phase diffusion generation process. The first phase optimizes embeddings from random noise, guided by the text description, while the second phase corrects invalid SMILES strings to form valid molecular representations. We demonstrate that TGM-DLM outperforms MolT5-Base, an autoregressive model, without the need for additional data resources. Our findings underscore the remarkable effectiveness of TGM-DLM in generating coherent and precise molecules with specific properties, opening new avenues in drug discovery and related scientific domains. Code will be released at: https://github.com/Deno-V/tgm-dlm.

Motivation & Objective

• To address the limitations of autoregressive models in text-guided molecule generation, which struggle with global constraints due to fixed generation order.
• To explore the use of diffusion models in SMILES-based generation, enabling holistic and iterative optimization of molecular representations.
• To improve molecular validity and alignment with textual descriptions by introducing a dedicated correction phase for invalid SMILES strings.
• To demonstrate that diffusion-based generation can surpass autoregressive baselines in generating coherent, precise, and valid molecules.

Proposed method

• TGM-DLM employs a two-phase diffusion generation process: phase one optimizes SMILES token embeddings from random noise using text description as guidance.
• Phase two corrects invalid SMILES strings by iteratively refining embeddings without text input, focusing on structural validity.
• The model is trained with two objectives: denoising embeddings using text guidance in phase one, and recovering uncorrupted SMILES from corrupted inputs in phase two.
• A shared transformer-based architecture is used for both phases, with separate training objectives to balance initial text alignment and structural correction.
• The correction phase uses a separate network head to refine molecular representations while preserving semantic consistency.
• The method leverages SMILES as a sequence of tokens and applies diffusion to optimize their embeddings collectively, not autoregressively.

Experimental results

Research questions

• RQ1 Can a diffusion-based approach outperform autoregressive models in text-guided SMILES generation by enabling better handling of global constraints?
• RQ2 How effective is a two-phase diffusion process—first guided by text, then correcting structural errors—in improving molecular validity and quality?
• RQ3 Does incorporating text input during the correction phase enhance the model’s ability to generate valid and relevant molecules?
• RQ4 To what extent can a diffusion model improve molecular generation without requiring additional pre-training data or fine-tuning on external datasets?

Key findings

• TGM-DLM achieves a validity rate of 88.3% on the ChEBI-20 dataset, significantly improving from 78.9% in the no-correction variant.
• The model improves the exact match score by tripling the performance of MolT5-Base, indicating higher precision in generating correct molecular strings.
• TGM-DLM achieves 18% to 36% higher fingerprint similarity metrics (MACCS FTS) compared to MolT5-Base, demonstrating better molecular structure alignment.
• The correction phase increases validity from 78.9% to 87.1% when using one phase-two step count, showing the effectiveness of structural refinement.
• Using text input during correction does not improve validity beyond 82.4%, suggesting it may hinder the model’s ability to fix structural errors.
• Joint training of both phases in a single network underperforms compared to separate training, indicating that distinct objectives benefit from independent optimization.