An Exploratory Study of Bayesian Prompt Optimization for Test-Driven Code Generation with Large Language Models

TL;DR

This paper introduces BODE-GEN, a Bayesian optimization method that adaptively searches for optimal prompts in continuous embedding space to improve the accuracy of code generated by large language models, demonstrating efficiency and effectiveness.

Contribution

It proposes a novel BO approach for prompt optimization in code generation, utilizing embedding space and auxiliary LLMs to enhance performance over fixed prompts.

Findings

BODE-GEN improves code accuracy over fixed prompts.

It is sample-efficient, requiring few iterations.

Effective across multiple LLMs and benchmarks.

Abstract

We consider the task of generating functionally correct code using large language models (LLMs). The correctness of generated code is influenced by the prompt used to query the given base LLM. We formulate the problem of finding the appropriate prompt as combinatorial search process and propose a Bayesian optimization (BO) approach referred to as {\em BO for Code GENeration (BODE-GEN)}. BODE-GEN performs an adaptive data-driven search over prompts guided by training data in the form of prompts tried and the functional accuracy of the generated code over a set of given test cases. The key insight is to perform BO in continuous embedding space by using an auxiliary LLM to bridge the gap between discrete prompt space and continuous embedding space. We leverage two synergistic ideas, namely, random projections and dimensionality scaled priors, to build effective Gaussian process based surrogate models over the high-dimensional embedding space. Our experiments on the HumanEval+ benchmark using multiple base LLMs show that BODE-GEN can improve performance in terms of code generation accuracy compared to fixed prompts and manual prompt engineering. Additionally, we demonstrate that BODE-GEN is sample-efficient, requiring relatively few iterations of BO to demonstrate improvements in code accuracy.