# Retrieval-based Controllable Molecule Generation

arXiv: https://arxiv.org/abs/2208.11126 · HTML: https://mlsanity.com/paper/2208.11126/

## Key facts

- Cited by 44 papers in the map (data through 14 Oct 2026).
- Builds on 43 of its 100 references that are in the map.
- First posted on arXiv 23 Aug 2022.

## Builds on (most cited first, 30 of 43)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension](https://mlsanity.com/paper/1910.13461/) (2019)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [Billion-scale similarity search with GPUs](https://mlsanity.com/paper/1702.08734/) (2017)
- [REALM: Retrieval-Augmented Language Model Pre-Training](https://mlsanity.com/paper/2002.08909/) (2020)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [Improving language models by retrieving from trillions of tokens](https://mlsanity.com/paper/2112.04426/) (2021)
- [What Makes Good In-Context Examples for GPT-3?](https://mlsanity.com/paper/2101.06804/) (2021)
- [Junction Tree Variational Autoencoder for Molecular Graph Generation](https://mlsanity.com/paper/1802.04364/) (2018)
- [Learning To Retrieve Prompts for In-Context Learning](https://mlsanity.com/paper/2112.08633/) (2021)
- [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://mlsanity.com/paper/2104.04473/) (2021)
- [Graph Convolutional Policy Network for Goal-Directed Molecular Graph Generation](https://mlsanity.com/paper/1806.02473/) (2018)
- [GuacaMol: Benchmarking Models for De Novo Molecular Design](https://mlsanity.com/paper/1811.09621/) (2018)
- [Grammar Variational Autoencoder](https://mlsanity.com/paper/1703.01925/) (2017)
- [Molecular De Novo Design through Deep Reinforcement Learning](https://mlsanity.com/paper/1704.07555/) (2017)
- [Memorizing Transformers](https://mlsanity.com/paper/2203.08913/) (2022)
- [Optimization of Molecules via Deep Reinforcement Learning](https://mlsanity.com/paper/1810.08678/) (2018)
- [GreaseLM: Graph REASoning Enhanced Language Models for Question Answering](https://mlsanity.com/paper/2201.08860/) (2022)
- [Constrained Language Models Yield Few-Shot Semantic Parsers](https://mlsanity.com/paper/2104.08768/) (2021)
- [Deep Reinforcement Learning for De-Novo Drug Design](https://mlsanity.com/paper/1711.10907/) (2017)
- [A Survey on Retrieval-Augmented Text Generation](https://mlsanity.com/paper/2202.01110/) (2022)
- [MARS: Markov Molecular Sampling for Multi-objective Drug Discovery](https://mlsanity.com/paper/2103.10432/) (2021)
- [Instance-Conditioned GAN](https://mlsanity.com/paper/2109.05070/) (2021)
- [Case-based Reasoning for Natural Language Queries over Knowledge Bases](https://mlsanity.com/paper/2104.08762/) (2021)
- [Multi-Objective De Novo Drug Design with Conditional Graph Generative Model](https://mlsanity.com/paper/1801.07299/) (2018)
- [Learning Multimodal Graph-to-Graph Translation for Molecular Optimization](https://mlsanity.com/paper/1812.01070/) (2018)
- [MEGATRON-CNTRL: Controllable Story Generation with External Knowledge Using Large-Scale Language Models](https://mlsanity.com/paper/2010.00840/) (2020)
- [Molecular generative model based on conditional variational autoencoder for de novo molecular design](https://mlsanity.com/paper/1806.05805/) (2018)
- [Learning to Extend Molecular Scaffolds with Structural Motifs](https://mlsanity.com/paper/2103.03864/) (2021)

## Built on by (most cited first, 30 of 44)

- [Retrieval-Augmented Generation for AI-Generated Content: A Survey](https://mlsanity.com/paper/2402.19473/) (2024)
- [Retrieval-Augmented Generation with Graphs (GraphRAG)](https://mlsanity.com/paper/2501.00309/) (2024)
- [Multi-modal Molecule Structure-text Model for Text-based Retrieval and Editing](https://mlsanity.com/paper/2212.10789/) (2022)
- [MiDi: Mixed Graph and 3D Denoising Diffusion for Molecule Generation](https://mlsanity.com/paper/2302.09048/) (2023)
- [Buffer of Thoughts: Thought-Augmented Reasoning with Large Language Models](https://mlsanity.com/paper/2406.04271/) (2024)
- [Reliable, Adaptable, and Attributable Language Models with Retrieval](https://mlsanity.com/paper/2403.03187/) (2024)
- [ReasonFlux: Hierarchical LLM Reasoning via Scaling Thought Templates](https://mlsanity.com/paper/2502.06772/) (2025)
- [GenMol: A Drug Discovery Generalist with Discrete Diffusion](https://mlsanity.com/paper/2501.06158/) (2025)
- [RealGen: Retrieval Augmented Generation for Controllable Traffic Scenarios](https://mlsanity.com/paper/2312.13303/) (2023)
- [Thought Propagation: An Analogical Approach to Complex Reasoning with Large Language Models](https://mlsanity.com/paper/2310.03965/) (2023)
- [Computer-Aided Multi-Objective Optimization in Small Molecule Discovery](https://mlsanity.com/paper/2210.07209/) (2022)
- [ChatGPT-powered Conversational Drug Editing Using Retrieval and Domain Feedback](https://mlsanity.com/paper/2305.18090/) (2023)
- [UniMoT: Unified Molecule-Text Language Model with Discrete Token Representation](https://mlsanity.com/paper/2408.00863/) (2024)
- [Molecule Generation with Fragment Retrieval Augmentation](https://mlsanity.com/paper/2411.12078/) (2024)
- [Domain-Agnostic Molecular Generation with Chemical Feedback](https://mlsanity.com/paper/2301.11259/) (2023)
- [Drug Discovery with Dynamic Goal-aware Fragments](https://mlsanity.com/paper/2310.00841/) (2023)
- [RANGE: Retrieval Augmented Neural Fields for Multi-Resolution Geo-Embeddings](https://mlsanity.com/paper/2502.19781/) (2025)
- [Retrieval Augmented Diffusion Model for Structure-informed Antibody Design and Optimization](https://mlsanity.com/paper/2410.15040/) (2024)
- [Retrieval-Retro: Retrieval-based Inorganic Retrosynthesis with Expert Knowledge](https://mlsanity.com/paper/2410.21341/) (2024)
- [3M-Diffusion: Latent Multi-Modal Diffusion for Language-Guided Molecular Structure Generation](https://mlsanity.com/paper/2403.07179/) (2024)
- [Source Code Foundation Models are Transferable Binary Analysis Knowledge Bases](https://mlsanity.com/paper/2405.19581/) (2024)
- [PrefixMol: Target- and Chemistry-aware Molecule Design via Prefix Embedding](https://mlsanity.com/paper/2302.07120/) (2023)
- [Small Molecule Optimization with Large Language Models](https://mlsanity.com/paper/2407.18897/) (2024)
- [Retrieval-augmented Generation for GenAI-enabled Semantic Communications](https://mlsanity.com/paper/2412.19494/) (2024)
- [GraphVF: Controllable Protein-Specific 3D Molecule Generation with Variational Flow](https://mlsanity.com/paper/2304.12825/) (2023)
- [Diffusion Twigs with Loop Guidance for Conditional Graph Generation](https://mlsanity.com/paper/2410.24012/) (2024)
- [BioMol-MQA: A Multi-Modal Question Answering Dataset For LLM Reasoning Over Bio-Molecular Interactions](https://mlsanity.com/paper/2506.05766/) (2025)
- [REIS: A High-Performance and Energy-Efficient Retrieval System with In-Storage Processing](https://mlsanity.com/paper/2506.16444/) (2025)
- [LLMs for Supply Chain Management](https://mlsanity.com/paper/2505.18597/) (2025)
- [Toward Closed-loop Molecular Discovery via Language Model, Property Alignment and Strategic Search](https://mlsanity.com/paper/2512.09566/) (2025)

## Abstract (arXiv)

Generating new molecules with specified chemical and biological properties via generative models has emerged as a promising direction for drug discovery. However, existing methods require extensive training/fine-tuning with a large dataset, often unavailable in real-world generation tasks. In this work, we propose a new retrieval-based framework for controllable molecule generation. We use a small set of exemplar molecules, i.e., those that (partially) satisfy the design criteria, to steer the pre-trained generative model towards synthesizing molecules that satisfy the given design criteria. We design a retrieval mechanism that retrieves and fuses the exemplar molecules with the input molecule, which is trained by a new self-supervised objective that predicts the nearest neighbor of the input molecule. We also propose an iterative refinement process to dynamically update the generated molecules and retrieval database for better generalization. Our approach is agnostic to the choice of generative models and requires no task-specific fine-tuning. On various tasks ranging from simple design criteria to a challenging real-world scenario for designing lead compounds that bind to the SARS-CoV-2 main protease, we demonstrate our approach extrapolates well beyond the retrieval database, and achieves better performance and wider applicability than previous methods. Code is available at https://github.com/NVlabs/RetMol.