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 paper proposes a latent prompt Transformer model for solving challenging optimization problems such as molecule design, where the goal is to find molecules with optimal values of a target chemical or biological property that can be computed by an existing software. Our proposed model consists of three components. (1) A latent vector whose prior distribution is modeled by a Unet transformation of a Gaussian white noise vector. (2) A molecule generation model that generates the string-based representation of molecule conditional on the latent vector in (1). We adopt the causal Transformer model that takes the latent vector in (1) as prompt. (3) A property prediction model that predicts the value of the target property of a molecule based on a non-linear regression on the latent vector in (1). We call the proposed model the latent prompt Transformer model. After initial training of the model on existing molecules and their property values, we then gradually shift the model distribution towards the region that supports desired values of the target property for the purpose of molecule design. Our experiments show that our proposed model achieves state of the art performances on several benchmark molecule design tasks.

Figures (4)

• Figure 1: Latent Prompt Transformer. $x$ is the string-based representation of molecule. $y$ is the value of a target property. $z$ is the latent vector. $z_0 \sim {\cal N}(0, I_d)$. (1) The prior distribution of $z$ is modeled by a Unet transformation of $z_0$, i.e., $z = U_\alpha(z_0)$. Given $z$, $x$ and $y$ are independent. (2) $p_\beta(x|z)$ is the generation model, parametrized by a causal Transformer with $z$ serving as the prompt. (3) $p_\gamma(y|z)$ is the property prediction model, which is a non-linear regression on $z$ parametrized by a multi-layer perceptron (MLP).
• Figure 2: Distribution shift of ACAA1 binding affinity across optimization iterations. For each shift iteration, we plot the densities of property values estimated from AutoDock-GPU.
• Figure 3: Molecules produced during the multi-objective optimization for ESR1. The legends denote $\mathrm{K_D}\downarrow$, SA$\downarrow$ and QED$\uparrow$.
• Figure 4: Molecules produced during the multi-objective optimization for ACAA1. The legends denote $\mathrm{K_D}\downarrow$, SA$\downarrow$ and QED$\uparrow$.