TagBO: LLM-Driven Task-Aware Graph Bayesian Optimization for Scientific Discovery

AI summary

Overview

Research area: Machine learning for scientific discovery — specifically Bayesian optimization (BO) for mixed-variable experimental design, combining large language model (LLM) priors with graph-based surrogate modeling.

Technical level: Intermediate. Readers should be comfortable with Gaussian processes, kernels, and the explore/exploit logic of Bayesian optimization; the LLM component is conceptually simple (ranking, not prediction).

Scope: The paper proposes TagBO, a method that turns LLM-derived ordinal rankings of categorical candidates into a task-aware similarity graph, fuses it with physicochemical descriptor similarity, and runs constrained Gaussian process BO over the resulting joint discrete–continuous space, evaluated on six scientific benchmarks plus one real formulation experiment.

What This Paper Is About

Scientific optimization often requires choosing among categorical options (reagents, solvents, molecular fragments) while also tuning continuous conditions (concentration, temperature), where every experiment is expensive. Standard BO methods measure similarity between categorical choices using structural descriptors such as molecular fingerprints, implicitly assuming that structurally similar candidates behave similarly on the task — an assumption the authors argue often fails in practice, for example when two nearly identical solvents differ drastically in electrochemical stability. TagBO's goal is to replace this task-agnostic notion of similarity with one that reflects how candidates actually relate under the target task, by extracting coarse ordinal relational knowledge from LLM reasoning and blending it with structural information.

Key Contributions

1. Task-aware similarity modeling for scientific optimization. The authors propose TagBO, which builds task-aware similarity structures over categorical candidates by incorporating LLM-derived relational knowledge, rather than relying solely on task-agnostic structural descriptors.
2. Ordinal relational modeling from LLM reasoning. Instead of using LLM outputs as predictive scores, TagBO treats them as a potentially noisy source of relational priors and extracts only coarse ordinal signals (rankings of candidates along LLM-identified task-relevant factors) to define similarity.
3. Principled fusion of two similarity sources. A multi-objective formulation combines a ranking-preservation loss (from LLM orderings) with a graph-consistency loss (geometric-mean fusion of the LLM graph and the structural graph), solved with the Multi-Gradient Descent Algorithm (MGDA) so no fixed trade-off coefficient must be hand-tuned.
4. Sample-efficiency gains under fixed budgets. Across six mixed-variable benchmarks, TagBO consistently lowers the number of iterations needed to reach 95% of the oracle optimum (Step@95%), and remains effective across LLMs of varying scale and capability.

Main Findings

• Consistent sample-efficiency improvement: On all six tasks (Organic Photovoltaics, Suzuki–Miyaura Reaction, Organic Laser Molecules, Perovskite Bandgap, Redox Flow Battery, Crystallization Process), TagBO reached high-quality candidates with fewer evaluations and lower Step@95% than One-hot Feature BO, Property Feature BO, COMBO, BOPRO, and LGBO. The gains are described as most notable on the challenging discrete tasks Organic Photovoltaics and Organic Laser Molecules, and are strongest in the early and middle stages of optimization. The paper's provided text describes these gains qualitatively and via the Step@95% curves; specific percentage improvements are not reported in the content available.
• Both signals are necessary (ablation): An LLM-only variant performed strongly early but saturated later, while a structural-only variant was consistently weaker throughout. The full TagBO combination achieved strong early exploration and better later-stage refinement.
• LLM quality matters mainly later: Comparing Intern-S1-241B, Llama-3.3-70B, Qwen3-30B (Instruct and Thinking), Qwen3-4B, and Llama-3.2-3B, stronger models generally gave lower Step@95%, with Intern-S1-241B Thinking reaching 95% of the oracle optimum in the fewest evaluations. Early-search trajectories were relatively close across backbones, which the authors read as evidence that the ordinal formulation is not highly sensitive to a specific LLM's numerical calibration.
• A more optimization-friendly landscape: A 3D UMAP visualization of the Suzuki reaction space showed TagBO's learned representation produces a more connected, smoother high-performing region, with high-yield candidates more locally clustered and better neighborhood continuity than under one-hot or property features.
• Real-world validation on an adhesive formulation: Optimizing a thermally conductive adhesive over 10 fillers chosen from 29 candidate materials plus 10 continuous composition variables, the best measured thermal conductivity rose from 0.1823 W/(m·K) at the first evaluation to 2.995 W/(m·K) after 34 physical evaluations, exceeding the 2.8 W/(m·K) qualification threshold. The final formulation reached 3.36 MPa tensile strength, 2.92 MPa shear strength, and 23.5 kV/mm dielectric strength, satisfying all three requirements (≥1 MPa, ≥2 MPa, ≥7.9 kV/mm). The authors note that the standard TuRBO initialization of 2d would require 40 evaluations for a 20-variable problem — more than the entire budget used here.
• No added search cost: Because the representation coincides with the original discrete candidate set, search-space cardinality and the number of acquisition evaluations are unchanged.

Methodology in Plain English

TagBO handles each type of categorical variable (e.g., the ligand choice, the solvent choice) separately, in three stages.

Stage 1 — Ask an LLM for rankings, not scores. For a candidate set of N options, the LLM first lists F key factors likely to affect the target objective, then ranks all N candidates under each factor. Each candidate becomes a vector of ranks — a coarse, relative description of where it sits in task-relevant terms. These rank vectors are normalized into embeddings, and an RBF kernel with bandwidth γ converts pairwise distances into a similarity graph (W^llm).

Stage 2 — Add structural grounding. From an external database such as PubChem, each candidate gets a physicochemical descriptor vector of K features. An RBF kernel with bandwidth γ_str produces a second, structural similarity graph (W^str) over the same nodes.

Stage 3 — Fuse the two graphs into a refined embedding. Two objectives are optimized jointly: a ranking loss that penalizes any violation of the LLM's pairwise orderings, and a graph-consistency loss whose weight for each pair is the geometric mean of the two graphs' edge weights. That multiplicative form means smoothness is enforced mainly between pairs supported by both sources — the LLM's task-aware view and structural reality — so noisy LLM signals do not dominate. MGDA balances the two objectives adaptively instead of using a hand-set coefficient.

Stage 4 — Optimize. The refined embeddings for all categorical variable types are concatenated with the continuous variables to form a joint space, and a Gaussian process with an RBF kernel performs BO there. Acquisition maximization is restricted to the feasible set of real configurations, so only valid experimental settings are proposed.

Theoretical framing. The authors show that a representation map w applied to a GP on the latent space induces a valid GP prior on the original input space with kernel K(w(x), w(x′)), so LLM orderings act as a prior on the surrogate's similarity structure rather than as predictions. They further argue, under stated assumptions, that latent reconstruction error and surrogate approximation bounds order as LLM-based ≤ structural ≤ one-hot, and that the MGDA solution is equivalent to an adaptively scalarized objective with coefficient λ*.

Why This Matters

Impact on research. The paper reframes a common assumption in scientific BO — that descriptor proximity implies task proximity — as a modeling choice that can be wrong, and offers a way to inject literature-derived, task-specific knowledge without trusting LLM numbers. Framing LLM knowledge as a prior over the surrogate's geometry, rather than as pseudo-observations, is a reusable idea for other sample-starved design problems.

Real-world applications:

• Reaction condition screening in synthetic chemistry, such as the Suzuki–Miyaura ligand/base/solvent selection evaluated here.
• Molecular and materials discovery where candidates are discrete fragments or building blocks, such as organic photovoltaics and organic laser molecules.
• Perovskite bandgap tuning and redoxmer selection for non-aqueous redox flow batteries.
• Formulation engineering for industrial products, as in the thermally conductive adhesive experiment where a qualified product was reached in 34 physical trials.

Industry relevance. Each physical experiment in the adhesive case takes approximately three days; cutting the number of trials directly reduces R&D cost and time-to-formulation. The method's compatibility with different LLM backbones (from Qwen3-4B to Intern-S1-241B) means organizations can trade cost against prior quality rather than being locked into one model.

Future Directions

• Robustness to noisy or adversarial LLM priors: The paper bounds LLM decoding error by a parameter σ_L² in its analysis but does not characterize empirically how badly misleading rankings degrade optimization — a natural stress test.
• Extending beyond rankings: TagBO deliberately discards LLM numeric scores in favor of orderings; whether richer LLM outputs (mechanistic explanations, pairwise comparisons instead of full rankings) can improve the graph without reintroducing noise is left open.
• Scaling the real-world validation: The physical experiment used 34 evaluations on a single formulation problem; broader multi-task or multi-property campaigns would test whether the gains hold under competing objectives and batch constraints.
• Automating factor selection: The number of factors F and the choice of descriptors are inputs to the pipeline; how sensitive results are to those choices, and whether they can be selected or learned automatically, is not resolved in the reported content.

Target Audience

Graduate students and researchers in AI for Science, Bayesian optimization, and automated experimental design; computational chemists and materials scientists who run expensive mixed-variable campaigns; and industrial R&D teams in specialty chemicals, batteries, and formulated products who want to reduce physical trial counts. Readers with a working knowledge of Gaussian processes will get the most from the methodology and theory sections, while the applied results are accessible to domain scientists.

Authors’ abstract

Mixed-variable optimization in scientific discovery often requires exploring large combinatorial spaces under costly experimental evaluations, making sample?efficient Bayesian optimization (BO) essential. However, existing BO methods typ?ically rely on structural descriptors to define similarity among categorical variables, implicitly assuming that structural similarity reflects similarity in task performance. In scientific optimization, this assumption is often violated, leading to surrogate representations that are poorly aligned with the optimization objective. To address this limitation, we propose TagBO (Task-aware Graph Bayesian Optimization), a method that integrates LLM-derived relational priors with descriptor-based struc?tural information for mixed-variable Bayesian optimization. Rather than directly using LLM outputs as predictive scores, TagBO treats LLM outputs as a potentially noisy source of task-aware relational priors and extracts coarse ordinal relational signals from LLM reasoning to construct task-aware similarity structures for scien?tific optimization. TagBO then performs constrained Gaussian process Bayesian optimization over the resulting joint discrete-continuous representation space. Ex?periments on diverse chemistry and materials formulation tasks demonstrate that TagBO consistently improves sample efficiency under fixed evaluation budgets across LLMs of varying capabilities.