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

Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 11:04:06

Abstract: In-context learning is an emergent 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 irregularly non-smooth, regardless of task linearity. 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.

Main paper: 
 introduction: 
 Preliminary work. Under review by the 1st In-context Learning Workshop at the International Conference on Machine Learning (ICML). Do not distribute.scale performance with increased training dataset size and model capacity through the simple objective of next token prediction (Kaplan et al., 2020) . A key emergent behavior of these transformer-based models is in-context learning, which allows the model to learn tasks by conditioning on a sequence of demonstrations without explicit training (Wei et al., 2022) . This unique capability allows LLMs to adapt seamlessly to new tasks, often achieving superior performance in few-shot settings (Brown et al., 2020) . Despite significant successes, the underlying mechanisms of how in-context learning works remain partially understood.Recent attempts to understand in-context learning have focused on various aspects. From a theoretical standpoint, studies by Von Oswald et al. (2023) and Dai et al. (2023) have linked the mechanisms of in-context learning to gradient descent, suggesting that transformers can emulate optimization processes. On the practical side, research has investigated the impact of different factors on in-context learning. Works by Min et al. (2022b) and Shi et al. (2023) reveal that accurate demonstrations are not essential for effective in-context learning. On the other hand, factors such as the prompt structure and model size (Wei et al., 2023; Webson and Pavlick, 2022) , or the order of in-context examples (Chen et al., 2024) greatly affect outcomes. More recently, with the development of LLMs supporting longer context lengths up to 10M (Reid et al., 2024) , studies have shown that in-context learning performance improves with significant number of demonstrations (Agarwal et al., 2024; Bertsch et al., 2024) , where the performance can be comparable to fine-tuning on the same amount of demonstrations. Additionally, works by (Garg et al., 2022; Nguyen and Grover, 2022) have demonstrated that small transformers trained from scratch can learn unseen function classes in-context from examples.In contrast to existing approaches, our study introduces a fresh perspective by viewing in-context learning in large language models (LLMs) as a unique machine learning algorithm. This conceptual framework enables us to leverage a classical tool from machine learning -analyzing decision boundaries in binary classification tasks. By visualizing these decision boundaries, both in linear and non-linear contexts, we gain invaluable insights into the performance and behavior of in-context learning. This method allows us to probe the inductive biases and generalization capabilities of LLMs and offers a unique assessment of the robustness of their in-context learning performance. Consequently, this approach provides a comprehensive means to qualitatively analyze the underlying mechanisms that govern in-context learning and suggest ways to improve its performance in LLMs.To our surprise, we found that the recent LLMs struggle to provide smooth decision boundaries in all the classification tasks we considered, regardless of the model size, the number and ordering of in-context examples, and semantics of the label format. This issue persists even for simple binary linear classification tasks, where classical methods such as SVM can easily achieve smooth boundaries with fewer examples as shown in Figure 1 . This observation raises questions about the factors that influence the decision boundaries of LLMs. To explore this, we experimented with a series of open-source LLMs including Llama2-7b, Llama2-13b, Llama3-8b (Touvron et al., 2023) , Mistral-7b (Jiang et al., 2023) , pruned Llama2-1.3b (Xia et al., 2023) , as well as state-of-the-art closed-source LLMs GPT-4o and GPT-3-Turbo. We then explore methods to smooth the decision boundary, including fine-tuning and adaptive prompting strategies. Our contributions can be summarized as follows:• Introduced a novel mechanism to probe and understand in-context learning in LLMs by visualizing and analyzing the decision boundaries on classification tasks.• Discovered that state-of-the-art LLMs exhibit nonsmooth, irregular decision boundaries even on simple linearly separable tasks, unlike classical ML models.  
