Controllable and Decomposed Diffusion Models for Structure-based Molecular Optimization

We follow the decomposition algorithm proposed by Guan et al. (2023b) to decompose ligand molecules into arms and a scaffold. We define the part of proteins that lies within 10Å of any atom of an arm as its corresponding subpocket.

Following DecompDiff (Guan et al., 2023b), we represent each protein atom with the following features: one-hot element indicator (H, C, N, O, S, Se), one-hot amino acid type indicator (20 dimension), one-dim flag indicating whether the atom is a backbone atom, and one-hot arm/scaffold region indicator. If the distance between the protein atom and any arm center is within 10Å, the protein atom will be labeled as belonging to an arm region and otherwise a scaffold region. The ligand atom is represented with following features: one-hot element indicator (C, N, O, F, P, S, Cl) and one-hot arm/scaffold indicator. Different from DecompDiff, the atom features are enhanced by concatenating an SE(3)-invariant feature of arms and their corresponding subpockets encoded by the condition encoder after the orignal features.

Two graphs are constructed for message passing in the protein-ligand complex: a k-nearest neighbors graph $\mathcal{G}_{K}$ upon ligand atoms and protein atoms (we choose $k=32$ in all experiments) and a fully-connected graph ${\mathcal{G}}_{L}$ upon ligand atoms. The edge features are the outer products of distance embedding and edge type. The distance embedding is obtained by expanding distance with radial basis functions (RBF) located at 20 centers between 0Å and 10Å. The edge type is a 4-dim one-hot vector indicating the edge is between ligand atoms, protein atoms, ligand-protein atoms or protein-ligand atoms. In the ligand graph, the ligand bond is represented with a one-hot bond type vector (non-bond, single, double, triple, aromatic), an additional feature indicating whether or not two ligand atoms are from the same arm/scaffold.

The controllable and decomposed diffusion model consist of two parts: a condition encoder and a diffusion-based decoder. The building block is an SE(3)-equivariant layer that is composed of three layers: atom update layer, bond update layer, and position update layer.

We denote the protein pocket as ${\mathcal{P}}=\{({\mathbf{x}}^{\mathcal{P}}_{i},{\textnormal{v}}^{\mathcal{P}}_{i})\}_{i\in\{1,\dots,N_{\mathcal{P}}\}}$ and the ligand molecule as $\mathcal{M}=\{({\bm{x}}_{i},{\bm{v}}_{i},{\bm{b}}_{ij})\}_{i,j\in\{1,\dots,N_{\mathcal{M}}\}}$, where ${\bm{x}}$ is the atom position, ${\bm{v}}$ is the atom type, and ${\bm{b}}_{ij}$ is the chemical bond type between the atom $i$ and the atom $j$. For brevity, we omit the superscript ${\mathcal{P}}$ or ${\mathcal{M}}$ in the following. We use ${\mathbf{h}}_{i}$ to denote the the SE(3)-invariant hidden state of $i$-th atom, ${\mathbf{x}}_{i}$ to denote the $i$-th atom’s coordinate, which is SE(3)-equivariant, and ${\mathbf{e}}_{ij}$ to denote the hidden state of the edge between the $i$-th atom and the $j$-th atom. They can be obtained as we described in the previous subsection. And we use $t$ to denote the time embedding as that in Ho et al. (2020).

Given a pair of protein and ligand molecule, we first decompose the molecule to get the arms. We add noise to the ligand molecules in the training set to get the perturbed molecules as the forward process of diffusion models (equation 1). The forward process is a Markov chain with fixed variance schedule $\{\beta_{t}\}_{t=1,\dots,T}$ (Ho et al., 2020). We denote $\alpha_{t}=1-\beta_{t}$ and $\bar{\alpha}_{t}=\prod_{s=1}^{t}\alpha_{s}$. More specifically, the noises at time $t$ are injected as follows:

where $K_{a}$ and $K_{b}$ are the number of atom classes and bond classes respectively.

Then the arms and subpockets are input to the condition encoder. The output of condition encoder and the perturbed ligand are further input to the diffusion-based decoder. Then the reconstruction loss $L_{t}$ at time $t$ is defined as follows:

