What Makes In-context Learning Effective for Mathematical Reasoning
Authors: Jiayu Liu
Zhenya Huang
Chaokun Wang
Xunpeng Huang
ChengXiang Zhai
Enhong Chen
Venue: ICML 2025 poster (2025)
Keywords: In context learning
Large Language Models
Primary area: deep_learning->large_language_models

Abstract: Owing to the capability of in-context learning, large language models (LLMs) have shown impressive performance across diverse mathematical reasoning benchmarks. However, we find that few-shot demonstrations can sometimes bring negative performance and their effectiveness on LLMs' reasoning abilities remains unreliable. To this end, in this paper, we aim to theoretically analyze the impact of in-context demonstrations on LLMs' reasoning performance. We prove that the reasoning efficacy (measured by empirical prediction loss) can be bounded by an \emph{LLM-oriented semantic similarity} and an \emph{inference stability of demonstrations}, which is general for both one-shot and few-shot scenarios. Based on this finding, we propose a straightforward, generalizable, and low-complexity demonstration selection method named LMS3. It facilitates to select the most pertinent samples for different LLMs and includes a novel demonstration rejection mechanism to automatically filter out samples that are unsuitable for few-shot learning. Through experiments on three representative benchmarks, two LLM backbones, and multiple few-shot settings, we verify that our LMS3 has superiority and achieves consistent improvements on all datasets, which existing methods have been unable to accomplish. Our code is available at \url{https://github.com/Ljyustc/LMS3}.

What Makes In-context Learning Effective for Mathematical Reasoning
Jiayu Liu' Zhenya Huang" Chaokun Wang'
Xumpeng Huang
ChengXiang Zhai Enhong Chen
Abstract
MATH
0SMSK
Owing to the capability of in-context learning, large language models (LLMs) have shown im- pressive performance across diverse mathemat- ical reasoning benchmarks. However, we find that few-shot demonstrations can sometimes bring negative performance and their effectiveness on LLMss reasoning abilities remains unreliable. To this end, in this paper, aim to theoretically analyze the impact of in-context demonstrations on LLMs' reasoning performance. We prove that the reasoning efficacy (measured by empiri- cal prediction loss) can be bounded by an LLM- oriented semantic similarity and an inference sta- bility of demonstrations, which is general for both one-shot and few-shot scenarios. Based this finding, we propose a straightforward, generalizable, and low-complexity demonstration selection method named LMS3. facilitates to select the most pertinent samples for differ- ent ILLMs and includes novel demonstration rejection mechanism to automatically filter samples that are unsuitable for few-shot leamn ing. Through experiments on three representative benchmarks, two LLM backbones, and multiple few-shot settings, we verify that our LMS3 has superiority and achieves consistent improvements on all datasets, which existing methods have been unable to accomplish. Our code is available https: httphuu...oo//ysss////////
1.0
Ar1.3%
0.
0,8
A-133 MIS
0.5
01
Figure I. Problem-solving Accuracy of zcro-shot and one-shot settings. The hatched areas represent that in the one-shot setting, the model answers incorrectly A proportion of problems that are answered correctly the zero-shot setting.
2020; Liu et al., 2023). Currently, many large language models (LL.Ms) have exhibited strong performance across various mathematical benchmarks (Hendrycks al., 2021; Cobbe et al., 2021). A key capability of these LLMs is in-context learning (ICL) (Dong et al., 2022), which enables them to learn from few examples to implement specific log- ical structures (Wei et al., 2022) or utilize codes (Chen et al., 2023) to improve reasoning accuracy. Based on this ability, they have adeptly addressed problems across different types and difficulty, ranging from elementary word problems college-level algebra (Brown et al., 2020; Achiam al., 2023; Xue et al., 2024; Liu et al., 2025).
However, it remains an unresolved issue whether in-context learning truly enhances L.LMs' mathematical reasoning abil- ities. To show this phenomenon, in Figure 1, we present the zero-shot and one-shot results of Llama3-8B (Meta, 2024), ChatGPT, and GPT-4 (Achiam al., 2023) on two repre- sentative benchmarks MATH (Hendrycks al., 2021) and GSM8K (Cobbe et al., 2021). Surprisingly,
1. Introduction
find that 1) When given an example, their problem-solving accuracy does not consistently improve, and sometimes even declines (e.g., ChatGPT on MATH dataset). 2) Further analysis re- veals that the one-shot setting, these LLMs even fail in 1.83%-34.99 34.96% of problems (marked with white hatching) that they have correctly solved in the zero-shot scenario. This raises an important research question: Is in-context learning always effective for mathematical reasoning, and under what conditions does it work?
Mathematical reasoning serves as a milestone in assessing the progress of natural language processing (Zhang al.,
State Key Laboratory of Cognitive Intelligence, University of Science and Technology of China 'Institute of Artificial Intelli- gence. Hefei Comprehensive National Science Center HHon Kong University of Science and Technology 'Department of Computer Science, University of Illinois at Urbana- Champaign. Correspon- dence to: Zhenya Huang <huangzhy ustc.edu.en>..
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
address this issue, existing literature primarily analyzes the impact of demonstrations from an empirical perspective.

---

What Makes In-context Learning Effeetive for Mathematical Reasoning
For instance, researchers have revealed several important factors, including the similarity to test samples (Liu al.,
We validate our method on three mathematical bench- marks using multiple L.MSs as backhones, demonstrat- ing improvements in problem-solving accuracy, gener- alization ability, and interpretability.
2022), the diversity (Gao & Das, 2024), complexity (An et al., 2023) and perplexity (Sorensen et al., 2022) of demon- strations, as well as the Inference I.LM used (Peng et al.,
2024; Ye et al., 2023). Regarding the theoretical foundations of ICL, existing work has tried to explain the introduction of demonstrations as form of meeta-graiient optimization (Dai et al., 2023), kernel regression (Han et al., 2023), and token reinforcement (Yan et al., 2024). However, these studies I) have not provided precise quantification of the impact of demonstrations on LLMs' reasoning performance, nor 2) have they offered theoretical conclusions on when demon- strations are beneficial.
We further test our method on commonsense reason- ing task, validating its applicability to more general reasoning tasks.
2. Related Work
Mathematical Reasoning. Mathematical reasoning is critical benchmark for assessing the level of artificial intelli- gence (Zhang et al., 2020; Liu et al., 2023). Early research in this area mainly focused on rule-based, template-based, and statistical machine learning methods for simple math word problems (Feigenbaum et al., 1963; Fletcher, 1985). With the development of large language models (LLMs), re- cent efforts have primarily explored two directions. On one hand, some approaches equip LLMs with chain-like (Kojima et al., 2022), tree-like (Yao et al., 2024), or graph-like (Besta et al., 2024) reasoning processes, or require LLMs to use codes (Chen et al., 2023; Gao et al., 2023) and tools (Ma et al., 2025) to address potential numerical computation errors. On the other hand, there are efforts that provide the model with certain examples in the prompts through aurreval-augmentee generation (Wei et al., 2022; Asai et al., 2024), allowing the model to solve problems based on simi- lar approaches using its contextual learning abilities.
In this paper, we first theoretically analyze the impact of demonstration on the reasoning performance in one-shot scenario. We find that a sufficient condition for one-shot to outperform zero-shot is that 1) the demonstration and test sample have similar semantics encoded by the inference LLM and that 2) the inference LLM exhibits sufficient sta- bility in reasoning the answer of the demonstration itself. The former goes beyond traditional methods that rely solely on the semantic similarity between demonstrations and test samples, pointing out the critical role of the inference LLM's encoding parameters, while also being generalizable to these methods. The latter introduces the concept and measure- ment of inference stability of demonstrations for the first time. It should be emphasized that our theory is general and we further extend it to the k-shot scenario.
Based on theory, we propose simple yet effective demonstration selection method, named LMS3, to bal- ance the LLM-oriented Semantic Similarity and inference Stability of demonstrations, allowing for the automatic selec- tion of approximately optimal samples tailored to different LLMs. Additionally, to ensure that the sufficient condi- tion of our theories is essentially satisfied, we introduce an innovative demonstration rejection mechanism that can adaptively identify when few-shot learning should not be used, which is the first attempt in the field. Our method has strong theoretical advantages, generalization ability, and low complexity. Experiments on three benchmarks demonstrate its consistent improvements in both one-shot and few-shot scenarios. The contributions of this paper are as follows:
In-context Learning. In-context Learning (ICL) focuses on making LLMs learn and reason based on existing exam- ples (Dong et al., 2022). Its advantage lies in the adaptability and flexibility for different tasks and scenarios (Meade et al., 2023). However, the selection of examples remains cen- tral challenge, where current researches have developed supervised methods and unsupervised methods. This paper focuses on unsupervised methods, which can be grouped into three main categories. The first and currently most prominent method is called Similar-ICLL ICL (Liu et al., 2022: Luo et al., 2023; Zhang et al., 2023; Fu et al., 2022), which aims to find examples with closest semantic representations to the test sample. The semantic representation approaches include TF-ID,, BM25 (Robertson et al., 2009). T5 encod- ing (Raffel et al., 2020), BGE-M3 (Chen et al., 2024b), Ope- nAI embedding, etc. The second line of methods calculates the impact of each demonstration on the test sample (Peng et al., 2024). Impact calculation approaches include influ- ence function (Van et al., 2024; Chang & Jia, 2023), mutual information (Sorensen et al., 2022). perplexity (Gonen et al., 2023), code-length (Wu et al., 2023), etc. The third cate- gory uses the feedback from LLMS to dynamically select demonstrations (Nguyen & Wong, 2023; Qin et al., 2023).
We theoretically quantify the effect of demonstrations on ICL reasoning performance under one/few-shot settings. We prove that it can be bounded by LLM- oriented semantic similarity and inference stability of demonstrations.
We propose a novel demonstration selection method, LMS3, which can generalize to various existing meth- ods and offer better scalability and complexity.
Explanation of ICI., Regarding the underlying mechanisms
2

---

What Makes In-context Learning Effeetive for Mathematical Reasoning
of ICL, most existing research explored the impact of em- pirical factors such as the number of examples, gold labels, diversity, and types of LLMs from an experimental perspec- tive (Pan et al., 2023; Peng et al., 2024; Min et al., 2022). Some theoretical explorations explain ICL from perspec- tives including meta meta-gradient updates (Dai et al., 2023), kernel regression (Han et al., 2023), and token reinforce- ment (Yan et al., 2024). In comparison, to the best of our knowledge, we are the first to theoretically quantify the impact of demonstrations on reasoning performance and identify when they are effective.
parameter optimization, the prediction of a test sample hist is F(brest) Wo -hesst + VVI Wo) =0 -htet--
Based on this idea, Eq. (3) can be interpreted as: 1) We have a lincar function F(z) with initialized parameters
Wv va
Wo
hrese (Wihtest)
(4)
2) We introduce training data 2n = WWgh h to optimize the parameters, with the gradient at (20, Wo) satisfying:
Wv h.
VFI Wo)
(5)
vd
3. Theoretical Analysis
Notations. In in-context learning (ICL) setup, we have demonstration pool D and a test set Drest, which con- tain M and N' mathematical problems respectively. The k-shot ICL is formulated as appending k demonstrations {(1,, yi).... (X,, } S 'D' with the test data Xtest E Diest in prompt to reason the solution
3) We finally apply the optimized parameters to calculate the result of test data hest E Dlest-
Under this setting, we aim to estimate the influence of the data 20 = W (corresponds to demonstration X E DD on loss L(F(heest)) )). Before detailed derivation, we first give three mathematical annotations:
def (X,, (............ Ut), Xtest), Vrest
W def argmin W
iPprel EnEDpe ((F())
where X,; Xtest represent the problem context and y rep- resents the labeled solution. The prediction loss on Xtest is denoted as L(X test; Brest). In the following, we omit sym- bol y and use X to express each demonstration for brevity.
We.00
def argmin W
(Dprel DxED pr L(F(z)) + L(F(z)))
Hw
def
[Dpeel Ppe Vi (,2, W),
To evaluate the influence of a demonstration X on inferenc- ing the answer of X Xtest, we use h, heest E Rd to denote the representation of problem X and Xtest- Then, the Trans- former attention mechanism in ICL setting is denoted as:
(6) where Drre is the data for pre- pre-training an LLM, and Hw is the Hessian matrix which is positive definite by assump- tion (Van et al., 2024). It is worth noting that the pre-trained parameters W are actually the initialized parameters in our above setting, i.e., W Wo. Taking E [Dppel the testing
FieI.huead) = (Vnn(V K, Q. hicst)
(Wx(h, heose)" Weghesst Vd
Wy h, hreat- softmax
loss on heet represented as L( W [Perel ) On this basis, we derive the following theorem:
(2) where Wo, K Wy are the projection matrices for com- puting the attention queries, keys, and values, respectively. Without loss of generality, we omit Wo in Fier (brest) because we can redefine WK- = WT We- As a result, we only keep WX E Rdxd, Wr E Rd''x xd in our setting, where d' is the output dimension of layer FICL- Following a common approach (Dai et al., 2023), we approximate the attention to a simplified linear attention by:
Theorem 1. Assume L is Lipschitz continuous w.r.t JF with constant μ. If inequality (7) holds true, then L'bret: W ( < L(hrest: Wo,zo), i.c., introducing the training sample 20 (i.e.. demonstration X) can reduce _1 I are the largest and small- the testing loss heeat- Aas' A, est eigenvalues of Hw. respectively.
( h, h hieat Vd
Aaa A1
Wv hll uC1) vdd
(hee))) Wv hreat]
Vw (hhia WIll > l zoll (ll
Wy
Wv Vd
Wv Vd
hiest (Wicheea) hust Vd
h- (Wxh)"
heeat
G=
heael
|| Weethe |
(3)
(7)
We refer the readers to Appendix A for the detailed proof, and present the sketch here.
Analogy to Linear Optimization. We start our analysis of Eq. (3) by considering linear function F(z) def W 2, W E [ xd 2 E Rd. Specifically, given F(z) with an initialized parameters Wo. assume we have a training data 20 E Rd, then the gradient of loss L(F) can be written as AW- = V,I(Zo, Wo) =5- Applying the gradient to
Proof. With W. W the influence of upweighting 2g on
Please note that W.... is conceptual and intermediate tool for theoretical analysis, rather than actually training L.MSs both Dere and the demonstration 2g.
3

---

What Makes In-context Learning Effeetive for Mathematical Reasoning
the empirical loss is (Ling, 1984; Koh & Liang, 2017):
According to Eq. (7), the right-hand side of Eq. (16) is greater than 0, which leads the conclusion. D
[2) = dL dh We We,s0)
de
Extension to k-shot setting. In Theorem 1, we only con- sider one demonstration X (i.e., the one-shot scenario). For the k-shot scenario, Eq (3) can be written as
1=
--V Lheest, WTT T H Vw L( W) =
Then, the testing loss L(hrest, W
DDraa 2o) can be evaluated
Wv vd
(hres)
heast (Wichtest) hreat
by Taylor approximation since prea is sufficiently small, ppmn)1200
(17)
i... L(brest, W
Wv 2 vd
h. (We-(W. -heats
L(hreat, Wo.so) + [Dprel
1
dL dL(hueass We,z))
where h1, .... h& are the representations of demonstrations X1,.., Xg. This formalization can be interpreted intro- ducing & training samples 21 = Wx.h,..., 2k = W hę to optimize the linear function F(z) simultancously. where the gradient at each training sample 2; satisfies
de c=0
-L(heeat, Wo,so)
Vw L(bru 177
[Dprel
-Hw' L(o, W)
Wv vd
V Wo)
h.
(18)
Therefore, now the question turns to evaluate
Li de/ LW-naaa W) HW Vw L(zo W) (Vw W) W(hear W)) Hw Vw (bbesse, W)
Similar to the proof of Theorem 1, we derive the following Theorem illustrate the condition of these samples to ensure a reduction the loss of testing data Xtest, where
1.11
+ Wireeheest WT H Ve L( W)
def argmin W.,a w
(Dprel DLED L(F()))
1
(19)
L12
(10)
+ Srt (
Since Hw is positive definite, we denote A1 2 A2 ...
Aaa > 0 are the cigenvalues of Hw and can prove that
Theorem 2. Assume V is Lipsehitz continuous w.r.t F with constant ns. If inequality (20) holds true, then T.hheest W
L1 2 - WIll (III-(IIVZLL L W)
_ < L(hrest, Wo,s.), i.e., intro-
ducing training samples {z1,..., 2&3 (i.e., demonstrations XI,..,, can reduce the testing loss on h
L=0 W|ll [heee L Will | zoll) (11)
Since F is Lipsehitz continuous, we get I,11 2
kdad' VW WI > S -2.l--
X
X1 VW WI (4]) (hiest 20)1l | |+
(20)
Wv (II
hill + +(i)
V>L(zo, WIll [heet - =oll)
Vd
(12)
Applying Eqs. (4) (5) to Eq. (12), we have:
Theorem 2 further indicates that, under our setup, the joint effect of different demonstrations follows an additive re- lationship. This implies that the selection of k different demonstrations can be approximately considered indepen- dently. We leave the exploration of more complex interac- tions among demonstrations to future work.
[[ 20)1 20)1
hresell
Weaal |heest 2oll
dd
(13)
Wy Vd
[L,,, WIll I[brest
zoll =
h|| | - zoll (14)
For L12, we similarly prove that:
4. LMS3: Method Design
L12 = 2e02 2 Aaa II Vi WwwLCheee (heest W1
(15)
Based on Section 3, an ideal demonstration X needs to max- imize the value of Lı (i.e., minimize the empirical testing loss L(hrest: W | =0 in Eq. (9)). This equivalent to minimize the right-hind side of Eq. (7) according to Eq. (16)
Combining Eqs. (12)-(15), finally get:
L Zdae || Vw L L(hraat, WI|l >1I/V L(breat, WIll (u-C
and can be further divided into: 1) minimize
Wv hll | - 20ll).
* | - 201l
[lb.se|| || +
Vd
(16)
Sim(X) def |
W Wohl/,
(21)

---

What Makes In-context Learning Effective for Mathematical Reasoning
Test Problem test
I.LM-oriented
Wy
K-sbot prompling
What is the dossain of the function the where log
Semantic Similarity
((x) logez loatx 20j"
Large Language Model
cest
you refer Here donain of
is the base 10 leganiitm finction?
Sim (X) WP Wohll
Demonstration Rejection
formula g(+)
Wo
||resse
WF Woh
Now it's your tam!
Demonstration
Take deep breat
Think exp by step
I:
Wy
Inference Stability of Demonstration
Iwill $200
the finction g(+) 244
(x)
Solution
stab(X)=
W3
Q: Whet is the domoin of the funotion leght-lgge--00 whane
(0-0000000000000 U0(UU(EE U(-8,05) 9.
v7
Figure Illustration our proposed LMS3 method.
(recall 2g = Wgh and WK is indeed w Wq in the aforementioned section), and 2) minimize
heoretical Dependent Generality Complexity
Method
Guarantee on LLM
Similar-ICL
O(M + N) O D( + +.
Stab(X) def Wv Vd
Influence
x
h//-
(22)
InfICL
X
/
X
O(DM + DV NN)
ToDK+MDL
X
V
O(M + KN)
ConK+
X
O(M KN)
Specifically, Sim(X) reflects an LLM-oriented Semantic Similarity between the demonstration X and the test data Xtest- It illustrates the need to find a demonstration similar to the test sample. However, goes beyond traditional meth- ods by taking into account 1) the whole reasoning path of demonstrations (recall X includes both the problem context and the solution) and 2) the characteristics of the inference LLM itself, which is more consistent with intuition. The value of Stab(X) is an evaluation of the Inference Stability of Demonstration X. Based on Eq (5), Stab(X ) is indeed the length of gradient of the loss function on X. If Stab(X² low, indicates that the LLM has reached a stable predic- tion with minimal loss on X, and the parameters will not be excessively altered due to the introduction of this sample. In other words, the LLM has been able to stably deduce the correct answer for demonstration X.
IDS
O(M + QN)
MI
X
J
X
D(MN)
SPELL
X
V
OIMR +N) O(M+N)
LMS3 (ours)
V
V
Table 1. Comparison of methods, including Similar-ICI. (Liu al., 2022: Zhang et al., 2023; Fu et al.. 2022: Chen et al., 2024b). Influence (Nguyen & Wong. 2023) (S the size subset used estimate influences. V is the size of validation set), Inf[CL (Van et al., 2024) (D is the number parameters external LLMs). TopK+MDL (Wu et al., 2023), TopK+ConE (Peng al., 2024) (CC is the number candidate demonstrations), IDS (Qin et al., 2023) (Q is the number of iterations), MI (Sorensen et al., 2022), SPELL (Gonen et al., 2023) (R is the number of samples for estimating perplexity). The generality setting X indicates that these works are more suitable for classification tasks and hard to implement for mathematical reasoning task.
Since it is hard to simultaneously achieve the minimum of (i and Stab(X), ), two naive approximations are to minimize demonstration scoring function that calculates their sum or product:
highest Score(X) still do not satisfy Eq. (20). In such cases, unlike existing methods that always select top k examples, we tend to refuse to provide any demonstration and instead use a zero-shot approach, because our theorems suggest that providing examples in this case will have negative effect. We control Sim ) to achieve this rejection mecha- nism, because if an example's Sim(X) ) is already too large, Sim(X): pC1 might have exceeded the left-hand side of Eq. (7). However, setting an absolute threshold for Sim(X) (X) is challenging since μu, C1 is unknown, and calculating the gradient norm L W)ll is costly. Therefore, we adopt a simplified relative threshold. We expect that the Sim(X) of ideal demonstration should be small as possible relative to all examples. Consequently, we rank Sim(X) X)) of all candidate examples. If a demonstration X ranked top-k in Score(X) does not have Sim((X value within the top smallest, we reject to select it.
Score( X) = Sim(X) X) on Stab(X);
(23)
Score(X)= = Sim(X): Stab(X),
(24)
However, Eq. (23) requires considering the scale differences between the two objectives and adjusting the hyperparame- ter a1 based on different LLMs and datasets, which is chal- lenging to apply in practice. Therefore, we prefer Eq. (24) as the basic form of our scoring function. To implement k-shoo in-context learning, with the highest Score(X) X) demonstrations according to Theorem 2, which can ensure that the most relevant and stable samples are used to enhance the LLM's performance.
select the top le samples
Furthermore, we design a demonstration rejection mecha- nism. For instance, it possible that the examples with the
Theoretically, to compute Score (X), we need to input the
5

---

What Makes In-context Learning Effective for Mathematical Reasoning
Llama2-13B
Llama3-8B
MAWPS
GSM8K
MATH
MAWPS
GSM8K
MATH
zero-shot
0,835 0.009
0.414# 0.004
0,096 0.005
0.951 0,004
0.820= 0.016
0.324 0.022
Random
0.816# 0.004
0.405 0.007
0.090 0.010
0.951 = 0.005
0.813# 0.003
0.330# 0.009
Best-validate
0.826 0.001
0.410 0.005
0.096 0.007
0.932 0.000
0.817 0.008
0.332 0.008
TF-IDF
0,826 0.021
0.424 0.007
0.099 0.006
0.945 0.009
0.803 0.007
0.344 0.005
BM25
0,815 0.008
0.416 0.014
0.098 0.007
0.932 0.003
0.805 0.002
0.334 0.004
T5
0,810 0.004
0.426 0.013
0.093 0.006
0.948 0.021
0.817 0.002
0.330 0.009
BGEM3
0.818 0.013
0.407
0.100 0.011
0.938 0.017
0.802 0.000
0.340. 0.005
OpenAl
0.805 0.014
0.416 0.005
0.101 0.002
0.965 0.01
0.809 0.008
0.346 0.002
SPELL
0.797 0.009
0.394
0.085 0.003
0.945 0.005
0.82 0.008
0.343 0.004
Influence
0.836 0.010
0.405 0.009
0.102 0.000
0.929 L 0.009
0.800 0.015
0.333 0.006
IDS
0.839 + 0.005
0.424 0.012
0.088 0.001
0.920 0.003
0.808 0.001
0,330 0.001
LMS33 (ours)
0.854 0.008
0.447 = 0.014
0.124* = 0.003
0.966 0.010
0.837 0.011
0.353* + 0.002
Table 2. One-shot Answer Accuracy, with the best/runner-up methods highlighted in bold/underlined.
concatenation of each "(demonstration, testing data)" pair Xtett) into the LLM to obtain their semantic represen- tations h, hreet- However, in practice, this process requires ((M)) complexity (measured by the number of LLM API calls) for testing, which significantly limits the effi- ciency. Therefore, we approximate by inputting each data individually into the LLM to obtain its semantic represen- tation, reducing the testing complexity to O(M + WV) (the representations of demonstrations can be pre-computed).
Chang & Jia, 2023), our method does not require additional external LLMs, repeated testing of demonstration effects on validation set, or the computation of Hessian matrix, which brings much lower complexity.
5. Experiments
5.1. Experimental Setup
Datasets. We use three datasets that cover a variety of types and difficulty levels. MAWPS Kedziol-e et al., 2016) consists of 2,373 elementary-level math word prob- lems. GSM8K (Cobbe et al., 2021) composed of 8,792 more challenging elementary problems with higher num- ber of steps. MATH (Hendrycks et al., 2021) is collected from high school math competition, containing 12,500 prob- lems across seven categories, and is currently one of the most widely used benchmarks. Dataset partition and statis- tics are presented in Appendix C.
We illustrate the workflow of our method, named LMS3, in Figure 2 and present the pseudo-code in Appendix B. LMS3 has several advantages summarized in Table 1. 1) Theoretical Guarantee: To the best of our knowledge, we are the first to theoretically quantify the impact of demon- strations on ICL reasoning performance and explain why and when they work. 2) Rational Dependency: Our analy- sis verifies that the optimal demonstration depends on the inference LLM (i.e., how the representations h, hrest are en- coded). This is reasonable because LLM's understanding of similar problems sets the upper limit on its ability to lever- age these problems (Peng et al., 2024). Consequently, the optimal demonstration should be selected adaptively for dif- ferent LLMs However, existing methods like Similar-ICL estimate semantic similarity independently of the inference LLM and the demonstration is the same for all LLMs. 3) Generalization Ability: If we set WT We I identity matrix and omit Stab(X), ) our method degenerates into finding the demonstration with the closest semantic representation to the test data. This perspective unifies the current approaches, summarizing their main differences in the setting of WT- WQ to obtain semantic representations. Besides, our analysis is based on the impact of demonstra- tions on the test loss, which is not dependent on the task type. In addition to mathematical reasoning, it is also applicable to generation tasks or classification tasks. 4) Low Complex- ity: Compared to methods based on impact estimation or LLMs' feedback (Van et al., 2024; Nguyen & Wong, 2023;
Baselines. We use Llama2-13B (Touvron et al., 2023) and Llama3-8B (Meta, 2024) as the backbones and take 10 representative and SOTA baselines including:
Random randomly samples a data from D.
Best-validate selects the data with the highest accuracy on a validation set,
and some typical Similar-ICL methods:
TF-IDF represents each problem as a TF-IDF vector, and selects the nearest sample.
BM25 (Robertson al., 2009) selects demonstrations by retrieval method BM25.
T5 (Raffel et al., 2020) encodes problems with T5- large and selects the most similar one.
6

---

What Makes In-context Learning Effeetive for Mathematical Reasoning
BGEM3 (Chen et al., 2024b) integrate multiple infor- mation retrieval functionalities in unified embedding.
Rasdemi
TF-DF
OpesAl
1D5
LXSS
UM35
Influnnce
SPIILL
MAWPS
OpenAI (Neelakantan et al., 2022) represents prob- lems with OpenAl E-meedding-3-maal model,
0.955
0.90
Two
Thece
shot
and methods that do not rely on problem similarity:
0.85
GISMHK
SPELL (Gonen et al., 2023) selects demonstrations by calculating their perplexity.
0.90
Twoo-bot
Thrre-hot MATEE
Poor-shot
0.40
Influence (Nguyen & Wong. 2023) divides D into multiple subsets. The preference of demonstration is calculated by the difference in validation accuracy between subsets that include and exclude it.
0.35
0.50
Two-stor
Thrreshor
Fooe-sbot
IDS (Qin et 2023) iteratively selects demonstrations based on reasoning path similarity.
Figure 3. Few-shot Answer Accuracy of Llama3-8B.
GSMIK
MATH
13.0
Implementation Details. When implementing our LMS3, we set À to 10% for Llama2-13B -3B and 1% for Llama3-8B. The temperature for both L.LMs is set to 0.8. For the baseline Influence, the size of the subset is set to 20. For the baseline IDS, the number Q of iterations is set to 3. All experiments are conducted on a server with six NVIDIA RTX 3090 GPUs.
Tuesity
5.0
2.5
0.0
0.0
0.5 SearetX)
1.0
0.0
6.5 SoweUX
1.0 00
0.5
1.0
Scan(X)
Figure 4. Distribution of Score(X) in Eq. (24).
5.2. Performance on One-shot Reasoning
Figure 3. Firstly, our method remains superior across differ- ent settings, which is consistent with our Theorem 2, demon- strating its high applicability to various scenarios. Secondly, as k increases, the trend of reasoning performance varies across different datasets. On MAWPS and MATH, the per- formances of most methods consistently improve with a higher k. However, on GSM8K, the accuracy for almost methods declines after k = 3. This suggests that an ex- cessive number of demonstrations does not necessarily lead to increased accuracy and need to balance the number and length of demonstrations. A dataset with longer prob- lem lengths (i.e., GSM8K as indicated in Appendix C) may require fewer examples to achieve optimal performance.
In Table we present the performance of all methods in the one-shot setting. Firstly, it can be seen that our LMS3 outperforms all baselines across all datasets, and this effect is statistically significant with < 0.05 (marked +). This directly indicates that the demonstrations chosen by our method better stimulate the L.M's contextual learning abil- ity. We also present some cases in Appendix for clearer comparison. Secondly, our LMS3 is the only one that con- sistently provides improvements over the zero-shot setting, while other methods exhibit certain fluctuations across dif- ferent datasets. This can be attributed to our method being designed based on a theoretical analysis of when one-shot learning is effective (i.e., Theorem 1). These experimental results validate the rationality, effectiveness, and strong ro- bustness of our theoretical findings. Thirdly, we observe that one-shot learning generally improves the backbone's 's performance the more challenging MATH dataset, but sometimes shows a decrease on other datasets. We believe this is because the problems in MAWPS and GSMSK are relatively simple, and the LLM itself already has the capa- bility to solve them. Introducing additional examples in this case might instead mislead the model's reasoning thought.
5.4. Analysis of Scoring Function
Figure 4 presents the distribution of Score(X) X) in Eq. (24) normalized by z-score, which verifies that our Score(X) has good discriminative power for different samples. More importantly, we observe that the variances of the distribu- tions for Llama2 on GSMSK and MATH, Llama3 on MATH, are relatively small. This indicates that the differences be- tween samples in these cases are not significant, which can explain why most other one-shot baselines also perform bet- ter than the zero-shot setting in Table 2. In contrast, in other cases (e.g., on MAWPS), the performance gap between dif- ferent samples is larger, and only our LMS3 can consistently achieve better results than zero-shot setting.
5.3. Performance on Few-shot Reasoning
Now we validate our LMS3 in the k-shoo scenario, with Llama3-8B's performances at = (2,3,4) 4) visualized in
7

---

What Makes In-context Learning Effeetive for Mathematical Reasoning
Llané-48
ChatGPT GSM8K
GPT-4
S0AR
0020
MAWPS
MATH
MAWPS
GSMBK
MATH
0/82
zero-shot
0.906
0,800
0.482
0.941
0.954
0.758
499
Random
0,858
0,839
0.503
0.976
0.946
0.702
Best-validate
0.831
0.832
0.519
0.979
0.951
0.715
4.40
TE-IDF
0.895
0.820
0.514
0.975
0.947
0.724
0.123
BM25
0.901
0.828
0.510
0.987
0.953
0.691
0.900
031
TS
0.893
0.840
0.508
0,973
0,950
0.718
PATSIPS 307
185
20%
35%
100%
BGEM3
0.896
0.838
0.504
0.986
0.955
0.705
OpenAl
0.898
0.829
0.513
0.979
0.945
0.699
Figure 5. Performance with varying A. The dashed line corre- sponds to the result of the zero-shot setting.
Influence
0.878
0.848
0.515
0.974
0.955
0.702
IDS
0,908
0,848
0.505
0.979
0.959
0.742
LMS3 (ours)
0.909
0.862
0.517
0.990
0.961
0.752
Method
MAWPS
GSM8K
MATH
Random
0.951
0.813
0.330
Table 4. Generalization performance on ChaGGPT/GPT4,
+ours
0,952
0,818
0.349
Best-validate
0.932
0.817
0.332
when to include a demonstration in ICL.
+ours
0.941
0.829
0.344
TF-IDF
0,945
0,803
0.344
+ours
0,946
0.818
0.351
5.6. Generalization Ability
BM25
0,932
0,805
0.334
+ours
0.934
0.812
0.335
One advantage of our method is that it reasonably considers the inference LLM when selecting demonstrations. How- ever, it may raise concerns about its extensibility, as re- quires access to the LLM's internal parameters. To address this issue, we offer the demonstrations selected by LMS3 (Llama3-8B) directly to the most advanced LLMs ChatGPT and GPT-4, compared with other methods that do not rely on the LLM's parameters. From Table 4, we can see that our I.MS3 still achieves nearly optimal results, demonstrating the excellent generalization and representativeness of our selected demonstrations. Besides, we observe that all meth- ods negatively impact the reasoning performance on MATH dataset when applied to GPT-4. We attribute this to that chain-of-thougle examples may no longer be significantly beneficial for GPT-4, and future examples for GPT-4 might need to focus more on code (Chen et al., 2023; Gao et al., 2023) or other formats.
T5
0,948
0,817
0.330
+ours
0.953
0.828
0.333
BGEM3
0.938
0.802
0.340
+ours
0.941
0.822
0.350
OpenAl
0,965
0,809
0,346
+ours
0.973
0,818
0.347
SPELL
0,945
0,821
0.343
+ours
0.946
0.826
0.345
Influence
0,929
0,800
0.333
+ours
0.935
0.810
0.340
IDS 0,920
0,808
0.330
+ours
0,932
0.823
0.346
Table 3. Performance of combining our rejection mechanism with other demonstration selection methods.
5.5. Ne..ssesy of Demonstration Rejection
To validate the necessity of our proposed demonstra- tion rejection mechanism, we test the effects of X (1%, 5%, 10%, 20%, 40%, 60%, 80%, 100%). It is note- worthy that À = 100% is equivalent to removing our re- jection mechanism. From Figure 5, we can first observe that when = 100%. the accuracy of LMS3 falls below that of the zero- shot results, which highlights the necessity of our rejection mechanism. Secondly, when À increases, the performance of Llama2 initially rises and then falls, while the performance of Llama3 consistently declines. On one hand, this indicates that the strength of À needs to be bal- anced differently for various LLMs, but this process is not complicated since the optimal On the other hand, this demonstrates that our Sim(X) can effectively approximate the conditions in Theorems and 2, using it to implement the rejection mechanism can improve the model's performance. Thirdly, in Table 3, we apply our rejection mechanism to all baselines (denoted as "+ours") with Llama3-88 8B as the backbone. The consistent performance gains reflect the general applicability of our mechanism and underscores the importance of considering
5.7. Applicability to More Tasks
To expand the applicability of our work and demonstrate the versatility of LMS3, we include an additional experiment on CommonsenseQAA large-scale benchmark designed to evaluate commonsense reasoning task and has been widely used in ICL research (Ye et al., 2023; Qin al., 2023; Min et al., 2022). This extension allows us to showcase LMS3's effectiveness not only in mathematical reasoning tasks but also in broader reasoning domains that involve more nuanced and diverse knowledge.
basically within 10%.
The results in Table 5 show that LMS3 outperforms all baseline methods on both Llama2-133 and Llama3-8B, with an accuracy of 0.444 and 0.650, respectively. The im- provement is particularly significant on Llama2-13B, where LMS3 surpasses the next best method by a larger margin. These results highlight the effectiveness of LMS3 in enhanc- ing model performance across different tasks, demonstrating its potential for real-world applications.
8

---

What Makes In-context Learning Effective for Mathematical Reasoning
Llama2-13B
Llama3-8B
References
zero-shot
0.428
0,628
Random
0.415
0,625
Achiam, J., Adler, S., Agarwal, S., Ahmad, L., Akkaya, I. et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774 2023.
Best-validate
0.375
0.605
TF-IDF
0.368
0.617
BM25
0.383
0.627
T5
0,419
0,617
An, S., Lin, Z.,, Fu, Q. Chen, B., Zheng, N., Lou, J.-G., and Zhang, D. How do in-context examples affect composi- tional generalization? In Proceedings of the 6l Annual Meeting of the Association for Computational Linguistics (Volume Long Pupers), pp. 11027-11052, 2023.
BGEM3
0.377
0.589
OpenAl
0.393
0.607
SPELL
0.409
0.637
Influence
0.374
0,605
IDS
0,408
0.640
LMS3 (ours)
0.444
0.650
Table 5. Performance of CommonsenseQA benchmark.
Asai, A.. Wu, Z., Wang, Y., Sil. A., and Hajishirzi, H. Self- rag: Learning to retrieve, generate, and critique through self-reflection. The Twelfth International Conference on Learning Representations, 2024.
6. Conclusion
Besta, M., Blach, N., Kubicek, A., et al. Graph of thoughts: Solving elaborate problems with large language models. AAAI, volume 38, pp. 17682-17690, 2024.
In this paper. we theoretically analyzed how demonstra- tions affected LLMs' mathematical reasoning performance. On this basis, we proposed LMS3 method that balanced LLM-oriented semantic similarity and inference stability of demonstrations, and introduced a demonstration rejection mechanism to filter out negative situations. Experiments showed that our method was the only one to consistently improve the reasoning accuracy of LLMs, and our demon- strations exhibited strong generalization ability and inter- pretability. In the future, we will extend our method to more NLP tasks and apply our theory to broader scenarios. Please refer to Appendix E for more discussions and details.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Chang, T.-Y. and Jia, R. Data curation alone can stabilize in-context learning. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pp.
2023.
Chen, J., Lin, H., Han, X., and Sun, L. Benchmarking large language models in aitrivaal-augmente generation. In Proceedings of the AAAI Conference Artificial Intelli- gence, volume 38, pp. 17754-17762, 2024a.
Impact Statement
This work provides valuable insights into the underlying mechanisms of LLMs. By theoretically analyzing how in- context demonstrations affect LLM performance, we offer deeper understanding of LLM behavior, which paves the way for more effective optimizations in the future. This could lead to reduced costs and improved efficiency in de- ploying LLMs for reasoning tasks. Additionally, our LMS3 method has practical implications for improving the relia- bility and relevance of L.LM outputs in various applications, especially in decision-making fields that rely complex reasoning and aita-driven insights.
Chen, J., Xiao, S., Zhang, P., al. Bge m3-embedding: multi- cuatiiggranlarityy text embeddings through self-knowledge distillation. arXiv 2024b.
Multi-lingual,
preprint
Chen, W., Ma, X., Wang, X., and Cohen, W. W. Program of thoughts prompting: Disentangling computation from reasoning for numerical reasoning tasks. Transactions on Machine Learning Research, 2023.
Acknowledgement
Cobbe, K., Kosaraju, V., Bavarian, M., Chen, M., Jun, H., Kaiser, L., Plappert, M., Tworek, J., Hilton, J., Nakano, R. et al. Training verifiers to solve math word problems. arXiv preprint arXiv:21 10.1468,, 2021.
This research was partially supported by the Na- tional Natural Science Foundation of China (Grants No.62477044), Anhui Provincial Natural Science Foun- dation (No. 2308085QF229), the Fundamental Research Funds for the Central Universities (No. Noo.KK2500 10038). Zhenya Huang gratefully acknowledges the support of the Young Elite Scientists Sponsorship Program by CAST (No. 2024QNRC001).
Dai, D., Sun, Y., Dong, I.., Hao, Y., Ma, S., Sui, 7.,, and Wei, F. Why can gpt learn in-context? language models secretly perform gradient descent as meta-optimizers. In Findings of the Association for Computational Linguis- tics: ACI. 2023, pp. 4005-4014 2023.

---

What Makes In-context Learning Effective for Mathematical Reasoning
Ding, Y., Fan, W., Ning, L.., Wang, S., Li, H., Yin, D., Chua, T.-.,, and Li, Q- A survey on rag meets I1ms: To- wards retrieval-augmented large language models. arXiv 2024.
Kojima, T., Gu, S. S., Reid, M., Matsuo, Y., and Iwasawa, Y. Large language models are zero-shot reasoners. Ad- vances in neural information processing systems, 35: 22199-22213, 2022.
preprint arXiv:2405.0000
Dong. Q.. Li, Dai, D., Zheng, C., Wu, Z., Chang, B., Sun, X., Xu, J., and Sui, 7. A survey on in-context learning- arXiv preprint arXi::2301.0025: 2022.
Koncel-Kedziorsie R., Roy, S., Amini, A., et al. Mawps A math word problem repository. In NAACL-HII, pp. 1 1152-1157, 2016.
Feigenbaum, E. A., Feldman, J., et al. Computers and thought. New York McGraw-Hill, 1963.
Ling, R. Residuals and influence in regression, 1984.
Liu, J., Shen, D., Zhang, Y., Dolan, W. B., Carin, I.,, and Chen, W. What makes good in-context examples for gpt- 3? In Proceedings of Deep Learning Inside Out (DeeLIO 2022): The 3rd Workshop on Knowledge Extraction and Integration for Deep Learning Architectures, pp. 100- 114, 2022.
Fletcher, C. R. Understanding and solving arithmetic word problems: A computer simulation. Behavior Research Methods, Instruments, Computers, 17(5):565-57)) 1985.
Fu. Y.. Peng, H., Sabharwal, A., Clark, P.. and Khot, T. Complexity- based prompting for multi-step reasoning. In The Eleventh International Conference on Learning Representations, 2022.
Liu, J., Huang, Z., Zhai, C., and Liu, Q. Learning by ap- plying: A general framework for mathematical reasoning via enhancing explicit knowledge learning. In Proceed- ings of the AAAI Conference on Artificial Intelligence, volume 37, pp.
Gao, L., Madaan, A., Zhou, S., Alon, U.. Liu, P., Yang, Y., Callan, J., and Neubig, G. Pal: Program-aided lan- guage models. In International Conference on Machine Learning, pp. 10764-10799, PMLR, 2023.
2023.
Liu, J., Huang, Z., Liu, Q., Ma, Z., Zhai, C., and Chen, E. Knowledge eenwledge-centered dual -pooess reasoning for math word problems with large language models. IEEE Trans- actions on Knowledge and Data Engineering, 2025.
Gao, X. and Das, K. Customizing language model responses with contrastive in-context learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pp. 18039-18046, 2024.
Luo, M., Xu, X., Dai, Z., Pasupat, P., Kazemi, M., Baral, C., Imbrasaite, V., and Zhao, V. Y. Dr. icl: Demonstration retrieved in-context learning. arXiv preprint arXiv:2305. 14/28, 2023.
Gonen, H., Iyer, S., Blevins, T., Smith, N. A., and Zettle- moyer, L. Demystifying prompts in language models via perplexity estimation. In Findings of the Associa- tion for Computational Linguistics: EMNLP 2023, pp. 10136-10148, 2023.
Ma, 7.., Huang, Z., Liu, J., Wang, M., Zhao H., and Li, X. Automated creation of reusable and diverse toolsets for enhancing I1m reasoning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 39, pp. 24821-24832 2025.
Han, C., Wang, Z., Zhao, H., and Ji, H. Explaining emergent in-context learning as kernel regression. arXiv preprint arXiv:2305. 2023.
Mcade, N., Gella, S., Hazarika, D., Gupta, P., Jin, D., Reddy, S., Liu, Y., and Hakkani- D. Using in-context learning to improve dialogue safety. In Findings of the Association for Computational Linguistics: EMNLP 2023, pp. 11882- 11910, 2023.
He, D., Feng, Z., Jin, D.,. Wang, X., and Zhang, W. Joint identification of network communities and semantics via integrative modeling of network topologies and node con- tents. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 31, 2017.
Meta, A. Introducing meta llama The most capable openly available Ilm to date. Meta A1, 2024.
Hendrycks, D., Burns, C., Kadavath, S., Arora, A., Basart, S., Tang, E., Song, D., and Steinhardt, J. Measuring mathematical problem solving with the math dataset. In fiftty-fithh Conference on Neural Information Process- ing Systems Datasets and Benchmarks Track (Round 2), 2021.
Min, S., Lyu, X., Holtzman, A., Artetxe, M., Lewis, M.. Hajishirzi, H.,, and Zettlemoyer, L. Rethinking the role of demonstrations: What makes in-context learning work? In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp. 11048- 11064, 2022.
Koh, P. W. and Liang, P. Understanding black-box predic- tions via influence functions. In International conference on machine learning, pp. 1885-1894, PMLR, 2017.
Neelakantan, A., Xu, T., Puri, R., Radford, A., Han, J. M., Tworek, J., Yuan, Q., Tezak, N., Kim, J. W., Hallacy,
10

---

What Makes In-context Learning Effeetive for Mathematical Reasoning
C., al. Text and code embeddings by contrastive pre- training. arXiv preprint arXiv:2201. 2022.
Bhosale, S., et al. Llama 2: Open foundation and fine- tuned chat models. arXiv preprint arXiv:2307.09288 2023.
Nguyen, T. and Wong, E. In-context example selection with influences, arXiv preprint arXiv:2302.//002, 2023.
Van, M.-H., Wu, X., et al. In-context learning demon- stration selection via influence analysis. arXiv preprint arXiv:2402. 11750, 2024.
Pan, J., Gao, T., Chen, H., and Chen, D. What in-context learning" learns" in-context: Disentangling task recogni- tion and task learning. In The 6lst Annual Meeting Of The Association For Computational Linguistics, 2023.
Wei, J., Wang. X., Schuurmans, D., Bosma, M., Xia, F., Chi, E., Le, Q. V., Zhou, D., et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in Neural Information Processing Systems, 35: 24824 24837, 2022.
Pei, H., Yang, B., Liu, J., and Chang, K. C.-C. Active surveillance via group sparse bayesian learning. IEEE Transactions on Pattern Analysis and Machine Intelli- 2020.
Wei, J., Wei, J., Tay, Y., et al. Larger language mod- els do in-context learning differently. arXiv preprint arXiv:2303.m 03846, 2023.
gence,
Peng, K., Ding. L., Yuan, Y., Liu, X., Zhang, M., Ouyang, Y., and Tao, D. Revisiting demonstration selection strategies in in-context learning. arXiv preprint arXiv:2401.12087, 2024.
Wu, Z., Wang, Y., Ye, J., and Kong, L. Self-adaptive in- context learning: An information compression perspec- tive for in-context example selection and ordering. In Proceedings of the 61st Annual Meeting of the Associ- ation for Computational Linguistics (Volume 1: Long Papers), pp. 1423-1436, 2023.
Qin, C., Zhang, A., Dagar, A., and Ye, W. In-context learn- ing with iterative demonstration selection. arXiv preprint 10.00981, 2023.
Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. J. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21 (140):1667,2 2020.
Xiao, T., Liu, J., Huang. Z., Wu, J., Sha, J., Wang, S., and Chen, E. Learning to solve geometry problems via simulating human deal-neg process. In Proceedings of the Thirty-Thise Third International Joint Conference on Artificial Intelligence, pp. 6559-6568 2024.
Robertson, S., Zaragoza, H., et al. The probabilistic rele- vance framework: Bm25 and beyond. Foundations and Trends@ in Information Retrieval, 3(4):333 -0933338922000 2009.
Xue, S., Huang, Z., Liu, J., Lin, X., Ning, Y., Jin, B., Li. X., and Liu, Q. Decompose, analyze and rethink: Solv- ing intricate problems with human-like reasoning cycle. Advances in Neural Information Processing Systems, 37: 357-385, 2024.
Socher, R., Perelygin, A., Wu, J., Chuang, J., Manning, D., Ng, A. Y., and Potts, C. Recursive deep models for semantic compositionality over sentiment treebank. In Proceedings of the 2013 conference on empirical methods in natural language processing, pp. 1631-1642, 2013.
D.,
Yan, J., Xu, J., Song, C., Wu, C., Li, Y., and Zhang, Y. Understanding in-context learning from repetitions. In The Twelfth International Conference on Learning Repre- sentations, 2024.
Sorensen, T., Robinson, J., Rytting, C., Shaw, A., Rogers, K., Delorey, A., Khalil, M. Fulda, N., and Wingate, D. An ceermattion-thorettte approach to prompt engineering without ground truth labels. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 119-862, 2022.
Yao, S., Yu, D., Zhao, J., et al. Tree of thoughts: Deliberate problem solving with large language models. Advances in Neural Information Processing Systems, 36, 2024.
Ye, J.. Wu, Z., Feng, J., Yu, T., and Kong, L. Compositional exemplars for in-context learning. In International Con- ference on Machine Learning, pp. 99888-39833. PMLR, 2023.
Talmor, A., Herzig, J., Lourie, N., and Berant, J. Com- monsenseqa: A question answering challenge targeting commonsense knowledge. In Proceedings of the 2019 Conference of the North American Chapter of the Associ- ation for Computational Linguistics: Human Language Technologies, Volume (Long and Short Papers), pp. 4149 4158, 2019.
Zhang, D., Wang, I.., et al. The gap of semantic parsing: A survey on automatic math word problem solvers. IEEE Transactions on Pattern Analysis and Machine Intelli- gence, 42(9)):287-2323 2020.
Zhang, P., Xiao, S., Liu, Z., Dou, Z., and Nie, J.-Y. Re- trieve anything to augment large language models. arXiv preprint arxiv:2310.07554, 2023.
Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P.,

---

What Makes In-context Learning Effective for Mathematical Reasoning
A. Proof of Theorem 1
Therefore, now the question turns evaluate
We interpret the attention mechanism in ICL setting as: 1) We have a linear function F(2) with initialized parameters
L1 T.-eat W)" H W L/o, (31)
Wv heest : (Wxchestt" vd
Specifically, L1 can be represented as L1
W
(25)
(Vw W) Vw (heestt W)) Hw Ve w L (Ibreatss W)
2) We introduce a training data 20 = Wgh to optimize the parameters, with the gradient at (zo, Wo) satisfies
En
+ Ww-Leheass W)J w V Vw Libeest W)
Wv vd
V Wo)
(26)
L12
(32)
Since Hig- is positive definite, there exists an orthogonal
3) We finally apply the optimized parameters to calculate the result of test data htest-
matrix O and a diagonal matrix E = diag(A (>,, Add'
satisfying H OT E O, where A 2 A2 2 ... z
Theorem 1. Assume F is Lipschitz continuous w.r.t F with constant μ.. If inequality (7) holds true, then L(brest, W 200 < L(biest, Wo i.e., introducing the training sample zo (i.e., demonstration X) can reduce 1_ 1 are the largest and small- the testing loss on brest. dear
Aad 0 are the eigenvalues of H For L11 it equals
W
VOVVITaa W) Vw I(hrest, W))) din,
est eigenvalues of Hw, respectively.
Ow L(uest, W) = ndd daan 2 -v -V 2/d oz Idd (Arbs)²
Add A
Vw Lhheett Will > |lbeest zoll (Il
Wy
hll + (CI)
2 Vil Will I| V. W) - Vw I(zo, Will
Vd
- LeweL(hess WIll ||Flat (V=L(heest) W)- hif
Wy
Ci -
hest - ||Wgh [heaa
Vd
VzL(zo, =0"
(27)
Proof. With W. We..o. the influence of upweighting 2p on the parameters W is (Ling, 1984; Koh & Liang, 2017):
2 -Ail/V L(heat,, WIll (IFlat((V>L(beees W)
V/L(z0, W)) h + [[lat(V>>L(0), W) (brest 20)")1l)
dW..zo de
V Lhrest, Will (I W) V-I WIll
I Iparamater(22)
(28)
==
[hresll + [V WIll || - =20ll)
H- Vw L(o. W).
(33)
where Flat(W) is the operation that flatten matrix W E gd'xd xd into a vector of length Rd'' Since V is Lipschitz continuous, we get L11 2
and the influence on the loss function is:
Tiosa (20) = dLChtesst W
de
\1ll L(braat, WIll (nllWW 20)l |[brecttt+++
c=0
dW... c,50
(29)
|V WIll |[hrest (aaat-20) zoll)
Vw Libest, WTT.
de cma
(34)
Applying Eqs. (25)(26) to Eq. (34), we have:
= - Whieat WTT Hw Vw T(20, W)
| (heet 20)1 = V (heest - 20) WTW (heest - -20)
Then, the testing loss Lhbestt, W
Paprll 20 can be evaluated
Wv hecarl (WW (bicat - -00]
(35)
by Taylor approximation since
is sufficiently small,
Vd
12 prel
Wy
ie., [(hestt W
hicat II || (Wgheelll | - 201l
"prat .20
Vd
L(buaats Wo,zo) +
dh W
LV WIll
||ll/|/hsttt zoll
Wv vd
hll- 1ll-lheeet zoll (36)
Darel de
=I(hrest, Wo.
1
Vw (hrest, WTT
(30)
For L1z, we similarly have:
D
Vw - -H L(2,,
L12 = Edd'' - 2,02 2 daa Vw I. W)11 (37)
12

---

What Makes In-context Learning Effeetive for Mathematical Reasoning
Algorithm 1 Our LMS3
demonstrations the validation set alone, is insufficient. It is essential to balance similarity and the inference stability of demonstrations, LMS3 does, to consistently achieve better results compared to zero-shot setting. In case 3, we again observe that the two baselines incorrectly answer a problem that zero-shot got right. In contrast, LMS3's rejection mechanism determines that the best demonstration still has relatively large similarity distance Sim(X) from the test sample (ranked in the top 1.19% exceeding À 1%). By automatically rejecting this demonstration and adopting the zero-shot setting, LMS3 maintains the original performance, which verifies the necessity and effectiveness of our proposed demonstration rejection mechanism.
Input:
k-shot, Xtest, D, A
Output:
Selected demonstration set D S D Calculate Sim(X), (ttabX)))
1:
for E
based on (gss(22)((22(225))
2: Define Score: S D as the set of l samples with the smallest Score(X X) values.
3: Define Simy C D as the set of X samples with the smallest Sim (X) values.
4: D2 = =(--
5: for X E Score; do
6: If XX E Simy do
7: D2 = D& u (X}:
MAWPS
GSM8K
MATH
E. Discussion and Future Work
Num. Problems
2,373
8,792
12,500
[D]
1,898
7,473
7,500
Here, discuss several future directions of LMS3.
Piesal
475
1,319
5,000
First, in this paper, motivated by LLMs' unreliable perfor- mance on mathematical reasoning benchmarks, we derive our Theorems and 2 to illustrate when in-context learning is effective for mathematical reasoning and design a corre- sponding demonstration selection method LMS3. However, our theories and method be easily extended to other nat- ural language processing tasks, such commonsense rea- soning (Talmor et al., 2019) and sentiment analysis (Socher et al., 2013; He et al., 2017). Besides, it is also worthwhile to explore how our analysis can be further adapted to multi- modal tasks, such as geometric reasoning (Xiao et al., 2024) and graph learning (Pei et al., 2020). Therefore, in the fu- ture, are willing refine our theories by incorporating the characteristics of different tasks and to test our LMS3 more broadly.
Avg.
Problem Length
30.08
45,88
30.79
Avg.
Solution Length
5.90
48.31
84.50
Difficulty Level
Elementary Elementary High School
Table 6. Statistics of datasets.
Combining Eqs. (34)-(37), we finally get:
L1 Zdde || Vw L(hisst, Will - A1/ L(busst, WIIl (p
Wy
hresell ||Wcheesll- |[heest zoll -|leeettttlt
Vd
Wv
hll
lhest - zull).
Vd
(38)
According to Eq. (27), the right-hand side of Eq. (38) is greater than 0, which leads to the conclusion. D
Second, from our theories, we further find that the cor- rectness of the reasoning paths in the demonstrations may have less impact on the test performance compared with its semantics, which aligns with observations reported in previous studies (Wei al., 2023; Min et al., 2022). This phenomenon does not meet the goal of enabling LLMs to learn and imitate the reasoning paths of demonstrations. Therefore, we will explore ways to amplify the influence of the demonstration reasoning paths based on our theory for improving the interpretability and controllability of ICL.
B. Pseudo-Code of LMS3
The pseudo-code of our LMS3 is presented in Algorithm 1.
C. Dataset Partition and Statistics
For GSM8K and MATH, we follow their publicly available train/test splits D and Drest- For MAWPS, we randomly split the dataset into 8:2 ratio for D/Drest. We summarize the dataset statistics in Table 6. For each dataset, we also randomly select 200 problems from D as the validation set to support the implementation of some baselines.
Moreover, similar ICL, retrieval-augmenned generation (RAG) (Ding et al., 2024; Chen et al., 2024a) also enables LLMs to perform reasoning based on given information (e.g., retrieved knowledge). However, this information is often semantically unrelated to the test problem (e.g.. pro- viding external knowledge) or may not involve reasoning problems in itself (c.g., the knowledge is a declarative state- ment that does not require LLMs to reason). Therefore, we believe that how extend our theory to analyze when RAG is effective presents promising research direction.
D. Case Study
Figure 6 shows three cases to validate LMS3's interpretabil- ity. For brevity, we omit the solutions for demonstrations.
For cases and 2, the baselines OpenAl and Best-validate both made errors. This indicates that considering only the semantic similarity of demonstrations or the effectiveness of
13

---

What Makes In-context Learning Effeetive for Mathematical Reasoning
Case I: Find the vun of the skope and y-intcreep the line through the paints a7.5)and and (9.0).
Zero-shot: slage line ttruuhh (7.8) and (5,0) m- aerelljeee (0-00000-)/ (0-5/1977 2. The
Antercppt the
Tharefore, the sum the stope and intercopt 8-8mitbcaodde
X
OpenAl
Demonstration: What the sum of the coordinates of the midpoint of the line segment connecting points (5.4) (-9882
Answer: Lat's solie the posbion stop by steg. First, nead find the siuge fee)
Nou
we sced to fhd the y intorrate (b) i. So, the y-Jaaoreape is 26, S.Faialty, we need find the nus
of slope and Intercopt + 28 24. Therefore, the sum the sispe and intereept is 24. X
Rest-validaae
(areerarr axed, &ttestt yo12 1a-5 &'testt if] -2 Je vill li2s-b &itestt (f x s-22 aendjarra faarraa//rightt /Find arb the pieeowise function contineous (which means that graph can be drawn without lifting your pencill from the
Demonstration: Le (f6t)
Answer: Let's doep breath nosk nop stpp The slope of
The slope
nntreggt form of line tay m nox b .. So. the meereeer is 3n The nus of the slope and y murrrrgr to: w b of 36- 32. Therefore the wer of the sape and y-intoroppe is 32.
LMS33 (Ours)
Demonstration: Find the stope of the line that is perpendicular to the line 2x 3y
Answer: Zach fovak it down stop viep! First need in find the slope of the fior... We con the
So,, the stage
find the y-intereapp the line.
Hl can use the foreala.... So, the y-teereept and y interegpe..... ... Therefore, the sum of the stope and onterepp is 32.
Fiually we Fid the sure of the sloge
Case 2: Which type of conic section described by the equation ((x yltx 22 12 1imer "C" for circle. 'P" for parabola, "E" for ellipse. 'II" for hyperbols, and 'N" for none of the shove.
shot: The gives cquation cas be stopiiffid at fotions i [bggnoaalign)) sit- & 2y12 Millerr 2 &- 2722 fillline' 2 &- dillinge' 2) d- wnloom//algn 7 This the equatton a loporboia
X
Opennl
Demonstration: Which
conic section deserihed by equation 11=23 Zv 324 Enter "C" circle. "p" parabola," "g" for elligse "H" by byperbola, "N"
for none of the above.
Answer: Let's breah in down step by supt First we con expand the lof hand stde of the equation: 2y Sn she equattin bermmeer 2 y'3 -ly'2 - Nove, le'ss more all the terms squation ciede' Theofore
sisle 2'2 y22
Best-validate
Demonstration: Shown below rows 2,2, and of Puscal's triangle. V'beginfarran] (eccecec]&
&
l'end|array]]
Sta DVS 5(h_ 13S Stc_ 1yS he the sequence. from left right, elements the 2005th, 2006ch, and 2007th rows, respetfivel, with the leftmost element eccurring at Si U.S.
Computel /sum ii- (2006) Wreech Rie H- leam [i- 1i (2055) frucia Bitb B.3 chatlenge!
Anywer:
factoring the equation: 1ile-
2y"2 112 2.122 22 IiI) (Vor2-yy2 2-2 1) 22 2212 (002-y'2 2-1-2223 2) Dif
21 3y f)- G1)). This is quadrantc equation in 2 y22 and it's the form:
Ifar 2yy
where b and e are constonte. The equiin represente conic tection
Sanee the cofflecinn x (o) prutivee and the cosficiied y2 (A) pustive, the equation
regrnsants illynee
X
MS3 (Durs)
Demonstration Find the form of the parahala xy2
Answer: detonsine the oye
section deseriled the given equation, hand ride, we get, 2x02
stepriffingg equation (964)99 i -2y22 Dif Exponding
y"2 - 2y-2 f11] Rearrnnging the equation, we get: y22 - 1']. This is the equation
curlle so the correct osswer C
Cass 3: There me four hinds at the Caboose, Sally Two three older than Gimmn Rod. Sally Sally and Sally Thirytwo years what's the totil age
Gramy
times
some age Sally Thityywo, birds?
Zero-soot: Let's break dows
step I.. Solly Four also 8 years old, 2.,, Gramey years Nin,, let's up ages
Red ix 16 years old 3.. Sally
the four binds... The ioted age of the fowr binds 8+ 19 + 16 8- 51 yeary old
OpenAl
Demonstration: Jerry counts six birds nesting in the bushes, 2/3rd of that number of birds flying overhead, and 3 groups of eight birds each feeding, How many hirds does be count in tutal?
Answer: Lavks break down step supelt... Suily abn old.2.. Graany
yours 2... so Sutty Too ages of fow binds: 8 (Saily Fowry 16 /rraaay Rodi+19 (Sally Bwoj 44 (Calsosse) (Puietthee 47
years Finally wy odl up
X
Rest-validate
Demonstration: Natalia sold clips to 48 her friends April, and then she sold half clips
How elips Natalia altogether April and May?
Answer: break down stoy noppl. 10 Sally
also yeurs
... Grazny
Redit 2 x8 - 16 yoars old.3... so Saft Tro is 16 3- 19 years oid. Now lot't ndd up the ogos of the for bins, Sally Four (8), Salty (19) Grany Red (16), and wan there's one
find but we're
ite age' Ta's nemoe it's years old in baby hinl!) Then, the
age forr hirts years.
yorrs.
and age of the hinbs
X
LMS3 (Ours)
Demonstration: Sebastian plays drams for percusion in an orchestra, He sits next the brass settion where four people play tromhont, two play trumpet, and one plays Freneh strings section with three vielinists, cellist,
freat of them
quartet Bute players woodwinds section. The moestre conducts them all. How many people are in the orchestra." (Rejested, ranked 1.19% by Sim(())
contrabavilet, Nest to strings are
diarinets
step! I.. Sally Four also 8 years old 2... Granny yoars old le's add the ages of the four binls.. toial uge of the four lind & 19 16 51 years old.
Anywer: Ler's bresh down
years old Sully Thos is
Figure 6. Case Study.
14


Reviews:
[{"review_id": "hiyD9Q5mXj", "reviewer": "Reviewer_2APu", "summary": "In-context learning has been a key driver of LLM performance over the past few years. However, the performance of a model can vary (and sometimes even be negatively impacted) based on the content of the few-shot demonstrations provided in-context. This work provides a theoretical analysis of the conditions under which ICL is beneficial and find that performance depends on 2 keys factors :1) The semantic similarity of the demonstration to the test problem and 2) The inference stability of the demonstration which indicates how easily the LLM can solve the demonstration itself. Building on their theoretical insights, the authors present LMS3, a simple algorithm to select the demonstration(s) for a given problem by trading-off between the 2 objectives defined above. Results show that their proposed algorithm improves performance over existing ICL selection methods and surprisingly, they also find that for some questions, having no demonstrations is actually beneficial!", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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:\n\n$\\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.\n\n$\\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. \n\nTo 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.\n\nWe sincerely appreciate your thoughtful comment and hope this explanation addresses your concerns.\n\n$\\bf{Q2}$: In the K-shot setting, is it fair to treat the demonstrations as independent of each other?\n\n$\\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.\n\nMoreover, 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.\n\n$\\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.\n\n$\\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. \n\nFurther 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.\n\nWe greatly appreciate your thought-provoking idea and will supplement the discussion it in our revised version. \n\n$\\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?\n\n$\\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. \n\nWe appreciate your thoughtful question and hope this clarifies our approach.", "claims_and_evidence": "Yes, theoretical claims are supported by appropriate proofs - however I have not checked the math carefully. Their proposed algorithm is validated on 3 standard benchmarks for mathematical reasoning and shows strong performance.", "methods_and_evaluation_criteria": "Yes, standard evaluation criteria (accuracy) are chosen and appropriate benchmarks are selected. Confidence intervals are also provided for the main set of results.", "theoretical_claims": "I did not check the correctness of proofs.", "experimental_designs_or_analyses": "Yes, experimental design and analysis is valid.", "relation_to_broader_scientific_literature": "This paper provides theoretical grounding to empirical findings that have been observed in prior work.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "- This a very-written paper with theoretical analysis backed by a practical algorithm. The paper also does a good job of building intuition and the final findings align well with prior work in the area. \n- The algorithm presented requires white box access to the LLM. However, given the strong generalization performance across LLMs, this might not be an issue.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "- In the K-shot setting, is it fair to treat the demonstrations as independent of each other? Aren't there potentially significant interaction effects that need to be accounted for? \n-  Can insights from this work potentially be used to design a method that generates effective demonstrations? That is, use some form of optimization to generate a X that maximizes the score function. This might result in better performance compared to a fixed offline demonstration set. \n- By how much would the inference runtime for a query be affected (compared to random demonstration) if LMS3 is used using a fixed demonstration dataset. If I understand currently, since the embeddings of all demonstrations are computed independently and stored, the overhead should be minimal?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "EeDPSnV4Sv", "reviewer": "Reviewer_spYJ", "summary": "This paper aims to explore the underlying mechanism of in-context learning (ICL). To this end, the authors first theoretically analyze the influence of the demonstrations on inference performance, where they prove that the performance is bounded by an LLM-oriented semantic similarity and the demonstration stability. Then, based on this finding, they proposed a LMS3 method for demonstration selection. From extensive experiments on two widely-used LLM backbones and multiple few-shot settings, they validate the superiority of LMS3.\n\n## update after rebuttal\nThe authors have addressed my concerns, and l  would like to retain my positive score.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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:\n\n$\\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).\n\n$\\bf{A1}$: Thanks for your recognition of our theoretical analysis and this valuable suggestion. We appreciate the opportunity to discuss its broader applicability.\n\nAlthough 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.\n\nTo 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. \n\nWe 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. \n\n[1] Compositional Exemplars for In-context Learning. \n\n[2] In-Context Learning with Iterative Demonstration Selection. \n\n[3] Rethinking the Role of Demonstrations: What Makes In-Context Learning Work?\n\n$\\bf{Q2}$: Combine the rejection mechanism with other demonstration selection method to further validate its necessity.\n\n$\\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. \n|LLaMA3-8B| MAWPS | GSM8K | MATH |\n|-|-|-|-|\n| Random|0.951|0.813|0.330|\n| +Our|0.952|0.818|0.349|\n| Best-validate|0.932|0.817|0.332|\n| +Our|0.941|0.829|0.344|\n|TF-IDF|0.945|0.803|0.344|\n| +Our|0.946|0.818|0.351|\n| BM25|0.932|0.805|0.334|\n| +Our|0.934|0.812|0.335|\n| T5|0.948|0.817|0.330|\n| +Our|0.953|0.828|0.333|\n| BGEM3|0.938|0.802|0.340|\n| +Our|0.941|0.822|0.350|\n| OpenAI|0.965|0.809|0.346|\n| +Our|0.973|0.818|0.347|\n| SPELL|0.945|0.821|0.343|\n| +Our|0.946|0.826|0.345|\n| Influence|0.929|0.800|0.333|\n| +Our|0.935|0.810|0.340|\n| IDS|0.920|0.808|0.330|\n| +Our|0.932|0.823|0.346|\n\nThe 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.\n\nThank you again for your valuable suggestion! We will incorporate this experiment and its analysis into the revised version to further support our findings.\n\n$\\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)?\n\n$\\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. \n\nHowever, 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.\n\n$\\bf{Q4}$: I found some typos.\n\n$\\bf{A4}$: Thanks for your meticulous review and pointing out these typos! We will carefully correct them in the revised version.", "claims_and_evidence": "I believe that the claims are supported by clear and convincing evidence. On one hand, the authors provide fundamental theoretical analyses to reveal the impact of demonstrations on ICL performance, offering a solid foundation for the proposed method. On the other hand, they conduct experiments with two representative LLMs and compare them against 10 baselines under various few-shot settings. Therefore, I believe the experiments are sufficient to validate the effectiveness of the proposed method.", "methods_and_evaluation_criteria": "The proposed method has strong theoretical foundation and is well designed for the problem. This paper employs answer accuracy as the evaluation metric, which is a widely used and reasonable setting. This choice ensures the reliability of the results.", "theoretical_claims": "I have checked the proofs of all theoretical claims (i.e., Theorem 1 and Theorem 2) and verified their correctness. Additionally, I have ensured that all assumptions and derivations are logically consistent and properly justified.", "experimental_designs_or_analyses": "Yes, I have carefully checked the soundness and validity of the experimental designs and analyses. Specifically, I examined the correctness of the experimental setups to ensure that they align with standard practices in the field. I also examined the selection and preprocessing of the datasets used for evaluation. Additionally, the experiments have been run multiple times to ensure the stability of the findings. No significant issues were identified, and the experimental results are consistent with theoretical expectations.", "relation_to_broader_scientific_literature": "This paper builds on prior work in in-context learning by providing a theoretical analysis of how demonstrations influence LLM reasoning performance. Unlike prior heuristic-based or semantics-based selection methods, the proposed LMS3 is theoretically grounded, generalizable, and introduces a novel demonstration rejection mechanism. The empirical results further strengthen its contribution by demonstrating consistent improvements across multiple benchmarks and LLMs, addressing a key limitation of previous methods that lacked robustness across settings.", "essential_references_not_discussed": "I think the related work has been discussed sufficiently.", "other_strengths_and_weaknesses": "Strengths:\n\n1. The paper provides a rigorous theoretical analysis of ICL, revealing the importance of LLM-oriented semantic similarity and inference stability for reasoning performance. These theoretical findings offer a deeper understanding of when and why demonstrations help or hurt ICL.\n\n2. The proposed LMS3 method is simple but practical and efficient. Notably, the demonstration rejection mechanism is a novel contribution. I think it is the first exploration of when ICL should not be used. This perspective fills an important gap in existing research.\n\n3. The empirical evaluation is thorough, covering two LLM backbones, 10 baselines, and several few-shot settings. Additionally, the experiments have been run multiple times and the authors provide the confidence intervals. These results consistently support the theoretical claims and demonstrate the robustness of LMS3.\n\n4. The paper is clearly written, well-structured, and easy to reproduce.\n\nThere are some minor issues with this paper:\n\n1. I think the theoretical analyses are well-founded, but I hope to see if the authors could give more discussions about how to extend them to other domains or tasks (please see Question 1 below).\n2. More experiments could be conducted to validate the necessity of the rejection mechanism (please see Question 2 below).\n3. There exist some typos.", "other_comments_or_suggestions": "I found some typos, including:\n\n--Line 87, “remain”->”remains”\n\n--Line 116, “satisfies”->”satisfying”\n\n--Line 265, “suggests”->”suggest”", "questions_for_authors": "1. 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).\n\n2. Apart from the experiments in Section 5.5., I think the authors can combine the rejection mechanism with other demonstration selection method to further validate its necessity. \n\n3. 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)?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "GcFHjBVheW", "reviewer": "Reviewer_voNS", "summary": "In this paper, the authors investigate the theoretical explanation of in-context learning (ICL). They prove that the influence of the demonstrations can be estimated by two factors: LLM-oriented semantic similarity and inference stability of demonstrations. Based on it, they propose a LMS3 method, and the experiments on Llama2 and Llama3 validate the effectiveness of LMS3 under both one-shot and few-shot settings, as well as for both mathematical reasoning and commonsense reasoning tasks.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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. \n\n$\\bf{Q1}$: The generalizability of the theory and methods in this paper could be further discussed. \n\n$\\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. \n\nAs 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.\n\nFollowing your suggestions, we will supplement the above discussions in the revised version.\n\n[1] Compositional Exemplars for In-context Learning. \n\n[2] In-Context Learning with Iterative Demonstration Selection. \n\n[3] Rethinking the Role of Demonstrations: What Makes In-Context Learning Work?\n\n$\\bf{Q2}$: In Eq.(6), I am unsure how the pre-trained data $D_{pre}$ and the demonstration $z_0$ are optimized simultaneously.\n\n$\\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}$. \n\nBased 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.\n\n$\\bf{Q3}$: In Eq.(23), $\\lambda_1$ is misused. There are some typos.\n\n$\\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.", "claims_and_evidence": "I have reviewed all the theoretical proofs in this paper and find the claims and results both reasonable and correct. Additionally, the authors have conducted experiments across one-shot to four-shot settings, providing convincing evidence of the proposed method's effectiveness.", "methods_and_evaluation_criteria": "The proposed method is grounded in the theoretical findings, which makes sense and offers a novel and deeper understanding of ICL. Additionally, the experiments are conducted on three mathematical reasoning datasets and one commonsense reasoning dataset. Therefore, this paper provides strong empirical validation.", "theoretical_claims": "Yes, I have checked the correctness of the proofs in Section 3. They are self-contained and reasonable, providing a good perspective for discussing the influence of demonstrations.", "experimental_designs_or_analyses": "Yes, the experiments are conducted using two LLMs, four datasets, and four few-shot settings. Therefore, I think the effectiveness is validated.", "relation_to_broader_scientific_literature": "The authors expand the existing understanding for ICL based on a widely-used theoretical setting. Although the simplification of the softmax function in Eq.(3) may deviate from reality, it has been widely adopted in previous work.", "essential_references_not_discussed": "No. The references are cited and discussed sufficiently", "other_strengths_and_weaknesses": "Strengths:\nFrist, building on existing analyses of ICL, this work innovatively derives the theoretical relationship between test loss and demonstrations. Since no prior work has reached a similar conclusion, the novelty of this paper is well justified.\n\nSecond, this paper explains how the effectiveness of ICL is determined by both the LLM-oriented Semantic Similarity and Inference Stability of Demonstration. I think this conclusion makes sense and it provides insights for practical applications of ICL.\n\nThird, the experiments are thorough, with the authors using two LLMs as backbones and conducting evaluations across 1-shot to 4-shot settings.\n\nLast, the writing is clear, making the paper easy to understand and reproduce.\n\nWeakness:\nThe generalizability of the theory and methods in this paper could be further discussed (Q1 below). Besides, I still have some questions regarding the theoretical analyses (Q2 below).", "other_comments_or_suggestions": "In Eq.(23), lambda_1 is misused, as it was previously introduced to represent an eigenvalue\nLine 105, “and offers” should be “and offer”\nLine 75, “generate code” should be “generate codes”", "questions_for_authors": "Q1: In this paper, the author emphasizes the analysis of situations where ICL is effective for mathematical reasoning. I believe the analysis is sound, but the conclusion could also be applied to other tasks, as the author has demonstrated its effectiveness on CommonsenseQA as well. Therefore, I think the author could further discuss the generalizability of this theory and explain whether it could be effective for other tasks.\n\nQ2: In Eq.(6), I am unsure how the pre-trained data D_pre and the demonstration z0 are optimized simultaneously, because in my understanding, D_pre is for pretraining and z0 is for the gradient update in inference phase?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper investigates the effectiveness of in-context learning from a theoretical perspective and derives a lower-bound (Theorem 1) on the prediction error over the test example by its similarity to the in-context samples and the model stability to these examples. Using Theorem 1, the paper derives a scoring scheme, LMS3, on few-shot examples that could potentially be useful for improving the accuracy on the test examples. Experiments are provided using two LLMs on varied reasoning tasks, demonstrating promising results.

The paper is organized well and received overall positive reviews. All the reviewers appreciated the theoretical insights and the empirical results. There are minor concerns on the generalizability of the approach and inference time of LMS3 (2APu), extension of the insights to other domains beyond mathematical reasoning (spYJ), and extension of the LMS3 scheme to other demonstration schemes (spYJ).  Authors addressed these concerns and provided additional results demonstrating the generlizability of the scheme to an elaborate set of demonstration selection methods with promise. 

There are three key outstanding concerns in the paper, namely: 
1) the proposed approach needs a scheme to derive the stability of the in-context examples, which would ideally need access to the network weights, 
2) simplified analysis of the approach to a single layer transformer setting with linearized softmax. 
3) the applicability/empirical analysis of the approach only to mathematical reasoning problems, 

AC thinks these issues may not be critical given the theoretical focus of the paper.  That being said, AC makes a note of the concern raised by Reviewer voNS regarding the derivations in Theorem 1 that uses the annotations in (6)-(ii). Specifically, the theoretical approach in the paper is to investigate in-context learning through the lens of few-shot training and (6) considers a setting using hypothetical weights \hat{W}_\eps, z0 that could have been obtained, had the in-context examples been part of training (as few-shot examples). While the idea simplifies the theoretical derivations, the subsequent derivations may not necessarily be truthful to understanding in-context examples as standard in-context formulations do not assume any updates to the network weights and are usually test-time prompt samples. Thus, the paper may be improved by providing more clarity into these underlying assumptions.

All the reviewers recommend acceptance of this paper. AC concords with the reviewers' sentiment and see that the paper has value in terms of theoretical insights (that could be explored further) as well as elaborate empirical results. As such AC recommends acceptance. Authors are recommended to address reviewers' comments, especially on clarifications to (6), in the camera-ready.
Decision: Accept (poster)