🌍 引言
在智能的世界里，记忆就像是一个巨大的仓库，能够根据具体的上下文进行高效计算。尽管如此，基于变换器的语言模型（LLMs）却常常被其注意力机制的限制所束缚。本文介绍了一种新的高效方法，可以将基于变换器的LLMs扩展至无限长的输入，同时确保内存和计算的限制。这种方法的核心是一个名为“Infini-attention”的新型注意力技术，它将压缩内存引入传统的注意力机制，并在单一的变换器模块中结合了遮蔽的局部注意力和长期线性注意力机制。
📚 背景
在深度学习的早期，递归神经网络（RNNs）通过逐步读取每一个输入来处理序列信息。然而，RNNs在处理长序列时的局限性让人难以忍受。在此背景下，变换器模型的出现无疑是一次技术革命，它通过并行计算极大地提高了处理速度，但仍然受到注意力机制的计算复杂度（记忆和计算时间均为平方级别）所困扰。例如，一个500B的模型在2048的上下文长度下，其注意力键值（KV）状态的内存占用高达3TB。
为了应对这一挑战，压缩内存系统被提出，它承诺在处理极长序列时比传统的注意力机制更具可扩展性和效率。与随输入序列长度增长的数组不同，压缩内存主要保持固定数量的参数，以存储和检索信息，从而降低存储和计算成本。
graph TD;
A[输入序列] --> B[局部注意力计算]
B --> C[压缩内存]
C --> D[全局上下文输出]🔧 方法
1. Infini-attention
Infini-attention是我们提出的核心技术，它通过在每个变换器块中集成压缩内存和局部注意力，从而打破了上下文长度的限制。与传统的变换器模型不同，Infini-attention能够在处理新输入时，保留旧的KV状态，从而维护整个上下文历史。
这种设计使得Infini-attention不仅可以处理短期上下文，还能有效整合长期记忆，形成一个既高效又灵活的注意力机制。
2. 记忆与有效上下文窗口
我们的实验表明，Infini-attention在长上下文语言建模基准测试中表现出色，且在内存大小上实现了114倍的压缩比。通过在100K序列长度的训练下，模型的困惑度显著降低。我们展示了Infini-attention如何使得1B的LLM自然扩展至1M的序列长度，并在传递密钥检索任务中表现优异。
🧪 实验
在我们的实验中，Infini-Transformer模型在涉及极长输入序列的基准测试上取得了显著的成果，包括长上下文语言建模、1M长度的密钥上下文块检索和500K长度的书籍摘要任务。我们通过对现有LLMs进行连续预训练，突出显示了我们方法的插件式长上下文适应能力。
实验结果
| 模型 | 内存大小 (压缩比) | 上下文长度 | PG19 | Arxiv-math |
|---|---|---|---|---|
| Transformer-XL | 50M (3.7x) | 2048 | 11.88 | 2.42 |
| Infini-Transformer (线性) | 1.6M (114x) | 2048 | 9.65 | 2.24 |
我们的Infini-Transformer在PG19和Arxiv-math基准上均表现出了优越的困惑度，而其内存参数量却大幅低于其他模型。
📝 结论
在处理长上下文时，记忆系统的有效性至关重要。我们的研究展示了Infini-attention如何通过将压缩内存模块与传统的点积注意力层紧密结合，打破变换器的上下文限制。我们的方法不仅能够扩展至百万长度的输入序列，在长上下文语言建模基准和书籍摘要任务中也表现优异，展现了强大的长度泛化能力。
📑 参考文献
- Munkhdalai, T., Faruqui, M., & Gopal, S. (2024). Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention. arXiv:2404.07143v2.
- Vaswani, A. et al. (2017). Attention is All You Need. Advances in Neural Information Processing Systems.
- Rae, J. W. et al. (2019). Compressive Transformers for Long-Range Sequence Modelling. arXiv:1911.05507.