Context-Aware Kernel Search for Bayesian Optimization with Large Language Models

Authors: Richard Cornelius Suwandi, Feng Yin, Juntao Wang, Renjie Li, Tsung-Hui Chang
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
Keywords: bayesian optimization, gaussian processes, kernel design, large language models

Abstract:
The efficiency of Bayesian optimization (BO) relies on careful selection of the surrogate model to balance exploration and exploitation under limited budget. Traditional BO methods often struggle with sub-optimal kernel choices when using Gaussian processes (GPs) as the surrogate model. When the kernel is inadequately chosen, BO may converge slowly or even get stuck at an undesired local minimum. To address such drawback, we propose the novel Context-Aware Kernel Search (CAKES) to automate optimal kernel design in BO with large language models (LLMs). Concretely, CAKES exploits LLMs as crossover and mutation operators to adaptively generate and refine GP kernels based on the observed data. CAKES works entirely in-context and can be easily integrated into existing systems without requiring any fine-tuning. We further present a theoretical analysis demonstrating that our method achieves sub-linear regret relative to the budget for any input dimension. Experimental results demonstrate that CAKES outperforms various salient baseline methods in numerous synthetic and real-world optimization tasks. Notably, CAKES improves the overall performance on benchmark functions by roughly 36\%. In hyperparameter tuning tasks, CAKES can effectively leverage fewer data samples to quickly identify high-performing configurations and consistently ranks first across various datasets. As an encouraging real application, we successfully applied CAKES to design photonic chips,  achieving significant improvements in key performance indicators while speeding up the design cycle by a factor of ten compared to the baselines. Our code is accessible at https://github.com/cakes4bo/cakes.

Decision: withdrawn
Reviewer scores: 5, 6, 5, 8

Reviews and discussion:
<thread>
creation date:
1730107854492

thread_type:
official_review

<thread_object>
creator: reviewer_bsdq

<content>
summary:
Gaussian processes play a key role in Bayesian optimization, where the quality of the model heavily depends on the selection of the kernel type. This paper aims to address this selection issue with the assistance of large language models. Extensive experiments are conducted to verify the superior performance of their proposed framework.

strengths:
The topic is compelling, as kernel design is crucial in Gaussian process modeling. The paper is well-written and easy to follow. The method's performance is excellent, potentially reducing the selection cost for BO practitioners.

weaknesses:
- The presentation of the algorithm places too much emphasis on kernel search. However, BO is just one application of kernel design in Machine Learning. Therefore, I recommend the authors refine the presentation of this BO-specific method.
- The theoretical analysis is too simplistic and could apply to any adaptation method of kernel design, such as symbolic regression. LLM-related analysis may better describe the performance.

questions:
- I am curious whether the setting of weights in photonic chip design, where the last three objectives are much more important than the first two, is a commonly recognized practice or carefully designed. This is particularly relevant as CAKES performs outstandingly well on the fourth objective but only fairly on the other objectives.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730302977562

thread_type:
official_review

<thread_object>
creator: reviewer_bro2

<content>
summary:
This submission proposes a novel context-aware kernel search (CAKES) for constructing surrogate models in Bayesian Optimization. CAKES replies on LLM to operate the proposed crossover and mutation method for kernel design. Evolutionary approach is adapted to construct new kernels for the GP models. Some theoretical analysis is presented to show a sub-linear regret bound of the proposed method. Experiments on benchmarks and real-world applications are performed to show CAKES outperforms other surrogate models in BO.

strengths:
The paper is well structured and fluent to read. The idea of CAKES is well explained in this submission. Motivation for a novel surrogate model design method in BO is convincing. The usage of a developed and trained LLM without any fine-tuning in the few-shot learning setting of BO is new. Experiments consider both synthetic functions and real-world application that shows practical potential of CAKES.

