Target-conditioned GFlowNet for Structure-based Drug Design

We represent structure of the pocket $P$ as a standard $K$-nearest neighbor (KNN) residue graph $\mathcal{G}^{\mathcal{P}}=(\mathcal{V}^{\mathcal{P}},\mathcal{E}^{\mathcal{P}})$ - following previous work in protein representation (Jing et al., 2021). The $i$-th residue node $v_{i}^{\mathcal{P}}\in\mathcal{V}^{\mathcal{P}}$ is featurized using its geometric and chemical properties. These features include the type of residue, the dihedral angles of the atoms in the residue backbone, and the directional unit vectors. An edge $e_{ij}^{\mathcal{P}}\in\mathcal{E}^{\mathcal{P}}$ is formed if the $j$-th residue $v_{j}^{\mathcal{P}}$ is among the $K$-nearest neighbors of residue $v_{i}^{\mathcal{P}}$, as measured by the euclidean distance between their respective $C_{a}$ atoms. We set the number of neighbours $K=30$. An edge is featurized with the Euclidean distance, distance along the backbone and the direction vector between the two residues. These features sufficiently describe the features of the protein pocket.

We apply a graph neural network with geometric vector perceptrons (GVP) layers (Jing et al., 2021) to the KNN pocket graph $\mathcal{G}^{\mathcal{P}}$ to learn the node embedding $\mathbf{h}_{v_{i}^{\mathcal{P}}}$ for each residue. The node embeddings {$\mathbf{h}_{v_{i}^{\mathcal{P}}}\}$ are then averaged to obtain an embedding of the entire graph $\mathbf{h}_{\mathcal{G}^{\mathcal{P}}}$. We use GVP because it encodes the protein pocket into an embedding that is invariant to rotations and translations.

In this section, we discuss how to employ the pocket structure, more specifically its latent embedding, to condition the GFlowNet to generate molecules that interact with a given protein pocket. Furthermore, we describe our fragment-based molecular generation framework.

Molecule representation. During molecular generation, we represent ligands as a 2D molecular graph $\mathcal{G}^{L}=(\mathcal{V}^{L},\mathcal{E}^{L})$ with node $v^{L}_{i}\in\mathcal{V}^{L}$ representing a molecule fragment, and directional edges $e^{L}_{ij}\in\mathcal{E}^{L}$ indicating the attachment atom of fragment $v^{L}_{i}$ that connects to fragment $v^{L}_{j}$. Since a molecule’s reward (representing its desirability as a real-world drug) should be the same regardless of its predicted 3D conformation, we represent ligands as 2D graphs here. ¹¹1Docking score is dependent on a molecule’s 3D conformation. However, we consider the best docking score for a molecule as the reward here, which is invariant to its conformation in this context. Taking the best docking score is a valid approach, because when a compound is tested for binding in the wet lab, the compound simply binds the protein with the conformations which results in the strongest affinity. Furthermore, the predicted 3D position of generated molecules from existing diffusion-based SBDD models often changes significantly upon re-docking and likely do not reflect the true conformation. (Reidenbach, 2024).

Fragment vocabulary construction. TacoGFN generates molecules by adding one molecular fragment at a time. To create the vocabulary of fragments used, we extract common fragments from a chemical database in a data-driven and chemically valid way. To obtain a fragment vocabulary, we first apply BRICS decomposition (Degen et al., 2008) to 250k ZINC20 (Irwin et al., 2020) molecules. BRICS breaks molecules via retrosynthetic rules and provides synthetically accessible building blocks, i.e. molecules that are easy to prepare (Seo et al., 2023). To reduce the fragment vocabulary size, we further break all single bonds connecting a heavy atom to a ring structure. Next, we retain the fragments that occur in more than 50 (or 0.02%) of the molecules in our ZINC set. Only atoms (of a fragment) in the bonds decomposed by the BRICS can join to form new fragment connections during generation time; This reduces the occurrence of non-synthesizable attachment points forming bonds. This results in a set of 475 building block-like fragments. As our fragments are mined from a synthetically accessible virtual library and poor synthetic accessibility is penalized in the reward function, TacoGFN achieves superior synthetic accessibility compared to all existing SBDD methods.

