Prompt Optimization with EASE? Efficient Ordering-aware Automated Selection of Exemplars

Summary

The paper proposes EASE, an exemplar selection algorithm for in-context learning of large language models (LLMs). It takes into account the order of the exemplars and offers the possibility to jointly optimize the instruction and the exemplars. A neural network is iteratively trained on the embeddings of the current exemplar sequence to predict the score of the next exemplar. The search space is reduced using a technique based on optimal transport. The paper also empirically shows that the effect of exemplar selection is not the same for all tasks.

Strengths

- A wide range of baseline methods and tasks are covered in the experiments. - The paper provides insights on the impact of exemplar selection on some specific tasks.

Weaknesses

- The performance of EASE, which requires training a model for each task, on Instruction Induction (II) is very close to the random search baseline for most tasks. - Experiments are mainly conducted on gpt-3.5-turbo-1106, an API that might include unknown pre- and post-processing. Moreover, the availability and behaviour of these APIs is not guaranteed hindering reproducibility. Consequently, the algorithm should also be validated on open-weight models (not just when doing progressive finetuning) in addition to models only available through APIs. Minor issues: Concerning the presentation, although the main contribution is clearly stated, the paper is structured in a way that makes it a bit hard to read. For example, the main paper lacks a dedicated related work section that would better contextualize the proposed algorithm.

Questions

- “out-of-distribution” tasks is defined as as tasks on which the LLM is not already well trained. How are these tasks selected without access to the training data? - what is the effect of the length of exemplars k on performance (using a range instead of a fixed value as in D.5)?

Rating

6

Confidence

4

Soundness

2

Presentation

2

Contribution

2

Limitations

There is a section on limitations in the main paper and a section on broader social impact in the appendix.

Summary

The paper introduces EASE (Efficient ordering-aware Automated Selection of Exemplars), a new approach to boost in-context learning (ICL) in large language models (LLMs). EASE optimizes the selection and ordering of input-label exemplars without needing model fine-tuning or test-time retrieval. EASE trains a neural network to predict the prompt performance using hidden embeddings from a pre-trained language model. Using the network as the scorer, it then employs a bandit algorithm to search for the best order exemplars for improved performance efficiently. Additionally, EASE can jointly optimize exemplars and instructions. In Instruction Induction (II) benchmark tasks and four novel out-of-distribution tasks, EASE outperforms several basic baselines, especially when the LLM has limited knowledge of the task.

Strengths

- The problem addressed is significant, as finding a fixed set of examples and instructions that generalize well to testing time could be a valuable technique. - The proposed method is well-founded, incorporating a novel component that uses optimal transport to reduce the search space. - The performance improvement in the out-of-distribution (OOD) setting is clear.

Weaknesses

- The main framework is quite similar to prior work [1], with only minor differences in details. - The literature review is inadequate, failing to cover other prompt optimization approaches, including instructions and/or example ordering (see below). Including a comprehensive related work section is crucial to highlight the novelty of the approach and justify its preference over others. - Due to the lack of discussion on prominent related works, it is unclear why the authors did not include them in the baseline set. The current baseline set is weak, relying only on heuristic methods. It would be better to incorporate proper optimization-based methods using RL[2,3], GA [4], or even LLM [5,6] as optimizers. - The current presentation of the results is unconvincing. For instance, there is no comparison in terms of running time or efficiency, areas where the paper claims to excel. The second-best baseline, Best-of-N, shows competitive results in real benchmarks. If N increases, its performance might surpass EASE. Without efficiency comparison between methods, it is hard to determine if EASE is useful. - Another major concern is the method seems to only work well with synthetic OOD settings. More experiments with real datasets are preferred. [1] Xiaoqiang Lin, Zhaoxuan Wu, Zhongxiang Dai, Wenyang Hu, Yao Shu, See-Kiong Ng, Patrick Jaillet, and Bryan Kian Hsiang Low. Use your INSTINCT: Instruction optimization using neural bandits coupled with transformers. In NeurIPS Workshop on Instruction Tuning and Instruction Following, 2023. [2] Mingkai Deng, Jianyu Wang, Cheng-Ping Hsieh, Yihan Wang, Han Guo, Tianmin Shu, Meng Song, Eric Xing, and Zhiting Hu. 2022. Rlprompt: Optimizing discrete text prompts with reinforcement learning. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 3369–3391. [3] Tianjun Zhang, Xuezhi Wang, Denny Zhou, Dale Schuurmans, and Joseph E Gonzalez. 2022. Tempera: Test-time prompt editing via reinforcement learning. In The Eleventh International Conference on Learning Representations [4] Archiki Prasad, Peter Hase, Xiang Zhou, and Mohit Bansal. 2023. Grips: Gradient-free, editbased instruction search for prompting large language models. In Proceedings of the 17th Conference of the European Chapter of the Association for Computational Linguistics, pages 3827–3846. [5] Yongchao Zhou, Andrei Ioan Muresanu, Ziwen Han, Keiran Paster, Silviu Pitis, Harris Chan, and Jimmy Ba. 2022. Large language models are human-level prompt engineers. In The Eleventh International Conference on Learning Representations. [6] Chengrun Yang, Xuezhi Wang, Yifeng Lu, Hanxiao Liu, Quoc V Le, Denny Zhou, and Xinyun Chen. 2023. Large language models as optimizers. arXiv preprint arXiv:2309.03409.