where $({\mathbf{x}}_{t},{\mathbf{v}}_{t},{\mathbf{b}}_{t})$, $({\mathbf{x}}_{0},{\mathbf{v}}_{t},{\mathbf{b}}_{t})$, and $(\hat{{\mathbf{x}}}_{0},\hat{{\mathbf{v}}}_{0},\hat{{\mathbf{b}}}_{0})$ represents atom positions, atom types, and bond types of the perturbed molecule at time $t$, ground truth molecule, and the predicted molecule respectively, ${\bm{c}}({\mathbf{v}}_{t},{\mathbf{v}}_{0})=\bm{c}^{\star}/\sum_{k=1}^{K_{a}}c_{k}^{\star}$ and $\bm{c}^{\star}({\mathbf{v}}_{t},{\mathbf{v}}_{0})=[\alpha_{t}{\mathbf{v}}_{t}+(1-\alpha_{t})/K_{a}]\odot[\bar{\alpha}_{t-1}{\mathbf{v}}_{0}+(1-\bar{\alpha}_{t-1})/K_{a}]$, ${\bm{c}}({\mathbf{b}}_{t},{\mathbf{b}}_{0})=\bm{c}^{\star}/\sum_{k=1}^{K_{b}}c_{k}^{\star}$ and $\bm{c}^{\star}({\mathbf{b}}_{t},{\mathbf{b}}_{0})=[\alpha_{t}{\mathbf{b}}_{t}+(1-\alpha_{t})/K_{b}]\odot[\bar{\alpha}_{t-1}{\mathbf{b}}_{0}+(1-\bar{\alpha}_{t-1})/K_{b}]$. Note that the condition encoder and the diffusion-based decoder are jointly trained.

In practice, we set the loss weights as $\gamma_{v}=100$ and $\gamma_{b}=100$. Follwing the setting of Guan et al. (2023b), we set the number of diffusion steps as 1000. For this diffusion noise schedule, we choose to use a sigmoid $\beta$ schedule with $\beta_{1}=\texttt{1e-7}$ and $\beta_{T}=\texttt{2e-3}$ for atom coordinates, and a cosine $\beta$ schedule suggested in Nichol & Dhariwal (2021) with $s=0.01$ for atom types and bond types.

We use Adam Kingma & Ba (2014) with init_learning_rate=0.0005, betas=(0.95, 0.999) to train the model. And we set batch_size=16 and clip_gradient_norm=8. During the training phase, we add a small Gaussian noise with a standard deviation of 0.1 to protein atom coordinates as data augmentation. We also schedule to decay the learning rate exponentially with a factor of 0.6 and a minimum learning rate of 1e-6. The learning rate is decayed if there is no improvement for the validation loss in 10 consecutive evaluations. The evaluation is performed for every 1000 training steps. We trained our model on one NVIDIA GeForce GTX A100 GPU, and it could converge within 237k steps.

To sample molecules using the pre-trained controllable and decomposed diffusion model, assume that there are available arms as conditions, we can first sample a noisy molecule from the prior distribution and derive a molecule by iteratively denoising following the reverse process (equation 2). More specifically, the denoising step at time $t$ corresponds to sampling molecules from the following distributions:

where $\tilde{{\bm{\mu}}}_{t}({\mathbf{x}}_{t},\hat{{\mathbf{x}}}_{0})=\frac{\sqrt{\bar{\alpha}_{t-1}\beta_{t}}}{1-\bar{\alpha}_{t}}\hat{{\mathbf{x}}}_{0}+\frac{\sqrt{\alpha}_{t}(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_{t}}{\mathbf{x}}_{t}$, $\tilde{\beta}_{t}=\frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_{t}}\beta_{t}$, $\tilde{\bm{c}}({\mathbf{v}}_{t},\hat{{\mathbf{v}}}_{0})=\tilde{\bm{c}}^{\star}/\sum_{k=1}^{K_{a}}\tilde{c}_{k}^{\star}$ and $\tilde{\bm{c}}^{\star}({\mathbf{v}}_{t},\hat{{\mathbf{v}}}_{0})=[\alpha_{t}{\mathbf{v}}_{t}+(1-\alpha_{t})/K_{a}]\odot[\bar{\alpha}_{t-1}\hat{{\mathbf{v}}}_{0}+(1-\bar{\alpha}_{t-1})/K_{a}]$, $\tilde{\bm{c}}({\mathbf{b}}_{t},\hat{{\mathbf{b}}}_{0})=\tilde{\bm{c}}^{\star}/\sum_{k=1}^{K_{b}}\tilde{c}_{k}^{\star}$ and $\tilde{\bm{c}}^{\star}({\mathbf{b}}_{t},\hat{{\mathbf{b}}}_{0})=[\alpha_{t}{\mathbf{b}}_{t}+(1-\alpha_{t})/K_{b}]\odot[\bar{\alpha}_{t-1}\hat{{\mathbf{b}}}_{0}+(1-\bar{\alpha}_{t-1})/K_{b}]$. Here $(\hat{{\mathbf{x}}}_{0},\hat{{\mathbf{v}}}_{0},\hat{{\mathbf{b}}}_{0})$ is the molecule output by the diffusion-based decoder, whose input is the noisy molecule at time $t$ and the condition feature. The sampling step is illustrated as Figure 6. During sampling, we also apply validity guidance proposed by Guan et al. (2023b), which encourages the model to generate molecules with valid structures.