Probing the Decision Boundaries of In-context Learning in Large Language Models

In-context learning is a key paradigm in large language models (LLMs) that enables them to generalize to new tasks and domains by simply prompting these models with a few exemplars without explicit parameter updates. Many attempts have been made to understand in-context learning in LLMs as a function of model scale, pretraining data, and other factors. In this work, we propose a new mechanism to probe and understand in-context learning from the lens of decision boundaries for in-context binary classification. Decision boundaries are straightforward to visualize and provide important information about the qualitative behavior of the inductive biases of standard classifiers. To our surprise, we find that the decision boundaries learned by current LLMs in simple binary classification tasks are often irregular and non-smooth, regardless of linear separability in the underlying task. This paper investigates the factors influencing these decision boundaries and explores methods to enhance their generalizability. We assess various approaches, including training-free and fine-tuning methods for LLMs, the impact of model architecture, and the effectiveness of active prompting techniques for smoothing decision boundaries in a data-efficient manner. Our findings provide a deeper understanding of in-context learning dynamics and offer practical improvements for enhancing robustness and generalizability of in-context learning.

Paper

References (43)

Scroll for more · 31 remaining

Similar papers

Peer review

Reviewer AvnH6/10 · confidence 4/52024-07-12

Summary

This study investigate in-context learning in LLMs by examining their decision boundaries in binary classification tasks. The authors investigate the performance of several mainstream models on these tasks. Despite achieving high test accuracy, the decision boundaries of these LLMs are often irregularly non-smooth. Factors such as model size and prompts' influence on decision boundaries are explored. Additionally, the author examines fine-tuning and adaptive sampling methods, finding them effective in improving boundary smoothness.

Strengths

- This is the first study, to my knowledge, to explore the decision boundaries of in-context learning LLMs. - The experiments are thorough, and some findings, such as the use of uncertainty-aware active learning to help LLMs learn decision boundaries, are beneficial for future research.

Weaknesses

- The tests are mainly conducted on binary classification tasks, making it unclear if the findings can be generalized to other tasks. - Although different model sizes are covered, these models are from different series. The reason for using these models should be clarified in the paper. - The decision boundaries appear to be heavily influenced by the prompts, such as the design of labels and the order of in-context learning examples. Coule the authors explore the robustness of their results to variations in prompts, like different formats or synonym replacements? If the robustness to prompts is poor, applying these results to other fields could be challenging.

Questions

- Line 69: “Large Language Models (LLMs) are trained on vast corpora of text using unsupervised learning.” Shouldn't this be self-supervised learning? - Line 132-134: “For the open-source models, we use the approach described in 3.2 to get predictions. For the closed-source models, we use the next token generation as the prediction…” This is confusing since section 3.2 only mentions visualization methods. What are the differences in prompt design or experiments between these two approaches?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Overall, this is a good paper. I would raise my score if the authors could address my main concerns, especially regarding the robustness of their results to variations in prompts.

Authorsrebuttal2024-08-07

