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

📄 arXiv: 2505.01110v1 📥 PDF

作者: Murtadha Ahmed, Wenbo, Liu yunfeng

分类: cs.CL

发布日期: 2025-05-02

🔗 代码/项目: GITHUB

💡 一句话要点

MateICL：缓解大规模上下文学习中的注意力分散问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 上下文学习 注意力机制 大型语言模型 长文本处理 注意力分散

📋 核心要点

  1. 现有大语言模型上下文学习受限于固定长度，扩展上下文导致注意力分散，影响性能。
  2. MateICL将上下文分割成多个窗口分别处理，并引入额外层重新校准注意力权重。
  3. 实验表明，MateICL能有效利用更大上下文，提升ICL性能，且优于检索方法。

📝 摘要（中文）

大型语言模型（LLMs）在上下文学习（ICL）中展现了卓越的能力。然而，预训练模型中固定的位置长度约束限制了演示示例的数量。最近扩展上下文的尝试，随着演示数量的增加，会遭受注意力分散的问题。本文提出了缓解大规模ICL中注意力分散的MateICL，使LLMs能够在上下文大小增长时保持有效的自注意力。我们首先将上下文分割成多个窗口，每个窗口都填充到模型的上下文容量，并分别处理。然后，我们引入一个额外的层来重新校准注意力权重，随着演示数量的增加，优先考虑查询token。我们的实验结果表明，MateICL可以有效地利用更大的上下文来提高ICL性能。与基于检索的基线相比，MateICL始终能获得更好的性能，而无需外部训练的检索模型。尽管最近在推理策略方面取得了进展（例如，32k token上下文），但我们的结果表明，MateICL在计算资源受限的环境中仍然有益。代码已在https://github.com/amurtadha/MateICL上公开。

🔬 方法详解

问题定义：现有的大型语言模型在进行上下文学习时，由于预训练模型固有的长度限制，无法有效利用大量的演示示例。直接扩展上下文长度会导致注意力分散，使得模型难以关注到关键信息，从而降低了学习效果。因此，如何在大规模上下文中保持有效的注意力是亟待解决的问题。

核心思路：MateICL的核心思路是通过分而治之的策略来缓解注意力分散。首先，将长上下文分割成多个较小的窗口，每个窗口的大小都在模型的上下文处理能力范围内。然后，分别处理这些窗口，并在最后通过一个额外的注意力校准层来整合信息，从而使模型能够更好地关注查询token，提高学习效率。

技术框架：MateICL的技术框架主要包含两个阶段：上下文分割和注意力校准。首先，将输入的上下文分割成多个窗口，每个窗口包含一定数量的演示示例。然后，每个窗口分别通过语言模型进行处理，得到每个窗口的表示。最后，将这些表示输入到注意力校准层，该层通过学习权重来重新分配注意力，从而使模型更加关注查询token。

关键创新：MateICL的关键创新在于引入了注意力校准层，该层能够根据演示示例的数量动态调整注意力权重，从而缓解了注意力分散的问题。与传统的注意力机制不同，MateICL的注意力校准层更加关注查询token，从而提高了模型在长上下文中的学习能力。

关键设计：注意力校准层使用了一个可学习的权重矩阵，该矩阵用于重新分配注意力权重。具体来说，对于每个查询token，注意力校准层会计算一个权重，该权重表示该查询token与所有演示示例的相关性。然后，使用这些权重来重新加权演示示例的表示，从而使模型更加关注与查询token相关的示例。损失函数的设计目标是最大化查询token与相关示例之间的互信息，从而提高模型的学习能力。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，MateICL在多个上下文学习任务上取得了显著的性能提升。与基于检索的基线方法相比，MateICL在不需要额外训练检索模型的情况下，始终能够获得更好的性能。即使在32k token上下文的设置下，MateICL仍然能够带来性能提升，证明了其在计算资源受限场景下的有效性。

🎯 应用场景

MateICL可应用于各种需要利用大规模上下文信息的自然语言处理任务，例如问答系统、文本摘要、机器翻译等。该方法能够有效提高模型在长文本上的学习能力，从而提升任务的性能。此外，MateICL在计算资源受限的环境中也具有优势，可以降低对硬件的要求。

📄 摘要（原文）

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.