Source-linked AI summary
Multi-modal Molecule Structure-text Model for Text-based Retrieval and Editing
Shengchao Liu, Weili Nie, Chengpeng Wang, Jiarui Lu, Zhuoran Qiao, Ling Liu, Jian Tang, Chaowei Xiao, Anima Anandkumar
TL;DR
Existing molecule-learning studies largely use chemical structures while overlooking textual chemistry knowledge. MoleculeSTM jointly learns molecular structures and descriptions with contrastive learning, and its evaluations show improved performance across zero-shot retrieval, editing, and property prediction tasks. The authors report that it can retrieve novel drug-target relations and modify molecule substructures toward desired properties, while noting data insufficiency and limited chemical-structure-model expressiveness.
Problem
Existing studies mainly use molecular chemical structures and overlook textual chemistry knowledge needed for text-based design objectives and complex biological-activity prediction.
Method
MoleculeSTM jointly learns molecular chemical structures and textual descriptions through contrastive learning, using PubChemSTM with over 280,000 structure-text pairs.
Results
MoleculeSTM consistently improves performance over existing methods across two zero-shot tasks and one property-prediction benchmark, including retrieval and editing capabilities.
Takeaways & Limitations
MoleculeSTM can retrieve novel drug-target relations and modify molecule substructures toward desired properties, supporting downstream practices such as repurposing and multi-objective lead optimization.
Takeaways & Limitations
The work is limited by data insufficiency and by the expressiveness of its chemical structure models, including SMILES, GNN, and generative-model components.
Abstract
from arXiv · showhide
There is increasing adoption of artificial intelligence in drug discovery. However, existing studies use machine learning to mainly utilize the chemical structures of molecules but ignore the vast textual knowledge available in chemistry. Incorporating textual knowledge enables us to realize new drug design objectives, adapt to text-based instructions and predict complex biological activities. Here we present a multi-modal molecule structure-text model, MoleculeSTM, by jointly learning molecules' chemical structures and textual descriptions via a contrastive learning strategy. To train MoleculeSTM, we construct a large multi-modal dataset, namely, PubChemSTM, with over 280,000 chemical structure-text pairs. To demonstrate the effectiveness and utility of MoleculeSTM, we design two challenging zero-shot tasks based on text instructions, including structure-text retrieval and molecule editing. MoleculeSTM has two main properties: open vocabulary and compositionality via natural language. In experiments, MoleculeSTM obtains the state-of-the-art generalization ability to novel biochemical concepts across various benchmarks.
Results
MoleculeSTM supports zero-shot structure-text retrieval and text-based molecule editing by connecting molecular representations with natural-language descriptions. Across retrieval, editing, and property prediction, it achieves strong performance while exhibiting open-vocabulary and compositional behavior.
- Method: MoleculeSTM uses separate chemical-structure and textual-description branches, aligning their representations through contrastive learning.The structure branch can use SMILES Transformers or molecular-graph GNNs, while the text branch encodes high-level molecular functionality.
- Structure-text retrieval: MoleculeSTM achieves approximately 50%, 40%, and 15% accuracy improvements over the best baseline on three zero-shot retrieval datasets with T = 20.Both structure encoders outperform other methods across the three datasets and two retrieval directions; GNN representations are stronger than SMILES representations on Description and Pharmacodynamics.
- Structure-text retrieval: MoleculeSTM retrieves ground-truth ATC labels with high rankings in four drug re-purposing case studies.The evaluation retrieves the 10 most similar labels from 600 ATC labels given a molecule’s chemical structure.
- Text-based molecule editing: MoleculeSTM obtains the best satisfactory hit ratios across all 20 zero-shot editing tasks spanning four editing categories.The results cover single-objective, multi-objective, binding-affinity-based, and drug-relevance prompts, for both SMILES and molecular-graph encoders.
- Text-based molecule editing: Visual analyses show edits matching requested properties, including functional-group changes that alter solubility, permeability, hydrogen-bond acceptors, and donors.Multi-objective edits also combine polar and hydrophobic structural changes to target solubility and permeability jointly.
- Molecular property prediction: On eight MoleculeNet classification tasks, MoleculeSTM has the best overall average performance for both SMILES and molecular-graph representations.The SMILES model improves on six of eight tasks, while the graph model performs best on four and is comparable to the best baselines on the other four.
Discussion
MoleculeSTM consistently improved performance across two zero-shot tasks and a standard property-prediction benchmark, while supporting novel drug-target retrieval and molecule editing. The authors identify data insufficiency and limited chemical-structure-model expressiveness as remaining bottlenecks, and note terminology-related effects in some tasks.
- MoleculeSTM consistently improved performance over existing methods on two zero-shot tasks and one standard property-prediction benchmark.
- The model retrieved novel drug-target relations and modified molecule substructures to obtain desired properties.
- Downstream task outcomes were consistent with chemistry-expert feedback, reflecting domain-knowledge exploration ability.
- PubChemSTM remains data-insufficient despite being 28× larger than datasets used in existing work.
- Chemical-structure encoders and the SMILES-based generative model remain bottlenecks in expressiveness.
- Text tokenization requires care because rich terminologies, including DrugBank-ATC codes, affect overall performance.
Methods
MoleculeSTM combines pretrained chemical-structure and scientific-language encoders through contrastive learning, then applies the joint space to zero-shot retrieval, text-based editing, and property prediction. PubChemSTM supplies 281K structure-text pairs, while downstream editing aligns latent spaces before prompt-conditioned optimization and decoding.
- Dataset construction: PubChemSTM contains 250K molecules and 281K structure-text pairs constructed from PubChem textual fields.
- Pretraining: The structure branch encodes SMILES strings and 2D molecular graphs, using MegaMolBART pretrained on 500M ZINC molecules for SMILES.
- Pretraining: The textual branch uses BERT and adapts SciBERT pretrained on chemical and biological text to encode molecular functionality.
- Pretraining: MoleculeSTM uses separate chemical-structure and textual-description branches connected through contrastive learning in a joint representation space.
- Zero-shot structure-text retrieval: Zero-shot retrieval selects the textual description or chemical structure with highest similarity using frozen pretrained encoders and projectors.
- Zero-shot text-based molecule editing: Molecule editing first aligns the generative-model space with MoleculeSTM, then optimizes a latent code close to both input-molecule and text representations.
- Zero-shot text-based molecule editing: The optimized latent code is decoded into an output molecule, and satisfactory hit ratio measures whether it satisfies the text prompt.
- Molecular property prediction: Property prediction fine-tunes a pretrained structure encoder with a prediction head for categorical or scalar molecular properties.
Data Availability
The paper reports that licensing constraints hindered release of the PubChemSTM dataset.
- Licensing evaluation for each of the 280K structure-text pairs hindered release of PubChemSTM.
Code Availability
The release provides code, pretrained checkpoints, downstream-task scripts, and notebook demonstrations for using MoleculeSTM.
- The release includes source code, Zenodo materials, pretraining and downstream-task scripts, pretrained checkpoints, demos, and dataset-customization support.
Author Contributions Statement
The experiments were conceived and designed by S.L., W.N., C.W., Z.Q., C.X., and A.A., with S.L. performing them and S.L. and C.W. analyzing the data.
- S.L., W.N., C.W., Z.Q., C.X., and A.A. conceived and designed the experiments.
- S.L. performed the experiments.
- S.L. and C.W. analyzed the data.
- S.L., C.W., and J.L. contributed analysis tools.
- The paper was written by S.L., W.N., C.W., J.L., Z.Q., L.L., J. T., C.X., and A.A.
Supplementary Information
The supplementary information describes PubChemSTM construction, preprocessing variants, vocabulary comparison, and reproducibility provisions.
- PubChemSTM Construction: PubChemSTM is extracted from the PubChem database as a chemical structure-text pair dataset.
- PubChemSTM Construction: The construction yields 250K molecules with unique PubChem IDs and 281K chemical structure-text pairs.
- Pre-processing Details: PubChemSTM-raw retains raw annotations, whereas PubChemSTM-extracted replaces molecule names with text prompts.
- Pre-processing Details: Preprocessing manually removes common-name or IUPAC-name prefixes, the extra word “Pure,” and identified typos.
- Reproducibility: The authors provide preprocessed datasets and source code because PubChem is updated frequently.
- Comparison: The SciBERT tokenizer produces a small vocabulary-size difference between PubChemSTM-raw and PubChemSTM-extract.
A.2 Architecture Details
MoleculeSTM uses separate chemical-structure and textual-description branches aligned by contrastive learning, with structure inputs, text inputs, and pretraining specifications documented.
- Architecture: MoleculeSTM has chemical structure and textual description branches, denoted f_c and f_t.
- Chemical Structure Branch: The chemical structure branch represents molecules from SMILES strings or 2D molecular graphs using deep learning encoders.
- Textual Description Branch: The textual description branch encodes high-level molecular functionality with a BERT-based text encoder adapted from SciBERT.
- Specifications: The supplementary tables list model specifications, pretraining hyperparameters, and running times for SMILES and 2D graph inputs.
- Pretraining Objective: MoleculeSTM applies contrastive pretraining using EBM-NCE or InfoNCE to align chemical structures and textual descriptions.
- Pretraining Objective: The EBM-NCE objective uses structure-text pairs as positive samples and randomly sampled structure or text items as negatives.
B Design Principles for Downstream Tasks
The downstream tasks are designed for exact computational evaluation while accommodating the non-unique outputs possible in molecule editing.
- Applicable Evaluation: The design principles focus on downstream tasks that are computationally feasible to evaluate.
- Applicable Evaluation: Drug-discovery outputs such as edited-molecule properties can be evaluated exactly in silico or in vitro, unlike many vision-language tasks.
- Fuzzy Matching: Molecule editing uses fuzzy matching because multiple output molecules may satisfy a text prompt.
C.1 Dataset Construction
The retrieval experiments use three DrugBank text fields—Description, Pharmacodynamics, and ATC—and filter structure-text pairs for molecular and textual overlap. Dataset statistics are summarized in Table 7.
- DrugBank fields: Three DrugBank fields support retrieval: Description, Pharmacodynamics, and anatomical therapeutic chemical (ATC).Description covers chemical properties, history, and regulatory status; Pharmacodynamics covers organism-level effects; ATC categorizes molecules.
- Filtering: Chemical structures are filtered when a molecule with the same canonical SMILES already appears in PubChemSTM.
- Filtering: Textual data are filtered using a similarity based on text length and Levenshtein distance between DrugBank and PubChemSTM descriptions.
- Dataset statistics: Table 7 reports statistics for the three DrugBank fields after filtering.
C.2 Experiments
The experiments evaluate zero-shot retrieval with three baselines, including a randomly initialized proof-of-concept baseline. Results are reported for three datasets and their corresponding retrieval settings.
- Experimental design: Three baselines are introduced in the main experiments, alongside a randomly initialized baseline called Random.Random uses randomly initialized encoders for both modalities.
- Experimental design: Zero-shot retrieval results are reported on three datasets in Tables 8 to 10.
- Retrieval results: Table 8 reports accuracy for DrugBank-Description T-choose-one retrieval.
- Retrieval results: Table 9 reports accuracy for DrugBank-Pharmacodynamics T-choose-one retrieval.
- Retrieval results: Table 10 reports accuracy for molecule-ATC T-choose-one retrieval.
C.3 Ablation Study: Fixed Pretrained Encoders
The ablation study fixes pretrained structure and language encoders while optimizing only projection layers, then reports retrieval results across three datasets. The surrounding editing setup highlights challenges in controllable molecular generation.
- Fixed pretrained encoders: The fixed-encoder ablation optimizes only the projection layers while keeping the chemical and textual encoders fixed.The main experiment instead updates all model parameters during contrastive pretraining.
- Fixed pretrained encoders: The ablation initializes the branches from GraphMVP, MegaMolBART, and SciBERT checkpoints.
- Retrieval ablations: Fixed-encoder ablation results are reported for DrugBank-Description retrieval.
- Retrieval ablations: Fixed-encoder ablation results are reported for DrugBank-Pharmacodynamics retrieval.
- Retrieval ablations: Fixed-encoder ablation results are reported for DrugBank-ATC retrieval.
- Molecule editing context: Molecule editing is framed as controllable generation, but obtaining a suitably disentangled molecular generative backbone is nontrivial.
D.1 Experiment Set-up
The editing evaluation covers single-objective, multi-objective, binding-affinity, and drug-relevance tasks using randomly sampled ZINC molecules, with quantitative hit ratios and qualitative case studies. The downstream scope also includes property prediction, while performance is limited by the fixed 2D backbone.
- Experiment set-up: Four editing categories comprise eight single-objective, six multi-objective, six binding-affinity, and four drug-relevance tasks.Three neighborhood-searching case studies are also included.
- Experiment set-up: The editing evaluation uses 200 randomly selected ZINC molecules absent from PubChemSTM.The selection preserves property distributions consistent with the full dataset.
- Single-objective editing: Single-objective editing evaluates solubility, drug-likeness, permeability, hydrogen-bond acceptors, and hydrogen-bond donors using task-specific thresholds.LogP, QED, tPSA, HBA, and HBD provide the property measures.
- Quantitative evaluation: Table 14 measures single-objective editing with hit ratios of property changes for SMILES-string and molecular-graph latent optimization.
- Qualitative evaluation: The editing visualizations explain property-directed structural changes using LogP, tPSA, HBA, and HBD.
- Multi-objective editing: Multi-objective editing tests simultaneous satisfaction of property conditions under loose and strict thresholds.Examples combine solubility with permeability or HBA/HBD requirements.
- Case studies: Three case studies generate approved drugs from patented analogs, including text-consistent changes such as an Erlotinib tPSA reduction from 83 to 75.
- Property prediction: Molecular property prediction uses a fixed GIN 2D GNN backbone, so MoleculeSTM performance is limited by that backbone.The authors identify more advanced GNN backbones as future work.