Efficient Prompt Optimization Through the Lens of Best Arm Identification

The remarkable instruction-following capability of large language models (LLMs) has sparked a growing interest in automatically finding good prompts, i.e., prompt optimization. Most existing works follow the scheme of selecting from a pre-generated pool of candidate prompts. However, these designs mainly focus on the generation strategy, while limited attention has been paid to the selection method. Especially, the cost incurred during the selection (e.g., accessing LLM and evaluating the responses) is rarely explicitly considered. To overcome this limitation, this work provides a principled framework, TRIPLE, to efficiently perform prompt selection under an explicit budget constraint. TRIPLE is built on a novel connection established between prompt optimization and fixed-budget best arm identification (BAI-FB) in multi-armed bandits (MAB); thus, it is capable of leveraging the rich toolbox from BAI-FB systematically and also incorporating unique characteristics of prompt optimization. Extensive experiments on multiple well-adopted tasks using various LLMs demonstrate the remarkable performance improvement of TRIPLE over baselines while satisfying the limited budget constraints. As an extension, variants of TRIPLE are proposed to efficiently select examples for few-shot prompts, also achieving superior empirical performance.

Paper

Similar papers

Peer review

Reviewer gchr6/10 · confidence 4/52024-07-10

Summary

This study focuses on the efficient prompt design under budget constraints, where effective prompts may have to be designed without excessive evaluation of a very large number of candidate prompts. The paper aims to present a principled framework - called TRIPLE - for tackling this problem, which is achieved by making connections between the prompt optimization problem and the fixed-budget best arm identification problem (BAI-FB), thereby drawing from the rich literature on the multi-armed bandits (MAB) problems and methods. Based on prompt design experiments for various design tasks using various LLMs, the paper shows that TRIPLE can meaningfully improve the prompt design results when there are budget constraints.

Strengths

- It is a good and meaningful attempt to formulate the prompt optimization problem as a multi-armed bandit (MAB) problem. As shown in the paper, making connections between prompt optimization and best-arm identification in MAB allows one to draw from the rich MAB literature to devise effective prompt design methods. This is clearly demonstrated in the budget constrained case, which is the main focus of the current study, as one can leverage fixed-budget best-arm identification methods to improve prompt design under a limited budget for evaluating the (possibly large) prompt candidate pool. However, the benefits of the proposed framework are likely to go beyond the current problem setting, as MAB is an actively investigated field with rich outcomes and existing tools in MAB may be utilized for enhancing prompt optimization methods in the future. - The results (e.g., in Figure 2 and Tables 2-4) clearly demonstrate the efficacy of the proposed TRIPLE framework for prompt evaluation & selection and how it may also improve end-to-end prompt optimization by integrating the framework with popular prompt generation schemes. According to the results shown, the gains turn out to be fairly significant and also consistent across different tasks/settings.

Weaknesses