training large language models: 
 Large Language Models (LLMs) are trained on vast corpora of text using unsupervised learning. During training, these models learn to predict the next token in a sequence. Given a sequence of tokens (x 1 , x 2 , . . . , x t-1 ), the model predicts the next token x t by maximizing the likelihood P (x t |x 1 , x 2 , . . . , x t-1 ). The training objective typically involves minimizing the cross-entropy loss:EQUATIONwhere T i is the number of tokens in the i-th sequence and N is the total number of sequences in the corpus. During training, teacher forcing is often employed, where the model receives the ground truth token x t as input at each time step instead of its own prediction, enabling parallel training.After training, LLMs can generalize to new tasks through a mechanism known as in-context learning.Let S = {(x 1 , y 1 ), (x 2 , y 2 ), . .  
methodology: 
 We aim to better understand in-context learning in Large Language Models by investigating their decision boundaries on a series of binary classification tasks. To increase the generality of our framework, we evaluate several existing LLMs on different task distributions under different settings. We present the general framework here, and refer to Section 4 for specific experiment settings. We prompt the LLM with P and obtain its prediction for x test by choosing the most likely class in the next token distribution. Formally, let V denote the size of the LLM's vocabulary, and l ∈ R V be the vector of logit values for each of the tokens. To obtain a class prediction, we convert each class label i into a unique token id, say c(i) and choose the class with the maximum logit value as the predicted label for x query , i.e., ŷ = arg max i∈{1,...,K} l c(i) .To visualize the decision boundary of a model π, we generate a grid of points covering the feature space defined by the in-context examples set S. Let S = {(x 1 , y 1 ), (x 2 , y 2 ), . . . , (x k , y k )} represent the set of incontext examples, and x min , x max ∈ R d denote the minimum and maximum values of the features in S along each dimension. We create a uniform grid with G points along each dimension, resulting in a total of G d grid points. The grid points are denoted asX grid = {x query | x query ∈ [x min , x max ] d , x query = x min + i∆x, i ∈ {0, 1, . . . , G - 1}} where ∆x = 1 G-1 (x max -x min )is the grid spacing along each dimension. Each point x query ∈ X grid is a query input, and the model π is prompted with the sequence (x 1 , y 1 , . . . , x k , y k , x query ) to predict the corresponding class label ŷ. The decision boundary is then visualized by plotting the predicted labels ŷ over the grid X grid . 
experiments: 
 In this section, we examine existing LLMs through the lens of decision boundaries by conducting a series of binary classification tasks under varying conditions. Our experiments aim to address the following key questions:• How do existing pretrained LLMs perform on binary classification tasks? §4.1• How do different factors influence the decision boundaries of these models? §4.2• How can we improve the smoothness of decision boundaries through finetuning or prompting? §4.3Classification Tasks. We investigate the decision boundary of LLMs by prompting them with n in-context examples of binary classification tasks, with an equal number of examples for each class. We generate classification datasets using scikit-learn (Pedregosa et al., 2011) , creating three types of linear and non-linear classification tasks: linear, circle, and moon, each describing different shapes of ground-truth decision boundaries. Detailed information on the dataset generation can be found in Appendix H. In addition to the in-context examples, we calculate the in-context learning accuracy on a held-out test set of size 100. We sample in-context examples and test points from classification task and convert them into prompt, with an example shown in Appendix F.Obtaining Decision Boundaries of Language Models.We study an extensive range of models, with sizes ranging from 1.3B to 13B parameters, including open-source models such as Llama2-7B, Llama3-8B, Llama2-13B, Mistral-7B-v0.1, and sheared-Llama-1.3B. We also extend our analysis to state-of-the-art closed-source LLMs, including GPT-4o and GPT-3.5-turbo. We generate the decision boundaries of the open-source models with 8-bit quantization due to computational constraints. We choose a grid size scale of 50 x 50, resulting in 2500 queries for each decision boundary.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.4.1. Non-Smooth Decision Boundaries of LLMs. Class decision regions, which means small changes in the input features can result in different classifications. This raises concerns about the reliability of LLMs and their practical deployment, as even when test accuracy for classification is high (shown in Figure 3 ), the underlying decision boundary lacks generalization. In the following sections, we will explore factors that affect decision boundary smoothness and investigate methods to improve smoothness.Impact of Model Size on Decision Boundary and Accuracy From Figure 2 , model sizes increase from left to right, yet there is no clear correlation between model size and the smoothness of the decision boundary. Even the most powerful model, GPT-4o, demonstrates fragmented decision regions. This suggests that increasing model size alone is insufficient for improving decision boundary smoothness. However, as shown in Figure 3 , larger models tend to perform better in terms of test accuracy compared to smaller models, with Llama-1.3B often performing the worst.Increasing (Dettmers et al., 2022) . When transitioning from 8-bit to 4-bit quantization, the red regions around the red in-context learning examples turn blue. This indicates that the reduced precision from 4-bit quantization significantly affects points near the decision boundary or areas where the model is most uncertain. To further investigate, we plotted the probability prediction for class 1 (Figure 5b ). The white regions, indicating a 50% probability for both classes, highlight the areas most impacted by quantization. This provides a practical insight: varying quantization levels can flip or revert the LLM's decisions in the regions of highest uncertainty.Are Decision Boundaries Sensitive to the Prompt Format? Yes, decision boundaries are sensitive to the labels' names, as shown in Figure 6 . Using semantically unrelated labels, such as "Foo" and "Bar" as suggested in (Wei et al., 2023) , results in flipped predictions compared to using reversed class names like "Bar" and "Foo". This suggests that the LLM's prediction still depend on its semantic prior knowledge of the labels. Can We Finetune LLMs on a Dataset of Classification Tasks to Achieve Smoother Decision Boundaries? Previous works have shown that finetuning a pretrained LLM on a large collection of tasks improves its in-context learning performance on unseen tasks (Min et al., 2022a) . In this section, we investigate if the same paradigm helps improve the decision boundary smoothness of LLMs. To do this, we finetune a pretrained Llama model (Touvron et al., 2023) on a set of 1000 binary classification tasks generated from scikit-learn (Pedregosa et al., 2011) , where the ground-truth decision boundary is either linear, circleshaped, or moon-shaped, with equal probabilities. For each task, we sample randomly N = 256 data points x ∼ X grid and their corresponding label y ′ s. We then sample the number of context points m ∼ U[8, 128], and finetune the LLM to predict y i>m given x i>m and the preceding examples:L(π) = E N i=m+1 log p(y i | x i , x 1:i-1 , y 1:i-1 ) , (2)where the expectation is with respect to task, data points {(x i , y i )} N i=1 , and the number of context points m. After training, we evaluate the same finetuned model on various binary classification tasks with varying numbers of context points. To ensure the test tasks are unseen during training, we use different parameters in creating the datasets, such as the separateness between two classes and the scale between the inner and outer circles in the circle task. See Appendix H for more details.We consider several finetuning settings for ablation studies. 1) In the first setting, we finetune the pretrained LLM using LoRA (Hu et al., 2021) and finetune the attention layers. 2) We finetune only the token embedding layer of LLM. 3) We finetune only the linear head layer of LLM. Then we consider modifying the architecture of the LLM: In this setting, we keep the core transformer backbone of the LLM frozen, attach randomly initialized embedding layers and prediction head to the model, and train the entire model using objective (2). This stems from the intuition that taskspecific embedding and prediction layers allow the model to maximally utlize the general pattern-matching capabilities of the transformer backbone for the new task. We refer to this model as CustomLLM, and consider its three variants, which add 1) a new embedding layer for x, 2) a new prediction head for y, and 3) new embedding layers for x, y, and a new prediction head for y. The embedding layers and prediction head are MLPs with one hidden layer. We embed the raw numerical values instead of the text representation of x whenever a new embeddding layer for x is used (same for y), and predict directly the binary class values instead of text labels whenever the new prediction head is used. Results of Finetuning LLM and CustomLLM in Figure 8 and Figure 18 show that finetuning the intermediate and earlier embedding layers leads to smoother decision boundary compared to finetuning the top prediction head. Can LLMs finetuned on one in-context learning task generalize to more complex in-context learning tasks? We demonstrated that SFT on the dataset can smooth the decision boundary on that dataset. In this section, we further explore whether a LLM fine-tuned only on a linear task can achiever smoother decision boundaries on unseen and more complex tasks. As shown in Figure 9 , we compare the decision boundaries of Llama3-8b before and after SFT on the linear task only. Unexpectedly, we found it generalizes to unseen non-linear tasks as well as 3-class and 4-class classification tasks, despite only being trained on a binary linear task. The smoother decision boundaries observed in these unseen tasks suggest that fine-tuning on a synthetic incontext learning task can have downstream benefits for other tasks, enabling the model to be more robust in in-context learning.Can we train a transformer from scratch to learn smooth decision boundary in-context? One may wonder whether a small transformer trained from scratch can provide smooth decision boundaries. To answer this, we train TNPs (Nguyen and Grover, 2022) , a transformer-based model specifically designed for in-context learning. For each sequence of data points {(x i , y i )} N i=1 from a task C, TNPs learn to predict the query labels y i>m given the query inputs x i>m and the context pairs, assuming conditional independence among the queries given the context:L(θ) = E N i=m+1 log p(y i | x i , x 1:m , y 1:m ) , (3)where the expectation is with respect to task C, data points {(x i , y i )} N i=1 , and the number of context points m. TNPs employ a specialized mask to ensure the conditional independence assumption. We showed in Appendix D that transformers trained from sctrach can learn to in-context learn smooth decision boundary. Details are in Appendix D.How to Use Uncertainty-aware Active Learning to Learn Decision Boundaries We investigate whether the decision boundary can be smoothed by providing the LLM with labels of the most uncertain points on the grid as additional in-context examples. Uncertainty is measured as the entropy of the probability distribution of the two classes after softmax normalization of the logits. Our study focuses on an active learning scheme where new in-context examples are incrementally added based on the LLM's current uncertainty. Initially, we obtain the decision boundary conditioned on the existing in-context examples. To refine this boundary, we query the LLM over a grid and select the top-k most uncertain points, ensuring they are spatially distant from each other using a greedy sampling approach. For labeling these uncertain points, we use a logistic regression model well-trained on a larger dataset with perfect accuracy as the ground truth decision boundary. As shown in Figure 10  
related works: 
 Understanding in-context learning in transformers and LLMs is an active area of research, with existing works approaching this problem from both theoretical and practical perspectives.Theoretical understanding of in-context learning Recent works aim to establish a theoretical connection between incontext learning and gradient descent (GD). The pioneering work by Von Oswald et al. (2023) proves the equivalence between linear self-attention and GD on linear regression by construction. Similarly, Dai et al. (2023) shows that attention in transformers has a dual form of GD and views transformers as meta-optimizers. Subsequent works extend these ideas to characterize the global optimum of singlelayer linear transformers. Ahn et al. (2024) observe that with the optimal parameters, the transformer implements a single step of preconditioned gradient descent, while Zhang et al. (2023) shows that at the global optimum, the transformer achieves competitive prediction error with the best linear predictor on a new prediction task. In addition to theoretical connections to GD, a complementary direction aims to establish statistical complexity and generalization bounds of in-context learning in transformers (Bai et al., 2024; Li et al., 2023b; Wies et al., 2024; Wu et al., 2023) . The common limitation of these existing theoretical frameworks is the reliance on strong assumptions about the transformer architecture or the functional form of the in-context learning tasks which may not reflect real-world practices.Practical understanding of in-context learning More relevant to our paper is a line of works focusing on understanding the practical aspects of in-context learning in LLMs. Many existing works investigate the roles of incontext examples and prompts. Min et al. (2022b) show a surprising result that ground-truth demonstrations are not required for in-context learning, while other factors such as the label space, input text distribution, and overall sequence format play an important role. Shi et al. (2023) investigate the distractibility of LLMs and shows that their performance dramatically drops when irrelevant context is included. Subsequently, Wei et al. (2023) setting with hundreds or thousands of examples (Agarwal et al., 2024; Li et al., 2023a; Bertsch et al., 2024) .Learning to learn in-context In contrast to the emergent in-context capabilities of LLMs, existing works have also studied methods that learn to perform in-context learning explicitly. Min et al. (2022a) propose MetaICL, a meta-training framework for finetuning pretrained LLMs to perform incontext learning on a large and diverse collection of tasks. MetaICL outperforms several baselines including emergent in-context learning and multi-task learning followed by zeroshot transfer. Going beyond the text domain, TNP (Nguyen and Grover, 2022) and PFNs (Müller et al., 2021) are two concurrent works that propose to train transformer models to perform in-context prediction for a family of functions, which allows in-context generalization to unseen functions after training. Similarly, (Garg et al., 2022) show that autoregressive transformers can be trained from scratch to learn function classes such as linear functions and 2-layer ReLU networks. These works present an interesting set of baselines for our work to examine the in-context learning ability of LLMs. 
conclusion: 
 We propose a novel approach to understanding in-context learning in LLMs by probing their decision boundaries in in-context learning in binary classification tasks. Despite achieving high test accuracy, we observe that the decision boundaries of LLMs are often irregularly non-smooth.Through extensive experiments, we identify factors that affect this decision boundary. We also explore fine-tuning and adaptive sampling methods, finding them effective in improving boundary smoothness. Our findings provide new insights into the mechanics of in-context learning and suggest pathways for further research and optimization. 
 
 Appendix: Anonymous Institution, Anonymous City, Anonymous Region, Anonymous Country. Correspondence to: Anonymous Author <anon.email@domain.com>. 


Decision: This paper proposes a new way to interpret in-context learning for binary classification by analyzing the decision boundaries. The work shows that these boundaries frequently are not smooth, examines the reasons for this and proposes methods for improving the generalizability of in-context learning. Authors perform additional experiments on 6 NLP datasets/tasks based on reviewer suggestions. Rebuttals result in reviewers increasing their scores.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: **[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: 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.
limitations: limitations: Yes
paper_summary: paper_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.
strength_weakness: strength_weakness: None that I can see.
questions: questions: N/A.
limitations: limitations: N/A.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: See the weaknesses.
limitations: limitations: The main limitation is that the experiments are conducted only on toy datasets.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: 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?
limitations: 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.

Comments: title: Thank you for raising score!
Response to the author
comment: 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!
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.
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!
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!
We thank the reviewer for raising the score and are pleased that you find our work novel and our experiments solid.
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.
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.