Towards Compute-Optimal Many-Shot In-Context Learning

Thank you for recognizing the effectiveness and straightforward nature of our proposed strategies, as well as the wide range of experiments used in our study. We appreciate your insightful feedback and the opportunity to respond to your concerns in detail. We address each of your points one by one in the following, in the order they are presented.

The tested models are both from the Gemini family, but it doesn't show any results on other LLMs, especially the open-source models like llama.

Thank you for your input. We would like to highlight that our choice of models and datasets was guided by the computational resources available to us, a strategy followed by all prominent prior works on many-shot ICL, such as Agarwal et al. (2024) and Bertsch et al. (2025). For example, Agarwal et al. (2024) primarily used a single LLM (Gemini 1.5 Pro) in their experiments, while Bertsch et al. (2025) mostly focused on classification datasets. Informed by these precedents, we did our best to carefully balance the number of models and datasets, ensuring diversity in both model size and task type. Specifically, we selected two LLMs of significantly different sizes, and our dataset choices span a range of tasks, from complex classification to tool use, reflecting practical applications valued by the community.

As a result, given our available resources, we strongly believe that our study includes a sufficiently diverse set of models, datasets, and scenarios to effectively demonstrate the benefits of our proposed demonstration strategies in many-shot settings.

Some experiment results seem to contradict one of the motivations. In Line 84-87, it states that similarity-based selection effectively becomes equivalent to random selection. However, in Figure 4, the performance gap between these two methods can be as large as 10%+ accuracy (see TREC dataset). The statement in Line 230-232 becomes "performance for both models remains almost unchanged beyond the 50-shot setting".

Thank you for your input. We would like to clarify that our claim—that “similarity-based selection effectively becomes equivalent to random selection”—is always accompanied by the important qualifier “beyond a certain point.” This qualifier is explicitly included in the relevant parts of the paper, specifically in lines 45–49 and 84–87, to emphasize that this approximation occurs only after a certain threshold. Our results, shown in Figure 4, clearly support this claim, as do previous studies such as Agarwal et al. (2024) and Bertsch et al. (2025).

Regarding the TREC dataset, it can be seen in Figure 4 that the performance of the random selection strategy begins to converge to the similarity-based strategy as the number of demonstrations increases—particularly from 150 to 200. However, for this dataset, the point at which similarity-based selection starts to closely approximate random selection seems to occur beyond 200 demonstrations. This convergence point is not visible in our current results due to the experimental limit of 200 demonstrations. Nevertheless, this does not invalidate our claim, and the TREC dataset still falls within the scope of our assertion. In fact, as shown in the work of Bertsch et al. (2025), specifically in Figure 2b, further increasing the number of demonstrations leads to a close approximation between the two strategies for this dataset. Therefore, our claims in the paper are fully consistent with the reported results, and there is no contradiction between our assertion and the data.

---

Questions:

In Figure 5, why does hybrid similarity-random outperform hybrid similarity-k-means most of the time? Intuitively, involving k-means should be no worse than random selection, but your experiments give a reverse conclusion.

Thank you for your question. We would like to emphasize that the effectiveness of one selection strategy compared to another depends on several factors, such as the downstream task, pretraining/fine-tuning data, training recipe, model size, and more. For example, as shown in Figure 4, the hybrid similarity-random method performs better than the hybrid similarity-$k$-means method on the ANLI dataset using Gemini Flash. However, when using Gemini Pro on the same dataset, the results are reversed. A similar but opposite trend can also be seen with the MetaTool dataset. Therefore, there is no strict rule to guarantee that one specific method will consistently perform better than the other.

Thank you for highlighting the effectiveness, practicality, and straightforward nature of our proposed demonstration selection strategies. We appreciate your thoughtful feedback and the opportunity to respond to your concerns in detail. Below, we address each of your concerns in the order they are raised.

When the entire prompt is cached (i.e. random selection), the complexity of attention calculation on the demonstration segment should reduce to O(1), yet Figure 2 depicts a different scaling; moreover, Fig. 2’s X‑axis metric and Fig. 4’s Y‑axis metric appear inconsistent.

Thank you for your input. We would like to clarify that the complexity of attention computation, even when the prompt is cached, is not $\mathcal{O}(1)$. Although the cached portion does not need to be recomputed, each token from the downstream test sample still has to attend to all cached tokens as well as to itself.

More formally, when considering proportional inference cost, the complexity is given by $\mathcal{O}(t^2 + tc)$, where $t$ represents the number of tokens in the downstream test sample, and $c$ represents the number of cached tokens. Since $t$ is fixed across different $n$-shot ICL settings, the dominant term becomes $\mathcal{O}(tc)$. As $c$ increases, the dominant term simplifies to $\mathcal{O}(c)$. Therefore, complexity grows linearly with the number of cached tokens.

For a more detailed explanation, we would like to refer you to lines 207–215 of the paper, where this analysis is formally detailed.

Experiments evaluate only two models from the same Gemini family, so the claimed compute–accuracy benefits may not hold for other LLM architectures or vendors.

Thank you for your input. First, we would like to clarify that our experiments involve two LLMs and seven different datasets, tested under two main scenarios: a data-rich setting and a low-data regime. This results in a total of 14 different experimental settings.

Additionally, we would like to highlight that our choice of models and datasets was guided by the computational resources available to us, a strategy followed by all prominent prior works on many-shot ICL, such as Agarwal et al. (2024) and Bertsch et al. (2025). For example, Agarwal et al. (2024) primarily used a single LLM (Gemini 1.5 Pro) in their experiments, while Bertsch et al. (2025) mostly focused on classification datasets. Informed by these precedents, we did our best to carefully balance the number of models and datasets, ensuring diversity in both model size and task type. Specifically, we selected two LLMs of significantly different sizes, and our dataset choices span a range of tasks, from complex classification to tool use, reflecting practical applications valued by the community.

As a result, given our available resources, we strongly believe that our study includes a sufficiently diverse set of models, datasets, and scenarios to effectively demonstrate the benefits of our proposed demonstration strategies in many-shot settings.

Regarding the generalizability of our findings, we would like to examine the two core components that underpin our proposed strategies: (1) scaling the number of demonstrations, and (2) key-value (KV) caching.

For the first component—scaling demonstrations—we would like to refer to Bertsch et al. (2025), which showed that ICL performance consistently improves as the number of demonstrations increases, across various LLMs, such as LLaMA models. Since our method is also based on the same principle of increasing demonstrations, our results also generalize similarly. The second component, KV caching, is inherently model-agnostic, meaning the computational savings it provides are not tied to any specific LLM architecture. Therefore, as these two generalizable components form the foundation of our strategies, we are confident that our strategies generalize well to other LLM architectures as well.

Thank you for pointing out the effectiveness, practicality, and straightforward nature of our proposed strategies. We are grateful for your constructive input. We appreciate the opportunity to address your concerns. To do so, we have provided our responses in the same order in which your concerns are listed.

1. The study evaluated the methods on two models - Gemini Pro and Flash to further test the generalizability of the method. It would be helpful if more open-source LLMs could be tested on the datasets.

Thank you for your input. First, we would like to clarify that our experiments involve two LLMs and seven different datasets, tested under two main scenarios: a data-rich setting and a low-data regime. This results in a total of 14 different experimental settings.

Additionally, we would like to highlight that our choice of models and datasets was guided by the computational resources available to us, a strategy followed by all prominent prior works on many-shot ICL, such as Agarwal et al. (2024) and Bertsch et al. (2025). For example, Agarwal et al. (2024) primarily used a single LLM (Gemini 1.5 Pro) in their experiments, while Bertsch et al. (2025) mostly focused on classification datasets. Informed by these precedents, we did our best to carefully balance the number of models and datasets, ensuring diversity in both model size and task type. Specifically, we selected two LLMs of significantly different sizes, and our dataset choices span a range of tasks, from complex classification to tool use, reflecting practical applications valued by the community.

