LICO: Large Language Models for In-Context Molecular Optimization

We thank the reviewer for their constructive feedback and for appreciating our paper's novelty, competitive performance, and clear writing. We address each specific concern below.

In some places, the paper seems to indicate that ICL for Bayesian Opt has not been done in chemistry.

We did not intend to say that ICL for Bayesian Opt has not been done in chemistry. We have added a section in the related work to discuss LLMs for molecular optimization, including the two papers that the reviewer mentioned. We discuss them briefly below.

• [1] prompts GPT to perform ICL directly in the text space while LICO embeds each molecule to one token in the feature space of the LLM and performs ICL in this space. This allows LICO to use more context points because one molecule only occupies one token compared to several tokens in the text space. Moreover, [1] obtains uncertainty estimations by using token probabilities from the LLM outputs, while LICO learns a separate prediction head to output the mean and std predictions.
• [2] similarly prompts GPT in the text space to generate new material designs and/or predict the score/property of a design, while human experts and/or a verifier provide feedback to iteratively improve the generation and prediction.

[1] Ramos, Mayk Caldas, et al. "Bayesian optimization of catalysts with in-context learning." arXiv preprint arXiv:2304.05341 (2023).

[2] Völker, Christoph, et al. "LLMs can Design Sustainable Concrete–a Systematic Benchmark." (2024).

A bit related is the use of LLM-derived embeddings for Bayesian Opt in chemistry. This, for example, has been reported in https://openreview.net/forum?id=A1RVn1m3J3

Thank you for pointing us to this work. It proposes an orthogonal direction of using LLM embeddings as molecule/chemical representations for Bayesian Opt, while we leverage the pretrained transformer backbone of the LLM to perform in-context learning. We have cited and discussed this work in the paper.

How do you obtain the uncertainties do you directly use the probabilities returned by the model? Are those probabilities well-calibrated?

The new prediction head that we attach to the pretrained LLM outputs both the mean prediction and standard deviation, and we use the standard deviation to obtain uncertainties. Figure 2 compares the predictive performance of LICO with a Gaussian Process, and the Root-mean-squared calibration error metric (RMS Cal) shows that the predicted probabilities are indeed well calibrated.

The ordering of the examples seems important - what is the impact and how is this taken care of?

This was also our thought when developing the method. However, some preliminary experiments where we ordered the context points in different ways (random shuffling, sorting according to scores, etc.) indicated that the performance is very robust to different orderings. Therefore, in all of the experiments in the paper, we keep the order of context points the same as they appear in the optimization process, i.e., append new context points to the right of the sequence.