Molecular generation framework. We formulate molecular generation as a sequential decision process and implement it using a GFlowNet. At the $t$-th step, the forward action policy $P_{a}$ samples an action $a$ depending on the current molecule state $\mathcal{G}^{L}_{t}$ and pocket embedding $\mathbf{h}_{\mathcal{G}^{\mathcal{P}}}$. The transition function $\mathcal{T}$ is a deterministic function which applies the action to the molecular graph at step $t$ to produce a molecule graph at step $t+1$.

Previous autoregressive models often formulate action prediction as a supervised task, where the goal is to predict the correct ground truth actions obtained from masked molecules (Peng et al., 2022). Instead, our molecule generation policy $P_{a}$ aims to generate molecules with probability $P(\mathcal{G}^{L}|\mathbf{h}_{\mathcal{G}^{\mathcal{P}}})$ proportional to the reward.

Pocket conditioned molecular generation. Here, we adopt the architecture for molecular generation introduced in Multi-Objective GFlowNet (MO-GFN) (Jain et al., 2023). Instead of conditioning the GFlowNet on multi-objective preference, we condition the GFlowNet on pocket embeddings to learn a family of molecular distributions corresponding to a family of reward functions induced from the pocket structure diversity.

We use a graph transformer (Yun et al., 2020) to model the probability $P_{a}(\mathbf{a}|\mathcal{G}^{L}_{t},\mathbf{h}_{\mathcal{G}^{\mathcal{P}}})$, by taking the partially constructed molecular graph at the $t$-th time step $\mathcal{G}^{L}_{t}$ and the pocket embedding $\mathbf{h}_{\mathcal{G}^{\mathcal{P}}}$ as input. The input feature $\mathbf{h}^{L(0)}_{i}$ of molecular node $v^{L}_{i}$ is a one-hot encoding of the node’s fragment type. The molecular edge input feature $\mathbf{e}_{ij}^{L(0)}$ is a one-hot mapping of the attachment atom index in node $v^{L}_{i}$ which connects to node $v^{L}_{j}$. We add an additional virtual conditioning node $\mathbf{h}^{V(0)}$, featurized using pocket embedding $\mathbf{h}_{\mathcal{G}^{\mathcal{P}}}$, to the graph $\mathcal{G}^{L}_{t}$ (Pham et al., 2017). This virtual node is connected to all other nodes and serves as a graph-level node to provide pocket information. After $N$ Transformer layers, the set of final node embeddings $\{\mathbf{h}^{L(N)}_{i}\}$ and edge embeddings $\{\mathbf{e}^{L(N)}_{ij}\}$ are obtained. The final graph level embedding $\mathbf{g}^{L(N)}$ is obtained via the concatenation of the global average pooling of the node embeddings and the final virtual node embedding $\mathbf{h}^{V(N)}$, as seen in Equation 4.

Using these final molecular graph embeddings, we follow the actions defined in previous works on fragment-based molecular generation (Bengio et al., 2021; Jin et al., 2018). Full details can be found in Appendix C.1.1.

Reward function. A drug candidate must not only have a high affinity to the pocket but also satisfy drug-like properties and synthetizability requirements to be selected for experimental validation. We design a reward function by multiplying the normalized score from all three aspects, using QED (Bickerton et al., 2012) as a measure of drug-likeliness, SA (Ertl & Schuffenhauer, 2009) as a measure of ease-of-synthesizability, and predicted docking score as a measure of affinity between ligand and protein²²2Computing the reward is very fast, for example, computing it for a batch of 64 protein-ligand pairs took under 0.15 seconds.. In drug discovery, it is crucial to design molecules that simultaneously satisfy all of these relevant properties, despite the inherent trade-offs among these properties. By multiplying the scores, the reward function ensures that a low score in any one of the factors (QED, SA, or DS) will significantly reduce the overall reward. Therefore, multiplying the scores is a more appropriate choice than summing the scores here.

We just have to optimize QED and SA up to a certain threshold for a molecule to be suitable as a drug candidate - optimizing them further does not bring additional utility (Coley, 2021). Therefore, we clip the reward for the QED or SA component to 1 when they achieve their respective threshold $t_{QED}$ and $t_{SA}$. For example, the reward function will prioritize more on optimizing QED if $t_{QED}$ is higher, at the expense of other properties such as affinity (See details of reward function in Appendix C.1.2).

