CASE: Challenger Arm Sampling for Efficient In-Context Reasoning

Thank you for your insightful comments and highlighting the strengths of our paper. Please find below our responses:

                   TRANSLATION OF THEOREM 1 AND LEMMA 1 TO IN-CONTEXT LEARNING

• In the ICL (In-context learning) setting, each arm corresponds to a subset of exemplars. Our goal is to find top-$m$ arms (subsets of exemplars). At each iteration when an arm is pulled or sampled, we call the LLM with exemplars in that arm to get the true mean estimate (line 22 of Algorithm 1). We don't have any other particular assumptions for the ICL setting. The $\Delta_a$ (true gap, Theorem 1) depends on the dataset and the LLMs employed. $K$ is size of $U_T \cup N_T$.
• Theorem 1 helps derive an upper bound on the sample complexity of the multi-arm bandit algorithm employed for exemplar selection in this paper. Here, sample complexity refers to the expected total number of arms pulled made by the algorithm before it terminates. In the context of real-world complex reasoning tasks, this directly translates to the number of LLM calls required. As stated in Sections 3.1 and 3.2, we model a subset of exemplars as an arm, and pulling an arm involves constructing a prompt using the exemplars within that subset and then calling the LLM to get actual accuracy (reward for that arm) on a validation subset. We have also updated this information in the revised manuscript in lines 347-349.
• Lemma 1 aids in deriving the sample complexity bound for Theorem 1 by bounding the gap-index between $c$ and $b$ (the most confusing arms from $U_t$ and $N_t$), which is tied to the stopping criterion. Our proof of this lemma is based on the approach used by Reda et al., but it is specifically adapted to our CASE algorithm, which maintains a shortlist of challenger arms.
• In summary, Theorem 1 and Lemma 1 together provide an upper bound on the expected number of arm pulls required by the algorithm, which translates to the number of LLM calls needed when applying CASE to complex reasoning tasks. Thank you for this query, and we will make this discussion clear in the revised version of the manuscript. We hope this response addresses your concern.

                            COMPARISON TO LENS+MMR+SC AND LENS+KNN+SC

• Unlike existing instance-level selection methods, such as KNN and MMR, which select each exemplar independently, without considering the positive and negative interactions between the exemplars. CASE scores exemplar subsets as a whole rather than scoring exemplars independently. Similarly, LENS also selects exemplars independently, assigning a score to each exemplar without accounting for interactions among them.
• We would like to highlight that CASE+MMR+SC and CASE+KNN+SC represent a unique static/dynamic hybrid approach that preserves the order of exemplars by selecting an entire subset from the top $m$ subsets that has the highest average similarity (for KNN) or maximum marginal relevance (for MMR) to each test example.
• Following your suggestion, we extended this static/dynamic hybrid approach to the LENS method by first selecting the top 50 exemplars using LENS and then applying KNN+SC and MMR+SC over this set. The results of these experiments are provided in the table below. We observe from Table that CASE variants with KNN+SC and MMR+SC outperform corresponding LENS (KNN+SC, MMR+SC) variants. We posit that this is because, in CASE, the interaction of exemplars within a selected subset remains preserved. We observe this follows from a similar trend for vanilla CASE, LENS, and CASE + SC, LENS +SC where our proposed CASE variants outperform the LENS variants as observed from Table 1. We have also updated these results in Table 1 in the uploaded revised manuscript. We appreciate your feedback to improve the quality of the work and will include these results in the revised manuscript. We hope this response addresses your concern.

Thank you again for your valuable review. If we have misunderstood any aspect of your review, or if you have additional questions, please let us know.

Greetings Reviewer 6HR2,

Thank you for the response and the valuable review that helped improve the quality of our work. We have also added the above results to the revised manuscript in blue color. We are open to further discussions and are happy to clarify any remaining doubts. If there are still any outstanding concerns, we kindly request you to share them with us.