Questions

- What is the size of the dataset D? If D is small, a brute-force search could be feasible. Typically, LLMs are most beneficial when there is little to no data for downstream tasks, so a small D would be more realistic. - How was the NN’s uncertainty computed? - What method was used to tune the hyperparameters of NeuralUCB? - Please consider more real data such as those in the Tempera paper [3].

Rating

5

Confidence

4

Soundness

2

Presentation

2

Contribution

2

Limitations

The authors could also mention the cost of hyperparameter tuning

Summary

The paper introduces EASE, a method for optimizing ICL in LLMs by selecting and ordering exemplars efficiently. Unlike retrieval-based methods that incur extra computation and privacy risks, EASE uses a neural bandit algorithm and optimal transport techniques to find high-quality ordered exemplars without test-time computation. It also extends to jointly optimize exemplars and instructions for improved performance. Empirical evaluations demonstrate EASE's superiority over existing methods across various tasks, especially where the LLM lacks task-specific knowledge. EASE proves robust in noisy data scenarios, offering a comprehensive solution for enhancing ICL in practical applications.

Strengths

- The introduction of the EASE algorithm, which combines neural bandit optimization and optimal transport, is a unique approach to exemplar selection. - The paper includes extensive empirical evaluations, comparing EASE with a comprehensive suite of baseline methods across various tasks. The consistent outperformance of EASE underscores its effectiveness. - The paper provides clear explanations of the methodology, including the use of neural bandit algorithms and optimal transport, making the complex concepts accessible.

Weaknesses

- The requirement for on-the-fly computation of embeddings for ordered exemplar sequences is identified as a potential computational bottleneck. - EASE relies on the availability of a suitable validation set, which may not always be readily available in certain scenarios. This requirement could limit the method's applicability in some contexts. - The joint optimization of exemplars and instructions, while beneficial, adds complexity to the method. The paper could provide more detailed guidance on effectively implementing this joint optimization in practice.

Questions

- Can you provide more insights or potential solutions to mitigate the computational bottleneck caused by the on-the-fly computation of embeddings for ordered exemplar sequences? - While combining EASE with retrieval-based methods for large exemplar sets has shown better performance, could you elaborate on any potential limitations or challenges of this approach?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

- The need for a suitable validation set is recognized as a limitation, as such sets may not always be readily available. This could restrict the method's applicability in certain contexts where validation data is scarce or difficult to obtain. - The selection of exemplars could inadvertently reinforce biases present in the training data. The authors could address how EASE mitigates bias and ensures fairness in exemplar selection, especially in sensitive applications.

Summary

The authors propose EASE, a method for optimizing the selection of few-shot examples for prompting black-box LLMs. EASE is an iterative algorithm that combines NeuralUCB and Optimal Transport. In iteration t, EASE trains a neural network to map embeddings of strings (of few-shot examples) to their average score on a validation set. Then, they sample a subset of possible combinations of few-shot examples, filter them with optimal transport, and use NeuralUCB to select a set that maximizes their acquisition function. The selected set is evaluated and added to the pool of embedding -> score list. The authors evaluate on a set of existing tasks and a few synthetic tasks they created. Overall, the proposed method outperforms the baselines nearly across the board.

