Paper: Searching for High-Value Molecules Using Reinforcement Learning and Transformers (ICLR-2024)

Reviewer comment (Novelty):
Originality: While the paper offers a thorough numerical comparison of different design choices (problem representation, neural architecture, pretraining, etc.) in the context of string-based reinforcement learning for molecular discovery, its originality is somewhat constrained.

Author response:
We thank the reviewer for their feedback and summary of our work. We appreciate the reviewer’s assessment that the paper makes a significant contribution to molecular discovery and is technically sound and well-written. It seems that the main concern of the reviewer is the originality of the proposed algorithm. We agree that the basic components algorithms and design choices that we experimentally study have been used separately in prior work, ours is the first fair and extensive comparison at a large scale for string-based RL. To highlight prior work, we have included Table 1 in our paper, which contains relevant prior works that have used some of the methods that we incorporate.

Reviewer comment (Writing):
(Introduction) "...our own algorithm (MoLRL)" - what is MoLRL?

Author response:
This is a typo. We have substitute MoLRL -> ChemRLformer. Thank you for pointing this out.

Reviewer comment (Presentation):
(Table 1) What is the difference between "MoLRL" and "ChemRLformer (Ours)"? Is MolRL the general framework and ChemRLformer the best combination of design choices based on the ablation study? Please clarify.

Author response:
This is a typo as well, they are the same method. We have fixed this in our updated paper.

Reviewer comment (Writing):
(Section 4) "The molecular design space is complex but the benefit from finding improved options is great." This sentence reads a bit too colloquially. Please rephrase.

Author response:
We have changed this sentence to : “The space of drug-like molecules is vast, and reinforcement learning methods hold great promise in improving the speed and reducing the cost of drug discovery.”

Reviewer comment (Writing):
(Figure 2) "(a) Performance on SMILES-based molecular design with pertaining (left) and with pretrianing and RL (right)." Please correct the typos "pertaining" and "pretrianing".

Author response:
We have corrected these typos in the updates paper.

Reviewer comment (Presentation):
(Figure 2) "b) performance on SMILES-based molecular docking with pertaining (right) and with pretrainig and RL (right)." First (right) should be (left). Please correct.

Author response:
We have corrected this in the updated paper, thank you.

Reviewer comment (Reproducibility):
(Figure 2) How does the "pretraining only" results (left column) are computed? Is this a zero-shot evaluation? I can't find the answer in the text. Please clarify.

Author response:
The pretraining only scores are calculated by sampling molecules from the pretrained model, without any RL training. We have added a sentence to clarify this in the updated paper (section 5.2, 1st para) : “We compare the scores of the molecules generated by the policy after pretraining and after RL training.”

Reviewer comment (Novelty):
Some relevant works seem to be missing: MolGPT (https://www.ncbi.nlm.nih.gov/pmc/articles/PMC10232454/) also uses Transformer for molecular generation but without RL.

Author response:
We have added the first two papers (MolGPT, Taiga) in Table 1 of our updated paper, and can cite the third one as concurrent work after it is public.

Reviewer comment (Novelty):
Please also consider citing the relevant parallel work which applies offline RL for the similar problem: https://www.biorxiv.org/content/10.1101/2023.11.29.569328v1 (once it is public).

Author response:
We have added the first two papers (MolGPT, Taiga) in Table 1 of our updated paper, and can cite the third one as concurrent work after it is public.

Reviewer comment (Novelty):
The second work (Taiga) seems to be very similar in terms of approach and there is uncertainty about the novelty of the approach itself.

Author response:
Indeed, the second work is similar to ours in that both papers use a pre-trained transformer and train it using policy gradients, and we have added both papers in Table 1 of our updated paper. The difference is in the scope of our experiments and the large number of other design choices combined in a novel manner. For example, 
1) Both the papers mentioned by you only show results for 2 or 3 different tasks, which include QED or SA scores. These tasks have been categorized as substantially easier to optimize in the literature compared to the tasks in our experiments [1], which span 25 + tasks including docking scores. 
2) In previous papers, the size of the datasets used were around 1 million molecules. We conduct experiments with sizes up to 100 million molecules.
3) Moreover, we also study algorithmic components like regularization and the use of replay buffers which seem to have a huge impact on the results, but are not studied in the papers you mention.
4) Our experiments also include more architectural diversity, including fully-connected MLPs, RNNs and Transformer backbones.

Reviewer comment (Presentation):
From the plots I couldn't understand what "Div." and "Red." columns stand for.

Author response:
These mean diversity and redundancy respectively. We have added the clarification in Section 5.1, evaluation metrics paragraph and also in the caption of Figure 2.