Thank you for your insightful comments in the detailed review and for highlighting the strengths of our paper. Please find below our responses for the questions raised:

Q1: In no part of the algorithm, natural language generation is taken into account to choose the exemplar.

We would like to point out that calling the LLM/accounting for generation from LLM is taken into account in our algorithm. Specifically, the call to LLM is performed after the selection rule, where an arm (i.e., a set of exemplars) is sampled (line 22 in Algorithm 1). Here, $A(\pi_r(a(t)),\mathcal{V})$ refers to the computation of accuracy on the validation subset, as explained in lines 200-205, which involves LLM reasoning to generate predictions on this validation subset. Thus, the LLM's reasoning and generated outputs are explicitly accounted for within our algorithm. We hope this clarifies how the gap-index-based multi-arm bandit algorithm is adopted for selecting exemplars for complex reasoning. We further explain how the featurization of arms is estimated in the complex reasoning setting in response to Q6 below. Thank you for your query. We have incorporated this clarification into Section 3 (lines 300-305) in blue-colored text to make it more explicit in the revised manuscript.

Q2: CASE does not take into account the reasoning aspect of an LLM while selecting the exemplar. In the synthetic experiment vs the reasoning experiment, how do we actually understand the chosen exemplar leads to better reasoning? Better to show a table where for the same questions show what CASE vs LENS choose as exemplars/rationales

We think there is a slight misunderstanding of our setup and problem formulation, and we would like to clarify the same as follows:

• We would first like to clarify that an exemplar is defined as a triplet (input, rationale, output), as defined in line 42. The goal is to select few-shot exemplars from the training set to be used as in-context demonstrations for the generative model (LLM) during inference on the test set for complex reasoning tasks. Therefore, the selected exemplars may vary across different methods, depending on the algorithm's capability to identify the optimal exemplars. It is important to note that the rationales are already provided within the training examples for all the benchmarks utilized in our work.
• Secondly, the LLM reasoning and output generation on the validation subset is taken into account in our algorithm. Specifically, the call to LLM is performed after the selection rule, where an arm (i.e., a set of exemplars) is sampled (line 22 in Algorithm 1). Here, $A(\pi_r(a(t)),\mathcal{V})$ refers to the computation of accuracy on the validation subset, as explained in lines 200-205, which involves LLM reasoning to generate predictions on this validation subset. We have incorporated this clarification into Section 3 (lines 300-305) in blue-colored text to make it more explicit in the revised manuscript.
• Additionally, our results indicate that CASE selects high-quality exemplars, as demonstrated by the qualitative analyses in Tables 6-10. As explained in lines 470-475 and 968-1014, LENS often selects exemplars with redundancy, where multiple examples require similar reasoning steps and share thematic similarities. For instance, in Table 6 for AquaRat, exemplars 4 and 5 chosen by LENS are redundant, both centered on the theme of "work and time." Similarly, in Table 9 for TabMWP, exemplars 1 and 3 redundantly focus on calculating the median of a list of numbers, lacking diversity that would otherwise enrich the LLM’s inference capability.
• Unlike LENS, which assigns scores to exemplars independently without accounting for interactions, CASE captures the interactions between exemplars by scoring subsets as a whole. Furthermore, it learns weights ($\alpha$) for exemplars based on the subset's performance with respect to the LLM loss approximation error. Hence, this results in selecting more diverse and effective exemplars, leading to superior performance.

Q3: How do you choose $\epsilon$ in the algorithm? Should it be given as an input to the algorithm?

The stopping criterion, $\epsilon$, is a hyperparameter provided as input to the algorithm. We have set $\epsilon$ = 0 for synthetic experiments and $\epsilon$ = 0.1 for real-world experiments. These details are mentioned in Section 4.1, on lines 363 and 408.

Q4: LinGapE is similar to CASE from Figure 2(d). How do you differentiate between different bandit algorithms?

