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

The gist The proposed method, TagBO, integrates LLM-derived relational priors with structural information for mixed-variable Bayesian optimization to improve sample efficiency in scientific discovery;.

How it works

TagBO addresses the challenge of defining task-aware similarity structures over discrete variables by combining coarse-grained ordinal guidance from Large Language Models (LLMs) with fine-grained local structural information from external descriptors;. Rather than using raw LLM outputs as predictive scores, TagBO extracts coarse ordinal relational signals from LLM reasoning to represent task-aware relationships between categorical choices and combines these with structural information derived from physicochemical descriptors to form a joint representation space;.

The method constructs task-aware representations for discrete variables through several steps:

1. Coarse-grained LLM Guidance involves leveraging an LLM to extract a set of task-relevant factors and to characterize candidates along these dimensions for each discrete variable type, resulting in a ranking matrix R(m) which is then normalized into continuous embeddings z(m) that define a similarity matrix Wllm,(m) using an RBF kernel;.

2. Fine-grained Structured Modeling incorporates domain-specific descriptors to represent each candidate by a structural attribute vector a(m), and pairwise similarity in this space is measured using an RBF kernel to define the structural similarity graph Wstr,(m);.

3. Interpretation as Implicit Graph Fusion involves learning a refined representation that preserves the coarse relational ordering induced by the LLM while remaining locally smooth over structurally similar candidates, achieved by jointly optimizing ranking preservation loss and graph consistency loss;.

Joint Optimization Space and Constrained BO

The final optimization is performed in a unified representation space X˜, which is constructed by concatenating the refined representations of all variable types z˜ = [z(1),..., z(M)] with the continuous decision variables c, resulting in X˜ = n[z(1),..., z(M), c];. Gaussian process Bayesian optimization is then conducted over this space, with acquisition optimization restricted to the feasible set x˜t+1 = arg max x˜∈X˜ αt(x˜) ensuring that only valid configurations are considered during the search;.

Representation-Level Interpretation

The paper provides several interpretations for why task-aware representations improve performance, including:

** LLM relational priors as surrogate priors:**

The method is interpreted as transforming ordinal relational knowledge derived from LLM reasoning into a Gaussian process prior over the search space by defining f(x) = g(w(x)), inducing a valid Gaussian process on the original input space with kernel Kw(x, x′) = K w(x), w(x′);.

LLM ordinal priors and surrogate approximation:

The paper analyzes the effect of the LLM-induced representation relative to purely structural and one-hot representations, showing that if the LLM representation preserves task-relevant variation more effectively than task-agnostic structural descriptors, then DL ≤ DS ≤ DO;.

Adaptive trade-off interpretation:

The multi-objective formulation jointly optimizes ordinal relational consistency and structural graph smoothness, and under standard MGDA optimality conditions, the resulting Pareto-stationary solution z⋆ admits an equivalent adaptive scalarization ∇ (Lrank(z⋆) + λ⋆Lgraph(z⋆)) = 0, where the trade-off coefficient is determined adaptively by MGDA rather than manually specified;.

Experimental Results

Experiments on diverse scientific optimization benchmarks demonstrate that TagBO consistently improves sample efficiency under fixed experimental budgets across LLMs of varying capabilities;. In the real-world formulation experiment, TagBO successfully maximized the measured thermal conductivity to 2.995W/(m · K) after only 34 evaluations, exceeding the qualification threshold and satisfying all three performance requirements;. Furthermore, ablation studies show that combining coarse-grained relational guidance with fine-grained structural smoothness allows TagBO to achieve both strong early-stage exploration and improved later-stage refinement compared to using only LLM features or only physicochemical property features;.

Theoretical Guarantees

