提出新型注意力机制，让模型在长文本中高效访问远距离信息。
Efficient Length-Generalizable Attention via Causal Retrieval for Long-Context Language Modeling
- 用动态分块检索技术实现可泛化的长序列建模
- 训练长度仅16K时，可准确处理1600万上下文长度的测试任务
- 无需扩大注意力窗口，显著降低计算与内存开销
尽管Transformer取得成功，但处理长上下文仍面临长度泛化能力有限和自注意力二次复杂度的问题。现有方法通常需在训练后增大注意力窗口，大幅增加计算与内存成本。本文提出基于动态上下文的分组交叉注意力（GCA）机制，可在保持固定注意力窗口的前提下，将模型泛化能力扩展至预训练长度的1000倍，并有效获取远距离信息。对于输入序列，将其分块后，每块用于检索前序最相关的k个块以支持后续生成。不同于以往使用现成检索器的方法，本机制通过端到端学习，使检索器自动优化以最小化自回归生成损失。该设计使被检索块能以固定大小注意力窗口处理，极大降低训练与推理阶段的计算与内存开销。实验表明，基于GCA的模型在1600万上下文长度的密钥检索任务中达到接近完美的准确率，是训练长度的1000倍。
原文摘要 · Abstract (English)
Despite the success of Transformers, handling long contexts remains challenging due to the limited length generalization and quadratic complexity of self-attention. Thus Transformers often require post-training with a larger attention window, significantly increasing computational and memory costs. In this paper, we propose a novel attention mechanism based on dynamic context, Grouped Cross Attention (GCA), which can generalize to 1000 times the pre-training context length while maintaining the ability to access distant information with a constant attention window size. For a given input sequence, we split it into chunks and use each chunk to retrieve top-k relevant past chunks for subsequent text generation. Specifically, unlike most previous works that use an off-the-shelf retriever, our key innovation allows the retriever to learn how to retrieve past chunks that better minimize the auto-regressive loss of subsequent tokens in an end-to-end manner. Such a mechanism accommodates retrieved chunks with a fixed-size attention window to achieve long-range information access, significantly reducing computational and memory costs during training and inference. Experiments show that GCA-based models achieve near-perfect accuracy in passkey retrieval for 16M context lengths, which is 1000 times the training length.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。