What Makes In-context Learning Effective for Mathematical Reasoning: A Theoretical Analysis

TL;DR

The paper tackles whether in-context learning reliably enhances mathematical reasoning and provides a theoretical bound on the benefit of demonstrations via semantic similarity and inference stability. It introduces LMS3, a low-complexity demonstration-selection method with an adaptive rejection mechanism, and proves extensions to $k$-shot settings. Empirical results on MAWPS, GSM8K, and MATH across two backbones show LMS3 achieving consistent improvements over strong baselines, with evidence of generalization to large language models like ChatGPT and GPT-4. The work offers a principled understanding of when demonstrations help and a practical approach to harnessing them effectively in diverse reasoning tasks.

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 a LLM-oriented semantic similarity and an 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 can adaptively facilitate 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.

Figures (6)

• Figure 1: Problem-solving Accuracy of zero-shot and one-shot settings. The hatched areas represent that in the one-shot setting, the model answers incorrectly $\Delta$ proportion of problems that are answered correctly in the zero-shot setting.
• Figure 2: Illustration of our proposed LMS3 method.
• Figure 3: Few-shot Answer Accuracy of Llama3-8B.
• Figure 4: Distribution of $Score(X)$ in Eq. \ref{['main_score_mul']}.
• Figure 5: Performance with varying $\lambda$. The dashed line corresponds to the result of the zero-shot setting.