A Sober Look at LLMs for Material Discovery: Are They Actually Good for Bayesian Optimization Over Molecules?

TL;DR

This work critically examines whether large language models truly accelerate Bayesian optimization for molecular discovery. It compares two principled surrogate frameworks: (i) using pretrained LLM embeddings as fixed features for standard Bayesian surrogates, and (ii) fine-tuning LLMs with parameter-efficient methods and performing Bayesian inference via Laplace approximations. Across eight real-world chemistry tasks and multiple LLMs, domain-specific, chemistry-trained models (e.g., T5-Chem, MolFormer) consistently outperform general-purpose LLMs, and finetuning with PEFT generally improves BO performance, yielding reliable uncertainty estimates. The study also demonstrates that a principled, Bayes-based approach with domain-adapted LLMs is cost-effective compared to in-context learning baselines, and it provides an open-source library to facilitate principled BO over discrete molecular spaces.

Abstract

Automation is one of the cornerstones of contemporary material discovery. Bayesian optimization (BO) is an essential part of such workflows, enabling scientists to leverage prior domain knowledge into efficient exploration of a large molecular space. While such prior knowledge can take many forms, there has been significant fanfare around the ancillary scientific knowledge encapsulated in large language models (LLMs). However, existing work thus far has only explored LLMs for heuristic materials searches. Indeed, recent work obtains the uncertainty estimate -- an integral part of BO -- from point-estimated, non-Bayesian LLMs. In this work, we study the question of whether LLMs are actually useful to accelerate principled Bayesian optimization in the molecular space. We take a sober, dispassionate stance in answering this question. This is done by carefully (i) viewing LLMs as fixed feature extractors for standard but principled BO surrogate models and by (ii) leveraging parameter-efficient finetuning methods and Bayesian neural networks to obtain the posterior of the LLM surrogate. Our extensive experiments with real-world chemistry problems show that LLMs can be useful for BO over molecules, but only if they have been pretrained or finetuned with domain-specific data.

Figures (17)

• Figure 1: LLMs seem to "understand" chemistry. However, they often produce completely wrong answers while sounding very convincing. Both APIs were accessed on 2024-01-07.
• Figure 2: The surrogates we consider in this work. "PEFT" refers to parameter efficient finetuning which adds a (proportionally) few trainable weights $\boldsymbol{\omega}$ to the transformer. Grey denotes frozen weights that act as conditioning variables in the posterior over the surrogate $g_t$. Green denotes weights that are inferred in a Bayesian manner (e.g., to obtain $p({\bm{w}}, \boldsymbol{\omega} \mid \mathcal{D}_t)$) and then marginalized over to obtain the posterior predictive distribution $g_t$ (e.g., $\iint p(g_t(\,\cdot\,) \mid {\bm{w}}, \boldsymbol{\omega}; {\bm{W}}_*) \, p({\bm{w}}, \boldsymbol{\omega} \mid \mathcal{D}_t) \, d{\bm{w}} \, d\boldsymbol{\omega}$). Both models are principled Bayesian surrogates, in contrast to the in-context learning frameworks considered by prior works on BO with LLMs ramos2023bayesopticlanonymous2023llambo.
• Figure 3: LLM as fixed feature extractors in BO over molecules. See \ref{['fig:surrogates']}(a) for the model schematic.
• Figure 4: Summarized performance of the results in \ref{['fig:last_layer_bo']} in terms of the GAP metric. Chemistry-focused features (T5-Chem, MolFormer, and even fingerprints) are better than general-purposed LLM features.
• Figure 5: Multiobjective BO performance in terms of the standard hypervolume metric.