中文

基于因果检索的高效长度可泛化注意力用于长上下文语言建模

计算与语言 2025-06-13 v4 人工智能

摘要

尽管Transformer取得了成功,但处理长上下文仍具有挑战性,由于自注意力的有限长度泛化和二次复杂度。因此,Transformer通常需要在更大的注意力窗口下进行后训练,显著增加计算和内存成本。本文中,我们提出一种基于动态上下文、分组交叉注意力(GCA)的新型注意力机制,能够在保持常数注意力窗口大小的同时,泛化到训练上下文长度的1000倍,并能够访问远距离信息。对于给定的输入序列,我们将其分割为块,并使用每个块检索后续文本生成的top-k相关过去块。具体而言,与大多数先前方法不同,使用标准检索器,我们的关键创新允许检索器学习如何检索更能最小化后续标记自回归损失的过去块,以端到端方式实现。这种机制容纳检索块与固定大小的注意力窗口,以实现长程信息访问,显著降低训练和推理期间的计算和内存成本。实验表明,GCA-based模型在16M上下文长度的密钥检索中实现近乎完美的准确率,这是训练长度的1000倍。

关键词

引用

@article{arxiv.2410.01651,
  title  = {Efficient Length-Generalizable Attention via Causal Retrieval for Long-Context Language Modeling},
  author = {Xiang Hu and Zhihao Teng and Jun Zhao and Wei Wu and Kewei Tu},
  journal= {arXiv preprint arXiv:2410.01651},
  year   = {2025}
}

备注

accepted to ICML 2025