weaknesses:
1. This submission does not present adequate contributions in my opinion. The idea of kernel design using transformers and language models has been studied in the field. For example, Simpson et al. (https://proceedings.neurips.cc/paper/2021/hash/56c3b2c6ea3a83aaeeff35eeb45d700d-Abstract.html) proposed their transformer-based architecture for kernel design, and their model is specifically trained using a vocabulary of known kernels. The authors claim the "no need for fine tuning" as an advantage for their method, but feels more like the key weakness to me. Considering all the conditioning (line 193-194, 198-199) and numerical information (such as observations and kernel's characteristic) are fed to the LLM model purely through text prompt, CAKES's performance highly (or even entirely) depends on how familiar the chosen LLM model is with GP and kernel design. In addition, the crossover and mutation operators (line 227-234) utilize very standard kernel grammar and simple replacement, which are already well-developed techniques. 

2.  Experiments design is not very consistent. In section 6, the authors use 5 different metrics to measure the performance of methods in different experiments. Given the BO settings, I don't see a good reason for varying performance metric between different experiments when the authors could just simply select one. Simple regret is commonly used and more importantly consistent with the regret bound proven in section 4. The "average rank" metric presented in Table 2 seems to highlight the same information as the "average regret" plots in Figure 2.

questions:
1. In line 406-408, the authors mention SE and Matern-5/2 as the two most commonly-used kernels. Why is Matern-5/2 kernel not included in the four base kernels defined in line 203 then?

2.  How do the authors define the values of number of crossovers $n_c$, probability of mutation $p_m$ and population size $n_p$ (line 163-165 in Algorithm 1) in final implementation?

3.  In line 235-236, the authors mention a brief analysis will be returned by the LLM model. Can the authors provide one actual example of such analysis?

4.  I'm a bit confused by the weighting strategy explained in line 242-243. Does it mean the most ideal kernel with very low BIC (thus low weight) and high acquisition value might not be selected while some worse kernel (higher BIC, lower acquisition value) might? Could the authors further justify the intuition behind such strategy?

5.  In Figure 2, what is "Best" (the fourth method in legend)? Does it refer to the "Utility" method in line 411?

6.  I notice the authors set the maximum number of function evaluation $T$ and number of replications using different random seeds quite small for their experiments. For the benchmark functions (section 6.1), $T = 10 d \leq 50$ since $\max d = 5$ for the chosen test functions. For the hyperparameter tuning tasks (section 6.2), $T=25$. Usually the value of $T$ will be larger (e.g. $50$, $100$) in BO experiments. Similarly the number of random seeds for each experiment (10 for section 6.1, 5 for section 6.2 and not mentioned for section 6.3) is relatively less than what people usually use (e.g. $20$, $25$). How many random seeds did the author use for experiment in section 6.3? Is the reason for such not very generous setting due to long generation time of kernel design?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730645838671

thread_type:
official_review

<thread_object>
creator: reviewer_whhf

<content>
summary:
The paper presents Context-Aware Kernel Search (CAKES), a novel method for optimizing Bayesian optimization (BO) by leveraging large language models (LLMs) to automate the design of Gaussian process (GP) kernels. CAKES addresses the challenge of sub-optimal kernel selection, which can hinder the efficiency of BO, by using LLMs as genetic operators to adaptively generate and refine kernels based on observed data. Theoretical analysis shows that CAKES achieves sub-linear regret, and experimental results indicate it outperforms existing methods across various optimization tasks, including hyperparameter tuning and photonic chip design, significantly improving performance and reducing design cycle times.

strengths:
One of the key strengths of this paper is its innovative use of large language models (LLMs) for in-context learning to improve Bayesian optimization (BO). The Context-Aware Kernel Search (CAKES) method leverages LLMs' few-shot learning capabilities to adaptively generate and refine Gaussian process kernels based on observed data, without requiring any fine-tuning. 

Theoretical analysis indicates that CAKES achieves sub-linear regret relative to the budget for any input dimension. Experimental results demonstrate its superiority over baseline methods in various optimization tasks, including benchmark functions and hyperparameter tuning and photonic chip design.

weaknesses:
Limited testing in high-dimensional spaces: The paper does not demonstrate the effectiveness of CAKES on high-dimensional optimization problems. This leaves uncertainty about how well the method scales to more complex search spaces with many variables.

questions:
Same in weaknesses part
</content>
</thread_object>

</thread>
<thread>
creation date:
1730656309222

thread_type:
official_review

<thread_object>
creator: reviewer_4qnv

<content>
summary:
- This work proposes a large language model-based kernel design method for Bayesian optimization.
- Since the performance of Bayesian optimization depends on kernel choice, the kernel search is crucial.
- This work utilizes a large language model to select kernel design.
- The authors test their method in diverse benchmark functions.

strengths:
- It shows an ability of large language models to understand natural language description in the problem of kernel search.
- Paper is generally well-written.

weaknesses:
- More real-world benchmarks are missing.
- Some baseline results of photonic chip design might be missing.
- More recent baselines are missing.

questions:
- My concerns are all related to experiments.
- I think that hyperparameter tuning tasks are involved with small simple models.  If more sophisticated models are tuned, it would be better.
- I don't know why only two baseline methods are compared to CAKES in photonic chip design.
- Some baselines such as (Malkomes and Garnett, 2018) should be compared.
- I think that few-shot learning is not correct in this context.  The authors didn't fine-tune large language models.  It should be few-shot prompting.
- Why did you use gpt-4o-mini only?  If other large language models such as gpt-4o and llama-3 are used, can the empirical results change?  And why did you choose a temperature of 0.7 and a top_p of 0.95 specifically?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732426487881

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
New Ablation Study

comment:
We thank the reviewers for their time and effort in reviewing our work. In response to the reviewers' feedback on our prompt design and LLM-related analysis, **we have conducted a new ablation study in Section F.3** to systematically evaluate the impact of each component in our prompt. Specifically, **we have highlighted the key components of our prompt**, which include:
- **Domain information:** This includes the role assignment as well as description of kernel functions and their properties.
- **Problem context:** This includes the optimization history, which consists of the observations as well as the fitness values.
- **Task instruction:** This includes the guidelines on how to generate the kernels based on the kernel grammar.

Based on the above components, we conduct an ablation study using the following configurations:
- **CAKES:** This is the vanilla CAKES setup employed in our experiments, incorporating all three components.
- **CAKES-NoInfo:** This variant assesses the impact of domain-specific information on performance. Specifically, it omits the role assignment and kernel descriptions while keeping the other components (see Figure 15).
- **CAKES-NoContext:** This variant evaluates the importance of problem context by removing the optimization history (i.e., the observations and fitness values), while keeping the other components  (see Figure 16).
- **CAKES-NoInstruct:** This variant excludes task instructions and omits the explicit kernel generation guidelines (i.e., how to use the crossover and mutation operators), while keeping the other components (see Figure 17).

Our ablation revealed that:
- **Standard CAKES setup outperforms the other variants**, highlighting the importance of each prompt component in enhancing overall performance.
- **CAKES-NoInfo demonstrated similar optimization performance to CAKES** despite lacking domain-specific information (i.e., the role assignment and kernel descriptions in the prompt). This suggests that the **LLM may already possess valuable prior knowledge**, which is likely derived from its pre-training on a vast amount of statistical information and code.
- **CAKES-NoContext exhibited worse performance than CAKES**, indicating that **the optimization history is crucial** and that the LLM effectively leverages this information. 
- **CAKES-NoInstruct performed poorly compared to the other variants**. To investigate this, we examined the percentage of valid kernels, defined as the proportion of kernels that satisfy the kernel properties. We found that **CAKES-NoInstruct generated valid kernels only 68.72% ($\pm$ 0.45%) of the time, while the other variants consistently produced valid kernels**. This is due to **the absence of the generation guidelines from the kernel grammar**, which underscores its importance in the kernel generation process.

We hope this additional study enhances the quality of our paper and strengthens its case for acceptance. We remain open to any further feedback from the reviewers.


With thanks,

Authors of #5709
</content>
</thread_object>

</thread>
<thread>
creation date:
1732635498383

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
New LLM-Related Analysis

comment:
We thank the reviewers once again for taking time in reviewing our submission. In response to the reviewers’ comments and requests to justify our choice of LLM and provide more LLM-related analysis, **we have conducted an additional experiment in Appendix F.4 to analyze the impact of the LLM choice** on the performance of our method. Specifically, **we consider the following state-of-the-art models**: Meta's `llama-3.1-8b`, Google's `gemini-1.5-flash`, Anthropic's `claude-3.5-haiku`, and OpenAI's `gpt-4o-mini`. We evaluate these models on the benchmark functions considered in Section 6.1.

Based on the results presented in Table 4, we can observe that `gpt-4o-mini` **outperforms the other models** in 8 out of 15 test functions. `claude-3.5-haiku` also demonstrates **competitive performance**, securing the best results in 6 out of 15 test functions. Moreover, as of November 2024, `gpt-4o-mini` is priced at 0.15 USD per million input tokens and 0.60 USD per million output tokens, making it **significantly more affordable** than `claude-3.5-haiku`, which costs 1 USD per million input tokens and 5 USD per million output tokens. In terms of speed, `gpt-4o-mini` delivers an output speed of approximately 104.5 tokens per second, **outperforming** `claude-3.5-haiku`'s 64.3 tokens per second. 

Considering the above metrics, we choose `gpt-4o-mini` as the default model in CAKES as it offers **an excellent balance of cost, speed, and intelligence for our implementation**. We hope that this additional analysis enhances the comprehensiveness and overall quality of our work.

With thanks,

Authors of #5709


| Function                | llama-3.1-8b | gemini-1.5-flash | claude-3.5-haiku | gpt-4o-mini   |
|-------------------------|---------------|-------------------|-------------------|----------------|
| Ackley-2                | 0.1019        | 0.0916            | 0.0812        | **0.0783**     |
| Ackley-5                | 0.2106        | 0.1907            | **0.1621**        | 0.1732     |
| Beale                   | 0.2621        | 0.2595            | 0.2640            | **0.2565**     |
| Branin                  | 0.0131        | 0.0199            | 0.0091        | **0.0070**     |
| Dropwave                | 0.4941        | **0.4640**        | 0.4878            | 0.4690     |
| Eggholder               | 0.1323        | 0.1309            | **0.1156**        | 0.1241     |
| Griewank-2             | 0.0337        | 0.0360            | 0.0289       | **0.0267**     |
| Griewank-5             | 0.0301        | 0.0247            | **0.0158**        | 0.0185     |
| Hartmann                | 0.0012        | 0.0005        | **0.0001**        | **0.0001**     |
| Levy-2                  | 0.0459        | 0.0426            | 0.0368        | **0.0353**     |
| Levy-3                  | 0.0710        | 0.0667            | **0.0480**        | 0.0505     |
| Rastringin-2           | 0.3595        | 0.3611            | **0.3322**        | 0.3341     |
| Rastringin-4           | 0.3301        | 0.3147        | 0.3253            | **0.3128**     |
| Rosenbrock              | 0.0623        | 0.0501       | 0.0507            | **0.0483**     |
| Six-Hump Camel          | 0.1103        | **0.1002**        | 0.1073            | 0.1015     |
</content>
</thread_object>

</thread>
<thread>
creation date:
1733215258423

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Rebuttal Summary

comment:
We thank the reviewers for their constructive feedback and for recognizing the contributions
of our work. 

We are encouraged that the reviewers found CAKES to address a **compelling
topic** in Gaussian process modeling and Bayesian optimization (BO), with its kernel design
deemed **crucial** (`Bsdq`, `BRo2`, `4qNv`). Our integration of large language models (LLMs) for kernel
search was appreciated as **innovative**, with the usage of a pre-trained LLM without fine-tuning
considered **novel and impactful** in the BO setting (`BRo2`, `wHHf`). We are pleased that the reviewers
found the paper **well-written, easy to follow, and systematic** in its structure, with a **strong motivation** for
the proposed surrogate model design (`Bsdq`, `BRo2`, `4qNv`). Experiments demonstrating CAKES'
effectiveness across synthetic benchmarks and real-world applications, such as hyperparameter
tuning and photonic chip design, were highlighted as showcasing its **practical potential** (`BRo2`,
`wHHf`). Theoretical analysis showing sub-linear regret relative to the budget further reinforced
its **strength** (`wHHf`).

We have also taken the reviewers' feedback into account and made the following key changes
to improve the paper:

- **Provided LLM-related analysis**, including analysis of the kernels' fitness distribution (Appendix
F.1), sample outputs from the LLMs (Appendix F.2), prompt ablation study (Appendix
F.3), and comparison of different LLMs (Appendix F.4).
- **Standardized the experiment design** by replacing the GAP metric with the regret across
all experiments (Section 6, Appendix G).
- **Conducted additional experiments to include two variants of the Matern kernel**: Matern-
3/2 and Matern-5/2 (Appendix G.1).
- **Added the explicit formulation** for the weighting strategy (Section 3).
- **Increased the number of random seeds** for the benchmark function experiments to 20
(Section 6.1).
- **Conducted additional experiments** using the Ackley and Levy functions **with higher input
dimensions**, i.e., $d \in$ {10, 20, 30} (Appendix G.1)
- **Included two new real-world benchmarks** that involve tuning a controller for robot pushing
(Appendix G.4) and lunar lander (Appendix G.5).
- **Added comparisons against three recent BO baselines**, BOM, MCMC, and ABO (Appendix
G.1).
- **Included additional baseline methods from the literature** for the photonics chip design
experiment (Appendix G.3).

We hope these revisions adequately address the reviewers’ concerns. We remain open to further feedback.

With thanks,

Authors of #5709
</content>
</thread_object>

</thread>


Metareview: