ControllableGPT: A Ground-Up Designed Controllable GPT for Molecule Optimization

TL;DR

The work tackles the need for controllable, bidirectional generation in molecule optimization by introducing ControllableGPT, a ground-up GPT model built around the Causally Masked Seq2seq (CMS) objective that blends MLM, CLM, and seq2seq. CMS enables per-token control with expansion, contraction, or mutation at specified spans while preserving designated regions in SMILES strings, facilitated by a three-phase pretraining regimen and a novel generation process. The model is evaluated on viral and cancer drug benchmarks, where it outperforms eight baselines, achieving higher similarity to original drugs and improved property scores, including toxicity reduction in targeted cases. This approach advances controllable generative modeling for drug optimization and demonstrates practical benefits for editing existing therapeutics with precision and safety considerations, paving the way for broader applications in molecular design.

Abstract

Large Language Models (LLMs) employ three popular training approaches: Masked Language Models (MLM), Causal Language Models (CLM), and Sequence-to-Sequence Models (seq2seq). However, each approach has its strengths and limitations, and faces challenges in addressing specific tasks that require controllable and bidirectional generation, such as drug optimization. To address this challenge, inspired by the biological processes of growth and evolution, which involve the expansion, shrinking, and mutation of sequences, we introduce ControllableGPT. This initiative represents the first effort to combine the advantages of MLM, CLM, and seq2seq into a single unified, controllable GPT framework. It enables the precise management of specific locations and ranges within a sequence, allowing for expansion, reduction, or mutation over chosen or random lengths, while maintaining the integrity of any specified positions or subsequences. In this work, we designed ControllableGPT for drug optimization from the ground up, which included proposing the Causally Masked Seq2seq (CMS) objective, developing the training corpus, introducing a novel pre-training approach, and devising a unique generation process. We demonstrate the effectiveness and controllability of ControllableGPT by conducting experiments on drug optimization tasks for both viral and cancer benchmarks, surpassing competing baselines.

Figures (6)

• Figure 1: Penicillin in drug optimization. With adding a simple functional group NH2 (in red), Ampicillin has resolved the rash side effect bring about by Penicillin.
• Figure 2: The visual representation of our causal masked objective on a molecule features two mask spans $(n=2)$, each with a specific size hint. The first span, <$mask\_1:2$>, covers two tokens, and the second, <$mask\_2:7$>, covers seven tokens.
• Figure 3: The visual representation of building the training corpus with both masked and seq2seq spans for seq2seq causal masked objective.
• Figure 4: Modification of an original molecule. This figure illustrates the process of altering a molecule's structure. Key steps include replacing original segments with masked and sequence-to-sequence tokens (highlighted in red), generating new molecular segments (in green) by the model, and manually reintegrating these segments into the molecule.
• Figure 5: Expansion of an original molecule: Mask tokens (in red) are inserted into the SMILES string, prompting the generation of new segments (in green). These segments are then manually added to the molecule, showcasing the model's capability to expand molecular structures both creatively and precisely.