• Thank you for your question. We would like to point out that the main advantage of CASE is its efficiency compared to existing multi-arm bandit frameworks, as explained in Section 4.2. This is demonstrated in Figures 2(a) and 2(b), where CASE shows significantly lower runtime compared to LinGapE (by approximately 5.6x) and LinGIFA (by about 12x). This efficiency gain is due to fewer arm comparisons and reduced gap-index computations.
• In our synthetic experiments, we observed that CASE not only provides comparable performance guarantees in terms of error rates but also significantly outperforms existing gap-index frameworks in efficiency. Thus, for real-world complex reasoning tasks, where other gap-index frameworks become computationally expensive, we utilize CASE for exemplar selection. Consequently, we compare our efficient MAB algorithm, CASE, against state-of-the-art task-level exemplar selection methods like LENS.
• Additionally, to the best of our knowledge, we are the first to formulate exemplar selection for in-context learning (ICL) as a top-$m$ arm selection problem, which sets our approach apart from existing methods.

Q5. How do you learn the $\alpha_i$ in eq(3)? Is it pre-specified?

The parameter $\alpha$ is learned using a least squares estimator. Initially, the $\alpha$’s are initialized from a normal distribution (line 8 in Algorithm 1). Subsequently, when an arm is pulled, its true mean (i.e., the validation accuracy in real-world experiments) is obtained. The values of $\alpha$ are then updated using the least squares estimator as specified in line 24 of Algorithm 1.

Q6. While the $\phi_u(a)$ is estimated using a pre-trained transformer, how do you estimate the featurization $x_i \in \mathcal{X}$?

Thank you for your query. We compute features for each arm as follows:

First, we calculate the similarity between the exemplars within an arm and the validation samples, as outlined in line 226. This results in a tensor of shape (number of exemplars, number of validation samples). We then average across the validation samples to obtain a feature vector of dimension (1, number of exemplars) for each arm. Consequently, the feature matrix has dimensions (number of arms, number of exemplars), as shown in Equation 3. We hope this clarifies how the features are constructed. The number of exemplars within an arm determines the dimension of its corresponding feature vector.

Q7. How the featurization of the actions $\mathcal{X}$ and the $k$-subsets of $\mathcal{X}$ are obtained?

Figure 1 and Section 4.1 illustrates the process of obtaining the $k$-subsets (arms) from $\mathcal{X}$ and subsequently populating the lists $U_t$ and $N_t$. Additionally, in response to your previous question, we have clarified how features are computed for each arm. We hope this addresses your concerns, and we have made these details clearer in the revised manuscript. We appreciate your feedback to further enhance the quality of the paper.

Thank you again for your valuable review. If we have misunderstood any aspect of your review, or if you have additional questions, please let us know.

Q8. Consider the definition of the UCB, $W_t(i,j)$ definition below eq (3). It is not clear what are (|Ei|Σ^tλ+|Ej|Σ^tλ). Several such areas require more explanation and definitions need to be defined more clearly.

• Thank you for pointing this out. We apologize this was a typographical error and we have fixed this in the revised version of the uploaded manuscript (lines 238-242). We hope this clarifies the setup.

We have also clarified other concerns raised in the review. Thank you for the insightful questions and suggestions that helped improve the quality of our work. We hope we have sufficiently addressed your queries and request you to share any additional queries or concerns you may have. Thank you for your time and consideration.

Greetings Reviewer wmwx,

Thank you for your valuable feedback and constructive comments. In our rebuttal, we have provided clarifications regarding the adoption of bandit framework to reasoning setup by taking LLM reasoning into account. We have added explanations in the revised manuscript in blue color to clarify and address all the concerns raised. We are open to further discussions and are happy to clarify any remaining doubts. If there are still any outstanding concerns, we kindly request you to share them with us.