The paper establishes theoretical bounds showing that the latent reconstruction errors satisfy DL ≤ DS ≤ DO, which implies that the induced surrogate upper bounds satisfy UL ≤ US ≤ UO;. This ordering confirms that the LLM representation induces the smallest surrogate upper bound, while the one-hot representation induces the largest one;. Moreover, under Assumption 1 (Lipschitz task link), each Bayes risk is controlled by its corresponding surrogate upper bound, RW ≤ UW, W ∈ [L, S, O];. The adaptive scalarization induced by MGDA ensures that the learned representation emerges from balancing coarse-grained ordinal guidance with fine-grained local structural smoothness;.

The effectiveness of TagBO depends on the quality of the task-aware relational priors extracted from the LLM, and future work plans include exploring retrieval-augmented generation (RAG) to alleviate limitations in domains where relevant scientific knowledge is sparse;. The paper confirms that LLMs are core components of the proposed framework, as they are used to extract relational guidance, not just for writing or formatting purposes;. The research complies with the NeurIPS Code of Ethics;. All datasets and external codebases used in this work are properly cited, and their licenses and usage terms are respected;. The researchers provide code in supplemental materials to ensure reproducibility of the main experimental results;. The paper does not introduce new public assets such as datasets or models;. All experiments were conducted using 4 NVIDIA H200 GPUs, providing sufficient information on the compute resources required for reproduction;. The paper does not discuss both potential positive societal impacts and negative societal impacts of the work performed;. The proposed method does not introduce high-risk generative models or sensitive datasets requiring special safeguards;. The paper does not involve crowdsourcing or research with human subjects requiring IRB approval;. The LLM is used as an important, original, or non-standard component of the core methods in this research;. The paper does not use existing assets, and all datasets and external codebases used in this work are properly cited;. All reported analyses, evaluations, and conclusions were verified by the authors;. The paper reports error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments;. The paper provides sufficient information on the computer resources needed to reproduce the experiments;. The paper does not introduce new public assets such as datasets or models;. All reported analyses, evaluations, and conclusions were verified by the authors;. The LLM is used as an important, original, or non-standard component of the core methods in this research;. All reported analyses, evaluations, and conclusions were verified by the authors;. The paper does not use existing assets, and all datasets and external codebases used in this work are properly cited;. All reported analyses, evaluations, and conclusions were verified by the authors;. The LLM is used as an important, original, or non-standard component of the core methods in this research; <ref:2610.

1. Bold header: Task-Aware Similarity Modeling for Mixed-Variable BO. TagBO integrates LLM-derived relational priors with descriptor-based structural information for mixed-variable Bayesian optimization, allowing the system to construct a surrogate model that is task-aware rather than relying on task-agnostic descriptors which are not explicitly aligned with the task-specific structure of the objective.

2. Bold header: Ordinal Relational Modeling from LLM Reasoning. The system can extract coarse ordinal signals from LLM reasoning to represent task-aware relationships between categorical choices, enabling it to treat LLM outputs as a potentially noisy source of task-aware relational priors.

3. Bold header: Improved Sample Efficiency Under Fixed Budgets. The resulting optimization framework consistently improves sample efficiency under fixed experimental budgets across LLMs of varying capabilities, leading to faster convergence toward near-optimal solutions in expensive scientific discovery settings.

4. Bold header: Adaptive Coarse-to-Fine Trade-off. By using Multi-Gradient Descent Algorithm (MGDA), the system adaptively balances the two objectives, ensuring that the learned representation emerges from balancing coarse-grained ordinal relational guidance with fine-grained local structural smoothness.

5. Bold header: Representation Warping as a Gaussian Process Prior. The system can be interpreted as transforming ordinal relational knowledge derived from LLM reasoning into a Gaussian process prior over the search space, where the kernel is induced by a warping function, which induces a Gaussian process prior over functions on X.

6. Bold header: Representation-Level Surrogate Upper Bounds. The paper establishes that representations control prediction error via UW = L 2 hDW, proving that representations that better preserve the task-relevant latent variable Z⋆ induce tighter surrogate upper bounds, which in turn correspond to more accurate function models.

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.