[Paper Review] Can LLMs Generate Diverse Molecules? Towards Alignment with Structural Diversity

Hang Jea Jang, Yunhui Jang+2 more|arXiv (Cornell University)|Oct 4, 2024

Chemical Synthesis and AnalysisBiochemistry, Genetics and Molecular Biology3 citations

TL;DR

This paper proposes a two-stage fine-tuning method—supervised fine-tuning followed by reinforcement learning—to enable large language models (LLMs) to autoregressively generate structurally diverse molecules conditioned on previously generated ones. The approach significantly improves molecular diversity and quality over baseline decoding schemes and state-of-the-art LLMs, achieving higher diversity scores and better performance on unseen molecular property prompts.

ABSTRACT

Recent advancements in large language models (LLMs) have demonstrated impressive performance in molecular generation, which offers potential to accelerate drug discovery. However, the current LLMs overlook a critical requirement for drug discovery: proposing a diverse set of molecules. This diversity is essential for improving the chances of finding a viable drug, as it provides alternative molecules that may succeed where others fail in real-world validations. Nevertheless, the LLMs often output structurally similar molecules. While decoding schemes like diverse beam search may enhance textual diversity, this often does not align with molecular structural diversity. In response, we propose a new method for fine-tuning molecular generative LLMs to autoregressively generate a set of structurally diverse molecules, where each molecule is generated by conditioning on the previously generated molecules. Our approach consists of two stages: (1) supervised fine-tuning to adapt LLMs to autoregressively generate molecules in a sequence and (2) reinforcement learning to maximize structural diversity within the generated molecules. Our experiments show that the proposed approach enables LLMs to generate diverse molecules better than existing approaches for diverse sequence generation.

Motivation & Objective

• To address the critical gap in current LLMs for generating structurally diverse molecules, a key requirement for successful drug discovery.
• To overcome the limitation of existing decoding schemes, which enhance textual diversity but not molecular structural diversity.
• To develop a self-improving fine-tuning framework that enhances molecular diversity without relying on external diverse molecular datasets.
• To improve the robustness of LLM-based molecular generation by conditioning subsequent molecule generation on prior outputs.
• To demonstrate that reinforcement learning with molecule-wise diversity rewards leads to superior diversity compared to alternative training strategies.

Proposed method

• Supervised fine-tuning is applied to adapt a pre-trained LLM to autoregressively generate a sequence of multiple molecules from a single prompt.
• The training data is collected via iterative sampling from the LLM itself and filtered to improve quality and diversity.
• Reinforcement learning is used to maximize structural diversity by defining a reward function based on molecular similarity metrics such as Tanimoto similarity and NCircles.
• The method employs a multi-stage reinforcement learning setting, where each molecule generation is treated as a step with individual diversity rewards.
• The model is trained end-to-end using self-generated data, avoiding dependency on external diverse molecular datasets.
• The approach uses SMILES representations and evaluates diversity using metrics like Bertz complexity, QED, and molecular fingerprint-based similarity.

Experimental results

Research questions

• RQ1 Can LLMs be fine-tuned to generate a diverse set of molecules rather than repeating similar structures?
• RQ2 Does conditioning molecule generation on previously generated molecules improve structural diversity?
• RQ3 How does the proposed two-stage fine-tuning method compare to standard decoding schemes like beam search in generating diverse molecules?
• RQ4 Can reinforcement learning with molecule-wise rewards outperform supervised fine-tuning with hard filtering on diversity metrics?
• RQ5 Does a multi-stage reinforcement learning setup improve diversity compared to a single-stage return formulation?

Key findings

• The proposed method achieves significantly higher structural diversity, with NCircles scores of 21.98 at h=0.65, outperforming baseline methods.
• On the DrugAssist benchmark, the method consistently improves diversity and quality across all molecular property prompts, including unseen QED-based prompts.
• With 155 generated molecules, the method discovers 25.6 unique and diverse molecules, compared to 21.3 with beam search (BS=500), demonstrating superior scalability.
• The method reduces time cost by 80% compared to beam search (65s vs. 585s) while generating more diverse molecules.
• Multi-stage reinforcement learning outperforms single-stage RL and hard-filtered SFT, indicating better credit assignment and diversity optimization.
• The method maintains high performance even with unseen prompts, showing strong generalization and robustness in molecular generation.