• Thank you for your valuable feedback. We acknowledge that our approach focuses on selecting few-shot exemplars (I/p,o/p, rationales) in an efficient manner using CASE to enhance the performance of LLMs on complex reasoning tasks (as indicated in lines 24-25, Section 1, Section 3.1 and lines 300-306). We would like to humbly clarify that in prior works on few-shot in-context learning [1] , exemplar selection [2] , the process of employing rationales within few-shot examples to guide the LLM in generating both a rationale and the corresponding predictions is commonly referred to as LLM reasoning (as highlighted in [1, 2, 3, 4]). We adopt this terminology to maintain consistency with the existing literature. In our approach, rationale generation and output prediction by the LLM, conditioned on the selected few-shot examples, are explicitly considered when pulling an arm to update the parameters (lines 300-306). Furthermore, during inference, these exemplars are utilized to prompt the LLM to generate rationales and predictions for test samples. Our method selects exemplars that showcase diverse reasoning skills through rationales (analysis in Tables 6-10) for respective complex tasks.
• Our intention was not to overplay LLM reasoning, but to align with established terminology in the field. We are happy to reword any specific paragraphs that the reviewer thinks are overstating the reasoning aspect. We deeply value your feedback and appreciate the opportunity to refine our work further.

[1] Chain-of-Thought Prompting Elicits Reasoning in Large Language Models, Wei et al. NeurIPS 2022.

[2] COMPLEXITY-BASED PROMPTING FOR MULTI-STEP REASONING, Yao Fu et al. ICLR 2023.

[3] Self-Consistency Improves Chain of Thought Reasoning in Language Models, Wang et al. ICLR 2023.

[4] Large Language Models are Zero-Shot Reasoners, Kojima et al. NeurIPS 2022.

Greetings Reviewer wmwx,

Thank you for your feedback which helped improve the quality of our work. As the reviewer-author discussion period is nearing its conclusion (December 2nd at 11:59 pm AoE), we would like to request if we have addressed all your concerns in our previous responses.

We have tried to clarify queries regarding the different parameters of CASE and also regarding how CASE is adopted for LLM based rationale and answer generation for different benchmarks. In our approach, rationale generation and output prediction by the LLM (LLM generation), conditioned on the selected few-shot examples, are explicitly considered when pulling an arm to update the parameters (lines 300-306). Furthermore, during inference, these exemplars are utilized to prompt the LLM to generate rationales and predictions for test samples.

Further, we also clarified the use of terminology related to ``LLM reasoning”, in the previous response and we have also explicitly stated the core motivation and purpose of our work as efficient task-level exemplar selection to improve LLM performance (lines 72-98, lines 14-15, lines 24-25, Section 3.1). We are also happy to reword any instances to make this more explicit in the final version of the manuscript.

We hope that we have addressed all the main concerns raised in the initial review. If your concerns have been addressed, we politely and respectfully request you to consider revising your initial rating. Thank you so much for your time, consideration, and effort in reviewing our work.

Best regards,

The Authors

Thank you for your insightful comments and highlighting the strengths of our paper. Please find below our responses:

                                          PART A

a) Why $\pi$ needs $m$ sets of exemplars, rather than just a set of $k$ examples?

• We introduce $\pi$, a dynamic prompt generator, which operates on top of the task-level exemplars selected by CASE. This is equivalent to CASE+KNN and CASE+MMR, as outlined in lines 191-195, and corresponds to the CASE static-dynamic hybrid variants shown in Table 1 and further explained in Section 4.1, lines 417-421.
• CASE returns $U$, a set of $m$ subsets where each subset consists of $k$ exemplars. The dynamic prompt generator $\pi$ takes these $m$ subsets and applies methods like MMR or KNN to select a subset of exemplars to include in the prompt for each test instance. This approach reduces the search space at the sample level, as we search over these $m$ subsets instead of all training examples, making CASE a static/dynamic hybrid method.

b) Are S1…Sm allowed to overlap? Why?

