Searching for High-Value Molecules Using Reinforcement Learning and Transformers

Dear reviewer,

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.

Below, we will answer your questions in more detail, and incorporate your feedback by making updates to the paper.

(Introduction) "...our own algorithm (MoLRL)" - what is MoLRL?

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

Is MolRL the general framework and ChemRLformer

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

This sentence reads a bit too colloquially.

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.”

However, the corresponding transition function induced in the graph representation of molecules is more complex as it is determined by the encoding/decoding rules of the chosen text representation

These transition dynamics are determined by the encoding / decoding rules of the chosen text representation. But the underlying constraints occur due to valencies and other chemical constraints that come into play when combining atoms and bonds to form molecules.

Please correct the typos "pertaining" and "pretrianing".

We have corrected these typos in the updates paper.

Please correct.

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

How does the "pretraining only" results (left column) are computed? Is this a zero-shot evaluation?

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.”

Dear reviewer,

We thank the reviewer for their feedback and summary of our work. We appreciate the reviewer’s assessment that the paper proposes a promising approach for molecular generation with extensive, relevant ablations. It seems like the reviewer’s main concerns are missing relevant work, presentation of the results, and examples of generated and hacked molecules. We have modified our paper to incorporate these concerns:

1. In Table 1 we add the relevant citations and a conceptual comparison with our work.
2. In Table 6,7,8,9 we add individual scores of the pretrained and RL trained agent on all PyTDC and docking tasks.
3. In Figure 14, we include molecules that are sampled as a result of reward hacking.
4. In Figure 18 - 40, we add the top 5 molecules corresponding to all 21 PyTDC tasks.

Do these new updates and experiments, with our answers to specific questions, address all the reviewer's concerns?

Some relevant works seem to be missing:

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.

The second work seems to be very similar in terms of approach

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.

I haven't noticed the comparison against any other methods.

There are not many prior works that have been applied to text-based representations of molecules. Hence our goal was to extensively compare various design and algorithmic choices. However, many of our ablations can be recombined in different ways to recover prior algorithms. In Table 1, we show how many prior algorithms can be categorized by the design choices that we compare in our experiments.

What are "Div." and "Red." columns?

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.

I couldn't find reported scores for different tasks (e.g., similarity, QED or SA).

In Table 6,7,8,9, we have added the exact individual scores achieved by just the pretrained model and the benefits achieved from reinforcement learning on these pretrained models across all 26 docking and PyTDC tasks. These tables show that the RL algorithm improves the pretrained model by 69% and 35% on docking and PyTDC tasks respectively.

Reward hacking is mentioned by the authors but including resulting molecules is essential.

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.

Below, we will answer the specific reviewer questions and weaknesses. Do these revisions and answers address all the reviewer's concerns?

What are the principal differences between your approach and Taiga?

The largest difference is the scale of our experiments. Our experiments are much more extensive when compared to Taiga:

1. Num of tasks – ours : 25, Taiga : 2.
2. Docking scores – ours : Yes, Taiga : No.
3. Largest datasets used – ours : 100 million, Taiga : 1 million.
4. Architectures compared – ours : FC, Transformer, RNN. Taiga : Transformers.
5. Important Algorithmic choices — ours : hill climb replay buffers, log p regularization. Taiga : None.

We also have other contributions that differentiate us from prior work:

1. We show that current docking functions are not perfect, and can be easily hacked by ChemRLformer.
2. We show that when choosing a pre-training datasets, the alignment matters more than its size.
3. We show that transformers aren’t superior to RNNs for current string-based RL algorithms, at the scale of data currently available.

Dear reviewer,

We thank the reviewer for their feedback. We appreciate the reviewer’s assessment that the paper is of high importance and well-written with extensive experiments and ablation studies. It appears that the reviewer's main concerns are the choice of using the REINFORCE algorithm, additional experiments, details of the model and presentation of results. We address the concerns using new experiments and updates to the paper:

1. Figure 13 compares REINFORCE with PPO across 21 tasks. This result supports our choice of using REINFORCE which achieves superior performance while being simpler.

2. We add Table 6,7,8,9 which contain individual scores of the pretrained and RL trained agent on all PyTDC and docking tasks.

3. We update appendix A to contain all details related to the experiments, including details of the size and hyperparameters of all pretraining models in appendix A.3, para “pre training models”.

Do these new experiments, with our answers to specific questions, address all the reviewer's concerns?

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.

finetuning transformers is reinforcement learning with human feedback.

While we agree that incorporating feedback from humans is a promising future direction, one reason why this will be more difficult in the space of molecular optimisation is that the humans that would be needed to give feedback are expert chemists. Given the scarcity of experts, the RLHF framework may not scale well for this use case. Indirectly, those experts are performing this already by creating simulations (e.g. Autodock) which we use in our paper.

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.

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

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.

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.

In RLHF, the reward preference model

Learning a reward preference model needs a large amount of expert human feedback. Rewards are limited in real-world chemistry, unlike general language chatbots. We need methods that use minimal rewards. Hence, we focus on more sample efficiency [1].

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]

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.

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.

@article{christodoulou2019soft,
  title={Soft actor-critic for discrete action settings},
  author={Christodoulou, Petros},
  journal={arXiv preprint arXiv:1910.07207},
  year={2019}
}

Dear reviewer,

We thank the reviewer for their feedback. We appreciate the reviewer’s assessment that the paper addresses an important problem with a new algorithm, is well-written and provides an extensive analysis. It appears that the reviewer's main concerns are regarding the discount factor, the choice of REINFORCE and a shortage of domain experts' perspective on the effectiveness of ChemRLformer. We have incorporated new updates Section 4 (footnote on page 5 explaining the use of discount factor 1) and experiments :

1. Figure 13 compares REINFORCE with PPO across 21 tasks PyTDC. This result supports our choice of using REINFORCE which achieves superior performance while being simpler.

2. In Appendix D, we added remarks by domain experts on ChemRLformer.

Do these new experiments, with our answers to specific questions, address all the reviewer's concerns? We address all of these concerns in detail below:

It is confusing why the authors assume a discount rate of 1 in Section 4.1

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.

REINFORCE is a very classic yet old algorithm.

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.

It would be more convincing if domain experts can help to judge the effectiveness of ChemRLformer from molecular's perspective.

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.

[1] Ciepliński, Tobiasz, et al. "Generative Models Should at Least Be Able to Design Molecules That Dock Well: A New Benchmark." Journal of Chemical Information and Modeling (2023).

[2] Gao, Wenhao, et al. "Sample efficiency matters: a benchmark for practical molecular optimization." Advances in Neural Information Processing Systems 35 (2022): 21342-21357.