Reviewer comment (Evaluation):
There are no samples of generated molecules. RL could hack reward functions and generate molecules which make no sense. Reward hacking is mentioned by the authors but including resulting molecules is essential.

Author response:
In Figure 14 (Page 25) we have added examples of molecules sampled due to reward hacking. We can see that ChemRLformer agents are able to obtain unusually high docking scores by stacking together long chains and rings of sulfur, phosphorus or carbon atoms.

Reviewer comment (Experiments):
Some of the design choices seem to be dated. For example, it is not clear why REINFORCE is used as the policy learning algorithm, as opposed for instance to PPO.

Author response:
> It is not clear why the authors use Reinforce instead of more popular PPO for example

Although PPO has been shown to generally be a better algorithm for control problems, the comparison is not clear in the case of molecular optimization ([1,2]) . We perform an experiment across 21 PyTDC objectives (see Figure 13, page 22) to compare REINFORCE with PPO. Our results indicate that vanilla policy gradient algorithms achieve higher performance on all metrics and are more stable when compared to actor critic algorithms like PPO. This experiment guided the choice of using REINFORCE over PPO.

Reviewer comment (Experiments):
Some of the design choices seem to be dated. For example, there is no automatic entropy tuning for exploration.

Author response:
> Re: entropy tuning. 