• Yes, the subsets are allowed to overlap since we sample with replacement. This design aligns with our motivation to explore different combinations of exemplars, inspired by literature where in-context learning (ICL) samples have been shown to exhibit both positive and negative interactions with each other [1]. The key objective is for each subset to serve as a high-quality prompt, covering diverse aspects of the task. Properly capturing interactions among exemplars is critical, as redundant or negatively interacting exemplars can degrade performance. Our qualitative analysis in Tables 6-10 demonstrates that exemplars in each subset contribute complementary capabilities.
• For instance, in Table 6, for AquaRat we observe that exemplars 4 and 5 in the set of selected exemplars by LENS are redundant; they are highly similar problems that require the same reasoning steps and share a common theme of "work and time." Similarly, in Table 9, for TabMWP, exemplars 1 and 3 selected by LENS redundantly focus on computing the median of a list of numbers, failing to introduce new reasoning concepts that would aid the model in solving diverse problems during inference.
• In contrast, CASE consistently selects exemplars that are both thematically and conceptually distinct, ensuring a broader coverage of problem types. Thus, it is evident that the CASE algorithm leads to overall performance improvements, as further demonstrated by its comparisons with existing baselines.

[1] Voronov2024, Mind Your Format: Towards Consistent Evaluation of In-Context Learning Improvements.

c) How does the objective in (2) turn into the top-m problem?...

• Thank you for pointing this out. Yes, we assume that the top-$m$ exemplar subsets (prompts), which have the highest average accuracy (equivalent to the subsets with the highest individual accuracy), are also those that maximize the objective function in Equation (2).
• Please note that the transformation $\pi$ calculates the best subset among the top-$m$ subsets for a given validation example (static-dynamic hybrid approach as described earlier). For the total accuracy of the prompt generated by $\pi$ (Equation 2) to be high, each subset should maximize the accuracy across most validation examples in $\mathcal{V}$. The top-$m$ objective function assumes that the exemplar subset also performs well on other validation examples, so that the average accuracy of each subset over the considered validation subset is the highest. This assumption inherently filters out subsets that are highly accurate for a few validation examples but highly inaccurate for the majority.
• This assumption is critical due to the complexity involved in the search process. The search space is large—first over all possible sets of exemplars to form subsets, and then over all subsets to identify the top-$m$ subsets. This complexity is discussed in Section 3.2, specifically in lines 209-221.

We would like to answer rest of the questions as follows

                                                      PART B

c) As a comment, is the reduction to linear bandits necessary? You are trying to find the top $m$ sets of size $k$. It feels like you could do this in a more direct manner. In particular, if $a_1$, $a_2$, …., $a_n$ were an ordering of the individual arms by rewards as given in 3), you just need to return $a_1$, .. $a_{mk}$ appropriately segmented (if there are no overlaps - otherwise it would be $a_{k+m}$). Could you employ existing works on coarse ranking, eg Katariya '18?

We thank the reviewer for their valuable comments.

• Here, $k$ refers to the number of exemplars in a subset rather than the number of “arms.” We cannot construct a prompt using a single exemplar, and hence cannot score a single exemplar using the LLM loss. Therefore, we believe that the setting described in coarse ranking is not directly applicable to our problem. Our contribution is primarily an efficient version of gap-index framework applied to ICL setup for LLMs.

• The linear bandits framework is necessary for defining the confidence $W_{ij}$, which depends on the design matrix $V$.

• Additionally, the algorithm in Katariya '18 is linear in the number of arms, which is computationally impractical in our setting due to the large number of arms (subsets of exemplars).

                                                      PART C

a) I don’t really see the connection to SETC. I went and read this section of the bandit book again, and the comparison is very unclear. Please clarify.

• We mention SETC as it provides an algorithm with a sub-linear regret bound for linear bandits with an exponential number of arms. A key similarity between SETC and CASE is that both employ a uniformly random exploration followed by a selection strategy.
• In case of SETC, arms are deselected implicitly by freezing certain dimensions. For CASE, arms are selected within the set $U_T \cup N_T$ by implicitly updating the $\alpha$’s.