As a result, given our available resources, we strongly believe that our study includes a sufficiently diverse set of models, datasets, and scenarios to effectively demonstrate the benefits of our proposed demonstration strategies in many-shot settings.

Regarding the generalizability of our findings, we would like to examine the two core components that underpin our proposed strategies: (1) scaling the number of demonstrations, and (2) key-value (KV) caching.

For the first component—scaling demonstrations—we would like to refer to Bertsch et al. (2025), which showed that ICL performance consistently improves as the number of demonstrations increases, across various LLMs, including open-weight LLaMA models. Since our method is also based on the same principle of increasing demonstrations, our results also generalize similarly. The second component, KV caching, is inherently model-agnostic, meaning the computational savings it provides are not tied to any specific LLM architecture. Therefore, as these two generalizable components form the foundation of our strategies, we are confident that our strategies generalize well to other LLM architectures as well.

2. Besides the testing of the current 4 datasets, it might be helpful if more datasets, especially the domain-specific ones, could be included for further experiments.

Thank you for your input. We addressed this before. Please find our response above.

---

Lastly, we hope our responses have effectively addressed your concerns. If you have any further questions or need additional clarification, please let us know. We would be more than happy to discuss any aspect in greater detail.

---

References:

[1] Many-Shot In-Context Learning (Agarwal et al., NeurIPS 2024)

[2] In-Context Learning with Long-Context Models: An In-Depth Exploration (Bertsch et al., NAACL 2025)

Thank you for recognizing the practicality and simplicity of our method, and we are happy you found our cost analysis clear. We also appreciate your thoughtful and detailed feedback. We are grateful for the opportunity to address your concerns. Below, our responses are provided in the same order as your comments.

Limited Model and Task Generalization: The evaluation is restricted to two models from the Gemini family and four datasets. This narrow scope raises concerns about the generalizability of the findings to other LLM architectures (e.g., LLaMA) and task types, particularly domain-specific or diverse NLP benchmarks.

Thank you for your input. First, we would like to clarify that our experiments involve two LLMs and seven different datasets, tested under two main scenarios: a data-rich setting and a low-data regime. This results in a total of 14 different experimental settings.

Additionally, we would like to highlight that our choice of models and datasets was guided by the computational resources available to us, a strategy followed by all prominent prior works on many-shot ICL, such as Agarwal et al. (2024) and Bertsch et al. (2025). For example, Agarwal et al. (2024) primarily used a single LLM (Gemini 1.5 Pro) in their experiments, while Bertsch et al. (2025) mostly focused on classification datasets. Informed by these precedents, we did our best to carefully balance the number of models and datasets, ensuring diversity in both model size and task type. Specifically, we selected two LLMs of significantly different sizes, and our dataset choices span a range of tasks, from complex classification to tool use, reflecting practical applications valued by the community.

As a result, given our available resources, we strongly believe that our study includes a sufficiently diverse set of models, datasets, and scenarios to effectively demonstrate the benefits of our proposed demonstration strategies in many-shot settings.

Regarding the generalizability of our findings, we would like to examine the two core components that underpin our proposed strategies: (1) scaling the number of demonstrations, and (2) key-value (KV) caching.

For the first component—scaling demonstrations—we would like to refer to Bertsch et al. (2025), which showed that ICL performance consistently improves as the number of demonstrations increases, across various LLMs, including LLaMA models. Since our method is also based on the same principle of increasing demonstrations, our results also generalize similarly. The second component, KV caching, is inherently model-agnostic, meaning the computational savings it provides are not tied to any specific LLM architecture. Therefore, as these two generalizable components form the foundation of our strategies, we are confident that our strategies generalize well to other LLM architectures as well.

Dear Reviewer KMMK,

We would like to thank you for taking the time to review our paper. As the discussion period is coming to an end, we would like to kindly remind you that we have carefully and thoroughly responded to all your concerns. We value your feedback and would appreciate it if you could provide any additional input or confirm whether our responses have addressed your concerns.

Thank you for your time and attention.

Best regards,
The Authors