The approach is shown for a causal LM. But it seems that a masked approach (similar to https://www.nature.com/articles/s42256-023-00639-z) might be more effective in learning from the entire sequence

While it’s true that we’re using a causal pretrained LLM, we note that the model always learns from the entire sequence when predicting a target point, because we always place the unknown target points at the rightmost position of the sequence. Moreover, we stick to causal LLMs mainly because the strongest open-source LLMs right now such as Llama are causal and we want to finetune them in the same way they were pretrained to retain the existing capabilities of the model.

We thank the reviewer again for the constructive feedback and continued support. We believe your suggestion has significantly improved the paper, and we sincerely hope that the reviewer can take our discussion into account in the final assessment of the paper.

The ablation and baseline should be more comprehensive: there are several concurrent works for LLM for molecular optimization, the authors should also add them as a baseline, if applicable

We agree that MOLLEO and ChatDrug are interesting and relevant works and we have updated the paper to discuss them. However, they are completely different approaches and are not comparable:

• MOLLEO proposed to incorporate LLMs to perform crossover and mutation operations in a standard graph genetic algorithm. As shown in Appendix E of their paper, the performance of MOLLEO largely depends on the alignment of the prompt with the model's knowledge, and the optimal prompt varies between tasks. In contrast, LICO is a general-purpose model that can optimize any objective after training without any finetuning or ad-hoc optimization.
• ChatDrug was proposed for drug editing. ChatDrug requires a specific prompt design for each task, and requires the task/objective to be described in natural language. In contrast, LICO does not require such domain expertise, and one LICO model can optimize any molecular objective without changing any component. Moreover, the ReDF in ChatDrug requires a retrieval database of molecules that is used to inject domain feedback into the LLM. To perform the injection, the method also assumes it is cheap and easy to evaluate the validity of the molecules in this database for the editing task. These assumptions do not hold in optimization tasks where we do not have a fixed database, and there is not a clear notion of validity to retrieve from such a database for knowledge injection.

Two more straightforward baselines are: (1) drop the input embedding layers, simply extract the text embedding from prompts and train an MLP layer for the surrogate. (2) drop both the embedding and prediction layers, use the pretrained model to do in context learning only.

We thank the reviewer for this suggestion. We indeed considered these alternatives before coming up with LICO. To demonstrate (2), we conducted an experiment where we prompted GPT-4o to perform in-context property prediction in the text space. The figure at https://imgur.com/a/VosqY6o compares the predictive performance of GPT-4o with LICO on 3 tasks – median1, ranolazine_mpo, and troglitazone_rediscovery, similar to the paper. We vary the context length from 32 to 512, and for each context length average the mean squared error across 128 target molecules. The result shows that prompting GPT4-o directly in the text space performs poorly, while LICO works much better and the performance improves as we increase the context length.

In addition to the empirical results, there are three strong reasons why LICO is better than (1) and (2). First, since we train separate embedding and prediction layers, LICO can be easily adapted to any new, even non-textual domain beyond molecules, while using the pretrained text embedding cannot. Moreover, training a separate embedding layer that embeds each molecule as a single token reduces the sequence length significantly. If we represent each molecule by its SMILE string, and assume each string takes 10 tokens on average, then this approach results in 10 times longer sequences, and hence can use 10 times fewer in-context examples, compared to LICO. Last but not least, using a separate prediction head allows LICO to quantify uncertainties which is essential to the optimization algorithm. We have motivated these design choices in lines 211-215 in the main text.

The description of the experimental details is a bit lacking: e.g. how the embedding and prediction layers are integrated into the existing LLM backbone?

Figure 1 in the paper depicts how we integrate the new embedding and prediction layers into the pretrained LLM backbone. The input sequence is formatted as (task prompt, “x”, x_1, “y”, y_1, “x”, x_2, “y”, y_2, …, “x”, x_n, “y”, y_n), in which x_i is the fingerprint (a 2048-dimensional vector) representation of molecule i, and y_i is the scalar score of molecule i. Given this sequence, we embed the task prompt as well as “x” and “y” using the pretrained text embedding layer of the LLM, while embedding x_i and y_i using their respective MLP embedding layers to the same dimension of the LLM. This produces a sequence of hidden vectors with the same dimension, and we feed this sequence through the transformer layers of the LLM. Finally, the prediction layer takes the feature vector of the last transformer layer and outputs prediction y^_i, which consists of the mean and standard deviation for each molecule x_i.

We thank the reviewer for the constructive feedback and for recognizing the novelty and efficacy of our proposed method. We address each specific concern below.

The authors selectively show the numbers with a different sampling budget (1k instead of 10k in the original PMO setting)

We ran LICO with 10K budget during the rebuttal and updated the paper with the new results. We also named PMO with 1K budget PMO-1K to clearly distinguish the two settings. LICO outperforms the PMO baselines but slightly underperforms Genetic GFN and Augmented Memory. We note that it was highly non-trivial to scale an in-context learning method like LICO to 10K context examples. To do this, we adopted Llama-2-7B-32K-Instruct, the long-context checkpoint of Llama-2 provided by TogetherAI (https://www.together.ai/blog/llama-2-7b-32k-instruct), and utilized the Liger Kernel (https://github.com/linkedin/Liger-Kernel) with gradient checkpointing for finetuning and inferencing with long context. We hope the reviewer will appreciate our efforts in providing this result.

We also wanted to emphasize that we performed hyperparameter tuning with grid search for each of the baselines on the first 5 tasks in PMO-1K, then used the optimal hyperparameters for the remaining tasks. The table below specifies our grid search and the optimal set of hyperparameters for each method, which differs from what was reported in the original PMO paper. We did our best to make sure the baselines achieved their best performance in this benchmark. We have included the hyperparameter tuning details in the updated submission. We hope the reviewer will recognize our efforts in ensuring a fair and honest comparison.

Different sampling budgets to confirm the generalization of the proposed approach

We believe generalization should not depend on the sampling budget but on how many tasks a model can generalize to. In the paper, we show that a single LICO model can generalize to and solve 23 different optimization tasks in the PMO benchmark, which strongly indicates the generalization of our method.

To clarify, I was asking for which type of molecular language representation, x, was used in the LLM. But now it becomes clear to me, that you are using the 2048-dimension vector as x there (also appears in line 350). Is that ECFP/Morgan fingerprints? I am not sure if I missed the details somewhere. Can you also point me to the actual location in your paper that explain the fingerprint details? right now it's unclear what the 2048-dimension vector is, but I assume it's ECFP. Please correct me if I am wrong.

Overall, thank you for your detailed responses on your thoughts of the MOLLEO and other baselines, and running new experiments for 10k-optimization budget, and adding GPT4-o in-context learning as a baseline.

However, several of my major concerns still exist: primarily of the soundness of the designed LLM approach here, as well as the effectiveness compared with the state-of-the-art methods.

We believe generalization should not depend on the sampling budget but on how many tasks a model can generalize to. In the paper, we show that a single LICO model can generalize to and solve 23 different optimization tasks in the PMO benchmark, which strongly indicates the generalization of our method.

I totally agree with your point that generalization to multiple tasks matters a lot, but I don't think testing on sampling budget is meaningless. For example, right now when extending to 10k budget as the default setting of PMO, the advantage of LICO shrinks than other methods. Furthermore, the PMO results challenge the effectiveness of the method that appears in PMO-1k. (I don't object that your method is good overall, but the trend does change.)

We agree that MOLLEO and ChatDrug are interesting and relevant works and we have updated the paper to discuss them. However, they are completely different approaches and are not comparable:

ChatDrug may be less related but in terms of MOLLEO, although it's different approach, I don't think it's not comparable if our goal is to push the frontier of PMO (MOLLEO is tested on exactly the same benchmark). Let's take a step back, if it's not comparable since it's a completely different approach, what's the point of comparing any new methods with REINVENT, GFN?

If you do compare with MOLLEO (and/or other more recent methods) and show the advantages of LICO, I believe it will be more convincing.

To demonstrate (2), we conducted an experiment where we prompted GPT-4o to perform in-context property prediction in the text space.

Thank you very much for conducting this additional experiment. Can you please add the experimentation details and settings into your paper? as well as the results.

In addition to the empirical results, there are three strong reasons why LICO is better than (1) and (2). [...]

For (1), what if we use a SMILES representation of a molecule, extract the embedding from the LLM, and then add a prediction head for the target properties? This is a simple and scalable baseline. To keep it simple, you can freeze the LLM weights and only fine-tune the prediction head. I personally think this is a must-have baseline in your work.

[..] The input sequence is formatted as (task prompt, “x”, x_1, “y”, y_1, “x”, x_2, “y”, y_2, …, “x”, x_n, “y”, y_n), in which x_i is the fingerprint (a 2048-dimensional vector) representation of molecule i, and y_i is the scalar score of molecule i. Given this sequence, [...]

Sorry I missed this information in line 359-360. Thanks for clarifying this. Your answer actually reminds me of a more fundamental question: Why do we want to combine the molecular feature vector into the LLM architecture, is that due to the flexibility of adding the task-specific prompt ahead of it? It's not very clear to me why we can benefit from pretrained LLM if the molecule representation is based on chemistry-informed feature vectors. Have you tried training random forests with these features directly for the surrogating modeling, is that bad?

I look forward to your thoughts on these.

Dear Reviewer 8R1V,

We thank you for your constructive feedback and suggestions. It's been a few days since we posted our rebuttals and we sincerely hope the reviewer has had time to read it. Per the reviewer's request, we have included two more baselines: 1) MLP/GP with LLM embeddings, and 2) MOLLEO, which we believe have resolved the two major concerns previously raised: soundness of LICO and its effectiveness compared to state-of-the-art methods.