Thank you again for your valuable review. If we have misunderstood any aspect of your review, or if you have additional questions, please let us know.

Thank you for your insightful reviews. Below we discuss the remaining questions

                                        PART A

d) In practice, why not just use similarity between validation set and exemplars instead of m exemplar sets?

• The pairwise similarity between training prompts and validation examples is a baseline for instance-level (or dynamic) prompt selection. However, this approach often fails to capture the interactions between exemplars within a prompt. Apart from its additional computational overhead during inference, we observe that it does not consistently result in stable performance gains (Table 1). The fundamental problem tackled in our work is task-level exemplar selection.
• In contrast, our work offers a principled exploration over the space of exemplar subsets (arms), where the accuracy on a validation set is treated as the true reward for each arm. Through careful exploration, we model the mean estimates for the arm along with confidence information $W_t(i,j)$.
• We also perform ablations as shown in Table 2, where we fit the model only once on the validation set without exploration (similar to only considering similarity between validation samples and exemplars), and another ablation with exhaustive evaluation. We find that empirically, such methods perform worse, with reasons outlined in Section 4.5.

                                        PART B

a) Discussion of differences to other pure exploration works and relevant citations.

• We appreciate the reviewer’s suggestion regarding optimal exploration techniques in linear bandits. According to Li et al. [2], the main problem tackled is to optimally explore the parameter space of a linear bandit by maintaining a distribution $\lambda$ over the sample set $Z$, which could also represent the set of arms $X$.
• However, in our specific scenario, the set of arms is large (i.e., all possible subsets of training data points), making it impractical to maintain a distribution over such an extensive set. This problem is common across algorithms in the papers mentioned. Hence, this class of techniques is not directly applicable and is also computationally expensive compared to the proposed gap index-based approach.
• Primarily we focus on designing an efficient approach for gap-index-based MAB algorithms building upon GIFA framework for the LLM complex reasoning setup. However, we acknowledge that the proof techniques used for analyzing sample complexity could be applicable to the present problem. We will explore the application of these techniques to improve algorithm design and derive tighter sample complexity bounds in future work. Thank you for highlighting these connections. We have cited and discussed the relevant works in Section 2.2, lines 145-146 (marked in blue in the revised manuscript).

[2] Optimal exploration is no harder than Thompson Sampling, Zhaoqi Li et al. AISTATS 2024.

b) I didn’t really understand what Theorem 1 says. What does epsilon-delta PAC mean here? How does this sample complexity scale with a function of $m$, $k$, $n$? What is $K$?

We thank the reviewer for pointing out these issues.

• Clarification on Theorem 1: Theorem 1 adapts the result from Theorem 2 from Reda 2021 et al. [3] to our setting, where we restrict ourselves to arms in $U_T \cup N_T$. It states that the $\epsilon$-optimal top-$m$ arms from $U_T \cup N_T$ are present in $U_T$ with probability $1-\delta$, if $T> \tau_{\delta}$. We will include this clarification in the revised manuscript.
• Definition of $K$: $K$ refers to the size of $U_T\cup N_T$. This clarification has been added to lines 338-340 in the updated manuscript.
• Sample Complexity: We derive a high-probability upper bound on sample complexity, assuming that the average regret of $U_T \cup N_T$ is less than $\epsilon$. Under this assumption, Theorem 1, describes the condition on stopping time $\tau_{\delta}$, which ensures that $U_T$ contains the $\epsilon$-optimal top-$m$ arms. This clarification has been included in the revised manuscript (lines 319–323).
• Furthermore, the connection between Lemma 1 and the sample complexity term in Theorem 1 is elaborated in Appendix A and B.

[3] Top-m identification for linear bandits, Clemence Reda et al. AISTATS 2021.

Greetings Reviewer ZiMF,

