Retrieval-Augmented Foundation Models for Matched Molecular Pair Transformations to Recapitulate Medicinal Chemistry Intuition

Bo Pan; Peter Zhiping Zhang; Hao-Wei Pang; Alex Zhu; Xiang Yu; Liying Zhang; Liang Zhao

cs.LGFeb 18, 2026

Retrieval-Augmented Foundation Models for Matched Molecular Pair Transformations to Recapitulate Medicinal Chemistry Intuition

Authors: Bo Pan, Peter Zhiping Zhang, Hao-Wei Pang, Alex Zhu, Xiang Yu, Liying Zhang, Liang Zhao

Organizations: Department of Computer Science, Emory University, Atlanta, GA, USA · Merck & Co., Inc., Rahway, NJ, USA

Abstract

Matched molecular pairs (MMPs) capture the local chemical edits that medicinal chemists routinely use to design analogs, but existing ML approaches either operate at the whole-molecule level with limited edit controllability or learn MMP-style edits from restricted settings and small models. We propose a variable-to-variable formulation of analog generation and train a foundation model on large-scale MMP transformations (MMPTs) to generate diverse variables conditioned on an input variable. To enable practical control, we develop prompting mechanisms that let the users specify preferred transformation patterns during generation. We further introduce MMPT-RAG, a retrieval-augmented framework that uses external reference analogs as contextual guidance to steer generation and generalize from project-specific series. Experiments on general chemical corpora and patent-specific datasets demonstrate improved diversity, novelty, and controllability, and show that our method recovers realistic analog structures in practical discovery scenarios.

Figures & tables

Figure 1. An example of (a) Matched Molecular Pairs (MMP); (b) Matched Molecular Pair Transformation (MMPT) and its textual representation.

Figure 2. Overview of the proposed MMPT framework. (a) The foundation model (MMPT-FM) is trained on large-scale MMPT data. (b) MMPT-FM supports controllable generation via masked template prompting. (c) MMPT-RAG augments generation with retrieval, clustering, and MCS-based template extraction to guide context-aware transformation generation.

Table 1. MMPT generation performance on three tasks.

Panel of Figure 3 — Figure 3. Visualizations of the chemical space explored by our foundation model MMPT-FM (blue) versus Database Retrieval (red) on (a) ChEMBL and (b) PMV17 datasets.

Beam	400	600	800	1000	1200
Avg Validity	0.9992	0.9991	0.9989	0.9988	0.9986

Table 2. Effect of beam size on average validity of MMPT-FM, averaged on the ChEMBL MMPT held-out test set.

	ChEMBL			PMV17
$K$	Valid (%)	Recall (%)	#Uniq	Valid (%)	Recall (%)	#Uniq
1	100.0	58.0	1.00	98.0	46.0	0.98
10	86.0	96.0	8.60	79.2	90.0	7.92
20	74.1	100.0	14.82	62.9	96.0	12.58
50	48.0	100.0	24.00	39.24	98.0	19.62
100	31.7	100.0	31.70	24.84	98.0	24.84

Table 3. Prompted generation on ChEMBL and PMV17 at different numbers of generations (

K

). We report Validity, Recall of ground truth, and numbers of generated unique and valid variables (#Unique).

Figure 4. UMAP visualization of MMPT-FM and MMPT-RAG’s chemical landscape on PMV17. The grey shaded areas represent the reference dataset’s distribution. Compared to FM inference (blue), the MMP-RAG framework (red) populates structural voids where the foundation model is sparse or absent.

Figure 5. Hyperparameter Study. (a) the number of clusters to expand, (b) the number of variables to generate for each cluster, (c) the range of mask length to fill.

Figure 6. Examples of MMPT-FM generations. In each row, the left structure is the input variable, and the structures on the right are generated outputs.

Figure 7. Examples of MMPT-RAG generations. Retrieved variables are clustered, an MCS is extracted per cluster, and generation is conditioned on each cluster’s template (MCS).

Molecular Generation

Molecular generation develops computational methods to create chemical structures and conformations for discovery, optimization, and structure elucidation.

See all 234 papers

PairIT: Autoregressive Transformers for Low-Data Molecule Optimization
Oct 13, 2026Hannah Lawrence, Bodhi P. Vani, Ji Won Park +5Autoregressive Transformer Molecular Optimization
Genentech, South San Francisco, CA, USA · Massachusetts Institute of Technology, Cambridge, MA, USA · Prescient Design, Genentech, Basel, Switzerland +1
MolRGen: A Training and Evaluation Setting for De Novo Molecular Generation with Reasonning Models
Mar 18, 2026Philippe Formont, Maxime Darrin, Ismail Ben Ayed +1Molecular Generation Molecule Design
Universit´e Paris-Saclay · ´ETS Montr´eal · ILLS – International Laboratory on Learning Systems +4
Controllable Molecular Generative Foundation Models
May 14, 2026Yihan Zhu, Yuhan Liu, Weijiang Li +2RL for Diffusion Models Molecular Graph Generation
University of Notre Dame

Oct 13, 2026J/K move · Enter open · S save

PairIT: Autoregressive Transformers for Low-Data Molecule Optimization

Hannah Lawrence, Bodhi P. Vani, Ji Won Park, Natasa Tagasovska+4

Genentech, South San Francisco, CA, USA · Massachusetts Institute of Technology, Cambridge, MA, USA · Prescient Design, Genentech, Basel, Switzerland+1

Autoregressive Transformer Molecular Optimization Drug Discovery Transformer

Open paper AskPDF arXiv

Generating novel molecules with desired properties is a central challenge in drug discovery. Although foundational work focused on de novo generation, practical use cases require a more complex setting: conditioning on an initial seed molecule, optimizing multiple properties, and operating in data-sparse regimes. We introduce PairIT, a flexible training framework for autoregressive transformers over SMILES sequences that enables seeded generation by learning from pairs of structurally similar molecules. By leveraging matched pairs, the model learns to generate local, chemically meaningful edits that improve a property of the seed molecule. Prompt tokens extend this framework by specifying the target property, desired direction, and approximate magnitude of change. Multiple prompt tokens can be concatenated to enable multi-property optimization at inference time. For properties where pair supervision alone is insufficient, reinforcement learning post-training can further improve optimization, providing a tunable trade-off between proximity to the seed and the optimization objective. PairIT naturally supports pretraining on a large dataset and fine-tuning on a smaller dataset, which is critical for performance in data-sparse regimes. Surprisingly, we observe diminishing returns to the scale of the pretraining dataset. We demonstrate empirically on MoleculeNet and Therapeutics Data Commons datasets that, with appropriate paired-data curation, a standard transformer model of modest size (68M parameters) can achieve robust seeded molecular optimization in low-data regimes.

Abstract

Molecular Generation

Explore similar work

PairIT: Autoregressive Transformers for Low-Data Molecule Optimization

MolRGen: A Training and Evaluation Setting for De Novo Molecular Generation with Reasonning Models

Controllable Molecular Generative Foundation Models