The exploration of chemical space with a GFlowNet requires evaluating binding affinities for millions of molecules sampled during training. However, using molecular docking for evaluation is computationally expensive. Here, we propose a fast ML-based docking score predictor that generalizes well across molecule and protein pocket distributions, as described in Figure 2 and Equation 5. (See details in Appendix C.2)

We leverage PharmacoNet (Seo & Kim, 2023), a recent deep learning method to obtain the pharmacophore $\mathbb{P}$³³3 Pharmacophore is a point set, where each point describes the desirable motif properties (such as being aromatic) a ligand should possess at this geometric position to form energetically or entropically favourable interactions (e.g. $\pi$–$\pi$ stacking) with the protein target. (Wermuth et al., 1998; Yang, 2010). from a pocket structure $P$.To enrich the representation, we obtained the embedding of the pocket $\mathbf{z}^{P}$ and the embeddings of pharmacophore points $\{\hat{\mathbf{h}}^{\mathbb{P}}_{i}\}$ from PharmacoNet. The 1D representation vectors of the pharmacophore $\mathbf{z}^{\mathbb{P}}$ is computed by the global pooling of $\{\hat{\mathbf{h}}^{\mathbb{P}}_{i}\}$ .

During docking score prediction, we represent a ligand as an atom-level 2D molecular graph. We apply Graph Isomorphism Network (Hu et al., 2019) to the atom graph to obtain the node embeddings $\{\mathbf{h}^{L}_{j}\}$. $\mathbf{z}^{L}$ is a 1D representation vector of the ligand, obtained by the global pooling of node embeddings $\{\mathbf{h}^{L}_{j}\}$. Then, we incorporate a pairwise interaction map $\mathbf{I}$ - computed from the outer product of $\{\hat{\mathbf{h}}^{\mathbb{P}}_{i}\}$ and $\{\mathbf{h}^{L}_{j}\}$. This preserves the structural topology and binding interaction details. Notably, our docking score predictor uses the pharmacophore points involved in binding instead of atoms or amino acids to obtain the interaction map. It improves the generalization ability at a reduced computational cost through coarse-grained modelling of the pocket at the pharmacophore level.

We train and evaluate TacoGFN on the commonly used CrossDocked benchmark (Francoeur et al., 2020). We first apply the splitting and processing protocol on the CrossDocked dataset to obtain the same train and test split of the 100k protein-ligand pairs as previous methods (Luo et al., 2021; Peng et al., 2022; Guan et al., 2023b). (See details in Appendix D.1). We then train our docking score predictor on the training split of the protein-ligand pairs to predict their corresponding Vina Dock scores. Since TacoGFN is trained on the same set of protein-ligand pairs and evaluated on the same unseen pockets, our experimental results can be directly compared against the published results of existing structure-based generative models.

We train one TacoGFN model to generate molecules conditional to any protein target structures using predicted affinity, Synthetic Accessibility (SA), and drug-likeness (QED) as the reward. We also adopt Double GFN (Lau et al., 2023) to improve exploration in sparse reward domains and high-dimensional states, by initializing two networks to model the action policy: an online network and a target network.

For each training trajectory, one protein pocket is randomly drawn from the CrossDocked training set first. Then, a molecule is generated through sampling a sequence of actions using our target network which models the forward action policy $P_{target}$. The reward is then computed for that molecule with respect to the protein target. The Trajectory Balance loss (Equation 1) is used to train the online network modelling the action policy $P_{online}$. This loss has the objective for our model to generate objects with probabilities proportional to the rewards (consisting of predicted docking score, QED and SA). Periodically, the weights of the target network are updated by the weights of the online network using a delayed strategy. This strategy reduces training instability and promotes explorations of the larger chemical space.

In all evaluations, each structure-based generative model is tasked to produce 100 molecules (ligands) for each of the 100 unseen protein pockets from the CrossDock-100k test set.

