What Makes In-context Learning Effective for Mathematical Reasoning

We sincerely appreciate your recognition of our reasonable, correct, and self-contained theoretical analysis, the novelty and effectiveness of our method, and our convincing and strong empirical validation.

$\bf{Q1}$: The generalizability of the theory and methods in this paper could be further discussed.

$\bf{A1}$:Thanks for your constructive suggestion! In this paper, we are motivated by the observation that, on several math datasets, LLMs may perform worse in one-shot than in zero-shot setting. Therefore, in our experiments, we use these math datasets for evaluation, which can provide direct evidence of our method’s advantage.

As highlighted in Appendix E, our theories and method can also generalize to other datasets and tasks. This is because they are built upon a general setting of transformer attention layer and the relationship between demonstrations and test samples, which is also suitable to other domains and tasks. For instance, in Section 5.6, we conduct experiments on CommonsenseQA dataset, which is a widely used large-scale commonsense benchmark in ICL research[1,2,3]. From Table 4, our LMS3 still achieves the best performance, which confirms its effectiveness and highlight its generalizability on a broader range of datasets/applications.

Following your suggestions, we will supplement the above discussions in the revised version.

[1] Compositional Exemplars for In-context Learning.

[2] In-Context Learning with Iterative Demonstration Selection.

[3] Rethinking the Role of Demonstrations: What Makes In-Context Learning Work?

$\bf{Q2}$: In Eq.(6), I am unsure how the pre-trained data $D_{pre}$ and the demonstration $z_0$ are optimized simultaneously.

$\bf{A2}$: Thanks for your valuable question! We sincerely apologize for the confusion and would like to clarify as follows. In the “Analogy to Linear Optimization” section, we interpret the influence of adding a demonstration $x$ as follows: we start with a linear function $F$, whose parameters are initialized on the pretraining dataset $D_{pre}$. Introducing $x$ is essentially equivalent to adding a training example $z_0$ to further optimize $F$, after which the optimized $F$ is used to reason the test sample $x_{test}$.

Based on this idea, our goal is to quantify the change in test loss $L$ for $x_{test}$ resulting from adding $z_0$. To achieve this, inspired by the influence function, we define Eq.(6) to denote he parameters after training with $z_0$, which can be obtained by setting $\epsilon=\frac{1}{|D_{pre}|}$ in Eq.(6). On this basis, we further quantify the testing loss $L$ leveraging Taylor approximation as shown in Eq.(8) and (9), which serve as the foundation for subsequent theoretical analysis. Thus, in fact, Eq.(6) is a conceptual and intermediate tool for theoretical analysis rather than representing actual training on $D_{pre}$ and the demonstration $x$ simultaneously. In response to your comments, we will incorporate the above discussion and clarification to make our paper clearer.

$\bf{Q3}$: In Eq.(23), $\lambda_1$ is misused. There are some typos.

$\bf{A3}$: Thanks for your meticulous review and pointing out these issues! We will carefully correct the misuse of $\lambda_1$ and the typos in the revised version.

We sincerely appreciate your recognition of our theoretical analysis, the clarity and good writing of our paper, and the strong performance of our method. As for your concerns:

$\bf{Q1}$: The theoretical analyses in this paper are highly general. Therefore, I suggest the authors to discuss the possibility to extend them to other domains or tasks (e.g., multi-choice QA).

$\bf{A1}$: Thanks for your recognition of our theoretical analysis and this valuable suggestion. We appreciate the opportunity to discuss its broader applicability.

Although the motivation of our paper stems from observations in mathematical reasoning task, the underlying principles can be applied more broadly. Indeed, as discussed in Appendix E, our conclusions can be extended to other tasks beyond those explored in this work. This is because our theoretical analyses are based on a general setup of transformer architecture and the relationship between demonstrations and test samples. As long as a task can benefit from demonstration-based prompting (e.g., multi-choice QA as you mentioned), our theoretical conclusions about LLM-oriented Semantic Similarity and Inference Stability of Demonstration in Eqs. (21) and (22) remain applicable.

To validate this, we applied our method to CommonsenseQA dataset in Section 5.7, which is a large-scale commonsense benchmark and has been widely used in ICL research[1,2,3]. As shown in Table 4, our method still achieves the best performance, further demonstrating its general applicability.

We sincerely appreciate your insightful comment and are very willing to explore the performance of our method on more tasks for future research. We will also enrich our discussion section in the revised version.

[1] Compositional Exemplars for In-context Learning.

[2] In-Context Learning with Iterative Demonstration Selection.

[3] Rethinking the Role of Demonstrations: What Makes In-Context Learning Work?

$\bf{Q2}$: Combine the rejection mechanism with other demonstration selection method to further validate its necessity.

$\bf{A2}$: Thanks for your constructive suggestion! Following your suggestion, we supplement additional experiments using LLaMA3-8B as the backbone as follows. Specifically, we apply our demonstration rejection mechanism to all baselines, denoted as "+Our" in the table below.

