MateICL: Mitigating Attention Dispersion in Large-Scale In-Context Learning

TL;DR

MateICL addresses attention dispersion when scaling in-context demonstrations by splitting the context into parallel windows and adding an AtBias attention-calibration layer that prioritizes query tokens. This retrieval-free approach maintains effective self-attention as context grows and reduces reliance on external retrievers, achieving competitive or superior results across classification, multiple-choice, and MRC tasks, even on resource-constrained GPUs. The method demonstrates robust gains with moderate bias settings and larger models, while also highlighting trade-offs between context size, compute, and model scale. Overall, MateICL offers a practical pathway to leverage larger demonstrations in ICL without finetuning or heavy retrieval architectures, enhancing stability and performance in real-world settings.

Abstract

Large Language Models (LLMs) have demonstrated remarkable capabilities in In-Context Learning (ICL). However, the fixed position length constraints in pre-trained models limit the number of demonstration examples. Recent efforts to extend context suffer from attention dispersion as the number of demonstrations increases. In this paper, we introduce Mitigating Attention Dispersion in large-scale ICL (MateICL) that enables LLMs to maintain effective self-attention as the context size grows. We first split the context into multiple windows, each filled to the model's context capacity, which are processed separately. Then, we introduce an additional layer to recalibrate the attention weights, prioritizing the query tokens as the number of demonstrations increases. Our empirical results show that MateICL can effectively leverage larger contexts to improve ICL performance. Compared to retrieval-based baselines, MateICL consistently achieves better performance without requiring an externally trained retrieval model. Despite recent advances in inference strategies (e.g., 32k token contexts), our results demonstrate that MateICL remains beneficial in computationally resource-constrained settings. The code is publicly available at https://github.com/amurtadha/MateICL.

Figures (4)

• Figure 1: Visualization of attention weights for a context window ($\text{W}=9$) on the TREC Dataset using LLaMA-7B, comparing PCW (A) pcw and MateICL (B), where tokens $> 50$ represent the query.
• Figure 2: An example of our MateICL. The context is divided into multiple windows to match the LLM's capacity (e.g., up to 1024 tokens for GPT-2). Each window is independently encoded, restricting windows to only attend to their own tokens. MateICL adds an additional "AtBias" layer to ensure that the task tokens can still attend to themselves as the number of examples grows.
• Figure 3: Averaged accuracy and standard deviation across text classification datasets of ($1 \leq b \leq 7$) of Eq.\ref{['eq:bias_value']} within context windows ($W=6$) and ($W=9$).
• Figure 4: (A) Averaged accuracy and error bars for scaled demonstrations and (B) model size correlation with average text classification improvements.