Evaluation metrics. We adopt the following commonly used metrics from Guan et al. (2023a) and Reidenbach (2024): (1) Validity is the percentage of unique generated molecules free of reconstruction errors and disconnections as determined by RDKit. (2) Vina Dock approximates the binding energy between a generated molecule and a protein pocket, where a lower docking score indicates a higher binding affinity. (3) High Affinity measures the percentage of generated molecules with higher affinity than the reference molecule. (4) QED is a measure of drug-likeness, estimating a molecule’s suitability as an oral drug based on its properties (Bickerton et al., 2012). (5) Synthetic Accessibility (SA) estimates difficulty of synthesizing the molecule (Ertl & Schuffenhauer, 2009). The score is normalized between 0 and 1 using the formula $(10-SA)/9$. (6) Diversity is calculated as the average pairwise fingerprint Tanimoto distance between molecules generated for a pocket. (7) Success Rate is the percentage of molecules which pass the same criteria (QED > 0.25, SA > 0.59, Vina Dock < -8.18) as in Long et al. (2022); Guan et al. (2023b); Zhou et al. (2024); Reidenbach (2024). (8) Time is the average runtime (in seconds) for generating 100 unique and valid molecules for a pocket.

Baselines and problem settings. Similar to Zhou et al. (2024) and (Reidenbach, 2024), we separate existing methods by their problem definitions:

1) Generative methods are expected to generalize for pocket structures unseen during training. They generate molecules for test pockets at inference time without optimization loops or access to docking programs. Therefore, methods under this setting are only allowed to generate 100 molecules for each pocket in one-shot. We compare TacoGFN trained on CrossDocked-100k against the following generative models: liGAN (Ragoza et al., 2022), GraphBP (Liu et al., 2022), AR (Luo et al., 2021), Pocket2Mol (Peng et al., 2022), TargetDiff (Guan et al., 2023a), DiffSBDD (Schneuing et al., 2023), DecompDiff (Guan et al., 2023b). For consistency with existing baselines, molecules are docked using QVina in this problem setting (Trott & Olson, 2010; Alhossary et al., 2015). For details on the docking protocol for the generative setting, please see Appendix D.2.

2) Optimization methods are able to leverage docking on the target pocket. Unlike the generative setting, these methods typically iteratively optimize the candidate pool and select the top 100 molecules. We compare against the following methods: RGA (Fu et al., 2022) and EvoSBDD (Reidenbach, 2024) use a black-box algorithms to conduct rounds of the optimization process, with molecule fitness based on docking score to the target pocket. DecompOpt and TargetDiff+Opt (Zhou et al., 2024) optimizes molecules generated by pre-trained SBDD models via rounds of optimization process involving re-docking to the target pocket. TacoGFN+FT fine-tunes a pre-trained TacoGFN to tailor the model to the target pocket using docking as the reward.

Table 1 compares TacoGFN against other baselines in the generative problem setting for SBDD. We examine the docking score performances for individual pockets in Figure 4 and 5. We also show examples of generated molecules in Figure 3 and analyze their average physical properties in Table 2. We then compare TacoGFN+FT against existing methods for the optimization problem setting in Table 3 and 4, where docking is used in optimization rounds. Lastly, we conduct ablation studies to show the benefits of using a larger docking score dataset in Table 5 and validate the utility of the pocket conditioning in Table 6.

Generative setting. Table 1 highlights the strong performance of TacoGFN in the generative problem setting for the CrossDocked test set pockets. Notably, TacoGFN boasts a significant improvement in success rate at 56.0%, more than doubling the previous best of 24.5% achieved by DecompDiff. This improvement is due to TacoGFN’s ability to generate high-affinity molecules that simultaneously satisfy the drug-likeness and ease-of-synthesis requirements. In fact, TacoGFN records the best QED, SA, and High Affinity simultaneously among all generative methods. In Figure 3, we show examples of molecules generated by TacoGFN and show its ability to generate molecules with significantly improved docking scores compared to native ligands.

It is difficult to discover molecules that simultaneously exhibit better QED and Vina Dock compared to known binders and molecules from existing baselines due to the trade-off between Vina Dock and QED. Molecules with higher molecular weight are more likely to have strong Vina Dock because of the presence of more interacting atoms, but they tend to have worse drug-like properties (QED). ⁴⁴4This trade-off is clearly observed with Pocket2Mol and DecompDiff (see Table 1 and 2), where DecompDiff has strong Vina Dock but lower QED, while Pocket2Mol exhibits the opposite trend. In Figure 4, we show the performance breakdown for the 100 test protein pockets. TacoGFN achieves better average Top-10 Vina Dock than DecompDiff in $57\%$ of the test pockets. Notably, the top molecules generated by TacoGFN also consistently demonstrate higher QED values. As shown in Appendix Figure 5, in the pockets where DecompDiff achieves a lower Top-10 Vina Dock, the molecules often exhibit a molecular mass larger than 500 daltons. These heavier molecules violate the ideal properties of orally active drugs according to the Rule of 5 (Lipinski et al., 1997). In addition, heavy molecules with high docking scores are more likely to be false positives (Pan et al., 2002). TacoGFN is able to find molecular spaces that improve both QED and Vina Dock requirements not only by modelling the reward distribution but also by exploring a broader chemical space. This is achieved through learning from generated examples using our online policy, rather than being limited to the training data.

We note exploring millions of molecules online could not be easily achieved with existing SBDD baseline. This is because the generation process of TacoGFN is a few orders of magnitude faster than existing autoregressive or diffusion-based generative methods. Additionally, TacoGFN achieves 100% in validity and uniqueness, demonstrating the efficiency of our fragment-based 2D generation framework. As shown in Table 1, TacoGFN achieves a significant improvement in both time and validity over previous methods.

Molecules generated from TacoGFN have more ideal molecular weight and drug-likeness properties, in addition to achieving better Vina Dock; Therefore, they are more suitable as drug candidates. TacoGFN samples at reward temperature $\beta$ of 64 at inference time, resulting in the modelling of the probability $p(x|c)\propto R(x,c)^{64}$. This policy focuses more on the modes of reward function. Therefore the molecules sampled have higher quality but lower diversity than those generated by the methods that do not attempt to satisfy multiple objectives. By changing $\beta$ or reward function we can trade off rewards with diversity as shown in Appendix D.4.2.

We conduct additional ablation studies using the base TacoGFN under the generative setting to examine the effect of docking score prediction accuracy and pocket conditioning on our method.

Effects of using higher quality docking score predictor. Here, we study the effect of using a more accurate docking score predictor, which is trained on a larger dataset, for reward. Since training a docking score predictor does not require high-quality protein-ligand structural data such as the CrossDock-100k set, we can introduce a second, larger dataset for docking score prediction called ZINCDock-15M. It consists of about 15M docking simulation data - from docking 1,000 random ZINC20 (Irwin et al., 2020) molecules into each of the 15,207 unique pockets from CrossDock-100k training split using QVina. We then train our pharmacophore-based docking score predictor on this larger dataset. Please see Appendix Table 11 for a comparison of docking score accuracy between using CrossDock-100k and ZINCDock-15M. As shown in Table 5, TacoGFN using the docking score predictor trained on the larger ZINCDock-15M dataset demonstrates improvements in average Vina Dock, high-affinity rate and success rate. This confirms it is possible to leverage the easily generated large-scale docking score data to generate more novel and higher affinity molecules.

Effects of pocket conditioning. To examine the effect of the proposed pocket conditioning for GFlowNet, we train a molecular generation policy unconditioned on pocket information. The docking score predictor is unchanged - meaning it still predicts a docking score for a molecule with pocket information. The results are shown in Table 6. We observe that the pocket-conditioned GFlowNet achieves higher docking scores compared to the GFlowNet without pocket conditioning. We further measure the number of non-covalent interactions of a molecule to respective pocket, by obtaining the binding pose using QVina (See Appendix D.4.1). We demonstrate generated molecules by pocket-conditioned GFlowNet result in more non-covalent interactions of for all categories (Hydrophobic interactions, Van der Waals contacts, Hydrogen binding) in Table 9. More non-covalent interactions are indicative of the molecule having better specificity to the protein target. This ablation validates that our method is indeed leveraging pocket conditioning to learn a family of molecular distribution across different pocket structures.