From Few to Many: Self-Improving Many-Shot Reasoners Through Iterative Optimization and Generation

Authors: Xingchen Wan, Han Zhou, Ruoxi Sun, Sercan O Arik
Primary area: foundation or frontier models, including LLMs
Keywords: many-shot, in-context learning, large language models

Abstract:
Recent advances in long-context large language models (LLMs) have led to the emerging paradigm of many-shot in-context learning (ICL), where it is observed that scaling many more demonstrating examples beyond the conventional few-shot setup in the context can lead to performance benefits. However, despite its promise, it is unclear what aspects dominate the benefits and whether simply scaling to more examples is the most effective way of improving many-shot ICL. In this work, we first provide an analysis on the factors driving many-shot ICL, and we find that 1) many-shot performance can still be attributed to often a few disproportionately influential examples and 2) identifying such influential examples ("optimize") and using them as demonstrations to regenerate new examples ("generate") can lead to further improvements. Inspired by the findings, we propose BRIDGE, an algorithm that alternates between the optimize step with Bayesian optimization to discover the influential sets of examples and the generate step to reuse this set to expand the reasoning paths of the examples back to the many-shot regime automatically. On Gemini, Claude, and Mistral LLMs of different sizes, we show BRIDGE led to significant improvements across a diverse set of tasks including symbolic reasoning, numerical reasoning and code generation.

Decision: accept (poster)
Reviewer scores: 5, 6, 8, 6

Reviews and discussion:
<thread>
creation date:
1730328865625

thread_type:
official_review

<thread_object>
creator: reviewer_qv5c

<content>
summary:
This paper presents a study on the performance of many-shot in-context learning on NLP tasks that require reasoning process (the reinforced ICL setup). The authors first define a importantce scoring function of input examples based on estimated input gradient, and observe that 1) many-shot ICL could be driven by a few high-performing examples, 2) scaling demonstration examples could still be beneficial with improved quality of reasoning. Motivated by these observation, the authors propose a method called BRIDGE to improve many-shot ICL in an iterative manner. Experiments of Gemini across a diverse set of tasks show the effectiveness of BRIDGE.

strengths:
1. Overall, this paper is well writtern and clear to follow.
2. The proposed method is well-motivated through experiment results (Figure 1 and Figure 2). The observation that many-shot ICL could be driven by few high-performing examples is interesting.
3. The proposed method shows decent performance compared to the baseline.
4. The method is evaluated on various  challenging datasets

weaknesses:
1. The analysis is limited to reinforced ICL setting that only considers reasoning-heavy tasks. It's not clear what many-shot ICL would perform on simpler tasks that do not heavily relies on reasoning. Also, could BRIDGE anyway be helpful on simpler tasks?
2. The experiments are only evaluated with Gemini. What would be the performance of proposed method on open-source methods like Llama-3.1? More importantly, for open-source model, would there exists a better way of importance attribution? Also, for base model where the reasoning ability is limited, would the proposed method also be helpful?
3. Lack of the analysis on the computation cost. What's the computational cost of optimize and generate step of BRIDGE, respectively? 
4. The definition of importance score function is well-justified. In section 2 Line 153, the importance of each example is estimated by the averate performance of randomly sampled demonstrations on the validation set. The variance would be huge considering different choices of subsets and orders. Have the author performed analysis on the varaince of importance score?
5. Importance attribution is not well justified. In section 2 Line 153, it's not clear how importance score is actually estimated and why it would make sense.
6. The description of "optimzie" step is not clear (Line 261). The text in page 6 is generally unclear to me. For example, what does "TCM" stand for in algorithm 2 step 6? How does the problem become a "bi-objective" problem and how the optimize step is connected to "Bayes optmization"?

Minor issue: ")" is missing in Line 149

questions:
Please refer to the "weakness".
</content>
</thread_object>

</thread>
<thread>
creation date:
1730557374729

thread_type:
official_review

<thread_object>
creator: reviewer_yxzh

<content>
summary:
The paper studies the many-shot in-context learning using Gemini family of models on BigBenchHard (BBH). Specifically, the authors demonstrate that identifying a subset of influential examples reinforced with model-generated reasoning can significantly outperform naively using all available examples. They also propose BRIDGE, a simple approach that instantiates this finding to iteratively seeking for the subset of examples and refining model-generated reasoning for these examples. BRIDGE is evaluated on the subset of BBH tasks and the substantial gains over the considered baselines are demonstrated.

strengths:
* Paper studies many-shot in-context learning, the important novel prompting technique that is used to significantly improve LLMs' predictions. Given that, this technique was made available only recently with the rise of long-context LLMs, studying its behaviour is important topic that can be relevant to broad community.

* The proposed approach is simple and algorithmically applicable to both close-source and open-source models.

weaknesses:
* The important baseline might be missing to ensure the strength of the claims, please see Questions section.

