Research map / Large language models / Language, safety and interpretability
Probing the Decision Boundaries of In-context Learning in Large Language Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Probes LLM in-context learning via decision-boundary visualization on synthetic binary classification, finds non-smooth boundaries, and tests fine-tuning, custom architectures, and uncertainty-based active sampling to smooth them.
Problem
In-context learning in LLMs is poorly understood; it is unclear whether LLM decision boundaries on simple in-context classification tasks are smooth and generalizable.
Approach
Treat LLM in-context learning as a classifier and visualize its decision boundaries over a 2D grid of queries on scikit-learn linear, circle, and moon tasks; then test fine-tuning on synthetic classification tasks, a CustomLLM with new embedding/head layers, and entropy-based active sampling to smooth boundaries.
Key results
- Across LLMs (1.3B to GPT-4o), decision boundaries on linearly separable tasks are non-smooth and fragmented even with 128 in-context examples.
- Fine-tuning Llama3-8B on a single linear task generalized to unseen non-linear and 3- and 4-class tasks with smoother boundaries.
- Uncertainty-aware active sampling yielded smoother boundaries and higher test accuracy with fewer in-context examples than random sampling.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Increase model size | Pretrained LLMs · Linear, circle, and moon classification tasks | decision-boundary smoothness | No clear effect |
| Increase the number of in-context examples | Pretrained LLMs · Linear, circle, and moon classification tasks | decision-boundary smoothness | No clear effect |
| Fine-tune on the in-context examples | Llama3-8B · Binary classification tasks | decision-boundary smoothness | No clear effect |
| Fine-tune on synthetic classification tasks | Llama series models · Linear, circle, and moon classification tasks | decision-boundary smoothness | Helped |
| Fine-tune on a single binary linear task | Llama3-8B · Unseen nonlinear, 3-class, and 4-class classification tasks | decision-boundary smoothness | Helped |
| Use uncertainty-aware active sampling rather than random sampling | LLM in-context classification · Binary classification tasks | decision-boundary smoothness | Helped |
| Use uncertainty-aware active sampling rather than random sampling | LLM in-context classification · Binary classification tasks | test accuracy | Helped |
| Train a transformer from scratch with the TNP objective | TNP · Moon classification task | decision-boundary smoothness | Helped |
| Fine-tune attention layers rather than only the linear head (ablation) | Llama3-8B · Binary classification tasks | decision-boundary smoothness | Helped |
| Fine-tune the token embedding layer rather than only the linear head (ablation) | Llama3-8B · Binary classification tasks | decision-boundary smoothness | Helped |
| Add a new input embedding layer to CustomLLM (ablation) | CustomLLM · Binary classification tasks | decision-boundary smoothness | Helped |
| Add a new prediction head to CustomLLM (ablation) | CustomLLM · Binary classification tasks | decision-boundary smoothness | No clear effect |
Surprising findings
- Fine-tuning Llama3-8B on in-context examples alone did not smooth decision boundaries.
- Larger model size did not correlate with smoother decision boundaries, even GPT-4o showed fragmented regions.
Limitations
- Open-source models were evaluated with 8-bit quantization due to compute constraints
- Active sampling labels come from a logistic regression ground-truth proxy rather than true labels
- Results are mostly qualitative visualizations with limited error bars
Experimental setup
Datasets: scikit-learn synthetic linear, circle, and moon classification tasks SUBJ financial phrasebank (FP) RTE ETHOS SST-2 AG_NEWS
Benchmarks: scikit-learn synthetic linear, circle, moon tasks SUBJ FP RTE ETHOS SST-2 AG_NEWS
Metrics: test accuracy on held-out set of 100 decision boundary smoothness (qualitative)
Baselines: SVM with polynomial kernel MLP Decision tree K-nearest neighbors TNP
Builds on: LoRA TNP MetaICL LLaMA
Closest papers by meaning
- In-Context Learning Dynamics with Random Binary Sequences (ICLR 2024 · similarity 0.90)
- Supervised Knowledge Makes Large Language Models Better In-context Learners (ICLR 2024 · similarity 0.89)
- The mechanistic basis of data dependence and abrupt learning in an in-context classification task (ICLR 2024 · similarity 0.89)
- Dual Operating Modes of In-Context Learning (ICML 2024 · similarity 0.89)
- Why Larger Language Models Do In-context Learning Differently? (ICML 2024 · similarity 0.89)
- Is attention required for ICL? Exploring the Relationship Between Model Architecture and In-Context Learning Ability (ICLR 2024 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.