Small Molecule Optimization with Large Language Models

Thank you for the review, which raised concerns about baselines, modeling decisions, data quality, and the proposed method’s robustness. We address the points in turn below.

Inclusion of Baseline Comparison with Similar Methods

The reason we omitted a comparison with "Efficient Evolutionary Search over Chemical Space with Large Language Models” in our work is that the paper did not provide an evaluation of all of the tasks from the practical molecular optimization benchmark, making a direct and fair comparison with the remaining methods difficult. We provide below a direct comparison for the tasks that they provided in their analysis:

Link to the table

Our largest model Chemma-2B is on par with the MolLEO based on the closed GPT-4 model (the difference in the total score is less than the standard deviation), while even our smallest model Chemlactica-125M is significantly stronger than MolLEO based on the open models. With a closer look we see that our approach is dominating in multi-property optimization tasks, while MolLEO on GPT-4 is better at similarity/rediscovery and SMARTS based tasks.

Model Selection and Justification

We acknowledge that the robustness of the method’s performance to the underlying model was not extensively evaluated, primarily due to the computational cost of optimizing, tuning, and executing the continued pretraining protocol for the models, which were trained on dozens of billions of tokens. This method shows strong performance in both Gemma and Galactica architectures. We selected the models Galactica and Gemma as a basis of our work based on a combination of domain-specific knowledge, generally strong performance on standard benchmarks, as well as availability of variants with under 7B parameters. We describe this reasoning at the start of Section 4. The pretraining of the Galactica models included a lot of domain-specific data of interest in that it had been trained on a scientific corpus and had already seen SMILES representations enclosed in special tags, which we leverage in our optimization algorithm. Additionally, the Galactica models were implemented using the OPT architecture, which is virtually identical to GPT3. Thus, the Galactica models primarily differ from GPT in training data composition and pretraining methodologies, which abates concerns regarding the generalizability of our findings with respect to model architecture suitability. The Gemma models were not pretrained on a corpus of particular relevance to our work but demonstrated the strongest results on general-purpose language modeling benchmarks for their size. Given that LLama3 was not available at the time of our pretraining, and the Llama2 variant with the fewest parameters was LLama-2-7b, the Gemma-2b model was the strongest model available within our compute budget, (as we train on ~40B tokens). If there are other questions or points of clarification, we would be happy to elaborate on our reasoning.

Risk of Data Leakage and Benchmark Validity

The model can generate molecules that it has already been exposed to in its training data during optimization, and we recognize this as a limitation of our evaluation. However, despite the fact that PubChem is biased towards molecules of general pharmaceutical interest, the model did not, in its pretraining data, observe evaluations of the oracle scoring functions that are used throughout the molecular optimization tasks. This means that this mapping is learned during optimization.

In Figure 2, we demonstrate that our method generates molecules with increasing distance from our training dataset (PubChem) as the molecular optimization process progresses, demonstrating that in order to reach a high reward, the model must generate molecules that do not exist within the pretraining data. Current benchmarks do not explicitly measure the novelty of generations with respect to the training data. While we believe it would be a nice addition (along with many other improvements of the existing benchmarks), we leave the design of more sophisticated benchmarks to future work.

Task Selection and Generalizability

To clarify, we perform evaluation against all 23 tasks within the PMO benchmark, with the sum total comparison with prior SOTA methods presented in the rightmost column of Table 3. Table 10 in Appendix Section A.9 presents the results of all three models across all 23 tasks in detail. We additionally demonstrate results on two suites of docking multi-property molecular optimization tasks.

We made an effort to clarify our results, provide additional comparisons and underline our motivations. If this response has been to your standard, please consider revising your review score. We are also open to additional discussion if there are outstanding questions, points of unclarity or criticism. Thanks again for the review.

We greatly appreciate the thorough review. We provide our responses to your numbered inquiries in the corresponding order.