We apologize for the misunderstanding. Indeed we could regularize the entropy of the policy with a tuned coefficient. In Figure 41, we compared the entropy regularization with an automatically tuned coefficient with the Log P regulariser across all 21 PyTDC tasks.  We see that adding automatic tuning does not lead to much improvement. For the reasons explained in Section 5.4, from line 333, we believe that Log P regularization is significantly better for molecular optimization (It helps the agent get out of local optima. 

We note that Soft Actor Critic [Haarnoja et al 2018] proposed the automatic entropy tuning for continuous action spaces, while our tasks have discrete actions. Automatic entropy tuning needs to select a target entropy. We follow the Soft Actor Critc and use $-0.98 * \log(1 / |A|)$, where |A| is equal to the number of tokens in the vocabulary. We have also added this in the caption of Figure 41.

Reviewer comment (Reproducibility):
The details for the backbone language model are scarce. It would be also interesting to discuss the design choices for language modeling, e.g., pretraining architecture (BERT vs GPT), tokenizer etc.

Author response:
> The details for the backbone language model are scarce.

In appendix A.3, in paragraph ‘pretraining models’ we have included the number of parameters and architecture specification of all the models used.

Reviewer comment (Experiments):
It would be good to see if increasing the transformer size would lead to performance improvement. If I am not mistaken increasing the RNN size could lead to issues while training and therefore this could potentially become the strength of the paper.

Author response:
> It would be good to see if increasing the transformer size helps

We agree that this would be promising. But more importantly, our work uncovers an immediate problem of reward hacking of the docking functions that needs to be solved by domain experts. This is a problem that cannot be solved by increasing the transformer size alone. In addition, to make good use of a larger tranformer, we would also need larger, more diverse data, which is currently limited in the community.

Reviewer comment (Experiments):
A most popular approach nowadays to finetuning transformers is reinforcement learning with human feedback. In RLHF, the reward preference model is trained to choose between several samples of the transformer and then the fine-tuning is performed using PPO algorithm. Have the authors considered this approach?

Author response:
> I recommend adding a quick discussion (perhaps in the appendix) regarding RLHF

In Appendix E (green color), we have added a new paragraph discussing the use of RLHF in chemistry. We have also added a line (see line 347 in green) in the future work section that mentions learning a reward model in a sample efficient way. Like we said in the initial response, learning a reward model requires a large number of labeled samples. And although unlimited samples could be obtained from the simulators, we know that these simulators aren’t perfect [3]. Our experiment showing reward hacking (Figure 4 for results and Figure 14  for images of hacked molecules) also supports this. Real chemists will only be able to give feedback on a limited number of molecules and hence we should focus on RL algorithms that are sample efficient.

Reviewer comment (Experiments):
It is not clear why the authors use Reinforce instead of more popular PPO for example. I appreciate that the authors did an ablation study on the KL constraint, but PPO has other benefits in comparison with Reinforce.

Author response:
> It is not clear why the authors use Reinforce instead of more popular PPO for example.

We add a new Figure 13 (page 22) to compare REINFORCE with PPO across 21 molecular optimization tasks.  We found that for these tasks, more complicated approaches were more unstable than REINFORCE, which seems to perform better. These results resonate with other work in the molecular optimisation community [1,2]

Reviewer comment (Reproducibility):
Details of the backbone model are not given. Is it a standard GPT model? Was the tokenizer standard or also trained?

Author response:
> The details for the backbone language model 

In appendix A, we have included Table 4 and Table 5, which contains all the details on the pretraining models.

Reviewer comment (Experiments):
Haarnoja et al 2018 proposed an automatic tuning procedure of entropy that can explicitly state the target level of entropy. It would be interesting to have this as ablation as well.

Author response:
> automatic tuning procedure of entropy that can explicitly state the target level of entropy

The logp regularization we applied is similar to entropy regularization in that both regularize the log probability of the policy. Entropy regularization reduces log(prob), while logp regularisation increases 1 / log(prob). We add this clarification in the last paragraph of section 5.4

Reviewer comment (Theory):
I am confused about referencing Mnih et al 2013 regarding the use of replay buffer for on-policy algorithms. Don’t Mnih et al use an off-policy algorithm?

Author response:
> I am confused about referencing Mnih et al 2013 regarding the use of replay buffer for on-policy algorithms

You are correct in that Mnih et al 2013 use an off-policy algorithm, while we use REINFORCE which is on policy. We use replay buffers [3] to store high-scoring molecules and re-sample them to make REINFORCE updates, making the algorithm off-policy.  In appendix A.5, we add details in our paper about how this is done in practice.

Reviewer comment (Writing):
I am not quite sure what a replay buffer for on-policy algorithm means. Does it mean that we use samples for several previous iterations not just the most recent one?

Author response:
> Does it mean that we use samples for several previous iterations, not just the most recent one?

Yes. We store a small number of previously seen high-scoring molecules and apply REINFORCE updates on them. Although this makes the algorithm slightly off policy, it vastly improves the performance (Figure 5). In appendix A.5, we add details in our paper about how this is done in practice.

Reviewer comment (Presentation):
In Figure 2 CheMBL offers high diversity and high redundancy results, which seems confusing at first. Please comment in the figure caption.

Author response:
> In Figure 2 CheMBL offers high diversity and high redundancy results, which seems confusing at first. Please see the comment in the figure caption.

We have linked the explanation in the caption of Figure 2 and added it in the main paper.  This is because diversity is measured using average pairwise Tanimoto similarity between Morgan fingerprints of the molecules, which can be high even if a small fraction of molecules are vastly different from the rest.

Reviewer comment (Theory):
It is confusing why the authors assume a discount rate of 1 in Section 4.1. Even with a finite trajectory, a discount rate smaller than 1 is not obligatory, it can also help. A discount rate of 1 can prevent the agent from learning and executing long-term tasks, as it won't appropriately discount long-term returns, leading it to prioritize immediate rewards and neglect long-term benefits. It is suggested to justify this setting.

Author response:
The primary reason why we set the discount factor as 1 is related to the problem setting. During search, it is common to set a maximum length of the molecule (100 in our experiments) that the agent can sample. Hence, the MDP is a finite-length MDP. Rewards are only obtained once in an episode, when the agent samples the stop action or the length of the molecule hits the maximum length. Having a discount factor of 1 is an intentional choice because discount factors less than one bias the agent towards molecules with shorter lengths. This bias is undesirable in practice.

Reviewer comment (Experiments):
REINFORCE is a very classic yet old algorithm. It is highly recommend to try more recent algorithms, e.g., TROP or PPO. I am not an exert in RL for molecular optimization, but I doubt whether it is the state-of-the-art RL algorithm in this field. For example, (i) due to its reliance on stochastic policies and simple optimization methods, REINFORCE can be more susceptible to getting stuck in local optima, making it less effective in complex and high-dimensional environments; (ii) due to its on-policy mechanism, REINFORCE requires a large number of samples to estimate gradients accurately. This can make it computationally expensive and slow for complex tasks and environments.

Author response:
REINFORCE is a very classic yet old algorithm. It is highly recommended to try more recent algorithms, e.g., TROP or PPO. Although PPO has generally been shown to be a better algorithm for control problems, the comparison is not clear in the case of molecular optimization ([1,2]) . We perform an experiment across 21 PyTDC objectives (see Figure 13) to compare REINFORCE with PPO. Our results indicate that vanilla policy gradient algorithms achieve higher performance on all metrics and are more stable when compared to actor-critic algorithms like PPO. This experiment guided the choice of using REINFORCE over PPO.

Reviewer comment (Evaluation):
It would be more convincing if domain experts can help to judge the effectiveness of ChemRLformer from molecular's perspective.

Author response:
We include comments of a Domain expert who has read the paper: They find the versatility of ChemRLformer in solving multiple tasks, especially expensive tasks like docking, useful for early-stage drug discovery. The reward-hacking insights are particularly useful as they show the shortcomings of docking score evaluations and the need for better evaluation, both in simulation and real-world experiments.

We will add a deeper discussion on domain expert and application perspectives in the camera-ready version.