For completeness, we have also trained random forests for surrogate modeling to compare with LICO as the reviewer suggested. The figure at https://imgur.com/a/D4aKJHx shows the mean-squared-error predictive performance of LICO vs a 100-tree RF across the first 15 objective functions in PMO with varying numbers of observations. LICO clearly outperforms RF on 12/15 objectives.

As we approach the end of the discussion period, we would greatly appreciate your feedback on our rebuttals. We are happy to address any remaining concerns or questions to improve the paper further. Thank you so much!

Kind regards, The authors of paper 12932

We thank the reviewer for their constructive feedback and for recognizing the importance, novelty, and strong performance of our work. We address each specific concern below.

It would be valuable to understand which particular characteristics of LLMs contribute to the success of this molecular optimization task.

We believe the general in-context capabilities of the pretrained LLM are what contribute to its strong performance in a new domain like molecules. In Section 5.2.3, we compared a pretrained LLM with a model of the same size (7B) but trained from scratch, and the results showed that the pretrained LLM significantly outperforms the scratch model. Another piece of evidence is Figure 3, which demonstrates that larger models with stronger capabilities also perform better in molecular optimization. These results indicate that LICO successfully transfers general in-context learning capabilities of pretrained LLMs in the text domain to a completely new domain such as molecules.

Would domain-adaptive training on chemistry corpora further enhance results?