Dear Reviewer AvnH, Thank you for your feedback. We hope to address your questions and comments below. > Q1:"The tests are mainly conducted on binary classification tasks, making it unclear if the findings can be generalized to other tasks." 1. Our primary motivation for using synthetic 2D classification datasets is their amenability to decision boundary visualization, which is central to our paper's focus as decision boundaries are typically studied in 2D space, and our proposed mechanism offers novel insights into the generalization abilities of in-context learning. NLP classification tasks, with their discrete token nature and variable input lengths, pose significant challenges for visualization of decision boundaries along x and y axes. 2. Our use of synthetic function datasets is firmly grounded in established practices for studying in-context learning. Several recent high-profile works have exclusively used synthetic data to investigate ICL/learning algos: - NeurIPS 2022 paper (https://arxiv.org/pdf/2208.01066) focuses entirely on linear functions for ICL. - ICLR 2023 paper (https://arxiv.org/pdf/2211.15661) explores learning algorithms for ICL and primarily conducts experiments on synthetic linear functions. - Google's research (https://arxiv.org/abs/2303.03846) employs synthetic linear function classes (see Section 6). - NeurIPS 2019 paper (https://arxiv.org/pdf/1905.11604) utilizes similar scikit-learn datasets to study decision boundaries for SDE algorithms (see Figure 1). These papers demonstrate that our approach is quite prevalent in the field. 3. To extend beyond binary tasks, we conduct new experiments on 6 real-world NLP classification datasets. Note, to handle the high dimensionality of text embeddings, we projected them onto 2D space using t-sne. (Any dimensionality reduction technique will introduce confounders to the analysis, but this price is inevitable to extend our analysis). We experiment with a total of 6 widely-used NLP classification tasks, including both binary and multi-class settings. These include sentiment analysis and Textual Entailment Recognition, providing a broader perspective on the applicability of our approach. Our results, presented in **Figure 1 of the rebuttal PDF**, demonstrate similar non-smooth decision boundary characteristics in NLP tasks as observed in our synthetic datasets. > Q2: "Although different model sizes are covered, these models are from different series. The reason for using these models should be clarified in the paper." We use these LLMs to provide a comprehensive analysis across different sizes and architectures, representing current SoTA open-source LLMs. For the sizes, due to computational constraints in an academic setting and the expensive nature of querying decision boundaries with large grid points, we limited our analysis to models no larger than 13B parameters. Therefore to complete the size series and gain insights into smaller-scale models, we included the pruned llama 1.3B model. > Q3: "The decision boundaries appear to be heavily influenced by the prompts, such as the design of labels and the order of in-context learning examples. ... explore the robustness of their results to variations in prompts, like different formats or synonym replacements?" To address this concern, we have conducted experiments to explore how decision boundaries are affected by prompt formats, using 4 different synonyms for key terms. As shown in **Figure 2 of the rebuttal PDF**, we find that the LLM's decision boundary is indeed affected by the prompt format, which aligns with the importance of prompt engineering in ICL. However, the overall non-smoothness level of the decision boundaries remained consistent. We view prompt formats as another factor influencing the decision boundary rather than undermining our results, since our central observations regarding the non-smoothness of decision boundaries show similar patterns across different prompts. We will add this as another influencing factor in our paper. > Q4: “Large Language Models are trained on vast corpora of text using unsupervised learning.” Shouldn't this be self-supervised learning? Sorry for the typo! We will revise this. > Q5: “For the open-source models, we use the approach described in 3.2 to get predictions. For the closed-source models, we use the next token generation as the prediction…” This is confusing since section 3.2 only mentions visualization methods. What are the differences in prompt design or experiments between these two approaches? We apologize for the confusion. For both open and closed source models, we used the same prompts. For open source models, we used logits to get predictions for the class labels. For closed source models, we looked at the predictions since the logits were not available to us. We hope this addresses your concerns. Please let us know if you have any questions.

Reviewer AvnH2024-08-12

Response to the author

Thank you. Considering this is the first study on decision boundaries in LLMs (as far as I know) and the experiments are solid, I have raised my score.

Authorsrebuttal2024-08-12

Thank you for raising score!

We thank the reviewer for raising the score and are pleased that you find our work novel and our experiments solid.

Reviewer PktE3/10 · confidence 3/52024-07-12

Summary

This paper investigates the decision boundary of in-context learning of Transformers. The paper shows that for three toy tasks, the decision boundaries of in-context learning of various pretrained models are not smooth. The paper then explores the method for improving the smoothness of decision boundaries and finds that supervised finetuning can mitigate this non-smoothness. Furthermore, finetuning on one task can benefit the smoothness of decision boundaries on other tasks.

Strengths

1. The decision boundary of in-context learning is an interesting research topic, which seems not investigated by previous work. 2. The paper provides multiple empirical results and finds a method for improving the smoothness of decision boundary of in-context learning. 3. The writing of this paper is clear.

Weaknesses

1. The empirical results are limited to toy datasets, which may reduce the impact of the proposed smoothness-improving method. 2. The paper does not empirically show that smoother decision boundary smoothness leads to higher generalization accuracy. On the contrary, the paper finds that gpt-4o has high accuracy along with not smooth decision boundary, which makes improving smoothness less motivated. 3. No quantitative evaluation of smoothness is present in the papers, making smoothness difficult to compare.

Questions

See the weaknesses.

Rating

3

Confidence

3

Soundness

2

Presentation

3

Contribution

2

Limitations

The main limitation is that the experiments are conducted only on toy datasets.

Authorsrebuttal2024-08-12

Dear Reviewer PktE, Thank you again for your helpful feedback! Your suggestions have helped us solidify our experiments and analysis by: 1. Conducting additional experiments on NLP tasks spanning 6 NLP classification datasets, 2. Empirically demonstrating that smoothness on synthetic datasets can lead to improvements on downstream tasks, on average, across 7 ICL NLP tasks. These first two experiments address your main concern that "the main limitation is that the experiments are conducted only on toy datasets." 3. Using quantitative metrics, such as nearest neighbor entropy, as a measure of smoothness. We hope this addresses the concerns you listed in your review. As the discussion period is nearing its end, could you please let us know if you have any further questions for us? We are happy to address any questions or concerns you may have. Thank you for your time!

Reviewer WmaA10/10 · confidence 4/52024-07-13

Summary

The authors study the decision boundaries of LLMs in binary in-context learning tasks, finding that decision boundaries can be non-smooth despite high test accuracy and linear separability of the task itself. They examine numerous methods for smoothing the decision boundaries, including SFT and training a transformer that intentionally learns a smooth decision boundary.

Strengths

Fascinating, insightful paper. Well-written, good visualizations. Robustly evaluated across model families, model sizes, number of examples. Even examines the effects of quantization on the decision boundary, which might be one of the most important points in the paper (4-bit quantization affects the boundary quite a bit). Includes sections on fine-tuning to smooth decision boundary as well as training from scratch to do so. Great RW section, up to date with the most recent work. No-brainer accept.

Weaknesses

None that I can see.

Questions

N/A.

Rating

10

Confidence

4

Soundness

4

Presentation

4

Contribution

4

Limitations

N/A.

Reviewer PtXB5/10 · confidence 3/52024-07-25

Summary

This paper conducts a wide range of experiments exploring the smoothness of decision boundary generated by LLMs. Synthetic datasets from scikit-learn are used for the experiments, and experimental results are showing that various factors affect decision boundaries of LLMs.

Strengths

- There are analyses from diverse perspectives, in terms of model size, in-context examples, quantization, prompt format, example order. Extensive experimental results help interpret the targeted phenomenon. Moreover, changes when fine-tuning LLMs with in-context examples are observed.

Weaknesses

**[W1] Justification of task and dataset**: - There have been works trying to understand the mechanism of in-context learning by conducting experiments on NLP classification tasks. However, it seems unclear why the authors choose scikit-learn-generated dataset. Is this dataset proper to investigate the effectiveness of in-context learning? In addition, can experimental results using this dataset be generalized to NLP classification tasks? - Moreover, there are concerns on comparison groups. Traditional ML algorithms (e.g., Decition Tree, K-NN, etc.) map each datum into vector spaces, while LLMs map the entire context into vectors. Can comparing these two mechanism be regarded as a fair process? **[W2] Analysis**: I have concerns that there are sets of experimental reporting without interpretations of probable reasons. There is a core question: How the smoothness of decision boundary helps understand in-context learning mechanism? **(Update) After the author's rebuttal, most of these concerns are clarified, thus raising the score 3 to 5.**

Questions

- It would be better to define the concept of smoothness of decision boundry more formally since expecting a perfect smoothness of LLMs' decision boundary is rather unrealistic.

Rating

5

Confidence

3

Soundness

2

Presentation

3

Contribution

2

Limitations

Yes

Authorsrebuttal2024-08-12

Dear Reviewer PtXB, Thank you again for your helpful feedback! Your suggestions have helped us solidify our experiments and analysis by: 1. Conducting additional experiments on NLP tasks spanning 6 NLP classification datasets, 2. Demonstrating that smoothness on synthetic datasets can lead to improvements on downstream tasks, on average, across 7 ICL NLP tasks, and 3. Using quantitative metrics: nearest neighbor entropy, as a measure of smoothness. We hope this addresses the concerns you listed in your review. As the discussion period is nearing its end, could you please let us know if you have any further questions for us? We are happy to address any questions or concerns you may have. Thank you for you time!

Reviewer PtXB2024-08-13

Thank you for clarifying most of my concerns. I understand synthetic benchmarks are also used in this analytic field. Furthermore, I would appreciate for including NLP experiment results. It clearly helps my understanding. I believe refining smoothness can be a key to understand and improve LLMs, but I'm still doubtful on the hypothesis that smoothness has a causal relationship with model performance by relying on observed case samples. Thus, I would like to raise my score from 3 to 5.

Authorsrebuttal2024-08-13

Thank you for raising the score and for your valuable feedback! We appreciate your insights on smoothness and are glad the NLP experiments clarified your understanding. We will incorporate the analysis of these results into our paper. Thank you!

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC