Retrieval-based Controllable Molecule Generation
Authors: Zichao Wang
Weili Nie
Zhuoran Qiao
Chaowei Xiao
Richard Baraniuk
Anima Anandkumar
Venue: ICLR 2023 notable top 25% (2023)
Keywords: controllable molecule generation
retrieval mechanism
exemplar molecules
drug discovery

Abstract: Generating new molecules with specified chemical and biological properties via generative models has emerged as a promising direction for drug discovery. However, existing methods require extensive training/fine-tuning with a large dataset, often unavailable in real-world generation tasks. In this work, we propose a new retrieval-based framework for controllable molecule generation. We use a small set of exemplar molecules,  i.e., those that (partially) satisfy the design criteria, to steer the pre-trained generative model towards synthesizing molecules that satisfy the given design criteria. We design a retrieval mechanism that retrieves and fuses the exemplar molecules with the input molecule, which is trained by a new self-supervised objective that predicts the nearest neighbor of the input molecule. We also propose an iterative refinement process to dynamically update the generated molecules and retrieval database for better generalization. Our approach is agnostic to the choice of generative models and requires no task-specific fine-tuning. On various tasks ranging from simple design criteria to a challenging real-world scenario for designing lead compounds that bind to the SARS-CoV-2 main protease, we demonstrate our approach extrapolates well beyond the retrieval database, and achieves better performance and wider applicability than previous methods.