- As the main contribution of this paper seems to lie in making systematic connections between prompt optimization and best arm identification in MAB, after establishing these connections, the improvements achieved in prompt optimization are direct outcomes of adopting existing FB-BAI algorithms in the MAB literature. As a result, novel methodological contributions that go beyond the utilization of existing FB-BAI methods proposed and well-studied in MAB for prompt optimization purpose are somewhat limited. - While the paper discusses the issues that arise when the prompt pool is huge (hence a large number of arms) and also proposes practical schemes to address these issues, there is no in-depth discussion of how the pool size affects the overall prompt optimization performance. There should be further discussion on this issue and the paper should include at least some empirical results (e.g., similar to Figure 3) that shows how the gain may be affected by the prompt pool size. - In Figure 2, it is unclear when the "normalized evaluation score" exceeds 1 (i.e., outperforming the uniform scheme"). Please add a horizontal line to show which bars are above/below 1. - The comparison with BO-EI is interesting but as BO performance tends to vary widely depending on the acquisition function used, it would be meaningful to provide additional comparison based on at least one or two additional acquisition functions. Especially, using BO acquisition functions aimed at different aspects (e.g., uncertainty vs diversity) may be helpful.

Questions

- Please provide some discussion on how TRIPLE may be used when the prompt optimization task has multiple objectives. How can connections to best-arm identification be leveraged in case of multi-objective prompt optimization?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The paper discusses the limitations of the proposed approach in the Appendix (section B.2) as well as some future research directions to address these limitations.

Reviewer cpR66/10 · confidence 4/52024-07-11

Summary

This paper studies prompt optimization using BAI-FB. There are several contributions: 1. Draw a connection between prompt optimization with the BAI-FB 2. Benchmark different acquisition function for prompt optimization 3. Extensive experiments are done to show the effectiveness of the proposed prompt optimization and extend the setting to example selection.

Strengths

1. The presentation of this work is clear 2. The formulation of prompt optimization with limited budget as BAI-FB is clean with strong justification 3. The experiments are extensive with a good coverage of baselines, LLMs, and datasets

Weaknesses

1. The novelty of this work is limited, since there are already a lot of prompt optimization work using bandits as the authors have discussed in the paper 2. A larger domain of prompts need to be considered where exploration vs exploitation is more important. For example in [43], the domain of prompt is 10k and the experiments in [43] shows that NeuralUCB is a SOTA selection strategy in this case. A comparison of MAB-FB methods and regret minimization methods in a larger domain of prompt optimization can be more helpful.

Questions

What's the hyper-parameter used to run the regret minimization algorithm? e.g., UCB and NeuralUCB? Since these algorithm are highly sensitive to the selection of hyper-parameters.

Rating

6

Confidence

4

Soundness

2

Presentation

3

Contribution

2

Limitations

See my weakness

Reviewer 7VZU5/10 · confidence 4/52024-07-13

Summary

This work proposes an algorithm that adopts the fixed-budget best arm (BAI-FB) identification to search for the best prompt. The authors have considered two variants of BAI-FB including sequential halving (SH) and continuously reject (CR). They then utilize prompt embeddings to enhance the BAI-FB methods via clustering and function approximation.

Strengths

- The idea of using BAI-FB other than naive online MAB is interesting in prompt optimization. - The empirical results suggest TRIPLE is an effective method to identify good prompts.

Weaknesses

- The presentation could be substantially improved. For example, Section 3 seems redundant as most of its content is not critical and can be covered by Section 2. Besides, the empirical results show TRIPLE-CLST/GSE are generally better but their justifications and descriptions (including the clustering property which is critical for TRIPLE-CLST and the embedding in function approximation) are not enough in the main text. I strongly suggest authors to reduce Section 3 and extend Section 4 by adding explanations and moving some materials from the appendix into the main text. - The evaluation setting is not described in detail. For example, are the evaluation budgets per prompt the same for all considered baselines? If so, the exact amount of evaluation budgets per prompt only means how accurate this estimation is. - The considered prompt candidates for search are very limited only containing 150 at the maximum. However, in the literature (e.g., Lin et al. [43]), they have considered over 10k candidates. Generally, it is reasonable to consider a large prompt space for the search method to really find the optimal prompt. Otherwise, the result could be biased. - Although the paper emphasizes efficiency, it is very hard to see if the method is efficient in the empirical section, where only the performances under the specific budget setting are considered. - Recent baselines with/without embeddings like ZOPO and, OPRO should also be considered. - The tasks evaluated in this paper are mostly instruction induction tasks. It would be great to show that the method can work on GLUE tasks and reasoning tasks (e.g., GSM8K, MATH).

Questions

- Some literature adopts the deterministic function $f(\cdot)$ for prompt optimization and this is practical as some LLMs use greedy sampling. Will your method still work under the deterministic function $f(\cdot)$? See other weaknesses above.

Rating

5

Confidence

4

Soundness

2

Presentation

1

Contribution

2

Limitations

Yes.

Reviewer 55GB7/10 · confidence 4/52024-07-30

Summary

The authors study prompt optimization, with a focus on finding the best prompt from a pool of proposed prompts under highly limited budgets. They establish a connection to fixed-budget best arm identification in multi-armed bandits (MAB) and explore several algorithms from that problem, applied to prompt selection for evaluation. They then make additional improvements by observing that prompts are not independent and that information, through embedding the prompts, can be shared to make the exploration more effective. They test methods based on clustering and DNN reward functions, based off of off-the-shelf embeddings of the prompts.

Strengths

1. The paper is well-argued and well-presented. The authors argue that existing work on exploring prompt candidates under limited budgets apply a regret minimization framework, which is less well-suited than best arm identification, and then go on to explore highly-effective simple extensions that share information across prompts. 2. The results are compelling: the methods are compared against a convincing set of baselines, from UCB to Bayesian optimization (BO) with expected improvement. The outcomes consistently show the value of the author's selection of framework (BAI) and their information sharing.

Weaknesses

1. The authors evaluate on tasks that to my understanding are extremely limited in scope. Most are narrow reasoning puzzles and may not reflect the type or complexity of typical prompts people use in the increasingly elaborate open-ended LM systems out there. How would such open-endedness interact with the findings? 2. The results are somewhat hard to parse from the tables and especially so from the figures, though it is clear that they are overall quite positive. The budgets considered for evaluation are perhaps very strict, which is worth calling out, e.g. <1 evaluation per prompt, 5 evaluations per prompt. The method does show gains over baselines up to 30 evalutions per prompt, so it is very compelling nonetheless. 3. The claims at the start of the paper that existing work just rarely considers budget of optimization is a bit overblown and overlL unnecessary for the claims of the paper, which are compelling on their own. The authors compare with several baselines that *do* consider budget, just not as effectively or thoughtfully as the authors do.

Questions

Do TRIPLE-SAR and CSAR take ordering into account?

Rating

7

Confidence

4

Soundness

3

Presentation

4

Contribution

3

Limitations

n/a

Reviewer 55GB2024-08-07

Thank you for the detailed response. It helps me maintain my high score.

Authorsrebuttal2024-08-12

Thank you

We greatly appreciate your recognition of the contributions of this work! We will revise the paper accordingly to incorporate your suggestions.

Reviewer gchr2024-08-12

Response to authors

I would like to thank the authors for their careful response. The authors' additional clarifications have addressed many of my previous concerns and I am increasing the rating as a result.

Authorsrebuttal2024-08-12

Thank you

Thank you for recognizing the contributions of this work! We will carefully incorporate your suggestions in the revised paper.

Authorsrebuttal2024-08-13

It is our pleasure to hear that the responses are helpful in addressing your concerns. We will carefully incorporate your comments in the revised paper. Thank you for recognizing the contributions of this work!

Authorsrebuttal2024-08-13

Looking forward to Discussions

Dear Reviewer 7VZU, We would like to first thank you again for the valuable comments and suggestions on our work. As the discussion period is concluding, we would be grateful if you could share any further concerns or feedback you might have. If our responses have sufficiently addressed your concerns, we hope that you could consider raising the score of your evaluation. Thank you for your consideration. Best, Authors of Submission 13481

Reviewer 7VZU2024-08-14

Thanks for the additional results and some clarifications. >More descriptions will be added to Section 4 to highlight the empirically remarkable TRIPLE-CLST/GSE As the justification and description of the method are lacking here, I am not convinced by the novelty of this work. I hope the authors can improve this part in their revision. > This philosophy is very similar to that of APO [54], adopted as one pipeline in Section 5.2. Due to such a similarity, we mainly perform experiments with APO, with Table 3 demonstrating the superiority of TRIPLE. OPRO should be a much stronger baseline than APO. I think the authors should consider the strong and representative works in this area. --- Some of my concerns are addressed, although the comparison and clarity remain an issue. I will increase my score accordingly based on the additional results provided.

Authorsrebuttal2024-08-14

We are glad to hear that the responses are helpful in addressing your concerns. Thank you for providing further valuable suggestions on the clarity and comparisons! We will carefully incorporate them into the revised paper, in particular additional illustrations on TRIPLE-CLST/GSE and empirical comparisons with OPRO.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC