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

TL;DR

MateICL introduces a method to mitigate attention dispersion in large-scale in-context learning, enabling models to effectively utilize larger contexts and improve performance without external retrieval models.

Contribution

The paper proposes a novel approach that splits context into multiple windows and recalibrates attention, enhancing large language models' ability to handle bigger contexts in ICL.

Findings

MateICL improves ICL performance with larger contexts.

It outperforms retrieval-based baselines without external models.

It remains effective in resource-constrained 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.