* I would also ask for including similar analysis for other LLMs beside close-source ones like Gemini. At least on the subset of the considered tasks, to see whether the conclusions generalize to other models.

questions:
Can you please provide the baseline when a model conditions on all available examples but with all examples using model generated reasoning? This baseline can decouple the importance of subselecting a subset of examples and influence of having model-generated reasoning, and strengthen your claims if, indeed, you need both rather than only model-generated reasoning. The recent works [1, 2] can be used to ensure that a model is constrained to generate reasoning given the input and the ground truth answer.

[1] Hu et al. Amortizing intractable inference in large language models. ICLR 2024

[2] Zhao et al. Probabilistic Inference in Language Models via Twisted Sequential Monte Carlo. ICML 2024
</content>
</thread_object>

</thread>
<thread>
creation date:
1730594516803

thread_type:
official_review

<thread_object>
creator: reviewer_dzro

<content>
summary:
This paper studies the problem of how to improve many-shot in-context learning (ICL). The first observation is that the quality of ICL examples can be more important than the quantity -- a carefully selected subset of in-context learning examples can outperform the whole set. 

Based on this observation, this paper naturally moves to the question of how to improve the quality of ICL examples. The key idea is to iteratively select helpful examples from both the seed examples and those generated by the model itself, similar in spirit to STaR but applied to ICL instead. Specifically, in each iteration, the model is prompted to generate both reasoning steps and the answer. Only those with correct answers are kept. Then a Bayesian Optimizer is used to select the best ICL examples in this iteration, which are used as ICL examples for generation in the next iteration.

Results on BBH show that the proposed method largely outperforms baseline methods. Small improvements are achieved on MATH and GSM8K. A large improvement is observed on the BIRD-Challenging subset.

strengths:
- Many-shot ICL is an important problem from both the capability and safety angles. 
- This paper shows that even for many-shot ICL, good examples are still necessary. 
- The proposed method is a natural extension of the STaR framework to ICL. Given the hardness of optimizing ICL example selection, the proposed optimizer looks like a good contribution.
- Experimental results on BBH look strong, and small but non-trivial improvement are shown on math (MATH and GSM8K) and text-to-sql (BIRD) datasets.

weaknesses:
- The experiments are done only with Gemini models on text tasks. It would be interesting to see the performance on GPT-4o / o1 models and multimodal tasks. 
- Variance is not reported for MATH, GSM8K, and BIRD. Given the relatively small improvement on these datasets, it's unclear if the gain is significant.
- As the method optimize the performance on a dataset/task, it would be interesting to see the performance under distributional shifts. For example, does the gain on GSM8K transfer to GSM-symbolic?

questions:
I don't fully understand the claim that BRIDGE is tailored to many-shot ICL. I feel it also applies to few-shot ICL?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730711163702

thread_type:
official_review

<thread_object>
creator: reviewer_cswn

<content>
summary:
This paper addresses the many-shot in-context learning (ICL) paradigm. In contrast to few-shot learning, many-shot ICL leverages a larger number of demonstration examples to better exploit long-context capabilities. Initially, the paper demonstrates that selecting influential demonstration examples remains valuable, even in long-context settings. Building on this observation, the authors propose an algorithm that alternates between two key steps: an optimization step, which employs Bayesian optimization to identify influential example sets, and a generation step, which reuses these sets to automatically extend reasoning paths, thereby facilitating the transition back to the many-shot regime.

strengths:
- The proposed framework integrates example selection with LLM-generated examples to better align with real-world applications. 

- The paper validates this approach through extensive experiments conducted across a wide range of tasks.

weaknesses:
The validation of the findings appears overly simplistic. For instance, as model capabilities scale, the order and number of demonstration examples may have diminishing importance in affecting in-context learning (ICL) performance. The authors should reconsider their conclusions by re-evaluating their results on less advanced LLMs, providing a more comprehensive understanding of the interplay between long-context utilization and demonstration selection. Without this analysis, the conclusion may be over-claimed to specific models.

Lack of baselines. In the experiments, the authors did not compare their approach with existing learning-based strategies and learn-free ones for demonstration selection. 

The writing of the paper requires further refinement. Sections 3 and 4 are challenging to read and follow, despite the method itself being relatively straightforward. For instance, the description of the different variants could be more appropriately placed in the Experimental Setting section for better clarity.

questions:
Q1 In reference to Figure 1, does bottom_k refer to selecting examples that are less similar to the test case? Why does bottom_k achieve performance comparable to top_k as the number of demonstrations increases? Additionally, why does bottom_k outperform the all-examples approach?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732255321673

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Overall response

comment:
Overall Response
We thank all reviewers for their feedback and suggestions! In the common response below, we’d like to provide in-depth answers to some of the common questions, and we also summarize the key changes in the revised manuscript.

