PaperlayerView exact revision on arXiv

Source-linked AI summary

GenMol: A Drug Discovery Generalist with Discrete Diffusion

Seul Lee, Karsten Kreis, Srimukh Prasad Veccham, Meng Liu, Danny Reidenbach, Yuxing Peng, Saee Paliwal, Weili Nie, Arash Vahdat

arXiv:2501.06158v3cs.LG

TL;DR

Drug discovery requires molecular generation across multiple scenarios, while existing models often support only limited task coverage. GenMol uses one masked discrete diffusion framework with SAFE sequences, fragment remasking, and molecular context guidance, and reports state-of-the-art results across a wide range of tasks. Its authors note a quality–sampling-efficiency trade-off and potential misuse risks.

  • Problem

    Existing molecular generative models address only some drug discovery tasks, limiting their use across multifaceted drug discovery pipelines.

  • Method

    GenMol combines masked discrete diffusion, SAFE sequence generation, non-autoregressive bidirectional decoding, fragment remasking, and molecular context guidance in one framework.

  • Results

    GenMol achieves state-of-the-art results across a wide range of drug discovery tasks and outperforms existing methods in extensive experiments.

  • Takeaways & Limitations

    GenMol provides a unified and versatile tool for molecular design across diverse drug discovery scenarios.

  • Takeaways & Limitations

    Unmasking more than one token per step reduces generation quality despite improving sampling time, leaving a quality–efficiency trade-off for future work.

Abstract

from arXiv · show

Drug discovery is a complex process that involves multiple stages and tasks. However, existing molecular generative models can only tackle some of these tasks. We present Generalist Molecular generative model (GenMol), a versatile framework that uses only a single discrete diffusion model to handle diverse drug discovery scenarios. GenMol generates Sequential Attachment-based Fragment Embedding (SAFE) sequences through non-autoregressive bidirectional parallel decoding, thereby allowing the utilization of a molecular context that does not rely on the specific token ordering while having better sampling efficiency. GenMol uses fragments as basic building blocks for molecules and introduces fragment remasking, a strategy that optimizes molecules by regenerating masked fragments, enabling effective exploration of chemical space. We further propose molecular context guidance (MCG), a guidance method tailored for masked discrete diffusion of GenMol. GenMol significantly outperforms the previous GPT-based model in de novo generation and fragment-constrained generation, and achieves state-of-the-art performance in goal-directed hit generation and lead optimization. These results demonstrate that GenMol can tackle a wide range of drug discovery tasks, providing a unified and versatile approach for molecular design. Our code is available at https://github.com/NVIDIA-Digital-Bio/genmol.

1. Introduction

GenMol addresses the need for molecular generative models that support multiple drug discovery scenarios with one versatile framework. It combines masked discrete diffusion, SAFE sequences, fragment remasking, and molecular context guidance, and reports strong performance across the drug discovery pipeline.

  • Existing molecular generative models typically address only one or two drug discovery scenarios or require task-specific retraining.
  • GenMol combines masked discrete diffusion with the BERT architecture to generate SAFE molecular sequences for diverse drug discovery tasks.The framework targets de novo generation, fragment-constrained generation, hit generation, and lead optimization.
  • Non-autoregressive bidirectional decoding uses molecular context without relying on a specific token or fragment ordering and improves computational efficiency.
  • Fragment remasking replaces selected fragments with masked tokens so diffusion can generate new fragments for chemical-space exploration and molecule optimization.
  • GenMol outperforms existing methods across extensive experiments, while the best baseline results combine multiple task-specific models rather than one model.
  • The authors propose GenMol, fragment remasking, molecular context guidance, and validation across a wide range of drug discovery tasks.

2. Related Work

Related work spans discrete diffusion, molecular generation, and fragment-based drug discovery. Prior approaches provide task-specific capabilities, but the cited literature does not apply discrete diffusion to molecular sequence generation as a generalist foundation model.

  • Discrete diffusion: Discrete diffusion supports non-autoregressive generation, bidirectional dependency modeling, faster sampling, and flexible controllable-generation designs.
  • Discrete diffusion: Prior discrete-diffusion molecular models focus on molecular graphs or 3D structures rather than general molecular sequence generation.
  • Fragment-based drug discovery: Fragment-based methods use molecular substructures as building blocks, which can simplify generation and improve chemical validity.
  • Fragment-based drug discovery: Graph-based genetic algorithms have limited chemical-space exploration because they rely on random fragment combinations and low-probability local mutations.
  • Fragment-based drug discovery: f-RAG improves exploration–exploitation by retrieving fragments from a dynamically updated vocabulary but requires an information-fusion module for goal-oriented adaptation.