Published as a conference paper at ICLR 2023
RETRIEVAL-BASED CONTROLLABLE MOLECULE GENERATION
Zichao Wang Rice University
Weili Niet NVIDIA wnie@nvr com
Zhuoran Qiao Caltech zq aoêca tech. edu
jzwangeri ce.duu
Chaowei Xiao NVIDIA, ASU
Richard G. Baraniuk Rice University
Anima Anandkumar
NVIDIA, Caltech
chacweixenvidia.com
richberice reaebbriceeehh
Baraanckumarnnviddia..
ABSTRACT
Generating new molecules with specified chemical and biological properties via generative models has emerged promising direction for drug discovery. How- ever, existing methods require extensive training/fine-tunin
with large dataset, often unavailable in real-world generation tasks. In this work, propose a new retrieval- based framework for controllable molecule generation, We use a small set of exemplar molecules, i.c., those that (partially) satisfy the design criteria, steer the pre-trained generative model towards synthesizing molecules that satisfy the given design criteria. We design retrieval mechanism that retrieves and fuses the exemplar molecules with the input molecule, which is trained by new self- supervised objective that predicts the nearest neighbor the input molecule. We also propose iterative refinement process to dynamically update the generated molecules and retrieval database for better generalization. Our approach is agnostic the choice of generative models and requires task-specific fine-tuning. On various tasks ranging from simple design criteria to a challenging real-world sce- nario for designing lead compounds that bind to the SARS-CoVE CoV--CoV-2 main protease, we demonstrate our approach extrapolates well beyond the retrieval database, and achieves better performance and wider applicability than previous methods.
INTRODUCTION
Drug discovery a complex, multi-objective problem /Vamathevan 2019). For drug be safe and effective, the molecular entity must interact favorably with the desired target (Parenti and Rastelli 2012), possess favorable physicochemical properties such as solubility (Meanwell) [201]) and be readily synthesizable (Jiménez-Luna et al. (2021). Compounding the challenge is the massive search space (up to 1080 molecules Polishchuk et al. (2013)) Previous efforts address this challenge via high- throughput virtual screening (HTVS) techniques (Walters et al. 1998) by searching against existing molecular databases. Combinatorial approaches have also been proposed to enumerate molecules beyond the space of established drug-like molecule datasets. For example, genetic-algorinn (GA) based methods (Sliwoski er al.) 2013 [Jensen 2019: Yoshikawa et al.) 2018) explore potential new drug candidates via heuristics such hand-crafted rules and random mutations. Although widely adopted in practice, these methods tend to be inefficient and computationally expensive due to the vast chemical search space (Hoffman et al., [022I) The performance of these combinatorial approaches also heavily depends on the quality of generation rules, which often require task aaec-specific engineering expertise and may limit the diversity of the generated molecules.
To this end, recent research focuses on learning to controllably synthesize molecules with generative models (Tang et al.) 202T Chen 202T: Walters and Murcko, 2020p. It usually involves first training an unconditional generative model from millions of existing molecules (Winter et al., 2019a Irwin al. 2022) and then controlling the generative models to synthesize new desired molecules that
Work done during an internship at NVIDIA.
The first authors contributed equally this paper.

---

Published as a conference paper at ICLR 2023
Inpot molecule
Input embedding
Fused embedking
Ourpur molecule
Information
Encoder
Decoder
fusion
4.4 kenliinol
Retrieval databasc
Retriever
Shared weights
Retricval module
Encnder
Pre-trained module
-H0. Acailaaol
-HU. keauland
Retrieved exemplar molccules
Retrieved embeddings
Figure 1: An illustration of RetMol, a eetrieval-basee framework for controllable molecule generation. The framework incorporates a retrieval module (the molecule retriever and the information fusion) with pre-trained generative model (the encoder and decoder). The illustration shows example of optimizing the binding affinity (unit in kcal/mol; the lower the better) for an existing potential drug, Favipiravir, for better treating the COVID-19 virus (SARS-CoV-2 main protease, PDB ID: 7L.1 under various other design criteria.
satisfy one or more property constraints such as high drug- likg-likeness (Bickerton et al. 2012) and high synthesizability (Ertl and Schuffenhauer 2009). There are three main classes of learning- based molecule generation approaches: (i) reinforeement-lerrning (RL)-based methods (Olivecrona et al.
2017 Jin al. 2020a), (ii) supervised-leaminn (SL)-based methods (Lim al.
2018|Shin et al.
[2021] RL- and SL-based methods train fine-tune a pre-trained generative model using the desired properties as reward functions (RL) or using molecules with the desired properties as training data (SL). Such methods require heavy task specifse fine-tuning, making them not easily applicable to wide range of drug discovery tasks. Latent optimization- based methods, in contrast, learn to find latent codes that correspond to the desired molecules, based on property predictors trained on the generative model's latent space. However, training such latent-space property predictors can be challenging, especially in real-world scenarios where we only have a limited number of active molecules for training (Tin) al. 2020b Huynh et al., 2020). Moreover, such methods usually necessitate fixed- -xed-dimensiioal variil-v generative model with a compact and structured latent space (Xue et al. 2018 Xu
2021. and (iii) latent optimization- based methods (Winter et al.) [019b][aas et al.
et al. 2019), making them incompatible with other generative models with varying- diintngddmmeisiioii latent space, such as transformer-bas based architectures (Irwin et al., (022).
Our approach. In this work, we aim to overcome the aforementioned challenges of existing works and design controllable molecule generation method that (i) easily generalizes to various generation tasks; (ii) requires minimal training fine-tuning; and (iii) operates favorably in data-sparse regimes where active molecules are limited. We summarize our contributions as follows:
[1] We propose first-of-its-kind retrieval-basse based framework, termed RetMol, for controllable molecule generation. It uses a small set of exemplar molecules, which may partially satisfy the desired properties, from retrieval database guide generation towards satisfying all the desired properties.
[2] We design a retrieval mechanism that retrieves and fuses the exemplar molecules with the input molecule, a new self-supervised training with the molecule similarity as a proxy objective, and an iterative refinement process to dynamically update the generated molecules and retrieval database.
[3] We perform extensive evaluation of RetMol on a number of controllable molecule generation tasks ranging from simple molecule property control to challenging real-world drug design for treating the COVID-19 virus, and demonstrate RetMol's ' superior performance compared previous methods.
Specifically, as shown in Figure I the RetMol framework plugs a lightweight retrieval mechanism into a pre-trained, encoder-decoder generative model. For each task, we first construct a retrieval database consisting of exemplar molecules that (partially) satisfy the design criteria. Given an input molecule to be optimized, a retriever module uses it to retrieve a small number of exemplar molecules from the database, which are then converted into numerical embeddings, along with the input molecule, by the encoder of the pre-trained generative model. Next, an information fusion module fuses the embeddings of exemplar molecules with the input embedding to guide the generation (via the decoder in the pre-trained generative model) towards satisfying the desired properties. The fusion module is the only part in RetMol that requires training. For training. propose new
2

---

Published as a conference paper at ICLR 2023
self-supervised objective (i.e. predicting the nearest neighbor of the input molecule) to update the fusion module, which enables RetMol to generalize to various generation tasks without task-specific training/fine- -iuninn/nn-ttttnng.... For inference, we propose new inference process via iterative refinement that dynamically updates the generated molecules and retrieval database, which leads to improved generation quality and enables RetMol extrapolate well beyond the retrieval database.
RetMol enjoys several advantages. First, it requires only a handful of exemplar molecules to achieve strong controllable generation performance without task-specific fine-tuning. This makes it particularly appealing in real-world use cases where training or fine-tuning a task-specific model is challenging. Second, RetMol is flexible and easily adaptable. It is compatible with range of pre- trained generative models, including both fixed -ieedddimesiooall and varying- dimensional latenl-variable models, and requires training only lightweight, blug-and-play, task-agnostic retrieval module while freezing the base generative model. Once trained, it can be applied to many drug discovery tasks by simply replacing the retrieval database while keeping all other components unchanged.
In range of molecule controllable generation scenarios and benchmarks with diverse design criteria, our framework achieves state-of-the-art performance compared to latest methods. For example, on challenging four-property controllable generation task, our framework improves the success rate over the best method by 4.6% (96.9% vs. 92.3%) with better synthesis novelty and diversity. Using a retrieval database of only 23 molecules, we also demonstrate our real-world applicability drug discovery task of optimizing the binding affinity of eight existing, weakly-binding drugs for COVID-19 treatment under multiple design criteria. Compared to the best performing approach, our framework succeeds more often at generating new molecules with the given design criteria (62.5% vs. 37.5% success rate) and generates on average more potent optimized drug molecules (2.84 vs. 1.67 kcal/mmaaa /mol average binding aflinity improvement over the original drug molecules).
frontier
METHODOLOGY: THE RETMOL FRAMEWORK
We now detail the various components in RetMol and how we perform training and inference.
controllable generation setting in this paper. Con- cretely, let x E X be a molecule where X' denotes the set of all molecules, (e(x)) X + R a property predictor indexed by & [1..... L), and dz R desired threshold. Then, we formulate multi-property controllable generation as one of three problems below, differing in the control design criteria: (i) constraint satisfaction problem, where we identify set of new molecules such that {x X (I) 2 5τ, E), (ti) an unconstrained optimization problem, where we find a set of new molecules such that x' = argmax, s(x) where s(x) = Σt= wear(x) with the weighting coefficient We. and (iii) a constrained optimization problem that combines the objectives in (i) and (ii).
Problem setup. We focus the
2.1 RETMoL COMPONENTS
Encoder-decoder generative model backbone. The pre-trained molecule generative model forms the backbone of RetMol that interfaces between the continuous embedding and raw molecule representations. Specifically, the encoder encodes the incoming molecules into numerical embeddings and the decoder generates new molecules from an embedding, respectively. RetMol is agnostic to the choice of the underlying encoder and decoder architectures, enabling it to work with variety of generative models and molecule representations. In this work, we consider the SMILES string [Weininger] 1988) representation of molecules and the ChemFormer model (Irwin et al. 2022). which a variant of BART (Lewis et al. (020)) trained on the billion-scale ZINC dataset (Irwin and Shoichet 2004) and achieves state-of-the-art generation performance.
Retrieval database. The retrieval database XRR contains molecules that can potentially serve as exemplar molecules to steer the generation towards the design criteria and is thus vital for controllable generation. The construction of the retrieval database is task-specific: it usually contains molecules that at least partially satisfy the design criteria in a given task. The domain knowledge of what molecules meet the design criteria and how to select partially satisfied molecules can play an important role in our approach. Thus, our approach is essentially a hybrid system that combines the advantages of both the heuristic-based methods and baaning-based methods. Also, we find that database of only a handful of molecules (e.g., as few as 23) can already provide a strong control approach casily adapted to various tasks by quickly replacing the retrieval
signal. This makes database. Furthermore, the retrieval database can be dynamically updated during inference, i.e., newly
generated molecules can enrich the retrieval database for better generalization (see Section 2.3
3

---

Published as a conference paper at ICLR 2023
Molecule refriever. While the entire retrieval database can be used during generation, for com- putational reasons (e.g., memory and efficiency) it is more feasible to select a small portion of the most relevant exemplar molecules to provide a more accurate guidance. We design a simple heuristic-based retriever that retrieves the exemplar molecules most suitable for the given control design eriteria. Specifically, we first construct "feasible" set containing molecules that satisfy all the given constraints, i.c., = n { E XR ar(2)) 2 6t}. If this set is larger than K, i.e. the number of exemplar molecules that we wish retrieve, then we select K molecules with the best property scores, i.e., X, topx (A", s), where s(z) is the task-specific weighted average property score, defined in Section 2 Otherwise, we construct relaxed feasible set by removing the constraints one at a time until the relaxed feasible set is larger than K, at which point we select K molecules with the best scores of the most recently removed property constraints. In either case, the retriever retrieves exemplar molecules with more desirable properties than the input and guides the generation towards the given design criteria. We summarize this procedure in Algorithm|| I in Appendix We find that our simple retriever with K = 10 works well across a range of tasks. In general, more sophisticated retriever designs are possible and we leave them as the future work.
=
Information fusion. This module enables the retrieved exemplar molecules to modify the input molecule towards the targeted design criteria. It achieves this by merging the embeddings of the input and the retrieved exemplar molecules with lightweight, trainable, standard cross attention
mechanism similar to that in (Borgeaud et al. (022). Concretely, the fused embedding e is given by
(1) where fCA represents the cross attention function with parameters 0, and ein and E, are the input embedding and retrieved exemplar embeddings, respectively. The functions Attn, Query, Key, and Value compute the cross attention weights and the query, key, and value matrices, respectively. For our choice of the transformer- generative model (Irwin et al.
(Cm: 0) Attn Query(eia) Key(E,) Value (E,)
2022), we have that ein
Enn(((n) E TR LxD_ E, = ,e. e E R Lt)xD and et Enc(x,) E R where and La are the lengths of the tokenized input and the 6thh retrieved exemplar molecules, respectively, and D is the dimension of each token representation. Intuitively, the Attn function learns to weigh the retrieved exemplar molecules differently such that the more "important" retrieved molecules correspond to the higher weights. The fused embedding e E I I.x thus contains the information of desired properties extracted from the retrieved exemplar molecules, which serves as the input of the decoder to control its generation. More details of this module are available in Appendix A
2.2 TRAINING VIA PREDICTING THE INPUT MOLECULE'S NEAREST NEIGHBOR
The conventional training objective that reconstructs the input molecule (e.g.. in ChemFormer Irwin et al. 2022) is not appropriate in case, since perfectly reconstructing the input molecule does not rely on the retrieved exemplar modules (more details in Appendix (A). To enable RetMol to learn to use the exemplar molecules for controllable generation, we propose a new self-supervised training scheme, where the objective is to predict the nearest neighbor of the input molecule:
B
( = Σ CE (E(De((fca(e)) (0) EE ( FINN
())
(2)
where CF is the cross entropy loss function since we the BART model our encoder and decoder (see more details in Appendix [A]. ZINN represents the nearest neighbor of the input Tin. B is the batch size, and i indexes the input molecules. The set of retrieved exemplar molecules consists of the remaining K - nearest neighbors of the input molecule. During training, we freeze the parameters of the pre-trained encoder and decoder. Instead, we only update the parameters in the information fusion module fca. which makes our training lightweight and efficient. Furthermore, we use the full training dataset as the retrieval database and retrieve exemplar molecules by best similarity with the input molecule. Thus, our training is not task-specific yet forces the fusion module to be involved in the controllable generation with the similarity as a proxy criterion. We will show that during the inference time, the model trained with the above training objective and proxy criterion using similarity is able to generalize to different generation tasks with different design criteria, and performs better compared to training with the conventional reconstruction objective. For the efficient training, we pre-compute all the molecules' embeddings and their pairwise similarities with efficient approximate ANN algorithms (Guo et al.) 2020 Johnson et al. 2019).
Remarks. The above proposed training strategy that predicts the most similar molecule based on the input molecule and other K- similar molecules shares some similarity with masked language
4

---

Published as a conference paper at ICLR 2023
Table Under the similarity constraints, RetMol achieves higher success rate in the constrained QED optimiza- tion task and better score improvements the constrained penalized logP optimization task. Baseline results are reported from (Hoffman et al. 2021).
(a) Success rate of generated molecules that satisfy QED E [0.9, 1.0] under similarity con- straint 5 0.4.
(b) The average penalized logP improvements gen- erated molecules over inputs under similarity constraint
8 == f0.6, 0.4).
Method
Success (%)
Method
Improvement
MMPA Dalke al. 2018
32.9
= 0.6
5 0.4
JT-VAE
2013
8.8
JT-VAE Jin et
2018
0.28 =
1.03 + 1.30
GCPN You et ml. 2018
9.4
GCPN
You 2018
0.79 = 0.63
2.19 + 1.30
VSeq2Seq [Saipan
2015
58.5
MolDQN ZZou el al. al. 2019
1.86 4: 1.21
3.37 1: 1.62
VJINN+GAN (nn cuy
60.6
VSeq2Seg [Balsaana
2015
3.37 + 1.75
AtomG2G (Jin aral 2019
73.6
VJTNN etr 2019
2.33 = 1.24
3.55 + 1.67
HierG2G Uine 2019
76.9
Hie 2019
DES ES Marayaki el all 2020
77.8
GA (Nigan 2019)
3.44 = 1.09
5.93 : 1.41
QMO Hoffman at. 2021
92.8
[Hollman al. 2021
3.73 4: 2.85
7.71 1: 5.65
RetMol
94.5
RetMol
3.78 + 3.29
11.55 +
model pre-training in NLP (Devlin et al. 2019). However, there are several key differences: we
perform "masking" on a sequence level, I.e., a whole molecule, instead of on a token/word level: 2) "masking" strategy is to predict particular signal only (i.e., the most similar retrieved molecule to the imput) instead of randomly masked tokens: and 3) our training objective is used to only update the lightweight retrieval module instead of updating the whole backbone generative model.
2.3 INFERENCE VIA ITERATIVE REFINEMENT
We propose iterative refinement strategy to obtain an improved outcome when controlling genera- tion with implicit guidance derived from the exemplar molecules. The strategy works by replacing the input Tin and updating the retrieval database with the newly generated molecules. Such an iterative update is common in many other controllable generation approaches such as GA-based methods (Sli- woski et al. 2013 3013)00 Jensen 2019) and latent optimization-basee methods (Chenthamarakshan et al. 2020; Das et al. 2021). Furthermore, if the retrieval database Xx is fixed during inference, our method is constramed by the best property values of molecules in the retrieval database, which greatly limits the generalization ability of our approach and also limits our performance for certain generation scenarios such as unconstrained optimization. Thus, extrapolate beyond the database, we propose to dynamically update the retrieval database over iterations.
We consider the following iterative refinement process. We first randomly perturb the fused embed- ding M times by adding Guassian noises and greedily decode one molecule from each perturbed embedding to obtain M generated molecules. We score these molecules according to task-specific design criteria. Then, we replace the input molecule with the best molecule in this set if its score is better than that of the input molecule. At the same time, we add the remaining ones to the retrieval database if they have better scores than the lowest score in the retrieval database. If none of the generated molecules has a score better than the input molecule or the lowest in the retrieval database, then the input molecule and the retrieval database stays the same for the next iteration. We also add stop condition if the maximum allowable number iterations is achieved or a successful molecule that satisfies the desired criteria is generated. Algorithm in Appendix summarizes the procedure.
3 EXPERIMENTS
We conduct four sets of experiments, covering all controllable generation formulations (see "Problem setup" in Section 2 with increasing difficulty. For fair comparisons, in each experiment we faithfully follow the same setup as the baselines, including using the same number of optimization iterations and evaluating on the same set of input molecules to be optimized. Below, we summarize the main results and defer the detailed experiment setup and additional results to Appendix B and
3.1 IMPROVING QED AND PENALIZED LOGP UNDER SIMILARITY CONSTRAINT
These experiments aim to generate new molecules that improve upon the input molecule's intrinsic properties including QED (Bickerton et al. 2012) and penalized logP (defined by logP(z)) SA(z)) (Jin et al. (0188) where SA represents synthesizability (Ertl and Schuffenhauer (009)) under similarity constraints. For both experiments, we use the top 1k molecules with the best property values from the ZINC250k dataset (Jin et al. 2018) the retrieval database. In each iteration, we retrieve K 20
5

---

Published as a conference paper at ICLR 2023
exemplar molecules with the best property score that also satisfy the similarity constraint. The remaining configurations exactly follow (Hoftman et al. (021).
QED experiment setup and results. This is constraint satisfaction problem with the goal of generating new molecules z' such that asim (4', 2) 2 = 0.4 and AOED GOED(Z') 2 0.9 where æ is the input molecule, Ogim is the Tanimoto similarity function (Bajusz al. 2013), and AQED is the QED predictor. We select 800 molecules with QED in the range [0.7, 0.8] as inputs to be optimized. We measure performance by success rate, i.e., the percentage of input molecules that result in a generated molecule that satisfy both constraints. Table Ta shows the success rate of QED optimization by comparing RetMol with various baselines. We can see that RetMol achieves the best success rate, e.g., 94.5% versus 92.8% compared to the best existing approach.
Penalized logP setup and results. This constrained optimization problem with the goal to generate new molecules x' to maximize the penalized logP value (plogr(z') with similarity constraint (sim("', x) 2 8 where E {0.4, 0.6). We select 800 molecules that have the lowest penalized logP values in the ZINC250K dataset inputs to be optimized. We measure performance by the relative improvement in penalized logP between the generated and input molecules, averaged over all 800 input molecules. Table 1bb shows the results comparing RetMol to various baselines. RetMol outperforms the best existing method for both similarity constraint thresholds and, for the = 0.4 case, improves upon the best existing method by almost 50%. The large variance in RetMol is due to large penalized logp values of a few generated molecules, and is thus not indicative of poor statistical significance. We provide detailed analysis of this phenomenon in Appendix C
3.2 OPTIMIZING GSK33 AND JNK3 INIIIITIONN UNDER QED AND SA CONSTRAINTS
This experiment aims to generate novel, strong inhibitors for jointly inhibiting both GSK33 (Glycogen synthase ki- nase 3 beta) and JNK3 (-Jun N-terminal kinase-3) enzymes, which are relevant for potential treatment of Alzheimer's disease (Jin et al.] 2020a; Li et al. 2018). Following (Jin et al., 22000a). formulate this task constraint satisfaction problem with four property constraints: two positivity constraints (GSk3a(z") 2 0.5 and (INKa(œ') > 0.5, and two molecule property con- straints of QED and SA that aQED(x')> 2
Table 2: Success rate, novelty and diversity of generated
molecules in the task of optimizing four properties: QED, SA, and two binding affinities to GSK33 and JNK3 estimated by pre-trained models from (Jin al. 2020a). Baseline results
are reported from Din et al. 2202aa) Xie et al. (202T).
Method
Success
Novelty
Diversity
JT VAF (Ji al. 2018
1.3
GVAE-RI Unn et al 20220a
2.1
GCPN You
2018
4.0
REINVENT (livecrron et al. 2017
47.9
0.561
0.621
RationaleRL Din ctal. 20208
74.8
0.568
0.701
MARS (Xie et 20/4
92.3
0.824
0.719
MolEvol (Chen* at 2021
93.0
0.757
0.681
RetMol
96.9
0.862
0.732
0.6 and ASA (x') < Note that 4GSK38 and aINK3 are property predictors (Jin al. 2020a Li et al. 2018) for GSK3, 3 and JNK3, respectively, and higher values indicate better inhibition against the respective proteins. The retrieval database consists of all the molecules from the CheMBL. (Gaulton et al. 2016) dataset (approx. 700) that satisfy the above four constraints and we retrieve K = 10 exemplar molecules most similar to the input each time. For evaluation, we compute three metrics including success rate, novelty, and diversity according to (Jin et al. 2020a).
Table shows that RetMol outperforms all previous methods on the four-property molecule opti- mization task. In particular, framework achieves these results without task-specific fine-tuning required for RationaleRL and REINVENT. Besides, RetMol is computationally much more efficient than MARS, which requires 550 iterations model training whereas RetMol requires only 80 iterations (see more details in Appendix B.5). MARS also requires test-time, adaptive model training per sampling iteration which further increases the computational overhead during generation.
3.3 GUACAMOL BENCHMARK MULTIPLE PROPERTY OPTIMIZATION
This experiment evaluates RetMol on the seven multiple property optimization (MPO) tasks, a subset of tasks in the Guacamol benchmark (Brown et al. 2019). They are the unconstrained optimization problems to maximize a weighted sum of multiple diverse molecule properties. We choose these MPO tasks because 1) they represent a more challenging subset in the benchmark. as existing methods can achieve almost perfeet performance on most of the other tasks (Brown et al. 2019).. and 2) they the most relevant tasks to work, e.g., multi-property optimization. The retrieval database

---

Published as a conference paper at ICLR 2023
0.50
Couph MCTS
of ChUMBL
0.45
SMILES LSTM RetMol
4 4.5
Graph MCTS
Ckaph
4.0
SMILES GA
Gruph GA
3.5
Bos of CREMBL
SMILES LSTM
3.0
0.55 0.60
0.65
0,70
0.75
0,80
0.85
0.90
0.55
0.60
0.65
0,70
0.75
0,80
0.55
0.90
Benchmark Performance
Benchmark Performance
Figure 2: Comparison with the state-of-the-an methods in the multiple property optimization (MPO) tasks on the Guacamol benchmark. Left: QED (1) versus the averaged benchmark performance (t). Right: SA (L) versus the average benchmark performance (t). RetMol achieves best balance between improving the benchmark performance while maintaining the synthesizability (SA) and drug-likeness, (QED) of generated molecules.
Table 3: Quantitative results in the COVID-19 drug optimization task, where we aim to improve selected molecules' binding affinity (estimated via docking (Shidi al. (022 to the SARS-CoV-2 main protease under the QED, SA, and similarity constraints. Under stricter similarity condition, RetMol succeeds more cases (5/8 versus 3/8), Under milder similarity condition, RetMol achieves higher improvements (2.84 versus 1.67 average binding affinity improvements). Unit of numbers in the table is kcal /nol and lower is better.
0.6
5 0.4
Input molecule
Input score
RetMol
Graph GA (Jensen 2019
RetMol
Graph GA (Jensen
2019
Favipiravir
-4.93
-6.48
-7.10
-8.70
-7.10
Bromhexine
-9.64
-11.48
-11.20
-12.65
-11.83
PX-12
-6.13
-8.45
-8.07
-10,90
-8,31
Ebselen
-7.31
-10.82
-10.41
Disulfiram
-8.58
-9.09
-10.44
-10.00
Entecavir
-9.00
-12.34
Quercetin
-9.25
-9.84
9.81
Kaempferol
-8.45
-8.54
-10.35
10.19
Avg. Improvement
0.78
0.71
2.84
1.67
consists of 1k molecules with best scores for each task and we retrieve K 10 exemplar molecules with the highest scores each time.
We demonstrate that RetMol achieves the best results along the Pareto frontier of the molecular design space. Figure visualizes the benchmark performance averaged over the seven MPO tasks against QED, SA, two metrics that evaluate the drug- dingelikeneess and synthesizability of the optimized molecules and that are not part the Guacamol benchmark's 's optimization objective. Our framework achieves nice balance between optimizing benchmark performance and maintaining good QED and SA scores. In contrast, Graph GA @eensen 2019) achieves the best benchmark performance but suffers from low QED and high SA scores. These results demonstrate the advantage of retrieval-basee controllable generation: because the retrieval database usually contains drug-like molecules with desirable properties high QED and low SA, the generation is guided by these molecules in good way to not deviate too much from these desirable properties. Moreover, because the retrieval database updated with newly generated molecules with better benchmark performance, the generation can produce molecules with benchmark score beyond the best in the initial retrieval database. Additional results in Figure 2
in Appendix C.3 corroborate with those presented above.
3.4 OPTIMIZING EXISTING INHIBITORS FOR CARS-COV-2 MAIN PROTEASE
To demonstrate our framework's 's applicability the frontiers of drug discovery, we apply it the real-world task of improving the inhibition of existing weak inhibitors against the SARS-CoV-2 main protease (M PDB ID: 7L11), which is a promising target for treating COVID-19 19 by neutralizing the SARS-CoV-2 virus (Gao et al., 2022, Zhang et al., 202T). Because the the novel nature of this virus, there exist few high-potency inhibitors, making challenging for learning- based methods that require sizable training set not yet attainable in this case. However, the few existing inhibitors make excellent candidates for the retrieval database in our framework. We use a set of 23 known inhibitors (Jin et al. 2020b Huynh et al.) 2020; Hoffman et al.) 2021) as the retrieval database and select the 8 weakest inhibitors to M as input. We design an optimization task to maximize the binding affinity (estimated via docking (Shidi et al., 2022); ; see Appendix B for the detailed procedure) between the generated molecule and M while satisfying the following three constraints:
(q) > 0.6, aSa (x') < 4, and Osim (x',x) 2 8 where 5 E (0.4.0.6).
7

---

Published as a conference paper at ICLR 2023
coneact
combect
Original
Graph
(Durs)
(docking score=-9.64)
(docking scors=-11 (1.89)
(docking score--12.85)
Figure 3D visualizations that compare RetMol with Graph GA in optimizing the original inhibitor, Bromhexine, that binds to the SARS-CoV-2 main protease in the = 0.6 case. We can see the optimized inhibitor in RetMol has more polar contacts (red dotted lines) and also more disparate binding modes with the original compound than the Graph GA optimized inhibitor, which aligns with the quantitative results.
Table 4: Left: Comparing different training schemes in the unconditional generation setting. Right: Generation
performance with different retrieval database constructions based on the experiment in Section 3.2
Training objective
Validity Novelty Uuiqueness
Ret. database coustruction Success Novelty Diversity
RetMol (predict NN)
0.902
0.998
0.922
GSK38 + JNK3 QED + SA
96.9
0.862
0.732
Conventional (recon. input)
0.834
0.998
0.665
GSK33 + JNK3
84.7
0.736
0.700
GSK38 JNK3
44,1
0,571
0.708
Table the optimization results of comparing RetMol with graph GA, a competitive baseline. Under stricter (6- 0.6) similarity constraint, RetMol successfully optimizes the most molecules constraint (5 out of 8) while graph GA fails to optimize input molecules most of the time because it cannot satisfy the constraints or improve upon the input drugs' binding affinity. Under milder (8 = 0.4) similarity constraint, our framework achieves on average higher improvement (2.84 versus 1.67 binding affinity improvements) compared to the baseline. We have also tested QMO (Hoffman 202T). and find it unable to generate molecules that satisfy all the given constraints.
et al.
We produce 3D visualizations in Figure (3 (along with a video demo in httpt:://// http//ooorurrraro Emtys) to show the comparison of the RetMolloptimized inhibitors that bind to the SARS-CoV-2 main protease with the original and GA-optimized inhibitors. We can see that 1) there are more polar contacts (shown as red dotted lines around the molecule) in the RetMol-optimized compound, and 2) the binding mode of the Go-optimized molecule is much more similar to the original compound than the RetMooloptimmzze binding mode, implying RetMol optimizes the compound beyond local edits to the scaffold. These qualitative and quantitative results together demonstrate that RetMol has the potential effectively optimize inhibitors in real-world scenario. Besides, we provide extensive 2D graph visualizations of optimized molecules in Tables A3 and A4 in the Appendix. Finally, we also perform another of experiments on Antibacterial drug design for the MurD protein (Sangshetti et al. 2017) and observe similar results; see Appendix C.5 for more details.
3.5 ANALYSES
We analyze how the design choices in our framework impact generation performance. We summarize the main results and defer the detailed settings and additional results Appendix B and Unless otherwise stated, we perform all analyses on the experiment setting in Section 3.2
Training objectives. We evaluate how the proposed self-training objective in Sec. 2.2 affects
the quality of generated molecules in the unconditional generation setting. Table 4 (left) shows that, training with our proposed nearest neighbor objective (i.e., predict NN) indeed achieves better generation performance than the conventional input reconstruction objective (i.e., recon. input).
Types of retrieval database. We evaluate how the retrieval database construction impacts con- trollable generation performance, by comparing four different constructions: with molecules that satisfy all four constraints, i.e., GSK33 + JNK3 + QED + SA; or satisfy only GSK33 JNK3; or satisfy only GSK33 or JNK3 but not both. Table (right) shows that a retrieval database that better satisfies the design criteria and better aligns with the controllable generation task generally leads to better performance. Nevertheless, molecules that only partially satisfy the properties (c.g., GSK33 + JNK3), achicving comparable performance to RationaleRL. (Jin et al. 2020a).
note that RetMol performs reasonably well even with exemplar
Size of retrieval database. Figure (left) shows a larger retrieval database generally improves all metrics and reduces the variance. is particularly interesting that our framework achieves strong
8

---

Published as a conference paper at ICLR 2023
1.00
dess
0.35
Success rale
0.95
0,99
0.85
6.15
0.55
RetMal
0.75
0.80
ReMal
0.80
MARS
0.70
MARS
RationstOL
0.65
RationleRL
100
500
Size of retrieval database
#lterations
-12
32 30
60
Penalized logi values
Figure 4: Generation performance with varying retrieval database size (left), varying number of iterations (middle), and with or without dynamically updating the retrieval database (right). The left two plots are based on the experiment in Section 3.2 while the right plot hased on the penalized logP experiment in Section [3.
performance with small retrieval database, which already outperforms the best baseline on the success rate metric with only a 100-molecule retrieval database.
Number of optimization iterations. Figure 4 (middle) shows that RetMol outperforms the best existing methods with a small number of iterations (outperforms RationaleRL at 30 iterations and MARS 60 iterations); its performance continues improve with increasing iterations.
Dynamic retrieval database update. Figure 4 (right) shows that, for the penalized logP optimiza- tion experiment (Section 3.1 with dynamical update our framework generates more molecules with property values beyond the best in the data compared to the case without dynamic update. This comparison shows that dynamic update is crucial to generalizing beyond the retrieval database.
4 RELATED WORK
RetMol is most related to controllable molecule generation methods, which we have briefly reviewed and compared with in Section [ Another line of work in this direction leverages constraints based on explicit, pre-defined molecule scaffold structures to guide the generative process (Maziarz et al.) 2022: Langevin al. 2020). Our approach is fundamentally different in that the retrieved exemplar molecules implicitly guides the generative process through the information fusion module.
RetMol is also inspired by a recent line of work that integrates retrieval module in various NLP and vision tasks, such as language modeling (Borgeaud et al. [02TT; Liu et al., [2022] Wu et al.
2022), code generation (Hayati et al., 2018). question answering (Guu et al. |2200] Zhang al.
2022), and image generation (Tssng et al.) 2020 Casanova et al. [021] Blattmann et al. 2022 Chen
et al. 2022). Among them, the retrieval mechanism is mainly used for improving the generation quality either with smaller model (Borgeaud et al. [2021] Blattmann et al.) 2022) with very few data points (Casanova et al.) 2021: Chen et al.J (022). None of them has explicitly explored the controllable generation with retrieval. The retrieval module also appears in bio-related applications such as multiple sequence aligmment (MSA), which can be seen as a way to search and retrieve relevant protein sequences, and has been essential building block in MSA transformer (Rao et al., 202I) and AlphaFold (Jumper et al. 2021). However, the MSA methods focus on the pairwise interactions among a set of evolutionarily related (protein) sequences while RetMol considers the cross attention between the input and a set of retrieved examples.
There also exist priors work that study retrieval-based approaches for controllable text generation (Kim eet al.] [2020] Xu et al. (220). These methods require task-specific Maining//ine-tuning of the generative model for each controllable generation task whereas our framework can be applied many different tasks without it. While we focus on molecules, our framework is general and has the potential to achieve controllable generation for multiple other modalities beyond molecules.
5 CONCLUSIONS
We proposed RetMol, a new retrieval-based framework for controllable molecule generation. By incorporating a retrieval module with a pre-trained generative model, RetMol leverages exemplar molecules retrieved from a task-specific retrieval database to steer the generative model towards generating new molecules with the desired design criteria. RetMol versatile, requires no task- specific fine-tuning and is agnostic to the generative models (see Appendix C.7). We demonstrated the effectiveness RetMol on a variety of benchmark tasks and a real-world inhibitor design task for the SARS-CoV: virus, achieving state-of-the-an performances in each case comparing to existing methods. Since RetMol still requires exemplar molecules that at least partially satisfy the design criteria, it becomes challenging when those molecules unavailable at all. A valuable future work is to improve the retrieval mechanisms such that even weaker molecules, i.c., those that do not satisfy but are close to satisfying the design criteria, can be leveraged guide the generation process.

---

Published as a conference paper at ICLR 2023
ACKNOWLED
ZW and RGB are supported by by NSF grants CCF-1911094, IIS-1838177, and IIS-1730574; ONR grants N00014 4- 18- 12571, N00011-20-1-25530 and MURI 1000014-2--1278877 AFOSR grant FA9550-22-1-0060, and Vannevar Bush Faculty Fellowship, ONR grant N00014-18-1-204
ETHICS STATEMENT
Applications that involve molecule generation such drug discovery are high-stake in nature. These applications are highly regulated to prevent potential misuse (Hill and Richards 2022). RetMol as technology improve controllable molecule generation has the potential be subjected to malicious use. For example, one could change the retrieval database and the design criteria into harmful ones, such increased drug toxicity. However, we note that RetMol a computational tool useful for in silico experiments. As a result, although RetMol can suggest new molecules according to arbitrary design criteria, the properties of the generated molecules are estimations of the real chemical and biological properties and need to be further validated in lab experiments. Thus, while RetMol's real-world impact is limited to in silico experiments, it is also prevented from directly generating real drugs that can be readily used. In addition, controllable molecule generation is an active area of research; we hope that our work contribute to this ongoing line of research and make ML methods safe and reliable for molecule generation applications the real world.
REPRODUCIBILITY STATEMENT
To ensure the reproducibility of the empirical results, we provide the implementation details each task (i.e., experimental setups, hyperparameters, dataset specifications, etc.) in Appendix B The source code available at /sttps://githb... com httm/N/aabs//eeMee
REFERENCES
Jessica Vamathevan, Dominic Clark, Paul Czodrowski, Ian Dunham, Edgardo Ferran, George Lee, Bin Li, Anant Madabhushi, Parantu Shah, Michacla Spitzer, and Shanrong Zhao. Applications of machine lcarning in drug discovery and development. Nature Reviews Drug Discovery, 18(6):633 477, April 2019. doi: 10. 10.1088/44573-11--0000 0024-5. /ttpp:pp..//b...rgg/////// 1tt.:/1000000//11173--111010
Marco Daniele Parenti and Giulio Rastelli. Advances and applications of binding affinity prediction methods in drug discovery. Biotechnology Advances, 30(1):444-25 January 2012. doi: 10. 10.10166)biotechady..000 bitecchadv.2 1.08. 003.
http:///....ar///....////..i.aad http://...oottcchcccooo 20 003
Nicholas A. Mcanwell. Improving drug candidates by design: A focus on physieochemical properties a means of improving compound disposition and safety. Chemical Research in Toxicology, 1400):1420-1456 July 2011. doi: 10.1001/xxx0122 Iv. http::httts://.i..... .tr.:/ooorrr///o 1ttp:/10.0.00011001222000
José Jiménez-Luu Luna, Francesca Grisoni, Nils Weskamp. and Gisbert Schneider. Artificial intelligence in drug
discovery: recent advances and future perspectives. Expert Opinion Drug Discovery, 16(9);949-959, April 2021. doi: 10. (..1080/1746444 2021. 1909567. URL htt.httss:////// .rg. 1080 /141/144644411 .0021
67
Pavel G Polishchuk, Timur Madzhidov, and Alexandre Vamek. Estimation of the size of drug-like chemical space based gdb-17 data. Journal of computer-aided molecular design, 27(8):6756667, 2013.
W.Patrick Walters, Matthew T Stahl, and Mark A Murcko. Virtual screening an overview, Drug Discovery Today, (44):160-178 April 1998. doi: 0011016613599-644499701166-6--- URL https:///bbl.org//1
(016/81359-64469777
Gregory Sliwoski, Sandeepkumar Kothiwalc, Jens Meiler, and Edward W. Lowe. Computational methods in
drug discovery, Pharmacological Reviews, (6(1))334--99) December 2013. doi: 10.1 1111144::..1..077...11 12.000733.. 007336
Jan H. Jensen. A graph-based genetic algorithm and generative model/monte carlo tree search for the exploration
of chemical space. Chemical Science, 10 10112):567-3372 2019. doi: 10.103./c URL httpa
//doi /rr.iiorr/// httpa////....11..occcoco.o
Naruki Yoshikawa, Kei Terayama, Masato Sumita, Teruki Homma. Kenta Oono, and Koji Tsuda. Population- based de novo molecule generation, using grammatical evolution. Chemistry Letters, 47(11)14411-434, November 2018. doi: 10. /44//cl. 180665. Uttt/tttss:////t worg/10. 1246/c/11. 180665
10

---

Published as a conference paper at ICLR 2023
Samuel C. Hoffman, Vijil Chenthamarakshan, Kahini Wadhawan, Pin-Yu Chen, and Payel Das. Opti- mizing molecules using efficient queries from property evaluations. Nature Machine Intelligence, 4(1): 21-31, December 2021. doi: (1188/44225-0022-044222 URL 14422560021-005 021-00422-V
Bowen Tang, John Ewalt, and Ho-Leung Ng. Generative Al models for drig discovery. In Biophysical and Computational Tools in Drug Discovery, pages 221 243. Springer International Publishing, 2021. doi: 10. 01177/7335522211__11
[RL.stts:://dd .rgg / 0 1007/7355 2021 124
Hongming Chen. Can generative-moddl-bbsse drug design become new normal in drug discovery? Journal of Medicinal Chemistry, 6501):000-00,, December 2021. doi: 10.
URL
https: //cpi: http///ii..rrrree/ooo/o / 102 /ac..aacc... jmedchem. cc02042
W. Patrick Walters and Mark Mureko. Assessing the impact of generative Al on medicinal chemistry. Nature Biotechnology, 38(2):143314 145, January 2020, doi: (0.10388641587-0000044133 URL https::/dol
org/10.1038/s/// 00000/////41111 020-1-1---00022-2-----
Robin Winter, Floriane Montanari, Frank Noé, and Djork-Arné Clevert. Learning continuous and data-driven
molecular descriptors by translating equivalent chemical representations. Chemical Science, 1066)1192--771,
2019a. doi: 10. (039/c8cc04115)) URL http:httts///ttt //doi. wtt.:r...//.../////..... /000////////1555511
Ross Irwin, Spyridon Dimitriadis, Jiazhen He, and Eshen Jannik Bjerrum. Chemformer: a pre-trained transformer for computational chemistry. Machine Learning: Science and Technology, 3(1):015022, jan 2022. doi:
10. 22A-8266222215/a3ff1
hhttps://oo:
www./o..gg...... 088 2632 21 53///33//33
G. Richard Bickerton, Gaia V. Paolini, Jérémy Besnard, Sorel Muresan, and Andrew L. Hopkins. Quantifying the chemical beauty of drugs. Nature Chemistry, 4(2):90-98, January 2012. doi: 10. 10.1038/nchm..1255 URL
htttps://cc. http:/dci..rgg/l/// 10.
1243
Peter Ertl and Ansgar Schuffenhauer. Estimation of synthetic accessibility score of drug-like molecules based on molecular complexity and fragment contributions. Journal of Cheminformartics, 1(1), June 2009. doi:
10.1186/1758, 1111866158-22--111-1 1-8. URL /tttpp////bl..rrr/0...111/1111/3
Marcus Olivecrona, Thomas Blaschke, Ola Engkvist, and Hongming Chen. Molecular de-novo de- sign through deep reinforcement learning. Journal of Cheminformatiic, 9(1), September 2017. doi: 10. 1111666133111-1002225555
Wengong Jin, Re.igi Barzilay, and Tommi Jaakkola. Multi-objective molecule generation using interpretable substructures. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedingy of Machine Learning Research, pages 4849-48551 PMLR, 13- 18 Jul 2020a.
Jaechang Lim, Seongok Ryu, Jin Woo Kim, and Woo Youn Kim. Molecular generative model based on condi- tional variational autoencoder for de novo molecular design. Journal of Cheminforrmaticc. 10(1). July 2018. doi: 1 186/s 1311-33321--10222202 -01-0026-77 URL URL htt.::tttt//ttcttc doi /ttp.///........///1/332-------0
Bonggun Shin, Sungsoo Park, Jin Yeong Bak, and Joyce C. Ho. Controlled molecule generator for optimizing mul- tiple chemical properties. In Proceedings of the Conference on Health, Inference, and Learning, ACM, April 2021. doi: 10. 01111///4544993455555555
URL
Robin Winter, Floriane Montanari, Andreas Steffen, Hans Briem, Frank Noé, and Djork-Aréé Clevert. Efficient multi-objective molecular optimization in a continuous latent space. Chemical Science, 10(34):8016 8024, 2019b. doi: 10.1039/c9sc01923) URL httpp:::////eotto /tt::::///ooooooo org 10. 02539/999c0192/ 1988
Payel Das. Tom Sercu, Kahini Wadhawan, Inkit Padhi, Sebastian Gehrmann, Flaviu Cipcigan, Vijil Chenthama- rakshan, Hendrik Strobelt, Cicero dos Santos, Pin-Yu Chen, Yi Yan Yang. Jeremy P. K. Tan, James Hedrick, Jason Crain, and Aleksandra Mojsilovic. Accelerated antimicrobial discovery via deep generative models and molecular dynamics simulations. Nature Biomedical Engineering, (66)613 623, March 2021. doi: 10. 10.1008/4415111 0200000099900 x. URL
Zhenming al. Structure of mpro from SARS-CoVa and discovery of its inhibitors. Nature, 582 (781 1):289-2999 April 2020b. doi: 10. 15000//411586--0--22222
s4 14158 586 020 2223-y
Tien Huynh, Haoran Wang, and Binquan Luan. In Silico exploration of the molecular mechanism of clinically oriented drugs for possibly inhibiting SARS-CoVe 222 main protease. The Journal of Physical Chemistry Letters, ((11):443-44200 May 2020. doi: 10. hppl///.jpclltt..o. 0c00994,, URL httpp: / //o1 .org/10
1021/acs. jpc ett. 0c00994

---

Published as a conference paper at ICLR 2023
Dongyu Xue. Yukang Gong, Zhaoyi Yang, Guohui Chuai, Sheng Qu, Aizong Shen, Jing Yu, and Qi Liu. Advances and challenges in deep generative models for de novo molecule generation. WIREs Computational Molecular Science, 9(3), October 2018. doi: 10. 10.1002//em. 1395.
wcms. 1395
Youjun Xu, Kangjie Lin, Shiwei Wang. Lei Wang. Chenjing Cai, Chen Song, Luhua Lai, and Jianfeng Pei. Deep learning for molecular generation. Future Medicinal Chemistry, 11(6):567-55 -777 March 2019, doi: 10.4155/fmc-2011 10tef/mcc221----555 0-03888. UtL http:///.........1/0/.
2018 0356
David Weininger. SMILES, a chemical language and information system. 1. introduction to methodology and encoding rules. Journal of Chemical Information and Modeling, 28(1):31-36, February 1988. doi: 10. 10.102//000055:0055. URLihttps http:::://dcccoooo 10tp//11.1111/1111001001710000
Mike Lewis. Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, Abdelrahman Mohamed. Omer Levy. Veselin Stoyanov, and Luke Zettlemoyer. BART: Denoising sequence-to to-nneee----sseneeeeeeeeee pre-training for natural language generation, translation. and comprehension. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 787 7-880, Online, July 2020, Association for Computational I. Linguistics. doi: 10. 1/265//1120000.alllaaaao
703. U https: ac hnt logg wwtp:orgg/2o.....o acl mmain. 703
John J. Irwin and Brian K. Shoichet. ZINC a free database of commercially available compounds for virtual screening. Journal of Chemical Information and Modeling, 45(1):177-182, December 2004. doi: 10.1021/c:0497144. URL httpp:tt /doii wttp///00000000/14499/00111
Sebastian Borgeaud et al. Improving language models by retrieving from trillions of tokens. arXiv e-prints, art. 1eXiv:2112.04226. December 2021.
Ruiqi Guo, Philip Sun, Erik Lindgren, Quan Geng, David Simcha, Felix Chern, and Sanjiv Kumar. Accelerating large-scale inference with anisotropic vector quantization. In International Conference on Machine Learning, 2020. URL
aarxiv wwg. abs 908 10396
Jeff Johnson, Matthijs Douze, and Hervé Jégou. Billion-scale similarity search with GPUs. IEEE Transactions on Big Data, (3):533--577, 2019.
Jacob Devlin, Ming-Wer Mei Chang, Kenton Lee, and Kristina Toutanova, BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume I (Long and Short Papers). pages 4171-4186, Minneapolis. Minnesota, June 2019. Association for Computational Linguistics. doi: 10. 10.1863//1//119 1423. /RL/attss://ccaattoooo
BeggNNN org/ /N1 9 94123
Vijil Chenthamarakshan, Payel Das, Samuel C. Hoffman, Hendrik Strobelt, Inkit Padhi, Kar Wai Lim, Benjamin Hoover, Matteo Manica, Jannis Born, Teadoro Laino, and Aleksandra Mojsilovic. Cogmol: Target-specific and selective drug design for covid-19 using deep generative models. In Proceedings of the 34th International Conference on Neural Information Processing Systems, NIPS 20, Red Hook, NY, USA, 2020. Curran Associates Inc. ISBN 9781713829546.
Andrew Dalke, Jérôme Hert, and Christian Kramer. mmpdh: An open-source matched molecular pair platform for large multiproperty data sets. Journal of Chemical Information and Modeling, 58(5)9002-10,, May 2018. doi: 10.1022/acs. het.....ci...00011.... btt httttps://doi.opg/ 10. 1021/ / acs. jeim. 9b00173
Wengong Jin, Regina Barzilay, and Tommi Jaakkola. Junction tree variational autoencoder for molecular graph generation. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 2323-2332. PMLR, 10- 15 Jul 2018
Jiaxuan You, Bowen Liu, Rex Ying, Vijay Pande, and Jure Leskovec, Graph convolutional policy network for goal-directed molecular graph generation. In Proceedings of the 32nd International Conference on Neural Information Processing Systems, NIPS' 18. page 6412-6422, Red Hook, NY, USA, 2018. Curran Associates Inc.
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to aliga and translate. In Yoshua Bengio and Yann LeCun, editors, 3rd International Conference Learning Representations, JCLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015.
Wengong Jin, Kevin Yang, Regina Barzilay, and Tommi Jaakkola. Learning multimodal graph-to-grapp translation for molecule optimization. In International Conference on Learning Representations, 2019. URL htttps:/cpperevviet net/ forum? dtw.:mmmdddlxxssAA-- 115AA55577
Wengong Jin, Regina Barzilay, and Tommi Jaakkola. Hierarchical Graph-to-Graph Translation for Molecules. arXiv e-prints, art. arXiv:1907.11 1223, June 2019.
12

---

Published as a conference paper at ICLR 2023
Paul Maragakis, Hunter NisonofI, Brian Cole, and David E. Shaw. A deep-learning view of chemical space designed facilitate drug discovery. Journal of Chemical Information and Modeling, 60110) July 2020, doi: 10. 10.1021/acs.jcin Oc00321. URLhttps:://daa wttp:/pps//..l.lg//.......... http://....1///aaaaa/ jcin 0c00321
org/
Zhenpeng Zhou, Steven Keames, Li Li. Richard N. Zare, and Patrick Riley. Optimization of molecules via deep reinforcement learning. Scientific Reports, 9(1,, July 2019, doi: 10.1038/<411988011-4777
URL.
https:://ci.coo 10. 1038/ 4 598 01 9- 47-447144 148-x
AkshatKumar Nigam, Pascal Friederich, Mario Krenn, and Alán Aspuru- Guzik, Augmenting Genetic Algorithms with Deep Neural Networks for Exploring the Chemical Space. arXiv e-prints, art, 1rXiv:1999.11653 September 2019.
Dávid Bajusz, Anita Rácz, and Károly Héberger. Why is tanimoto index an appropriate choice for fingerprint- based similarity calculations? Journal of Chemingormattics, 7(1), May 2015. doi: 10.1186/113311-05--00
Yutong Xie, Chence Shi, Hao Zhou, Yuwei Yang. Weinan Zhang, Yong Yu, and Lei Li. [MARS]: Markov molec- ular sampling for multi-objectivve drug discovery. In International Conference on Learning Representations, 2021.
Binghong Chen", Tianzhe Wang*, Chengtao Li, Hanjun Dai, and Le Song. Molecule optimization by explainable evolution. International Conference on Learning Representations, 2021.
Yibo Li, Liangren Zhang, and Zhenming Liu. Multi-obeective de novo drug design with conditional graph generative model. Journal of Cheminformatiis, 10(1), July 2018. doi: 10.1 10.1186/s13221 33211-18-11108880000000000011111 URL
https://dci. crg/ 10. 1.0..11000
0287-6
Anna Gaulton, Anne Hersey, Michal Nowotka, A. Patrícia Bento, Jon Chambers, David Mendez, Prudence Mutowo, Francis Atkinson, Louisa J. Bellis, Elema Cibrián-Uhalte Mark Davies, Nathan Dedman, Anneli Karlsson,
María Paula Magariños, John P. Overington, George Papadatos, Ines Smit, and Andrew R. Leach. The ChEMBL database in 2017. Nucleic Acids Research, 45(D1):D945-D955 November 2016, doi: 10. 10.1099/arr/gkw 1074. URL https: http://ttoooog./////// 100 1ttp://..0000//////1gg/0
Nathan Brown, Marco Fiscato, Marwin H.S. Segler, and Alain C. Vaucher. GuacaMol: Benchmarking models de novo molecular design. Journal of Chemical Information and Modeling, 59(3): 1096-1 108, March 2019. doi: 10. 10211//cs.iii.Sb000033. URI. www.hhttsss/// //do org/ 0. 1021 /acs. jeim. 8bC0839
Tang Shidi, Chen Ruiqi, Lin Mengru, Lin Qingde, Zhu Yanxiang, Ding Ji, Wu Jiansheng, Hu Haifeng, and Ling Ming. Accelerating AutoDock VINA with GPUs. March 2022. doi: (0.26434/chemrxiv-222 2021-3qvn2--4
URLhhttps://de:: crg/ 10. 26434/ http://cheriiiiiieeee 2C21 202CCC13spp 3q
Wenhao Gao, Rocio Mercado, and Connor W. Coley. Amortized tree generation for bottom-up synthesis planning and synthesizabll molecular design. In International Conference Learning Representutions, 2022. URL
httpps://cpenrevieen net forum?
carrum?idee id= RxhHdnxt]
Chun-Hui Zhang, Elizabeth A. Stone, Maya Deshmukh, Joseph A. Ippolito, Mohammad M. Ghahremanpour, Julian Tivado-Rives, Krasimir A. Spasov, Shuo Zhang, Yuka Takeo, Shalley N. Kudalkar, Zhuobinn Liang, Farren Isaacs. Brett Lindenbach. Scott J. Miller, Karen S. Anderson, and William Jorgensen. Potent noncovalent inhibitors of the main protease of SARS-CoV from molecular sculpting of the drug perampanel guided by free energy perturbation calculations. ACS Central Science. 7(3):467-773 February 2021. doi: 10. 102 /aclsccccttssci cc00039. URL htttps://do/ orgg 1021 1ttp///c.111//ccccncncccccc 1c00039
Jaiprakash N. Sangshetti, Suyog S. Joshi, Rajendra H. Patil, Mark G. Moloney, and Devanand B. Shinde. Mur ligase inhibitors anti- bacterials: A comprehensive review. Current Pharmaceutical Design, 23 (21), August 2017. doi: Tt1F////338811222666666 8048
15048. URL httttp:://dii.orgggt/t
1381 616662226666 7021
Krzysztof Maziarz, Henry Richard Jackson-Flux, Pashmina Cameron, Finton Sirockin, Nadine Schneider, Nikolaus Stiefl, Marwin Segler, and Mare Brockschmidt. Learning to extend molecular scaffolds with structural motifs. In International Conference on Learning Representations, 2022. URL https:// openreview http:/oorumm.d=2233333355
Maxime Langevin, Hervé Minoux, Maximilicn Levesque, and Marc Bianciotto. Scaffold-constrained molecular generation. Journal of Chemical Information and Modeling, 60(12):5637-56 5646, 2020. doi: 10. 10.1021/acs.jeh 0c01015. URL
doi. wttee//g///oo 10. 1021/ acs jeim. 0c01015
PMID: 33301333.
Qi Liu, Dani Yogatama, and Phil Blunsom. Relational Memory Augmented Language Models. arXiv e-prints, arXiv:2201.0000 109680, January 2022.
13

---

Published as a conference paper at ICLR 2023
Yuhuai Wu, Markus Norman Rabe, DeLesley Hutchins, and Christian Szegedy. Memorizing transformers. In International Conference on Learning Representations, 2022. URL httpp:://opeeeeveee..oo nett
Shirley Anugrah Hayati, Raphael Olivier, Pravalika Avvaru, Pengcheng Yin, Anthony Tomasic, and Graham
Neubig. Retrieval-base neural code generation. Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pages 925-930, Brussels, Belgium, Octoher-Novemeee 2018 8. Association for Computational Linguisties. doi: 10. 10.1863//1/D18-1111 URL httpttts///aaaatthoott
org/D18-112
Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat, and Mingwei Chang. Retrieval augmented language model pre-training. In Proceedings of the 37th International Conference Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 3929-3938. PMIR, 13-18 Jul 2020.
Xikun Zhang, Antoine Bosselut, Michihiro Yasunaga, Hongyu Ren, Percy Liang, Christopher D Manning, and Jure Leskovec. GreaseL.M: Graph REASoning enhanced language models. International Conference on net. wwt.:eeeooommmm-111eeeeeeer
Learning Representations, 2022. URL
Hung- Yu Tseng, Hsin-Ying Lec, Lu Jiang, Ming-Hsuan Yang, and Weilong Yang. RetrieveGAN: Image synthesis via differentiable patch retrieval. In Computer Vision - ECCV 2020, pages 242-257. Springer International Publishing, 2020. doi: 10.1007998--3-33-5858858 15. URL https: / doi. org /00
30 58598- 15
Arantxa Casanova. Marlene Careil, Jakob Verbeek, Michal Drozdzal, and Adriana Romero Soriano. Instance- conditioned gan. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 27517-27529, Cur- ran Associates, Inc., 2021. URL http::ps:/pproceeeee www..n.....iiinoeeeeeeeeeee..eee neur ipsurips... 71 wttp:/aaddlaaaaaaaaaaaa
e 1lac288b0F22414445500
pdf
Andreas Blattmann, Robin Rombach, Kaan Oktay, and Björ Ommer. Retrieval- -ugmente- diffusion models. arXiv preprint orXiv:2204.18824 2022.
Wenhu Chen, Hexiang Hu, Chitwan Saharia, and William Cohen. Re-imagen: augmente text-to- image generator. arviv preprint arXiv:2209. 14491, 2022.
Roshan M Rao. Jason Liu, Robert Verkuil, Joshua Meier, John Canny, Pieter Abbeel, Tom Sercu, and Alexander Rives. Msa transformer. In Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 1844-8855 PMLR, 18 24 Jul 2021.
John Jumper et al. Highly accurate protein structure prediction with AlphaFold. Nature, 596(7873):583- July 2021.
Jihyeok Kim, Seungtaek Choi, Reinald Kim Amplayo, and Seung-won Hwang. Retrieval- controllable review generation. In Proceedings of the 28th International Conference Computational Linguistics, Barcelona, Spain (Online), December 2020. International Committee on Computational Linguistics.
Peng Xu, Mostofa Patwary, Mohammad Shoeybi, Raul Puri, Pascale Fung, Anima Anandkumar, and Bryan Catanzaro. MEGATRON-CNTES - Controllable story generation with external knowledge using large-scale language models. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), Online, November 2020. Association for Computational Linguistics.
Ray G Hill and Duncan Richards. Drug discovery and development: technology in transition. 2022. ISBN 9780702078040.
Matt J. Kusner, Brooks Paige, and José Miguel Hernández-Lobar Grammar variational autoencoder. Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 1945-1954. PMLR, 06-11 Aug 2017. URL https: http::///cccceeeeeeeeeee mwm...r................ wtw...igss...............
www.eseevoo. v kusner 7a html
Marwin H. S. Segler, Thierry Kogej, Christian Tyrchan, and Mark P. Waller. Generating Focussed Molecule Libraries for Drug Discovery with Recurrent Neural Networks. arXiy e-prints, art. 2017.
January
Liu, Y. Lin, X. Wen, R. N. Jorissen, and M. K. Gilson. Binding weh-accessible database of experimentally N. determined protein-ligand binding affinities. Nucleic Acids Research, 2500 January 2007. doi: 10. 1093//ar//k11999 URL sttps:/ wttp:/o..o/....0////. g
14

---

Published as a conference paper at ICLR 2023
Appendix
A FRAMEWORK DETAILS
More details of the information fusion module (Eq. (I) Recall that Eq. is
== JCA (ein: Er: Attn(Query(eaa)) Key(E,)) )) Value E,)
where SCA represents the cross attention function with parameters 0, and E RR and E, R xx)x are the input embedding and retrieved exemplar embeddings, respectively.
Concretely, Query, Key, and Value are affine functions that do not change the shape of the the input. The function input and output dimensions are as follows:
Query RLLD L2LD
(3)
Key: R
RE
LxIxD
R (EE Lx)XD
(4)
Value: RE L - R L.)XD
(5)
In particular, the Key and Value functions first apply to each of the k-th retrieved molecule embedding ef E R (which block of matrix in the retrieved molecule embedding E,) and then stack
these K output matrices horizontally to obtain the output matrices of shape (Σ*=1 Lx) x D.
The Attn function first computes the inner product between each slice of the query output matrix (of shape D) and the key output matrix (of shape (Σ*=) La) x D), followed by a softmax function which results in (E Lx) un-normalized weights. These weights are applied to the value output matrix (of shape (Σ k=1 Lx) X D) to obtain one slice in the output fused matrix The above (Ex=)
procedure is performed in parallel to obtain the full output fused matrix e, such that maintains the same dimensionality with the input embedding Cin i.e., R LxD
More details of the CE loss The cross entropy in Eq. (2) is the maximum likelihood objective generation tasks such as machine
commonly used in sequence modeling and sequence-to-sequemee translation. It takes two molecules as its inputs: one is the ground-truth molecule sequence, and the other one is the generated molecule sequence (i.c., the softmax output of decoder). Specifically, if we define y TINN and ý Dec(fcA(en)) E,); 0), then the "CE" in Eq. (2) is given as follows:
1-1V-1
CB(j, = -
yt log U1,
I=0
where L is the molecule sequence length, the vocabulary size, Yt is the ground-truth (one-hot) probability of vocabulary entry on the (-th token, and St, is the predicted probability (i.e., softmax output of decoder) the vocabulary entry on the /-th token.
More details of the RetMol training objective (Sec. 2.2 We first provide more detailed descrip- tions of our proposed objective: Given an input molecule, we use its nearest neighbor molecule from the retrieval database, based on the cosine similarity in the embedding space of the CDDD encoder, as the prediction target. The decoder takes in the fused embeddings from the information fusion module to generate new molecules. As shown in Eq. (2), we calculate the cross-entropy distance between the decoded output and the nearest neighbor molecule the training objective.
The motivation is that: Other similar molecules (i.e., the remaining the retrieval database, through the fusion module, can provide good guidance for transforming the input to its nearest neighbor (c.g., how to perturb the molecule and how large the perturbation would be). Accordingly, the fusion module can be effectively updated through this auxiliary task.
nearest neighbors) from
On the contrary, if we use the conventional encoder-decoder training objective (i.e., reconstructing the input molecule), the method actually does not need anything from the retrieval database to do well in the input reconstruction (as the input molecule itself contains all the required information already). As result, the information fusion module would not be effectively updated during training.
15

---

Published as a conference paper at ICLR 2023
Molecule retriever Algorithm describes the how the retriever retrieves exemplar molecules from a retrieval database given the design criteria.
Inference
A describes the inference procedure.
The total number of parameters in RetMol and in the base generative model (Irwin al., 2022) is 10471179 and 10010635, respectively. The additional 460544 parameters come exclusively from the information fusion model, which means that adds only less than 5% parameter overhead and very lightweight compared to the base generative model.
Parameters
Algorithm 1: Exemplar molecule retriever
Require: : Property predictors and desired property thresholds 8 for & E [l,..., for property constraints; scoring function for properties be optimized; retrieval database XR; number of retrieved exemplar molecules
Input :Input molecule Pin
Output :Retrieved exemplar molecules X,
f' L;
X" = far Nn while X2 < K do
/n [a(()) 2 5.):
I. := --
l' = L;
X' n (ir E XR (x())(()>> 5e):
end
X"
{r X' ar (x) og (zin): | K then
10
return X, = top (X X" 8);
else
12
return X, topx 8);
a end
B DETAILED EXPERIMENT SETUP
B. RETMOL TRAINING
We only train the information fusion model in our RetMol framework. The training dataset uses either ZINC250K (Irwin and Shoichet 2004) (for the experiments in Section 3.1 or CheMBL (Gaulton et al., 2016). The choice of the training data is to align with the setup for the baseline methods each experiment; The experiments in Sections 3.1 and 3.3 use the ZINC250K dataset while the experiments in Sections 3.2 and 3.4 use the CheMBL dataset. For the ZINC250K dataset, we follow the train/validation/test splits in (Tin al.) (019) and train the train split. For the ChemMBL dataset, we train on the entire dataset without splits. Training is distributed over four V100 NVIDIA GPUs, each with 16GB memory, with a batch size of 256 samples on each GPU, for 50k iterations.
The total training time is approximately hours for either the ZINC250K or the CheMBL dataset. Our training infrastructure largely follows the Megatron version of the molecule generative model in (Irwin et al. 2022), which uses DeepSpeed Apex and half precision for improved training efficiency. Note that, once RetMol is trained, we do not perform further task-specific fine-tuning.
http:///gittbbb.cco////////////
16

---

Published as a conference paper at ICLR 2023
Algorithm 2: Generation with adaptive input and retrieval database update
Require: Encoder Enc, decoder Dec, information fusion model SCA, exemplar molecule retriever Ret, property predictors ae(x) and desired property thresholds for l E [1,.... T.] for property constraints, retrieval database X', scoring function s(x) for properties to be optimized
Input
Input molecule Cin, number of retrieved exemplar molecules K, number of optimization iterations the number of molecules M sample each iteration Output :Opimizzed molecule
for t E [1..... [7) do
ein Enc(x).
X, Ret (Kn, fa,, 5e) 1. 8, 2in K):
E, Enc(At,)
= fCa (eins E,;0); X gen D; for I....,
M do ~~ N(0,1): 1); e+6 x' Dec(e;);
Xgen == Xgen U (ar):
11
end
Xgen Ret (Kgens (ar, dell 8, Fia K);
14
Fin == top (Xenn 8);
15
XR = X' U Exen \ Zin
16
end
17

---

Published as a conference paper at ICLR 2023
B.2 RETMOL INFERENCE
We use greedy sampling in the decoder throughout all experiments in this paper. When we need to sample more than one molecule from the decoder given the (fused) embedding, we first perturb the (fused) embedding with independent random isotropic Gaussian with standard deviation of I and then generate a sample from each perturbed (fused) embedding using the decoder. Inference uses a single V100 NVIDIA GPU with 16 GB memory.
Note that during inference, the multi multi-propent property design objective is provided with a general form s(r) Ei= (r).(r). In the experiments, we simply set all the weight coefficients uy to 1, i.c., the aggregated design criterion is the sum of individual property constraints.
B.3 BASELINES
We briefly overview the existing methods and baselines that we have used for all experiments. Some baselines are applicable for more than one experiment while some specialize in a certain experiment.
JT-VAE (Jin et al. 2018) The junction tree variational autoencoder (JT-VAE) reconstructs molecule its graph (2 dimensional) representation using its junction tree and molecular graph as inputs. JT-VAE learns a structured, exed-dimensional latent space. During inference, controllable generation is achieved by first performing optimization on the latent space via property predictors trained on the latent space and then generating a molecule via the decoder in the VAE.
MMPA DDlke al., 2018) The Matched Molecular Pair Analysis (MMPA) platform uses rules and heuristics, such as matched molecular pair, to perform various molecule operations such as search, transformations, and synthesizing.
GCPN You et al. 2018) Graph Convolutional Policy Network (GCPN) trains a graph convolutional neural network generate molecules in their 2D graph representations with policy gradient rein- forcement learning. The reward function consists of task-specific property predictors and adversarial loss.
Vseq2seq (Bahdanau et al. 2015) This is a basic sequence-to-sen to-sequence model borrowed from the machine translation literature that translates the input molecule to output molecule with more desired properties.
VJTNN JJin et al. 2019) VJTNN is graph-based method for generating molecule graphs based on junction tree variational autoencoders (Jin et al. 2018). The method formalizes the controllable molecule generation problem graph translation problem and is trained using adversarial loss to match the generation and data distribution.
HierG2G (Jín et al. 2019 the molecule's substructures, which are interleaved with the molecule's s atoms for more structured generation.
The Hierarchical Sraph-to-Graphe generation method takes into account
AtomG2G (Jin et al. 2019) The Atom Sraph-to-Graph generation method similar HierG2G but only takes into account of the molecule's atoms information without the structure and substructure information.
DESMILES (Maragakis et al. 2020) The DESMILES method aims to translate the molecule's fingerprint into its SMILES representation. Controllable generation is achieved by fine-tuning the model for task-specific properties and datasets.
MolDQN (Zhou et al., 2019) The Molecule Deep Q-Learning Network formalizes molecule generation and optimization as a sequence of Markov decision processes, which the authors use deep Q-learning and randomized reward function to optimize.
18

---

Published as a conference paper at ICLR 2023
GA (Nigam et al. 2019) This method augments the classic genetic algorithm with neural network which acts as a discriminator to improve the diversity of the generation during inference.
QMO (Hoffman al. 2021) The Query-based Molecule Optimization framework is a latent- optimization-based controllable generation method operated the SMILES representation of molecules, Instead finding a latent code the latent space via property predictor trained on the latent space, QMO uses property predictor the molecule space and performs latent optimization via zeroth order gradient estimation. Although this latent-optimization -aasdi method removes the need to train latent-space property predictors, find that it is sometimes challenging hyper-parameter to adapt QMO for different controllable generation tasks. For example, in the SARS-CoV-2 main protease inhibitor design task, we could not succeed generate an optimized molecule using QMO even with extensive hyper-parameter tuning.
the
GVAE-RL [Jin et al. 2020a) This method is the graph-based grammar VAE (Kusner et al. 2017) that learns expand molecule with of expansion rules, based variational autoencoder. GVAE-RL further fine-tunes GVAE with RL objective for controllable generation, using the property values the reward.
REINVENT (Olivecrona et al. 2017) techniques to generate new molecules.
REINVENT trains recurrent neural network using RL
RationaleRL (Jin et al. 2020a) The RationalRL method first assembles "rationale", a molecule fragment composed from different molecules with the desired properties. Then it trains decoder using the assembled collection of rationales as input by first randomly sampling from the decoder, scoring each sample with the property predictors, and use the positive and negative samles training data to tunetune the decoder.
MARS (ie et al. 2021) The MARS method models the molecule generation process as a Markov chain, where the transitions are the "edits" on the current molecule graph parametrized by graph neural network. The generation process proceeds by either accepting the newly edited molecule if has more desired attributes than the previous one. Otherwise, the previous molecule is kept and the Markov chain continues.
Graph MCTS and Graph GA (Jensen 2019) Graph MCTS and GA methods traverse the molecule space in its graph representation using genetic algorithm and Monte Carlo Tree Search algorithms, respectively.
SMILES LSTM (Segler et al. 2017) This method is a decoder (an LSTM) only method which generates molecules using seed input symbol. Controllable generation is achieved by fine-tuning it on a task-specific dataset with RL using an property scores as the reward function.
SMILES GA (Yoshikawa al., 2018) This method, similar to GA, applies various rules to the SMILES representation of molecules from starting population.
B.4 QED AND PENALIZED LOGP EXPERIMENTS
To ensure fair comparison with existing methods, rely using the same total number of calls to the property predictors. For example, an optimization process for an input molecule that runs for T iterations with M calls to the property predictors at each iteration will invoke total of x M property predictor calls. We use the same or less number property predictor calls for each molecule's optimization.
QED experiment setup. When running RetMol, for each input molecule, we set the maximum number of iterations to 1000 and sample 50 molecules at each iteration, resulting in a total of 1000 x 50 50, 000 property predictor calls. This number matches that in QMO (Hoffman et al.)
202I). where the maximum number of optimization iterations is 20 with 50 "restarts" and 50 samples each iteration for gradient estimation. Effeetivelyy, this results in total of 20 50 50 50, 000
19

---

Published as a conference paper at ICLR 2023
Table A1: Number of generated molecules (or number of calls to the property prediction) for the competitive methods in the task of optimizing four properties: QED, SA, and two binding affinities GSK388 and JNK3 estimated by pre-trained models from (Jin et al. 2020a).
Method
Number generated samples
RationaleRL Jin et al.
2020a
1,086,000
MARS Xie al. 2021
2,750,000
MolEvol Chenn Cheen"" et al. 2021
200,000
RetMol
296,000
calls to the scoring function, which is the same as in our setting. We evaluate performance by success rate. For each input molecule, if molecule that satisfy the design criteria (QED is above 0.9 and similarity with the input molecule is above 0.4), then we count it as a success, Success rate for all 800 input molecules thus
#successful imput molecules
success rate-
800
Penalized logP experiment setup. For each input molecule, we run the optimization for 80 iterations and sample 100 molecules at each iteration, which is exactly the same setting as in QMO (Hoffman et al., (02I). If optimization fails, i... no new molecules are generated that have higher penalized logP value than the input, then set the relative improvement to O. We evaluate performance by difference in penalized logP between the optimized and the input molecules, averaged over all 800 input molecules:
avg. improvement 1
Σ (aplogP (x)(1)) OplogI (1!11)
800
where OplogP is the penalized logP predictor, x' is the generated molecule and Iin is the input molecule.
B.5 GSK33 + JNK3 + QED SA EXPERIMENT SETUP
For fair comparison, we follow the same setup as in (Jin 2020a changes to RetMol's generative process. First, rather than starting with a known molecule, we first draw random latent embedding as the starting point of the molecule generation process. 2) In the iterative refinement process, we generate only one molecule and directly use it as the input molecule in the next iteration; that is, in this experiment, at each iteration, we not use the property predictors to select the best generated molecules because there is only one generated per iteration. For each input molecule, we run the optimization for 80 steps and generate one molecule at each step following (Jin) et al. 2018). Doing so results in total of 80 x x 3, 700 296, 000 generated molecules, which is an order of magnitude less than the 550 X 5, 000 = 2, 750, 000 number of generated molecules required for MARS (Xie et al. 2021). This number in MARS remains high even if change the number of samples at each iteration from 5, 000 to 3, 700 to align with the number of molecules evaluated in RetMol and the other baselines. In Table A1 we provide the complete comparison of the competitive methods regarding the number of generated samples, where we can see that our method sample efficient (with the best performance shown Table Ib).
2018) and make major
Unlike the previous tasks, there is no specification as to which molecule to use as the input molecule be optimized. To obtain the input molecules, first randomly select 3700 molecules from the CheMBL dataset (Gaulton et al,, 2016),, retrieve exemplar molecule randomly from the entire CheMBL dataset (Gaulton al. (016)), and greedily generate one molecule using the random input and random retrieved exemplar molecules. This one generated molecule is the input to RetMol. We choose 3700 molecules to be optimized because 3700 the number of molecules evaluated existing methods reported in (Jin et al. 2020a). Note that the inputs to RetMol differ from those in existing methods. However, we believe our choice of input does not put our method an advantage over existing methods and may even be at disadvantage compared to existing methods. For example, the inputs to RationaleRL (Jin et al. 2020a) are "rationales", i.e., molecule fragments that already satisfy either the 3 or the JNK3 property and that are pieced together. These rationales usually very close to the desired molecules. In contrast, the inputs to RetMol are mostly molecules with very low (close 0) GSK3,3 and JNK3 values, making the optimization challenging.
20

---

Published as a conference paper at ICLR 2023
Below, we show how to compute the metries when using RetMol, which largely follows the computa- tion in (Jin et al. 201 8 2020a). For success rate, we count number of input molecules that result in at least one successtul generated molecule (satisfy the four property constraints simultaneously) from all 80 molecules generated for that input molecule. If there exist more than one successful molecule for given input, choose the one with the least Tanimoto distance (Bajusz et al.] (013) with the input for the evaluation in the remaining metrics. For Novelty, we compute the percentage of the selected successful molecules that have less than 0.4 Tanimoto distance with any molecules in the retrieval database, i.e., those molecules that already satisfy the property constraints. For diversity, we compute the pairwise Tanimoto distance between each pair of the selected successful molecules.
B.6 GUACAMOL BENCHMARK EXPERIMENT SETUP
There are no known imput molecules to any tasks in the Guacamol benchmark. Therefore, we first randomly select population of molecules from the CheMBL dataset, which is the standard dataset in the benchmark. In this work, we choose five molecules as the starting population, although a larger population size is likely to yield better results the higher computational cost. We perform iterative refinement for each molecule in the input population for 1000 iterations. After the optimization ends, we collect all the generated molecules from all molecules in the population and select the best N molecules, based on which the benchmark score is computed for each benchmark task. In all the MPO tasks in the benchmark, N 100. The properties to be optimized and the scoring function details for each benchmark task are available in (Brown et al.) 2019) and at / htt/Benevelennt.ee ggaccmmlll The scoring functions and evaluation protocol are provided in the benchmark. The benchmark provide an API which takes population of input molecules and scoring functions to generated the optimized molecules. We implement RetMol using their API, which ensures fair comparison and evaluation using the same evaluation protocol defined in the Guacamol benchmark.
B.7 SARS-COV-2 MAIN PROTEASE INHIBITOR DESIGN EXPERIMENT SETUP
For RetMol, for each of the eight input inhibitors to be optimized, we run the iterative refinement for 10 iterations with 100 samples per iteration. For the graph GA baseline, for each of the eight inhibitors to be optimized, we use the same inhibitor as the initial population for crossover and mutation. We set the population and offspring size 100 and run for 10 iterations.
We use binding affinity between the generated inhibitor and the protein target. We computationally approximate binding affinity using Autodock-Gri -PU. which significantly speeds up docking compared to its CPU variants. The input files contains the receptor (target protein) of interest, the generated or identified ligands (inhibitor molecules), and a bounding box that indicates the docking site. For different docking experiments, the receptor and bounding box need to be prepared differently. Once the receptor and bounding box are prepared for a particular docking experiment, they are kept fixed and used for docking all ligands throughout the experiment. For the SARS-CoVa main protease docking experiments, we use the protein with PDB ID 7L11. We choose the bounding box by first docking the protein with a known inhibitor ligand, Compound 5 &Zhang al., [2021), and then extract the bounding box configurations from its best docking pose (Gao et al. 2022). The receptor, ligand, and bounding box preprocessing steps use the Autodock software suite
In addition, as we mentioned in the main paper, we also tested on QMO (Hoffman et al. 2021) with the following configurations: we set optimization iterations to 100, number of restarts to weights for the property (binding affinity) to be optimized in {0.005, 0.01, 0.1,, 0.25J, base learning rate in {0.05, 0.1). and the property constraints the same as in those in RetMol. In our experiments, these configurations did not succeed in generating an optimized molecule given the constraints. We follow the official QMO codebase in these experiments.
httpss//attchub.com////eeeeeeeee https://autodooa/ scri edu. https://github.cco wwt.:/ooo.///////. IMM/QMO
21

---

Published as a conference paper at ICLR 2023
B.8 ANALYSES EXPERIMENTS: TRAINING OBJECTIVES
The purpose of this experiment is two-fold: 1) We want to show if only updating the fusion module while fixing the weights of the base generative model during training will cause a degradation in the quality of generated molecules; and 2) we want check if our proposed nearest neighbor objective can achieve better generation quality than the conventional language modeling objective (i.e., predicting the next token of input).
To evaluate RetMol, the unconditional generation procedure is the same as that in training the information fusion module, i.e., the retrieval database is the training split of the ZINC250 dataset, and the retrieval criterion is molecule similarity. Besides, since we mainly focus on evaluating the quality of generated molecules, we thus use validity, novelty and uniqueness as our metrics, which are standard in literature for evaluating molecule generation models' performance. The results suggest that our training objective that only updates the fusion module does not sacrifice the unconditioned molecule generation performance and even results in slight improvement on the novelty metric.
We perform the evaluation on the test split of the ZINC250K dataset. For each molecule in the test set, we generate 10 molecules by first randomly perturb the input to the decoder, i.e., an encoded (or fused) embedding matrix, 10 independent times using a isotropic random Gaussian with standard deviation of 1. Then, for validity, we compute the percentage of the 10 generated molecules that are valid according to RDKit. averaged over all test molccules. For novelty, we compute the percentage of the 10 generated molecules that not in the training split of the ZINC250K dataset, averaged over all molecules in the test set. For uniqueness, we compute the percentage of the 10 generated molecules that are unique, averaged over all molecules in the test set.
B.9 REMARKS ON NUMBER OF ITERATIONS
In most experiments, we keep the number of iterative refinements RetMol the same as the baseline iterative methods. For example, in Section 3.1 we use 80 optimization iteration steps for RetMol, which is the same as QMO. In Section 3.22 we use 80 optimization iteration steps, which is the same as JT-VAE. Note that in this experiment, RationaleRL does not require such iterative refinement process but does require extensive task-specific training and fine-tuning. In Section 3.3] the benchmark results do not mention the number of iterations and therefore simply set the max number of iterations to 1k for RetMol. In Section 3.4, we use 20 iterations for both RetMol and Graph GA.
Our results also suggest that one does not need as many iterations as the basclines to achieve strong results. For example, we showed in Figure 3 (middle) that for experiments in Section 3.2 RetMol achieves better results than baselines with only 30 iterations.
C ADDITIONAL EXPERIMENT RESULTS AND ANALYSES
C.I QED AND LOGP EXPERIMENT
To visualize the property value improvements, we plot in Figure A1 (i) the QED of the optimized (generated) molecules and of the input molecules under similarity constraint 0.4 and (ii) the penalized logP improvement between the generated and the input molecules under similarity constraints 8 = (0.4, 0.6).
We can see that, for the penalized log? experiment, for similarity constraint 8 = 0.4, there are quite a few molecules with a large penalized logP value improvements. These molecules are the reason that the variance for RetMol in Table lb is large. But even if we remove those molecules with extreme values, i.e., penalized logP bigger than 20, we still obtain an average improvement 8.17 + 4.12, which is still better than the best existing method. Furthermore, we also compute the mode of the penalized logP values for our method (i.e., the x-axis value of the largest peak in Figure A1 Right), and we get 3.804 for values is still better than the mean in most of the baselines in Table Th. For QMO, their mode looks very similar to ours from Figure in their paper (Hoffman al. 202
0.6 and 6.389 for
0.4. We can see mode of the penalized logP
httpp::/ http///wwww......c www.:.dkkioorgggo
22

---

Published as a conference paper at ICLR 2023
6
Input QED
6 0.4
0.05
0.00
111
011
0.75
0.80
0.85
0.90
0.95
-10
20
30
60
QED
Penalized logP improvement
Figure A1: Left: distribution of QED values of the original and the optimized (generated) molecules under similarity constraint
0.4. Right: distribution of penalized logP improvement comparing similarity
constraints (0.1, 0.6).
1.00
Su0258 mie
0.35
0.85
0.9
0.70
0.90
0.00
0.85
0.75
0.70
BM
0.66
0.80
0.65
MARS
MARS
Raboralllti.
0.60
Ittioralekt.
Size retrieval data base
Size retrieval database
Size retrieval database
Figure A2: Generation performance with varying retrieval database size on the experiment Section 3.2. Our framework achieves strong performance with as few as 100 molecules in the retrieval database and performance generally improves with increasing retrieval database size on all metrics.
0:95
0.73
Divensiyy
0.85
0.30
0.7
0:75
RatMal
RetMel MARS RatioraloR
0.71
RecMal
MARS
MARS
RationdeR
05
RetioraleRt:
0.60
30
47
50
10 20
50
60
50
10
20
30
50
IEterations
#lterations
Alterations
Figure A3: Generation performance with varying number of optimization iterations on the experiments in Section 3.2. Dashed lines are the success rates and novelty scores of the two best baselines. We observe that all the metries improve increase the number of iterations.
C.2 GSK33 AND JNK3 EXPERIMENT
In Figure A4 we show 54 randomly chosen molecules generated by RetMol along with their highest similarity with the molecules in the retrieval database. These molecules demonstrate the diversity and novelty of the molecules that RetMol is capable of generating.
C.3 GUACAMOL, EXPERIMENT
We show Table A2 the detailed benchmark, SA, and QED results for all the MPO tasks in the Guacamol benchmark.
We additionally apply the functional group filters 3 to the optimized molecules of each method and compare the average number of molecules that pass the filters. Figure A5 visualizes the results, which align with those in Sec. 333 RetMol strikes the best balance between optimizing benchmark score and maintaining good functional group relevance. For example, Graph GA achieves the best benchmark performance but fails the filters more often than some of the methods under comparison. In contrast, RetMol achieves the second best benchmark performance and the highest number of
23

---

Published as a conference paper at ICLR 2023
ain-0287
sh1 219
sm-03366
sh100333
foff
sin-0389
Sir
sired
sin-3362
sim-0289
sin
sim:0323
sim=0300
sim=035t
sim:0343
sim=0284
sim=0
stm00378
skm-03111
skmw0
5im-0.778
sim=0284
5tm-0322
s1rr0310
séree
sin=03355
Fou oppoo
sim=03477
simed 381
sincd 264
sim=03077
sime:0312
sine0
sim-0200
$mm-0346
+lm-0323
5m--0341
sm=-0387
skm=0
4hnw0376
sim=0.393
sinn0.351
kimed3361
simrn 350
sirmm0288
Nim=0207
simm0.364
Figure A4: Visualizations of randomly chosen molecules generated by RetMol for the GSK33 JNK3 + QED + SA experiment. Below cach generated molecule, we show its highest similarity between each molecule in the retrieval database.
24

---

Published as a conference paper at ICLR 2023
RetMol
RDFitter #passed
SMILES LSTM
23
15
Graph MCTS
10
0.55
0.00
0.65
0,70
0.75
0.80
0.85
0.90
Benchmark Performance
Figure A5: Average number of optimized molecules that pass the functional group filters against Guacamol benchmark scores. The results align with those in Sec. 3.3 RetMol strikes the best balance between optimizing benchmark and maintaining gond functional group relevance.
Table A2: Detailed results from the Guacamol MPO results. The tables from the top the bottom the benchmark results, averaged SA values, averaged QED values, averaged numbers generated molccules that pass the functional group filters, respectively. SA and QED values and the number of filter-passing molecules are averaged over all the 100 molecules evaluated in each MPO task. Bold and underline represent the best and the second best in each metric in each benchmark task, respectively.
Benchmarks
SMILES GA
Graph MCTS
Graph GA
SMILES LSTM
Best Dataset
Ours
Osimertinib MPO
0.886
0.784
0.953
0,907
0.839
0.915
Fexofenadine MPO
0,931
0.695
0,998
0,959
0.817
0.969
Ranolazine MPO
0.881
0.616
0.920
0.855
0.792
0.931
Perindopril MPO
0.661
0.385
0.792
0.808
0.575
0.765
Amlodipine MPO
0.722
0.533
0.894
0.894
0.696
0.879
Sitagliptin MPO
0.689
0.458
0.891
0,545
0,509
0.735
Zaleplon MPO
0.413
0.488
0.754
0.669
0.547
0.713
Benchmarks SA
SMILES GA
Graph MCTS
Graph GA
SMILES LSTM
Best of Dataset
Ours
Osimertinib MPO
6.386
3.901
3.357
2.923
2.705
3.061
Fexofenadine MPO
3,590
4.671
3.897
3.171
3.097
3.546
Ranolazine MPO
6.071
4.110
4.111
2.900
3,226
3,456
Perindopril MPO
5.343
3.365
4,286
4.017
3.645
4.276
Amlodipine MPO
4.717
3.529
3.575
3.329
3.163
3.431
Sitagliptin MPO
6.743
5.183
6,804
2.794
2,886
3.715
Zaleplon MPO
3.244
3.216
2.899
2.387
2.294
2.622
Benchmarks QED
SMILES GA
Graph MCTS
Graph GA
SMILES LSTM
Best Dataset
Ours
Osimertinib MPO
0.256
0.443
0.197
0.240
0.478
0.264
Fexofenadine MPO
0.207
0.495
0.309
0,335
0,382
0,301
Ramolazine MPO
0.096
0.305
0.095
0.113
0.129
0.112
Perindopril MPO
0.481
0.477
0.365
0.465
0.421
0,546
Amlodipine MPO
0.146
0.582
0.351
0,386
0.472
0.365
Sitagliptin MPO
0.254
0.453
0.086
0.700
0.735
0.701
Zaleplon MPO
0.206
0.679
0.562
0.730
0.712
0.753
Benchmarks filters
SMILES GA
Graph MCTS
Graph GA
SMILES LSTM
Best of Dataset
Ours
Osimertinib MPO
0
23
0
0
19
Fexofenadine MPO
58
22
73
68
47
43
Ranolazine MPO
0
0
0
0
0
Perindopril MPO
12
29
42
60
40
90
Amlodipine MPO
4
31
56
78
49
58
Sitagliptin MPO
0
0
69
76
82
Zaleplon MPO
0
51
37
97
84
98
passed molecules. Please see Table A2 for the number of optimized molecules that pass the filters for each MPO task and for each method.
25

---

Published as a conference paper at ICLR 2023
Table A3: Visualizations of the original and optimized inhibitors from RetMol for the SARS-CoV. main protease. Similarity constraint is 0.6.
Inhibitor
Original
Properties (original)
Optimized
Properties (optimized)
Similarity map
Similarity
Docking
-1.93
Docking
-6.78
Favipiravir
QED 0.55
QED 0.77
0.60
SA 2,90
3,50
Docking -9.44
Docking
-11.48
QED
SA 2.94
SA 2.57
Docking
-6.13
Docking
-845
PX-12
QED
0.74
QED 0.64
SA 3.98
SA 3,80
Docking
Docking
Disulfiram
QED 0.57
QED 0.60
0.64
3.12
3.39
Docking -8.45
Docking
Knempferol
QED 0.55
QED 0.63
0.62
2.47
C.4 SARS-COV-2 MAIN PROTEASE INHIBITOR DESIGN EXPERIMENT
Tables A3 and A4 visualizes the original and the RetMol optimized inhibitors along with the similarity map: Tanimoto distance, QED, SA, and docking (unit kcal /mol) scores for both similarity constraint 5 {0.6, 0.4). We highlight the properties that the initial inhibitor do not satisfy red and the same properties the optimized inhibitor green. We can see that RetMol not only optimizes the docking score but also successfully improves the QED and SA scores for some inhibitors.
C.5 ANTIBACTERIAL DRUG DESIGN FOR THE MURD PROTEIN
addition to the experiments above and in the main paper, here we demonstrate another real-world case of RetMol for Antibacterial drug design. We choose the MurD protein (PDB ID: 3UAG) as the target. This a promising target for antibiotic design because it is necessary for the development of the cell wall, which is essential to the bacterial survival. Inhibiting this target thus has the potential to destroy the bacterial without harming humans because the cell wall and thus the target protein is absent in animals (Sangshetti et al. 2017
The design criteria in this controllable generation experiment is similar to those in the SARS-CoV-2 main protease inhibitor design experiment. In addition improving the binding affinity of selected weakly-binding molecules, we have several desired properties including small logP value (below 3), large QED value (above 0.6), small SA score (below 4), and molecule weight between 250-350 Da. Since the task encourages diverse generations, we set small similarity threshold to 0.2. We select 100 molecules from bindingDB (Liu al. (007)) that has experimental binding values
to MurD, and choose eight molecules with the lowest binding affinity input to be optimized. For
AAso see
here:
Overvies
26

---

Published as a conference paper at ICLR 2023
Table A4: Visualizations the original and optimized inhibitors from RetMol for the -AVS main protease. Similarity constraint here is 5 = 0.4.
Inhibitor name
Original
Properties (original)
Optimized
Properties (optimized)
Similarity map
Similarity
Docking
.493
Docking
-8.7
Favipiravir
QED
0.55
QED 0.62
0.41
SA 2.90
SA 3.25
Dacking
Docking
-12.65
Brombexinc
QED 0.78
QED 0.63
0.40
Docking
-6.13
Docking -10.90
PX-12
QED 0.74
QED 0.62
0.50
Docking
-7.31
Docking -1082
Ehselen
QED 0.63
0.61
0.40
2,05
SA 2,27
Docking
-8.88
Docking
-10.44
Disulfiram
QED
SA 3.12
SA 3.63
Dacking
.9.00
Dacking -12.34
Entecavir
QED 0.53
QED 0.64
0.41
SA 4.09
SA 3,76
Docking
Docking
Querectin
QED 0.43
QED 0.62
0.41
2,62
Docking -8445
Docking 10033
Kaempferol
QED
0.55
QED
0.67
0.41
Table A5: Antibacterial drug design with the MurD target comparing RetMol with Graph GA. RetMol optimizes input molecules better terms binding affinity, unit inkeal /mo than Graph GA under various property constraints.
Input
Input score
RetMol optimized
Graph GA (Jensen 2019
optimized
-7.33
-11.78
-9.76
-7.31
-12.82
-11.09
Oelc2500c2hml-eeecc@chchc
-7.53
-12.46
-9.31
-7.86
-13.50
-9.39
0e1c2SCCCc2ml
(C(F))
-14.72
-8.98
Oetc2S0Cc2ml1elleee
-7.63
-13.62
-8.81
Dele2SCCe2ml-cllee
-7.70
-14.00
-10.68
Delc2SCCe2mllelcorfff
-7.19
-13.05
-11.67
Average improvement
2.38
27

---

Published as a conference paper at ICLR 2023
RetMol, all the molecules resulted from bindingDB the retrieval database. The remaining experiment procedures, including docking, follows from those in the SARS-CoV-3 main protease inhibitor design experiment. Table compares the generation performance of RetMol with Graph GA. We can see that RetMol optimizes the input molecules better than Graph GA.
C.6 ANALYSES
With and without retrieval module. We test the base generative model in our framework on the experiment in Section 3.2 with the same setup our full framework. The base generative model without the retrieval module unable to generate any molecules that satisfy the given constraints demonstrating the importance of the retrieval module in achieving controllable generation.
Varying retrieval database size. Figure A2 shows how RetMol performs with varying retrieval database size with the novelty (middle) and diversity (right) metrics in addition to the success rate metric (left). We can observe that, similar success rate, RetMol can achieve reasonable novelty and diversity with a small retrieval database, albeit with larger variance. The performance continues to improve and the variance continues to decrease with a larger retrieval database. This experiment corroborate the analyses in the main paper that RetMol is efficient in the sense that small retrieval database can already provide a strong performance.
Varying refinement iterations. Figure A3 shows how RetMol performs with respect to the novelty (middle) and diversity (right) metrics in addition to success rate (left). These results are similar to and corroborate those presented in the main paper: RetMol achieves strong results and beats the best existing method with as little as 10 iterations on the diversity metric and with as little 20 iterations on the novelty metric. Performance also continues to improve with more iterations.
Comparing to parameter-free information fusion To show the effectiveness of our proposed information fusion module that involves additional trainable parameters from the cross entropy function, we introduce two parameter-free information fusion modules detailed below for comparison.
Both two methods need to calculate a property score for each of the retrieved molecules (and we take an average the property scores as the final score in the case of multi-property optimization), and then apply the softmax function on these scores obtain vector of normalized scores. The first method, termed softmax--ggregaation applies a weighted averaging of the retrieved embeddings the fused embedding, with the weights being set to the normalized scores. Note that since all retrieved embeddings have different lengths, we simply concatenate zero vectors to those embeddings with smaller length maintain the same dimensionality before the weighted averaging. The second method, termed softmax-sampling fused embedding, by treating the normalized scores from softmax as probabilities of a multinomial distribution. The rest procedures are kept the same with RetMol.
samples one retrieved molecule embedding as the
Table A6: We compare our proposed information fusion module with two parameter-fres fusion methods in the QED experiments Scction 3. where the results demonstrate the importance of our proposed information fusion module.
Method
Success (%)
Coftmax-aggregge
1.5
Softmax-sampling 53.6
RetMol
94.5
We summarize the results in Table A6 where we conduct the QED experiments in Section 3.1 Compared with RetMol (> 90% success rate), the softmax-aggregat method has nearly zero success rate. It implies that training the information fusion component is necessary and that simply aggregating the embeddings from the retrieved exemplar molecules results in undesirable results. The softmax- sampling method that does not directly aggregate the information also achieves subpar performance compared to RetMol, which also demonstrates the importance of our proposed information fusion module. These two comparisons together show that our proposed information fusion module with
28

---

Published as a conference paper at ICLR 2023
Average similartyy
0.10
Donsity
0.08
0.4
0.06
0.3
0.04
0.2
0.1
0.02
0.0
0.00
10
k-th nearest neighbor
Similarity
0.12
0.200
0.10
0.175
Density
0.08
Density
0.125
0.150
0.06
0.100
0.04
0.075
0.02
0.025
0.000
0.4
0.6
0.6
Similarity
Similarity
Figure A6: Analyses of the similarities between the retrieved molecules and the input molecules. Top left: the average similarity between the l--th most similar molecules to their corresponding input molecules; Top right, bottom left, and bottom right: the distribution of the similarities of the 1st, 2nd, and 10th most similar molecules to their corresponding input molecules, respectively.
trainable parameters that learns to dynamically aggregate retrieved information is critical for achieving good generation performance.
Computational cost For the memory cost, 1) the model size (in #parameters) does not increase much: The information fusion module contains 460,544 parameters. The base generative model in RetMol contains 10,010,635 parameters. These numbers indicate that the information module adds only less than 5% of the total parameters. 2) The memory cost in the fusion module scales linearly with the number of retrieved molecules K in order store the K extra embeddings of the retrieved molecules, i.e., O(K LD), where L is the average molecule length, and D is the embedding size. Since K is small (i.e., K 10 with T, 55, D 256), in experiments we find this additional memory cost is small.
For the time cost, we also observed infinitesimal difference in the runtime between the base generative model (i.e., without the fusion module) and RetMol (i.e., with the fusion module). To verify this, we ran a small experiment to compare their time cost the encoding step, as RetMol uses the same decoder as the base generative model. Specifically, generative model and RetMol with batch size of I and compute their average runtime (in seconds) NNVIDIA Quadro RTX 8000. The average encoding time for the base generative model and for RetMol is 0.00402 seconds and 0.00343 seconds, respectively. shows that the additional time cost of the fusion module in RetMol is indeed small.
gave 100 input molecules to both the base
Similarities of retrieved molecules We first plot in Figure A6 the similarities between the input molecule and each of its K nearest neighbors (measured by cosine similarity) where K = 10, respectively, and take average across all input molecules in the training set. We see that average, the most similar molecule has a similarity 0.66 to its corresponding input, and the second most similar molecule has a similarity of 0.57. The average similarities of the remaining 3rd to the 10th most similar molecules do not drop dramatically, implying the retrieved molecules are indeed similar of K 10. Second, we also plot in Figure A6 the distribution of
the input in the similarities for the 1st, 2nd, and 10th retrieved molecules to their input molecule, respectively, across the training set. We see that the majority of the top-k similar molecules have a similarity greater than 0.2 with their input, even for the distribution of the I 10th retrieved molecules. It implies the possibility of retrieving largely dissimilar molecule small. These results both imply that the noise present in the retrieved molecules is modest with small fixed K.
Attribute relaxation ordering in molecule retriever In molecule retriever, we may need to gradually relax attribute constraints construct relaxed feasible set. Our strategy is to relax the
29

---

Published as a conference paper at ICLR 2023
Table A7: Penalized logP experiment with HierVAE the base molecule generative model (encoder and decoder) RetMol framework. This result demonstrates that RetMol is flexible and compatible with models other than transformer-base ones and that the RetMol framework improves controllable molecule generation performance compared to the base model alone.
Method
6=0.6
8=0.4
HierG2G
2.49:11.09
3.98-11.44
RetMol (w/ HierG2G)
3.09
5.276
harder-to-satisfy constraints first and then the easier-to-satisfy ones. In other words, we first focus more on the simple attributes and then on the hard attributes (e.g.. similarity + QED + docking). which gradually increases the level of inference difficulty over iterations. The order of the constraints being removed is an important consideration during inference, especially when there are many of them. To show this, we perform sanity check on the impact of attribute relaxation ordering in the QED task in Sec where, for a quick experiment, we use subset 200 molecules and 20 iterations. We see that if we first relax QED and then similarity, we first relax similarity and then QED, we get 78.5% success rate. confirms our intuition that is preferable relax the harder-to-seee sorder-to-satiffy constraints first and then the easier-to-satisty ones.
get 89.5% success rate, and if
C.7 RETMOL WITH OTHER BASE GENERATIVE MODELS
We conduct all experiments in this paper until this point with pre-trained molecule generative model based on the transformer (i.e., BART) architecture as the encoder and decoder model in the RetMol framework. As further demonstration that RetMol is flexible and is compatible with other types of molecule generative models, and thanks to the reviewers' suggestions, we conduct the penalized logP experiment with HierG2G Jin et al. (2019) as the RetMol encoder and decoder models. The results in Table A7 suggest that the RetMol framework can also elevate the performance of graph-base based generative models, i.c., HicrG2G, on controllable molecule generation.
30


Reviews:
[{"review_id": "cRZlFWj5vT", "reviewer": "Reviewer_3hQB", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The work presents a retrieval-based framework for controllable molecule generation. The framework consists of:\n\n- A mechanism to retrieve a set of molecules that (partially) satisfy a set of desired properties the generated molecule should have (e.g. similarity to the input molecule, enhanced drug-like properties);\n- A fusion module that combines the retrieved molecules with the input molecule through cross-attention;\n- A novel self-supervised objective to train the fusion module, consisting in reconstructing the nearest neighbor molecule of the input molecule (as opposed to the classical self-reconstruction objective);\n- An iterative refinement process to update the retrieval database with the generated molecules.\n\nThe fusion module can be plugged-in into any pre-trained encoder-decoder architecture (in this work, they use the SMILES representation of molecules and the ChemFormer model), which is frozen (i.e. only the fusion model is trained for efficiency). \n\nThe proposed approach is evaluated in three controllable generation scenarios, one real-world drug design scenario (find novel inhibitors for SARS-COVID protease), and the GuacaMol benchmark. In every case, it obtains an improvement over a range of different metrics such as success rate or docking score.", "summary_of_the_review": "This appears to me as a strong submission. I am giving a score of 6 for the time being. Will be happy to raise my score when the authors answer the following set of questions/comments: \n\n- In section 2, you acknowledge the existence of weighting coefficients $w_{\\ell}$. However, they don't seem to be present in the rest of the work. What are they (or, are they just construction to explain the problem)?\n- What is the justification for having used ChemFormer instead of any other encoder-decoder architecture as the backbone model? It would be great to see RetMol applied to another backbone, to see if it works across backbones or if it is just a perfect match with ChemFormer.\n- In which order are the constraints removed to relax the feasible set (Sec. 2.1 - Molecule retriever)? It would be interesting to know the impact of using different ordering strategies.\n- In (Sec. 2.1 - Information fusion), I see that the set of retrieved exemplars is a matrix $\\in \\mathbb{R}^{(\\sum_{k=1}^K L_k) \\times D}$, which means that you concatenate all the embeddings of the retrieved molecules. Considering that they might have different dimensions, how do you handle the mismatch? Padding? Also, doesn't this introduce an unnecessary dependence on the order in which the molecules are retrieved (which could have been avoided, e.g. by a permutation invariant aggregation)?\n- I understand the model is very efficient to train. But what about preprocessing (i.e. computing all the embeddings and the pairwise similarities)? What is its impact?\n- This is for my understanding: do you confirm that the iterative refinement is applied only at inference, but not during training?\n- Due to the large standard deviation in Table 1b, I think it's better to also report the mode of the penalized logP values in addition to the mean.\n- How come uniqueness is sometimes not reported in the generation experiments (e.g. Table 2, Table 4-right)? \n\n\n**EDIT**\n\nAfter a successful rebuttal, I am raising my score to 8, since the method is sufficiently new, has good experimental results, and has wide applicability across different molecular generators.", "strength_and_weaknesses": "**Strenghts**\n\n- Strong performances in what appears to be a well-designed set of experiments;\n- Training is very efficient (just the retrieval module); inference is not affected too much since we are adding a small number of parameters with respect to the large size of the encoder-decoder backbone;\n- Works with a small set of retrieved molecules (up to 23 in the experiments), which is ideal when one wants to generate alternative molecules from limited data;\n- The nearest neighbor training objective is a simple but very clever idea that works experimentally.\n\n**Weaknesses**\n\n- Not sure it would be applicable in cases where no exemplar molecules to be retrieved are unavailable (as stated in Conclusions); this might limit its widespread applicability;\n- It poorly exploits the chemical space, i.e. it will most likely generate molecules that are very similar to the input molecule.\n- Even though it is claimed to work with any encoder-decoder architecture, performances are shown considering only a single backbone model (ChemFormer). I would have expected that its cross-applicability would be one of the main focuses of this paper, but unfortunately it is not.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "J9NKpJX0i3", "reviewer": "Reviewer_UNcN", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper studies molecule generation, proposing a new model using a retrieval mechanism, which has recently shown to be successful in NLP.\n\nThe method is evaluated on several established benchmarks and reaches performance comparable to state of the art", "summary_of_the_review": "good evaluation, decent results, limited novelty", "strength_and_weaknesses": "strengths\n- comprehensive evaluation\n- new, reasonably motivated model\n\nweaknesses\n- ML novelty somewhat limited\n- docking as oracle is not ideal", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "FirsSLi63qV", "reviewer": "Reviewer_PZ6g", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "Motivated by the need for sample-efficient molecule generation, this paper introduces a simple retrieval mechanism to retrieve and fuse the exemplar molecules for controllable molecule generation. Evaluated on a variety of synthetic molecule generation tasks, the proposed method outperforms other baselines. ", "summary_of_the_review": "In summary, this paper proposes a retrieval-based method for controllable molecule generation and extensively evaluates its effectiveness on benchmark datasets. The introduction of a retriever for molecule generation is new as far as I know, but it’s not new in other ML communities. However, the design of the retrieval components does not bring in any technical contributions to the community and is mostly based on ungrounded heuristics. ", "strength_and_weaknesses": "# Strength\n\n- The experiments are very extensive with lots of details, including both qualitative and quantitative materials.\n- It evaluates the effectiveness of simple retrieval modules for controllable molecule generation. \n\n# Weakness\n\n\n- Inference via iterative refinement introduced in Section 2.3 is quite heuristic, yet the heuristic is not designed carefully. The authors mention that iterative update is common in approaches such as GA methods, and there are many off-the-shelf GA tools. It seems to me that it’s a better choice to use those tools instead of relying on a design with much simpler heuristics.\n\n- On most benchmark tasks, RetMol does not significantly outperform other baselines, and those baselines methods are usually proposed in 2019. The authors should include recent stronger baselines to justify the value of retrieval modules.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "BTIQW-3RkM", "reviewer": "Reviewer_9XRZ", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposed a retrieval-based method for optimizing molecular properties under similarity constraints. The proposed method leverages a pre-trained encoder-decoder generative model and trains an information fusion module to steer the pre-trained model toward generating more similar molecules to the retrieved ones. During training time, the model is asked to predict the nearest molecule in the retrieved ones. And then, during inference time, an iterative refinement process is employed to perform multi-step optimization. The proposed method performed competitively on several molecular optimization tasks using only a small set of exemplar molecules.", "summary_of_the_review": "The paper proposed a retrieval-based method for molecular property optimization. Although the proposed method achieved competitive results, the methodology is less principled and needs more work.", "strength_and_weaknesses": "Strength:\n- Data efficiency is a practical issue in the optimization of molecular properties. The proposed method could be used to optimize molecules given a small number of positive examples.\n- To my knowledge, the self-supervised training strategy of predicting the nearest neighbors given the remaining neighbors is novel. The task-agnostic training makes it depend less on the number of task-specific samples.\n\nWeaknesses:\n- The paper did not justify why such self-supervised training strategy would work since the inference time objective is not aligned with the training time objective.\n  - During the training time, the information fusion is trained to predict the nearest molecule given the K-1 remaining nearest neighbors. An information fusion model that adds a small noise to the current molecule would yield a high objective score, which means the information fusion model learns to predict a similar molecule to the current one.\n  - However, during inference time, the model is asked to predict a similar molecule but in the direction of the positive examples. The inference time objective is quite different from the training time objective. This is a mismatch that needs further explanation and/or empirical analysis in the paper.\n- The design of the iterative refinement approach is ad-hoc, and the description in Section 2.3 seems to miss important details about random perturbing.\n  - From my understanding, the model is expected to interpolate between the current molecule and the positive examples in the axis of both structural similarity and property scores. It’s unrealistic to expect the model to extrapolate beyond the database, even if we do iterative refinement. The key mechanism to allow local exploration is random perturbing, mentioned briefly in Section 2.3. However, details are missing here.\n- The abstract and intro section gives the readers the wrong impression that the proposed method only requires a small number of samples for a property optimization task. However, the method still requires a large number of samples (number of calls to the property prediction), even in the SARS-COV-2 experiment where there are only 23 positive examples given. Since this paper focuses on sample efficiency, it would be more convincing if the authors could provide the exact numbers of samples used in the proposed method and all baselines for a full comparison. For example, in the multi-property optimization experiment in Section 3.2, the number of samples required by each methods are (copied from the Appendix): \n  - 296, 000 for the proposed method.\n  - 2, 750, 000 for MARS [1].\n  - 200, 000 for MolEvol [2] (missing in the baselines).\n  - xxx for RationaleRL[3].\n\n[1] Xie, Yutong, et al. \"MARS: Markov Molecular Sampling for Multi-objective Drug Discovery.\" International Conference on Learning Representations. 2021.\n[2] Chen, Binghong, et al. \"Molecule optimization by explainable evolution.\" International Conference on Learning Representation (ICLR). 2021.\n[3] Jin, Wengong, Regina Barzilay, and Tommi Jaakkola. \"Multi-objective molecule generation using interpretable substructures.\" International conference on machine learning. PMLR, 2020.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-25%