We thank you for your time and insightful review, which helped improve the quality of our work. We have tried to address the concerns raised in your review. Your suggestions have been thoughtfully incorporated into the revision, and we hope the updated manuscript addresses all the concerns you raised. If you have any further queries or require additional clarifications, we kindly request you to share them with us, and we will be happy to address them promptly. If your concerns have been addressed, we respectfully and politely request for a reconsideration of the score. Thank you once again for your time and consideration. We look forward to hearing from you.

We are pleased to have resolved all your queries and we politely and respectfully request you to consider revising your score. Thank you once again for your time and effort in reviewing our work.

Greetings Reviewer ZiMF,

Thank you for your insightful feedback and your effort in reviewing our work. We are pleased to have resolved all your queries. As the reviewer-author discussion period is nearing its conclusion (December 2nd at 11:59 pm AoE), we would like to request if you have any further concerns.

If we have addressed all your concerns, we politely and respectfully request you to consider revising your initial score. Thank you once again for your time and effort in reviewing our work.

Best regards,

The Authors

Thank you for your insightful comments and highlighting the strengths of our paper. Please find below our responses:

                 EVALUATIONS ARE FOR ICL, WILL BE NICE TO HAVE FULL FINETUNING

• Thank you for the question and feedback. We would like to humbly point out that our problem setup is efficient exemplar selection for in-context learning, not supervised fine-tuning, which falls outside the scope of our work. ICL allows LLMs to perform tasks by conditioning on a context that includes examples or instructions, without the need for additional fine-tuning, making it efficient and adaptable to new tasks. (lines 34-36). Notably, recent research has shown that many-shot ICL for complex reasoning can perform on par with, or even outperform, fine-tuning approaches [1,2,3].
• We would also like to humbly highlight that training a 70 billion parameter model is computationally intensive, requiring four clusters of eight A100 GPUs [4]. Even with optimizations like low-rank adaptation, fine-tuning still demands significant resources. Additionally, our work which is firmly centered on the ICL setting aims to improve efficiency and task performance through principle task-level exemplar selection.
• Our baselines are prior state-of-the-art exemplar selection methods that focus on the ICL setting. To compare our work with respect to fine-tuned approaches, we have included available LLM fine-tuning results from prior works as well as other supervised state-of-the-art results [4,5,6,7,8,9,10,11,12,13] for comparison on relevant benchmarks. We observe that CASE with gpt-3.5-turbo and gpt-4o-mini surpasses fine-tuning-based methods for both smaller and larger open models. Additionally, CASE with Llama2-7b, which does not require LLM fine-tuning, demonstrates competitive performance against versions of Llama2-7b that have been fine-tuned on specific tasks.

References

