DrugGen 2: A disease-aware language model for enhancing drug discovery

Abstract: Current computational approaches for drug design typically focus on generating molecules conditioned on specific targets or general molecular properties, often neglecting the influence of disease context on target behavior and therapeutic outcomes. To address this gap, we introduce DrugGen-2, a novel generative model that designs small molecules conditioned on both disease ontology and target protein sequences. DrugGen-2 was developed by fine-tuning a pre-trained GPT-2 model on a curated dataset of approved drugs linked to their diseases and targets, using a two-step strategy of supervised fine-tuning followed by reinforcement learning via group relative policy optimization (GRPO). This process was guided by reward functions optimizing for chemical validity, novelty, diversity, and high predicted binding affinity. When evaluated on five protein targets relevant to diabetic nephropathy, DrugGen-2 significantly outperformed baseline models (DrugGPT and DrugGen). It demonstrated a superior capacity to generate unique molecules, exhibited greater structural similarity to approved drugs, and achieved improved predicted binding affinities across all targets. Molecular docking analyses further supported these findings, identifying candidate ligands with strong binding potential, including compounds with predicted affinities (-9.917, -9.485, and -9.367) exceeding those of reference drugs such as enalapril for angiotensin-converting enzyme (-8.283). By integrating disease-specific context into molecular generation, DrugGen-2 advances AI-assisted drug discovery, offering a powerful tool for de novo design and drug repurposing that accounts for the complex interplay between diseases and molecular targets.

• The paper introduces DrugGen-2, an LLM-driven model that conditions molecular generation on disease ontologies and protein targets to enhance drug discovery.
• It employs a two-stage training approach, combining supervised fine-tuning on curated data with group-level reinforcement learning for efficient exploration.
• Results show significant improvements in output uniqueness, structural validity, and binding affinity compared to earlier models using quantitative benchmarks.

Disease-Aware Molecular Design: The DrugGen-2 Framework

Model Architecture and Training Paradigm

DrugGen-2 is an LLM-based generative model for disease-aware small molecule design, conditioned jointly on disease ontology (MeSH DAG) and protein target sequence. The architecture leverages a pre-trained DrugGPT (GPT-2 variant), fine-tuned via supervised learning on curated disease-target-drug data, and subsequently optimized using GRPO, a group-level reinforcement learning strategy combining reward functions for chemical validity, novelty, diversity, and PLAPT-predicted binding affinity. Model input is structured as <D><MeSH DAG><P><protein sequence><L><SMILES>, tokenized and truncated to ensure alignment for transformer-based processing.

Figure 1: Development and evaluation workflow for DrugGen-2, integrating supervised fine-tuning and GRPO with domain-specific rewards and multi-level assessment criteria.

The GRPO workflow circumvents the need for a value function by using baseline group rewards, achieving computational efficiency with stable advantage estimation. KL-regularization anchors learning to the supervised prior, mitigating pathological drift away from domain-relevant chemical structures.

Disease-Conditioned Molecular Generation: Scope and Performance

DrugGen-2's outputs were benchmarked against DrugGen and DrugGPT across five diabetic nephropathy-relevant targets: ACE, PPAR$\gamma$, NOS3, PAI-1, and TGF-$\beta$1. Disease context is injected via MeSH DAG identifiers, capturing hierarchical semantic information to enable nuanced generation conditioned on disease biology.

DrugGen-2 outperformed comparators with respect to:

• Unique molecule generation: Up to 444/500 unique outputs (DrugGen median 50, DrugGPT median 219), demonstrating marked enhancement in chemical space exploration capability.
• Structural validity: 99-100% valid outputs, confirming robust filtering of syntactic and valence errors through the RDKit pipeline.
• Similarity to approved drugs: Median Tanimoto similarity of 0.70 versus DrugGen (0.64) and DrugGPT (0.30), aligning generated molecules closer to clinically relevant scaffolds.
• Binding affinity: Median PLAPT-predicted affinities 9.26–9.97, substantially above DrugGen (7.15–8.49) and DrugGPT (5.86–6.22), with $\epsilon^2$ values indicating strong effect sizes and $P$ values $< 10^{-38}$ for all targets.

  Figure 2: DrugGen-2 achieves superior reward convergence, unique and valid generation, higher similarity to approved drugs, and elevated predicted binding affinity across targets; 3D protein-ligand docking highlights conformational alignment with reference binders and shows novel interactions.

Docking simulations validated the predicted affinities: DrugGen-2 ligands for ACE demonstrated docking scores (e.g., --9.917 for P12821-293) lower than reference drugs such as enalapril (--8.283), with binding pose similarity supporting functional relevance. For PPAR$\gamma$, designed ligands introduced novel pharmacophores without surpassing reference system affinity.

Reward Function Engineering and Reinforcement Strategy

Reward design integrates competitive priorities:

• PLAPT affinity: Quantitative assessment using pre-trained protein-ligand transformers, filtered by chemical validity.
• Novelty: Binary reward for SMILES absent from the approved drug dataset.
• Intra-batch diversity: Binary reward penalizing repetitive outputs to maximize local exploration of structure space.

GRPO's group-relative baseline ensures reward variance is robustly exploited, while KL divergence regularization prevents divergence from supervised chemical priors.

Methodological Implications and Scalability

DrugGen-2's modular architecture and reward strategy can be deployed to other disease ontologies and targets. Disease-conditioned generation enables context-aware drug design, potentially accounting for variable target-disease interaction semantics, such as ACE's differential roles in hypertension and neuroprotection or PPAR$\gamma$'s role in metabolic versus oncogenic pathways.

Docking protocol sensitivity was validated by RMSD calculations (Rosiglitazone/PPAR$\gamma$: 1.44Å), supporting reliability for in silico screening. Sequence truncation (≤768 residues) imposes constraints on capturing distal allosteric interactions; future work may leverage multimodal encoding strategies (protein structure+sequence) or hierarchical reward augmentation (ADMET, synthetic accessibility).

Theoretical and Practical Impact

DrugGen-2's disease-aware approach is a departure from prior target-only conditioning, enabling generation aligned with nuanced disease mechanisms. This has implications for both de novo design and drug repurposing, providing a computational framework for personalized pharmacology. The reinforcement learning integration expedites hit-to-lead workflows while maintaining clinical candidate relevance and chemical feasibility.

The paradigm is extensible: multi-objective reinforcement strategies and broader feature reward functions (toxicity, PK/PD properties) will further bridge the gap to translational medicine. Incorporation of cell-specific or tissue-specific context via additional input modalities may enable genuine precision therapeutics.

Conclusion

DrugGen-2 constitutes a scalable LLM-based framework for disease-aware molecular generation, combining hierarchical ontology encoding and protein sequence context to outperform existing models in uniqueness, validity, scaffold similarity, and binding affinity metrics. Reinforcement learning with group-level optimization and customized reward shaping enables efficient exploration of clinically relevant chemical space. The approach demonstrates the feasibility of context-sensitive AI-driven drug discovery, with future impact on precision medicine and adaptive therapeutic development.