1. .

   1. We acknowledge the challenge of disentangling the benefits of large-scale data, model architectures, and optimization methodology. The datasets involved — ZINC 250k, ChEMBL, and PubChem—differ significantly in size, diversity, and representational richness. For instance, prior analysis has shown that PubChem contains 1,281,788 unique Murcko ring with linker scaffolds, which exceeds the number of molecules in Zinc250K by an order of five times [1]. Supported by the observed positive relationship between training data scale and conditional generation performance in Table 1. (Chemma-2B trained on 2.1 billion vs 39 billion tokens), we hypothesize that this increased diversity contributes significantly to our results, although this effect is inextricably connected with our approach’s capacity to leverage this scale of data effectively.
   2. There are two ways we could entangle the data and the model architecture. We could try to (i) train Galactica and Gemma on a corpus produced from ZINC-250K, and (ii) replace ZINC-250K with PubChem in the baselines. For (i) we do not expect the models to learn any meaningful notion of similarity with ~500x less tokens, which is critical for the genetic algorithm. For (ii) we expect most of the small scale models will saturate quickly and will not be able to leverage 110M+ molecules. (e.g. Augmented Memory uses a 3-layer LSTM with less than 2M parameters). This implies that both “dimensions” play a critical role. A similar expectation comes from the intuition from the scaling laws that suggests proportional scaling of data and network size is needed to maintain optimality. This implies that scaling only one of the two (i.e. data or network size) by 2.5 orders of magnitude would lead to suboptimal results. Proper experimental verification of these expectations would require extensive tuning of hyperparameters, which we thought is beyond the scope of this work.
   3. To more directly address questions related to the optimization method, the comparison with the MolLEO baseline leveraging GPT-4 shown below demonstrates the strength of our approach in balancing model, data, and algorithmic scalability. Despite GPT-4 having a larger pre-training corpus and parameter count than our models, our method leveraging the 1.3B parameter Chemlactica model outperforms it on most tasks they report. Link to the table comparing MolLEO and our models
   4. In conclusion, thoroughly evaluating molecular optimization methods requires scaling pre-training datasets and adapting models to handle these scales. To this end, we encourage future work to expand on our findings and provide all pre-trained models and datasets for public use. We hope our work inspires further exploration of how different datasets, model architectures, and optimization algorithms interact.

2. We executed an equivalent hyperparameter search for conditional generation using the synthetic accessibility score (SAS) to compare to the results we presented on weight in the main text. The best-performing hyperparameters were the same as for weight. These results suggest that this hyperparameter configuration is at least near-optimal for conditional generation on simple molecular properties. Note that “DNF” is an acronym for did not finish, as the model either repeatedly generated the same substring or continued to generate invalid smiles such that the trial did not finish in a reasonable time. The remaining runs spanned three to four minutes, and the DNF trial exceeded 30 minutes of runtime when it was manually terminated.

   Link to the hyperparameter search table for SAS-conditioned generations

[1] Velkoborsky, J., & Hoksza, D. (2016). Scaffold analysis of PubChem database as background for hierarchical scaffold-based visualization. Journal of Cheminformatics, 8, 1-14.

Thank you for a considerate and thorough review, we address the limitations and questions asked in separate parts below.

Lack of Computational Efficiency Comparison

Despite the potentially high computational cost of LLM inference and the compute requirements of other molecular optimization methods, applications of computation methods in drug discovery are not typically constrained by computational cost. This is because the computation required to generate candidates is assumed to be less costly than the oracle, which would be a wet lab test in industrial use cases. Because oracles are prioritized over computational cost, the benchmarks included in the present work focus on sample efficiency rather than FLOP optimality or compute wall time.

To provide some concrete numbers, a single docking experiment (with one seed) using Chemlactica-125M on a single 40Gb A100 GPU takes approximately three hours. Notably, the docking simulation consumed 80% of this time, which is independent of our proposed method (assuming synchronous flow), meaning model inference and training alone required ~40 minutes. We do not provide a comparison with other methods as they did not make this information publicly available, and a comprehensive evaluation of the compute costs of molecular optimization methods (via our re-implementation) is beyond the scope of our work. We will add these numbers to the manuscript.

Finally, as was correctly identified in the concern addressed by the subsequent section, an advantage of our method is that it leverages large language models and thus inherits a broad range of techniques for more efficient training and inference, including different inference engines, parameter finetuning methods and quantization, among others. Thus, should computational cost become a bottleneck, our approach has an array of promising directions for further exploration to ameliorate such issues.

Exploration of Efficient Training Methods

As motivated in the previous section, this kind of analysis is generally beyond the bounds of this work, in which we primarily seek to demonstrate the effectiveness of a novel molecular optimization algorithm developed on continually pre-trained LLMs. Even so, we analyze the effects of running optimization with mixed precision, as it is one of the most ubiquitous methods used for more efficient training and inference. We show that lower precision training and inference can sometimes impede performance. Given our specific use case in which we have sufficient tokens and are not concerned with computational cost, it is preferable to do a full training of the entire model instead of using optimized methods, as they almost always incur a cost in task-specific metrics.

To explore this point further, we extended the experiment on the rat liver microsome (RLM) SFT, where we compared the performance and training time of the model with full finetuning and freezing 75% of the layers. The run time was decreased by 8.3% of the full finetuning. However, the model's performance was significantly worse, demonstrating poorer performance than the adme-fang-RandomForestRegressor, which all of our fully finetuned models outperformed. There appears to be a clear tradeoff between efficiency and performance in molecular-specific adaptation during finetuning. This result further supports the abovementioned analyses in the text on training and inference in half-precision.

Limited Exploration of LLM Capabilities for Multi-Property Optimization (MPO)

Our method should extend to the explicit optimization of multiple properties through their enumeration in prompts. We provide limited evidence in Appendix Section A.9, which shows that including individual learned properties in prompts during optimization improves performance on most task-model combinations we tried.

