Molecule Design by Latent Prompt Transformer

TL;DR

Experiments demonstrate that the Latent Prompt Transformer not only effectively discovers useful molecules across single-objective, multi-objective, and structure-constrained optimization tasks, but also exhibits strong sample efficiency.

Abstract

This work explores the challenging problem of molecule design by framing it as a conditional generative modeling task, where target biological properties or desired chemical constraints serve as conditioning variables. We propose the Latent Prompt Transformer (LPT), a novel generative model comprising three components: (1) a latent vector with a learnable prior distribution modeled by a neural transformation of Gaussian white noise; (2) a molecule generation model based on a causal Transformer, which uses the latent vector as a prompt; and (3) a property prediction model that predicts a molecule's target properties and/or constraint values using the latent prompt. LPT can be learned by maximum likelihood estimation on molecule-property pairs. During property optimization, the latent prompt is inferred from target properties and constraints through posterior sampling and then used to guide the autoregressive molecule generation. After initial training on existing molecules and their properties, we adopt an online learning algorithm to progressively shift the model distribution towards regions that support desired target properties. Experiments demonstrate that LPT not only effectively discovers useful molecules across single-objective, multi-objective, and structure-constrained optimization tasks, but also exhibits strong sample efficiency.

Figures (13)

• Figure 1: Left: Overview of Latent Prompt Transformer (LPT). The latent vector $z\in\mathbb{R}^d$ is a neural transformation of $z_0$, i.e., $z = U_\alpha(z_0)$, where $z_0 \sim {\cal N}(0, I_d)$. Given $z$, $x$ and $y$ are independent. $p_\beta(x|z)$ is the molecule generation model and $p_\gamma(y|z)$ predicts the property value or constraint based on $z$. Right: Illustration of molecule generation model $p_\beta(x|z)$. The latent vector $z$ is used as a prompt in the $p_\beta(x|z)$ via cross-attention.
• Figure 2: Illustration of online learning LPT. For each shift iteration, we plot the densities of docking scores $E$ using AutoDock-GPU. The increase of the docking scores indicates better binding affinity.
• Figure 3: Illustration of the development of PHGDH inhibitors spillier2021phosphoglycerate. Surface Plasmon Resonance (SPR) and AutoDock $K_D$ values are reported for each inhibitor. The trends observed between the experimental SPR values and the computational AutoDock values align well, validating the computational approach.
• Figure 4: (a) Structure-constrained Optimization. Conditionally generated compounds C2 and C3 closely resemble the human-designed compounds C2 and C3 shown in \ref{['fig:mol_one2three']}. Additionally, the right column also presents further optimized compounds that achieve improved $K_D$ scores. (b) Illustration of generated molecules binding to PHGDH with docking poses generated by AutoDock-GPU. The left panel visualizes the molecule generated through multi-objective optimization, while the right panel displays the molecule generated via structure-constrained optimization.
• Figure 5: PHGDH with NAD binding site.