The results consistently show that our mechanism enhances all baselines, regardless of whether they rely on retrieval-based similarity metrics (e.g., TF-IDF, BM25) or influence-based strategies. This suggests that our rejection mechanism can serve as a general enhancement technique that improves the robustness and effectiveness of various demonstration selection approaches. This also highlight the necessity of considering when to include a demonstration in in-context learning, rather than always providing demonstrations indiscriminately. Besides, our method leads to performance improvements across different datasets, demonstrating its broad applicability.

Thank you again for your valuable suggestion! We will incorporate this experiment and its analysis into the revised version to further support our findings.

$\bf{Q3}$: Just for curiosity, to measure the Inference Stability of Demonstration $X$, could we directly test the performance of the inference LLM on $X$ (e.g., calculate the accuracy)?

$\bf{A3}$: Thanks for your insightful question! Yes, we believe that directly testing the performance of demonstration $X$, such as calculating accuracy, can serve as a direct measure of the Inference Stability of $X$. This provides a straightforward way to assess its stability.

However, one potential challenge is that achieving a reliable measurement of stability in this way might require multiple calls to the LLM. This is because a single inference may not fully reflect the model's performance over time, while averaging results from multiple runs could give a more accurate measurement of the demonstration's stability. Consequently, this method could incur additional computational costs due to the need for repeated evaluations of the demonstration.

$\bf{Q4}$: I found some typos.

$\bf{A4}$: Thanks for your meticulous review and pointing out these typos! We will carefully correct them in the revised version.

We sincerely appreciate your affirmation of the effectiveness of our experiments, the good writing of our paper, and the significance of our work. As for your concerns:

$\bf{Q1}$: The algorithm presented requires white box access to the LLM. However, given the strong generalization performance across LLMs, this might not be an issue.

$\bf{A1}$: Thanks for your valuable comment! Yes, as our theoretical analysis suggests, the influence of demonstrations is determined by two key factors: LLM-oriented Semantic Similarity and Inference Stability of Demonstration. As defined in Eqs. (21) and (22), both factors are computed based on model representations. While they require white-box access to the model, this aligns with intuition, as different LLMs may benefit from different demonstrations depending on their own capabilities and characteristics.

To address the concern about generalization, as you pointed out, we examined this aspect in Section 5.6 by applying the demonstrations selected by our LMS3 method using Llama3-8B as backbone to ChatGPT and GPT-4. As shown in Table 3, these demonstrations still significantly improve their accuracy, demonstrating the strong generalization capability of our method. Under these conditions, LMS3 continues to achieve the best overall performance, highlighting its potential to provide valuable demonstrations even when working with closed-source LLMs in practical applications.

We sincerely appreciate your thoughtful comment and hope this explanation addresses your concerns.

$\bf{Q2}$: In the K-shot setting, is it fair to treat the demonstrations as independent of each other?

$\bf{A2}$: Thank you for the insightful question! In this work, our theoretical analysis considers the most fundamental case and starts with a single attention layer in transformers. Under this setup, as shown in Eq. (17), the influences of different demonstrations on the representation of the test sample $h_{test}$ follow an almost linear relationship, which allows us to treat them independently. In a full transformer architecture where deeper interactions occur (e.g., multi layers of cross-attention), the representations of different demonstrations will interact in more intricate ways, which may lead to complex effects on $h_{test}$ that could not be directly measurable. Therefore, our findings provide a foundational understanding that can offer insights into practical scenarios.

Moreover, even under this theoretical simplification, our method LMS3 consistently outperforms the baselines across all settings from 2-shot to 4-shot (Figure 3). This validates the feasibility of our theoretical results and the effectiveness of our method. Following your comment, we are very willing to further explore the impact of different demonstration combinations in the future.

$\bf{Q3}$: Can insights from this work potentially be used to design a method that generates effective demonstrations? This might result in better performance compared to a fixed offline demonstration set.

$\bf{A3}$: Thanks for your constructive suggestion! We fully agree the idea that our work can be used for effective demonstration generation. This is because our theoretical analysis sheds light on what characteristics contribute to their effectiveness and we can easily estimate the scores of demonstrations by Eq.(24) in our paper. This presents an exciting direction worth further exploration.

Further considering this idea, we think the only challenge in implementation is ensuring that the generated demonstrations have correct answers. While a fixed offline demonstration set allows for manual curation to guarantee correctness, dynamically generated demonstrations require additional mechanisms to verify their validity and reliability. Developing such mechanisms remains an open and important question.

We greatly appreciate your thought-provoking idea and will supplement the discussion it in our revised version.

$\bf{Q4}$: Inference runtime for a query if LMS3 is used using a fixed demonstration dataset. If I understand currently, the overhead should be minimal?

$\bf{A4}$: Thanks for your valuable question. Yes, your understanding is correct. Since the representations of all demonstrations can be precomputed in advance, the inference process in our LMS3 only involves encoding the test sample and retrieving the relevant precomputed information. This ensures that our method achieves the minimal computational complexity during inference as shown in Table 1.

We appreciate your thoughtful question and hope this clarifies our approach.