Strengths

1) The paper overall is easy to follow. 2) The results are very strong and the range of baselines is satisfactory. 3) The hypothesis and synthetic task evaluation is quite valuable and insightful. 4) The method underlying EASE is likely to spur future approaches in this space.

Weaknesses

1) The method does not really justify why we should expect off-the-shelf embeddings (into a single unlearned vector no less) of few-shot example strings to encode enough information for NeuralUCB to work effectively. In general, off-the-shelf encoders may or may not highlight the important aspects of the prompt; they're not trained to do that. The authors discuss that embeddings will be sensitive to the ordering. This is true, but in what way do we know that such sensitivity could be leveraged effectively from a frozen task-misaligned sematic encoder? 2) A lot of the discussions regarding "retrieval-based methods" are just forced and the paper doesn't even need them. The work shows that it's empirically more effective. That's sufficient. You don't need to argue that retrieval-based methods will "significant additional computational overheads especially when the number of test queries is large". That's plainly false: retrieval in the way discussed in this paper (e.g., cosine similarity) can standardly be done from billion-scale collections in 2-3 milliseconds in practice, far faster than the LLM invocation. 3) Many experimental details don't seem to have been controlled for: what's the effect of the LM choice? what's the effect of the precise budget? of the training set size? of the validation set size? How large are these currently in the first place? 4) The tasks selected are arguably overly simplistic and generally this means that the few-shot examples considered are all rather trivial. To my understanding, there are no reasoning chains or open-ended generation in any of the labels. There's no argument on why we should expect a method that works for simple direct classification-type labels to generalize to far more open-ended tasks. Such tasks exist plentifully; it's unclear why the authors do not consider any of them. 5) The authors say: "Our proposed Evo and Best-of-N are the most competitive baselines with best performance in 9 and 11 tasks, respectively." This seems to suggest that the authors believe they are the first to consider such approaches. While I think the current baselines are satisfactory, both Evo and Best-of-N approaches are components in numerous existing methods (which may involve other components), e.g. see PhaseEvo "Unified In-Context Prompt Optimization" for the former and DSPy's BootstrapFewShot with RandomSearch for the latter.

Questions

1) The authors simultaneously find best-of-N to be highly effective as a baseline and yet argue for the importance of Optimal Transport in their method. Can we gain more insight on what OT ultimately selects under the hood? Intuition is lacking. 2) What's the difference between Tables 7 and 8? What about ablations on real tasks? 3) The authors say: > filter the large pool of data exemplars to eliminate those that are less relevant to the task. To this end, we propose to first use retrieval-based methods to select the exemplars that are more relevant to the task, and then run our EASE using this refined smaller set of exemplars. Specifically, we use a cosine similarity retriever and perform exemplar selection on D with a size n as large as 100000. What does this mean? What are examples that are "more relevant" to a task? What's the precise retrieval formulation here, i.e. what's the query and what is the corpus of documents.

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

N/A

About NN uncertainty computation

Responding to *Reviewer NbA9*: Following [14] and [42], the uncertainty of the NN is the Gaussian process posterior standard deviation when the empirical neural tangent kernel (NTK) is used as the kernel $k(z_1,z_2) = \nabla_{\theta} m(z_1;\theta_{t})^\top \nabla_{\theta} m(z_2;\theta_{t}) $ where $z$ is the embedding and $\theta$ is the parameter of the NN $m$. Specifically, define $V_{t} = \sum_{\tau=1}^{t} \nabla_{\theta} m(z_{\tau};\theta_{\tau}) \nabla_{\theta} m(z_{\tau};\theta_{\tau})^{\top}$ where $z_\tau$ is the history of embedding at step $\tau$. Then, the uncertainty is calculated according to $\sigma_{t-1}(z;\theta_t) = \sqrt{\nabla_{\theta} m(z;\theta_{t})^{\top}V_{t-1}^{-1} \nabla_{\theta} m(z;\theta_{t})}$.