[1] Many-Shot In-Context Learning, Rishabh Agarwal et al. (https://arxiv.org/abs/2404.11018)

[2] Deeper Insights Without Updates: The Power of In-Context Learning Over Fine-Tuning, Qingyu Yin et al. EMNLP 2024.

[3] Is In-Context Learning Sufficient for Instruction Following in LLMs?, Hao Zhao et al. (https://arxiv.org/pdf/2405.19874)

[4] The ART of LLM Refinement: Ask, Refine, and Trust, Kumar Shridhar et al. NAACL 2024.

[5] MuggleMath: Assessing the Impact of Query and Response Augmentation on Math Reasoning, Chengpeng Li et al. ACL 2024.

[6] neuralmagic/Llama-2-7b-gsm8k-pruned_70

[7] Tora: A Tool-Integrated Reasoning Agent For Mathematical Problem Solving, Zhibhin et al. https://arxiv.org/pdf/2309.17452

[8] Large Language Models are Zero-Shot Reasoners, Kojima et al. NeurIPS 2022.

[9] Measuring and Improving BERT’s Mathematical Abilities by Predicting the Order of Reasoning, Piotr Piękos et al. ACL 2021.

[10] Dynamic Prompt Learning Via Policy Gradient For Semi-Structured Mathematical Reasoning, Pan Lu et al. ICLR 2023.

[11] Palm: Scaling language modeling with pathways, Aakanksha Chowdhery et al. JMLR 2023.

[12] SeCoKD: Aligning Large Language Models for In-Context Learning with Fewer Shots, Weixing Wang et al. (https://arxiv.org/pdf/2406.14208)

[13] The Unreasonable Effectiveness of Easy Training Data for Hard Tasks, Peter Hase et al. ACL 2024.

Below we answer the second question

                    DATA SETUP AND WHERE DO THE IMPROVEMENTS COME FROM? 

• Thank you for the feedback. We have provided comprehensive details regarding the data setup in Appendix C and Table 3 and we will highlight this further in Section 4. Our evaluation is based on widely recognized benchmarks for complex reasoning tasks, each of which has undergone a rigorous curation process as detailed in the respective papers. The dataset is divided into training, validation, and test sets, with no overlap between these subsets to the best of our knowledge. The authors of these benchmarks have taken careful measures to deduplicate examples, ensuring data integrity. For instance, strategyQA uses disjoint annotators for curating train and test sets and also performs deduplication after formation of the dataset.
• For our experiments, we adopt the exemplar selection setup used in prior works like Auto-CoT and Complex-CoT, along with other baselines. Based on your suggestion, we also reverified that the exemplars selected using CASE are not part of the test set. The performance gains we observe stem entirely from the algorithm itself, as evidenced by our thorough ablation studies in Table 2. These studies demonstrate a significant drop in performance when key components, particularly the exploration step integral to CASE, are removed.
• Furthermore, we have provided the selected exemplars in our open-source repository and have detailed them in Tables 6-10, with additional qualitative analyses in the Results section and Appendix G. We observe that exemplars selected by CASE are thematically diverse and better represent the task compared to other baselines like LENS, which tend to select conceptually similar exemplars, leading to redundancy (Section 4.3, lines 470-475).
• For instance, in Table 6, for AquaRat we observe that exemplars 4 and 5 in the set of selected exemplars by LENS are redundant; they are highly similar problems that require the same reasoning steps and share a common theme of "work and time." Similarly, in Table 9, for TabMWP, exemplars 1 and 3 selected by LENS redundantly focus on computing the median of a list of numbers, failing to introduce new reasoning concepts that would aid the model in solving diverse problems during inference.
• In contrast, CASE consistently selects exemplars that are both thematically and conceptually distinct, ensuring a broader coverage of problem types. Thus, it is evident that the CASE algorithm leads to overall performance improvements, as further demonstrated by its comparisons with existing baselines.

We will highlight the above discussion in the revised manuscript for clarity. Thank you again for your valuable review. We hope we have addressed your concerns. If we have misunderstood any aspect of your review, or if you have additional questions, please let us know.

Greetings Reviewer vb9k,

We thank you for your time and insightful review, which helped improve the quality of our work. As the reviewer-author discussion period is closing soon (November 26 at 11:59 pm AoE), we would like to gently remind you that we are eagerly awaiting your feedback on our response. We have tried to address the concerns raised in your review and clarified the details regarding the data setup and other concerns. If you have any further concerns, we kindly request that you share them with us. If your concerns have been addressed, we respectfully and politely request for a reconsideration of the score. Thank you so much for your time and consideration.

Greetings Reviewer vb9k,

Thank you for your time and insightful feedback. As the reviewer-author discussion period is nearing its conclusion (December 2nd at 11:59 pm AoE), we would like to request if we have addressed all your concerns.

We greatly value your insightful comments, which have been instrumental in improving our work. Your suggestions have been thoughtfully incorporated into the revision, and we hope that the author's response and the updated manuscript address all the concerns you raised.

If you have any further queries or require additional clarifications, we kindly request you to share them with us, and we will be happy to address them promptly. If your concerns have been addressed, we respectfully request you to consider revising your initial score. Thank you once again for your time and effort in reviewing our work.

Best regards,

The Authors