3. Background

Masked discrete diffusion progressively masks and unmasks molecular sequence tokens, while SAFE represents molecules as permutation-invariant sequences of fragment blocks. Together, these representations address ordering difficulties that complicate fragment-constrained generation with SMILES.

  • Masked diffusion: Masked diffusion models define a discrete diffusion process that follows MDLM and trains with masked-language-modeling losses across diffusion time steps.
  • Masked diffusion: The forward process interpolates each clean token toward a masking token, reaching an all-masked sequence at t = 1.
  • Masked diffusion: The reverse process independently infers less-masked tokens from more-masked data using a denoising network that predicts clean-sequence token probabilities.
  • SAFE molecular representation: SMILES relies on heuristic depth-first traversal, so structurally close atoms can be far apart in sequence and complicate fragment-constrained generation.
  • SAFE molecular representation: SAFE represents molecules as unordered sequences of fragment blocks, with consecutive fragment tokens and attachment points preserved.

4. Method

GenMol is a unified molecule-generation framework built on masked discrete diffusion over SAFE sequences. It combines parallel decoding, fragment-based goal-directed optimization with remasking, and molecular context guidance for diverse generation tasks.

  • Masked Diffusion for Molecule Generation: GenMol uses a BERT-based masked discrete diffusion model to generate SAFE sequences with bidirectional, non-autoregressive parallel decoding.This avoids reliance on a specific token ordering, considers the whole molecule, and improves sampling efficiency.
  • Masked Diffusion for Molecule Generation: Confidence-based sampling balances molecular quality and diversity through the softmax temperature and randomness parameters.All masked tokens are predicted in parallel, while only the most confident predictions are confirmed at each step.
  • Exploring Chemical Space with GenMol: Goal-directed optimization scores fragments, attaches two vocabulary fragments, and remasks one fragment before iteratively generating a new molecule.The generated molecule is decomposed and used to update the fragment vocabulary for subsequent iterations.
  • Exploring Chemical Space with GenMol: Fragment remasking enables exploration beyond the initial vocabulary by regenerating a selected fragment at a finer-grained level.The method masks one fragment of an initially attached molecule and iteratively unmasks the replacement.
  • Exploring Chemical Space with GenMol: Variable-length mask chunks adapt to the decomposition rule and provide controllability by generating fragments of varying lengths.Mask lengths are sampled from a predefined distribution, such as fragment lengths in the training set.
  • Molecular Context Guidance: Molecular context guidance compares predictions from the original input and a partially corrupted version to calibrate GenMol’s logits.The corrupted input is created by masking a fraction of tokens, allowing the model to use molecular context in fragment-constrained and goal-directed generation.

5. Experiments

Experiments evaluate one GenMol checkpoint across diverse drug discovery tasks, showing strong generation quality, efficient sampling, and effective fragment-based optimization. Ablations support fragment remasking and molecular context guidance as useful components.

  • Experimental scope: A single GenMol checkpoint is evaluated across de novo, fragment-constrained, hit-generation, and lead-optimization tasks without task-specific finetuning.The experiments use a SAFE dataset combining molecules from ZINC and UniChem.
  • De novo generation: GenMol significantly outperforms SAFE-GPT and its non-confidence-sampling variant in validity, quality, and sampling time while maintaining near-perfect uniqueness.With N=3, GenMol achieves higher quality than the comparisons with 2.5x shorter sampling time and comparable diversity.
  • Fragment-constrained generation: GenMol outperforms SAFE-GPT on most fragment-constrained generation metrics while preserving high diversity under fragment constraints.The benchmark covers linker design, scaffold morphing, motif extension, scaffold decoration, and superstructure generation.
  • Goal-directed hit generation: GenMol achieves the best performance in 19 of 23 goal-directed hit-generation tasks and significantly outperforms previous methods on summed AUC top-10.The optimization strategy uses fragment remasking to discover optimized hits.
  • Goal-directed lead optimization: GenMol successfully optimizes seed molecules in 26 of 30 lead-optimization tasks and outperforms baselines in most tasks.Baselines frequently fail under the harsher δ=0.6 similarity constraint, whereas GenMol satisfies the stated optimization constraints in successful cases.
  • Ablation study: Ablations show that fragment remasking improves exploration over attachment alone, while molecular context guidance improves performance beyond fragment remasking alone.Fragment-level remasking outperforms token-level and GPT-based remasking, and MCG reaches the best performance in 19 of 23 hit-generation tasks.

6. Conclusion