We thank the reviewer for bringing up this very interesting point. To answer this question, we compare the performance of LICO with two different base LLMs: T5-base [1] and Nach0-base [2], which finetunes T5-base on molecule corpora. Due to time constraints, we did not perform optimization with these new models, but compared their predictive performance instead. The figure at https://imgur.com/a/v4ih18n summarizes the results. There are two interesting observations from this figure.

• Nach0 works significantly better than plain T5, confirming the hypothesis that proper finetuning of a language model on molecule data helps boost its in-context property prediction in LICO.
• Llama-2 works better than Nach0. As we explained in the paper, because we perform in-context learning in the embedding space of the language model, we rely on the general pattern-matching capability of the model, i.e., the ability to extract the relationship between embeddings of x and y from examples. From this perspective, it is not surprising that Llama-2 works better than a T5-based model, since it is much larger and has been pretrained with a lot more data, leading to a superior pattern-matching capability. One more benefit of using general LLMs like Llama is that they are domain-agnostic, which means we can finetune them for other non-molecule domains as well. We have included this result in the updated version of the paper.

[1] Raffel, Colin, et al. "Exploring the limits of transfer learning with a unified text-to-text transformer." Journal of machine learning research 21.140 (2020): 1-67.

[2] Livne, Micha, et al. "nach0: Multimodal natural and chemical languages foundation model." Chemical Science 15.22 (2024): 8380-8389.

The study offers a limited exploration of prompt formats. Further investigation into how different prompt structures might influence model performance would be beneficial

Table 2 compares LICO with different language instructions, and it shows that a more structured (but simple) prompt boosts performance. We intentionally kept the prompt format simple because 1) we want to develop a general method that is independent of the prompt, and can leverage any existing LLMs without changing the prompt to suit the LLM, and 2) it is not the main focus of our work and the current simple prompt already shows the strong performance of LICO.

It would be better to discuss several recent works for using LLMs for molecule optimization

We thank the reviewer for pointing us to these related works. We have cited and discussed them in the updated submission. We discuss them briefly below.

• DrugAssist creates the MolOpt-Instructions dataset that contains pairs of molecules and their actual property values to finetune a pretrained LLM, while LICO only uses the unlabeled molecular data. Due to unsupervised finetuning, LICO can generalize to any objective function after training, while DrugAssist can only optimize what it was finetuned on.
• MOLGEN was an LLM explicitly pretrained and prefix-tuned on molecular corpora to synthesize valid molecules, whereas LICO leverages generic LLMs like Llama for molecular optimization. MOLGEN and LICO are two orthogonal but complementary directions.
• A Sober Look at LLMs for Material Discovery paper studies two ways to use LLMs for molecular Bayesian optimization: use a pretrained LLM as a feature extractor, or finetune an LLM to serve as the surrogate. Unlike LICO, they did not study synthetic pretraining or in-context learning. Therefore, they had to finetune the LLM for each specific task, while LICO generalizes to arbitrary objective functions after training.

