Transformers improve molecule design with few examples and multiple goals
PairIT: Autoregressive Transformers for Low-Data Molecule Optimization
Machine LearningArtificial Intelligence
Summary
Designing new molecules for drugs is hard because you need to start from an existing molecule and improve several properties at once, often without a lot of data. The authors created PairIT, a model that learns to make small, meaningful changes to molecules by studying pairs of similar molecules. It can be guided to improve specific properties and even multiple ones at the same time. When data is limited, it can be fine-tuned or improved using reinforcement learning. The model works well even when trained on a modest amount of data.
What this means in practice
- •For drug discovery teams: Create improved drug candidate molecules from known seeds with limited data and multiple property goals.
- •For chemical engineers: Design custom molecules by generating targeted chemical modifications using a transformer trained on molecule pairs.
Authors
Hannah Lawrence, Bodhi P. Vani, Ji Won Park, Natasa Tagasovska, Samuel Don Stanton, Andrew Martin Watkins, Michael Maser, Ewa M. Nowara
Abstract
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.