GenMol integrates discrete diffusion with SAFE to address diverse drug discovery scenarios. Fragment remasking supports chemical-space exploration, MCG further improves performance, and experiments report state-of-the-art results across varied tasks.

  • Conclusion: GenMol integrates discrete diffusion with SAFE into a molecule-generation framework for diverse drug discovery scenarios.The framework is designed to operate effectively and efficiently across multiple stages of drug discovery.
  • Conclusion: Fragment remasking enables chemical-space exploration, while molecular context guidance further improves GenMol’s performance.These components are presented as central mechanisms supporting the framework’s flexibility and effectiveness.
  • Conclusion: Experiments report state-of-the-art results across a wide range of drug discovery tasks, supporting GenMol as a unified and versatile tool.The conclusion characterizes GenMol’s potential within the scope of the reported experiments.

Impact Statement

GenMol addresses a broad spectrum of drug discovery challenges but could generate harmful drugs if used maliciously. Suggested safeguards include toxicity-aware targets, toxic-fragment exclusion, or toxicity filtering.

  • GenMol can generate harmful drugs if used maliciously, despite addressing a broad spectrum of drug discovery challenges.
  • Potential safeguards include toxicity-aware target properties, excluding toxic fragments from the vocabulary, or filtering candidates by predicted toxicity.

A. Limitations

GenMol improves sampling efficiency through confidence-based parallel unmasking, but unmasking multiple tokens per step degrades generation quality. Balancing quality and efficiency remains future work.

  • Confidence sampling selects tokens for unmasking using their prediction confidence scores at each sampling step.
  • GenMol predicts tokens simultaneously but confirms only the top-N most confident masked tokens, remasking the rest for later steps.
  • Confidence sampling exploits dependencies between tokens for better sampling quality than randomly and independently selecting tokens to unmask.

C. Derivation of MCG (Eq. (7))

MCG extends classifier-free guidance to GenMol’s masked discrete diffusion by combining rate matrices from high-quality and poor models. The derivation connects this formulation to denoiser predictions and masking corruption.

  • The rate matrix R(z,z′) describes the transition probability from z to z′.
  • Classifier-free guidance combines rate matrices from a high-quality model and a poor model to construct a guided transition process.
  • Because masked discrete diffusion can be interpreted as a continuous-time Markov chain, the CTMC guidance formulation can be applied to GenMol.
  • For GenMol, the high-quality and poor rate matrices come from denoiser predictions on the original and partially corrupted inputs.
  • The derivation uses one-hot z_s and the implication that a masked z_t remains masked after further corruption.

D.2. Training GenMol

GenMol’s experiments cover de novo, fragment-constrained, hit-generation, and lead-optimization settings using standardized datasets, baselines, and evaluation measures. Results show strong generation quality, sampling efficiency, and controllable QED and SA distributions.

  • Training GenMol: GenMol’s training uses BERT with the SAFE dataset and tokenizer, vocabulary size K=1880, batch size 2048, and learning rate 3e−4.
  • Experimental setup: The de novo experiments use RDKit and TDC to calculate fingerprints, diversity, QED, and SA, with mask-chunk lengths sampled from ZINC250k.
  • Fragment-constrained generation: The fragment-constrained benchmark extracts fragments and structural components from 10 known drugs for linker design, motif extension, scaffold morphing, and related tasks.
  • De novo generation: GenMol significantly outperforms JT-VAE and DiGress in molecular quality and sampling time.
  • De novo generation: The de novo evaluation reports means and standard deviations over 3 runs, with N, τ, and r controlling unmasking count, temperature, and randomness.
  • De novo generation: GenMol generates molecules with higher QED and lower SA than SAFE-GPT, while temperature and randomness control these distributions.

E.3. Analysis on Quality and Diversity in De Novo Generation

GenMol trades diversity for higher molecular quality under one sampling configuration and can adjust this balance through its sampling parameters.

  • 84.6% quality and 0.818 diversity are achieved by GenMol with N=1, τ=0.5, and r=0.5.
  • 38.2% quality and 0.897 diversity are obtained from 100k molecules randomly sampled from the test set.
  • GenMol’s higher-quality configuration sacrifices diversity by selecting a specific molecular mode.

E.4. Analysis on QED and SA Thresholds in Quality Metric

The quality metric remains robust when its QED and SA thresholds change, and GenMol consistently outperforms SAFE-GPT across the tested settings.

  • Quality summarizes GenMol’s ability to generate chemically plausible molecules.
  • The standard quality thresholds are QED ≥0.6 and SA ≤4, with softer and harsher alternatives also evaluated.
  • GenMol outperforms SAFE-GPT under softer, standard, and harsher QED and SA threshold settings.
Loading 2501.06158v3…