In the current benchmarks for molecular optimization, the reward scores are implemented as scalarized combinations of individual scores. Thus, we do not alter the design of our method to account for potential access to multiple properties. Additionally, we do not look towards the challenges in the so called “multi-objective optimization” setting, such as determining Pareto optimality. We acknowledge that better benchmarks in future might consider these aspects as well.

We thank you again for your review focused on the potential for computational efficiency and flexibility in the LLM paradigm. We trust that our responses have adequately addressed your concerns with reasoning, data and experiments, as necessary. We hope that our revisions merit a positive re-evaluation of our work.

Thanks for raising considerations regarding the scope of our work and concerns about method’s effectiveness and novelty. We’ll address your concerns point by point:

1. It is true that the model only considers the SMILES representation of molecules rather than any explicit modeling of the 3D conformation state of the molecule. There are two extenuating factors explaining this. Firstly, the primary goal of the method is to optimize for “better” molecules within chemical space rather than conformations. Secondly, the reason for this direction is that the currently available benchmarks for evaluating molecular optimization methods are all SMILES-based. To elaborate, none of the practical molecular optimization (PMO) benchmark tasks are conformation sensitive, and all docking simulations change the input conformation during pose estimation. This lack of consideration for conformations is a limitation that applies to the current state of this field, and we agree that better, 3D-aware benchmarks would help guide more practical research in this space. However, designing such benchmarks requires careful oracle design, which ensures accuracy, generalizability, and computational feasibility; thus, the creation of novel benchmarks for molecular optimization is beyond the scope of this work.
2. Further improvements in sample efficiency would make the proposed method more applicable to industrial use cases, where oracle scores are more computationally costly. We believe this work is a step in exactly that direction: it presents a novel method that achieves state-of-the-art results across several molecular optimization tasks. For certain benchmarks (e.g. docking on DRD2, ACHE, and MK2) we got 3-4x improvements of sample efficiency compared to the prior art.
3. In a drug-discovery campaign, an LLM in-the-loop method should be used with wet lab testing as the oracle. We recognize that lack of wet-lab validation is a limitation of the present work, but lack of infrastructure and funding prohibits it. In addition, this would preclude us from comparing with other methods, none of which performed this kind of evaluation and analysis.
4. We recognize concern about the sensitivity of models to hyperparameter selection and acknowledge it as a limitation of the work. This aspect of the method's performance is explicitly presented via multi-seed runs and visual representations of optimization trajectories. Additionally, note that hyperparameters selected based on the two PMO tasks (zaleplon_mpo and perindopril_mpo, consistent with the PMO paper) were then used for the remaining PMO tasks and the DRD2, ACHE, and MK2 docking MPO tasks. Similarly, the illustrative experiment from Saturn was used to select hyperparameters for the remaining tasks in the benchmark. These results demonstrate that although the proposed method's performance is sensitive to hyperparameters, hyperparameter selection transfers well across diverse property optimization tasks.
5. Despite the existence of various methods for molecular optimization using large language models, which we enumerate in our related work section, our model demonstrates novelty in that it is the only work to combine iterative LLM finetuning, genetic algorithms, and discrete prompt tuning. In addition, it is the first work (as far as we know) that combined iterative language model training with evolutionary methods. The algorithm's novelty stems from combining a suite of disparate methods into an end-to-end pipeline that demonstrates strong performance in a variety of scenarios.
6. We would be happy to amend our related work upon recommendation of any relevant literature that we have omitted. The paper currently includes an extensive enumeration of related literature, covering large language models for chemistry, adjacent methods in molecular optimization, and work at the intersection of these domains.

We hope we addressed all of the concerns you raised to your satisfaction. If that is the case, we would ask to adjust the review score accordingly. We are open for more questions and feedback. Thanks again.

Thank you for your response, which addressed part of my concerns and led me to increase my score. However, I still have some concerns regarding novelty and references.

While your work introduces a molecular optimization algorithm and claims to be the first to combine "iterative" language model training with evolutionary methods as the main contribution, I noticed that a recent study, LM-Design [1], also utilized LLM and iterative refinement for protein design. This overlap in approach raises questions about the novelty of your work. This LM-Design work is missing reference and discussion.

Furthermore, in the Molecule Optimization Techniques, I found that other recent studies which seems quite relevant but not referenced or discussed, such as DrugImprover [2] and ReLeaSE [3], a RL-based molecule optimization algorithm, and DrugAssist [4], a Large Language Model for molecule optimization algorithm.

These omissions limit the clarity of your work's contributions. If you address these above concerns and properly discuss the relevant studies in your updated version, I will consider to further increasing my score.

[1] Structure-informed Language Models Are Protein Designers, ICML 2024 [2] DrugImprover: Utilizing Reinforcement Learning for Multi-Objective Alignment in Drug Optimization, NeurIPS 2023 Drug Discovery Workshop [3] Deep reinforcement learning for de novo drug design, Science 2018 [4] DrugAssist: A Large Language Model for Molecule Optimization, arXiv 2023