### Experiments on Additional LLMs
We appreciate the reviewers’ suggestions on experiments on additional LLMs, and we agree that it is an important point. As such, we have added results on Mistral large (mistral-large-2407) and Claude 3.5 Sonnet. While **we refer the reviewer to App. C5 of the updated manuscript for detailed results**, we’d also like to summarize the key results below (the Gemini results are taken from Tables 1 and 3 in the paper). We find that in all LLMs BRIDGE led to improvements, which suggest that the gains in the paper are not model-specific.


Table S1: Aggregated BBH results in 4 different LLMs
|                      Model                     | Reinf ICL | Iterative Reinf. |       | BRIDGE (Ours) |       |           |           |           |
|:----------------------------------------------:|----------:|-----------------:|------:|--------------:|------:|----------:|----------:|----------:|
|                  # Iterations                  |         0 |                1 |     2 |            1o |    1g |        2o |        2g |        3o |
|       Gemini 1.5 Pro (gemini-1.5-pro-001)      |     79.61 |            81.61 | 82.37 |         82.11 | 84.61 |     85.77 | **87.13** |   _86.33_ |
|     Gemini 1.5 Flash (gemini-1.5-flash-001)    |     80.25 |            81.91 | 80.72 |         81.61 | 82.79 | **83.79** |   _83.77_ |     83.25 |
|       Mistral Large (mistral-large@2407)       |     84.82 |            83.22 | 87.08 |         86.65 | 83.70 |   _88.07_ |     82.80 | **88.52** |
| Claude 3.5 Sonnet (claude-3-5-sonnet@20240620) |     88.45 |            89.89 | 90.35 |         89.16 | 91.26 | _92.00_ |   **92.20** |     90.97 |


### Key Changes
Here we summarize the key changes in the latest revision. We highlight all changes and additions in $\textcolor{magenta}{\text{magenta}}$ for the convenience of the reviewers, and below we list all key changes.

- Added standard deviation for MATH and GSM-Hard results in Table 2.
- Added the infilling baseline utilizing all samples + model-generated reasoning given inputs and ground-truth labels in Table 1 & 3, and the implementation details are in Appendix B.2.
- Added results for Mistral Large and Claude 3.5 Sonnet in Appendix C.5
- Added results comparing BRIDGE against additional demonstration selection methods in Appendix C.1.
- Added transfer learning experiments (transferring from GSM-Hard to GSM-8K) in Appendix C.6.
- Corrected typo in Algorithm 2 .
- Updated captions of Fig 1 to better clarify the meaning of “top-K” and “bottom-K”.
- Added the cost analysis to Appendix D.
- Miscellaneous adjustments to Sec 3 for better clarity.
- Updated Conclusion to reflect the future direction of more efficient methods with white-box model access.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper investigates the many-shot in-context learning (ICL) paradigm, emphasizing that the quality of selected demonstration examples can significantly influence performance even when leveraging long-context capabilities. To address this, the authors propose BRIDGE, an iterative algorithm that combines Bayesian optimization for influential example selection with model-generated reasoning examples. The method alternates between optimization and generation steps to enhance the quality of demonstrations and reasoning paths. The effectiveness of BRIDGE is demonstrated through experiments on various tasks, particularly on reasoning-heavy benchmarks like BigBenchHard (BBH), MATH, GSM8K, and BIRD.

Strengths

+ The paper addresses the critical problem of enhancing many-shot ICL, a relatively new but impactful area due to the rise of long-context models.
+ The proposed framework integrates example selection and model-generated reasoning, making it applicable to both closed- and open-source models.
+ BRIDGE is well-motivated by the observation that a subset of influential examples can drive performance improvements, even in a many-shot regime.
+ Extensive experiments on reasoning-heavy tasks (e.g., BBH, MATH, GSM8K, and BIRD) validate the proposed method, with strong improvements on BBH and smaller yet non-trivial gains on other datasets.
+ The study provides novel insights into the importance of example selection over mere quantity, highlighting its practical implications for ICL.

Weaknesses

+ Experiments are limited to the Gemini family of models, with no evaluation on open-source models like LLaMA or other state-of-the-art LLMs (e.g., GPT-4o), limiting the generalizability of the findings.
+ Important baselines, such as using all examples with model-generated reasoning, are missing, which weakens claims about the necessity of example selection and iterative refinement.
+ Variance in experimental results, particularly for tasks like MATH and GSM8K, is not reported, making it unclear whether the observed gains are statistically significant.
+ The computational cost of the optimization and generation steps in BRIDGE is not analyzed, which could limit the method's practicality.
+ The paper lacks clarity in certain sections, particularly the explanation of the "optimize" step, Bayesian optimization setup, and the estimation of importance scores.
+ The analysis is restricted to reasoning-heavy tasks, leaving open questions about BRIDGE's applicability to simpler tasks or models with limited reasoning abilities.

Most concerns have been addressed by the authors during the rebuttal period.