We thank the reviewer again for the constructive feedback and continued support. Your suggestion has significantly improved the paper, and we sincerely hope that the reviewer will consider our discussion in the final assessment of the paper.

Three other papers that could be cited and discussed

We thank the reviewer for pointing us to these relevant works. We have cited and discussed them in the updated submission. We briefly discuss them below.

• As mentioned by the reviewer, Optformer does not leverage a pretrained LLM. Moreover, Optformer operates in the raw text space, which leads to a more limited context length than LICO and does not work for non-textual domains.
• Chemlactica/Chemma is trained on a massive molecule-related corpus and works as a genetic algorithm to sample new molecules, whereas LICO is used to score generated molecules from a genetic algorithm. Similarly to Optformer, Chemlactica/Chemma also works directly in the text space. Chemlactica/Chemma and LICO are orthogonal but complementary approaches for using LLMs for molecular optimization.
• MOLLEO proposed to prompt a pretrained LLM to perform crossover and mutation operations in a standard graph genetic algorithm. As shown in Appendix E, MOLLEO’s performance largely depends on the alignment of the prompt with the model's knowledge, and the optimal prompt varies between tasks. In contrast, LICO is a more general-purpose method that can optimize any objective after training without any ad-hoc optimization.

Details of the optimization algorithm

We agree with the reviewer that missing this detail can cause confusion. We have updated the paper to move the detail in A.3 to the main text.

A minor aspect that could be considered in the future iterations: use more realistic oracles, like molecular docking.

We agree that testing LICO on more realistic oracles is an interesting and important direction. However, we believe the diversity of tasks in PMO has showcased the effectiveness and generalization of our method. We will continue adopting and adapting LICO to more oracles in future work.

We thank the reviewer again for the constructive feedback and continued support. We believe your suggestion has significantly improved the paper, and sincerely hope that the reviewer can take our discussion into account in the final assessment of the paper.

We thank the reviewer for their constructive feedback, and for appreciating the novelty and detailed experiments of our work. We address each specific concern of the reviewer below:

PMO uses 10K budget of oracle calls (as mentioned by the authors).

We thank the reviewer for this suggestion. We ran LICO with 10K budget during the rebuttal and updated the paper with the new results. We also named PMO with 1K budget PMO-1K to clearly distinguish the two settings. LICO outperforms the PMO baselines but slightly underperforms Genetic GFN and Augmented Memory. We note that it was highly non-trivial to scale an in-context learning method like LICO to 10K context examples. To do this, we adopted Llama-2-7B-32K-Instruct, the long-context checkpoint of Llama-2 provided by TogetherAI (https://www.together.ai/blog/llama-2-7b-32k-instruct), and utilized the Liger Kernel (https://github.com/linkedin/Liger-Kernel) with gradient checkpointing for finetuning and inferencing with long context. We hope the reviewer will appreciate our efforts in providing this result.

All methods from PMO, and the subsequent methods like Genetic GFN have their hyperparameters tuned for the 10K budget.

We note that we performed hyperparameter tuning with grid search for each of the baselines on the first 5 tasks in PMO-1K, then used the optimal hyperparameters for the remaining tasks. The table below specifies our grid search and the optimal set of hyperparameters for each method, which differs from what was reported in the original PMO paper. We did our best to make sure the baselines achieved their best performance in this benchmark. We have included the hyperparameter tuning details in the updated submission. We hope the reviewer will recognize our efforts in ensuring a fair and honest comparison.

PMO has 23 tasks, not 21. jnk3 and gsk3 are missing.

We excluded GSK3B and JNK3 in the original submission because we had a version mismatch problem with scikit-learn when trying to load the machine learning oracles. We fixed it and included the results for GSK3B and JNK3 for both PMO and PMO-1K in the updated submission.