提出一种新型注意力机制，兼顾高效与长文本处理能力。
Online Vector Quantized Attention
- 采用稀疏记忆更新的向量量化策略，实现线性计算与常数内存。
- 在64k序列长度下性能接近自注意力，内存仅为后者的极小部分。
- 适合需要长上下文且资源受限的场景，如移动设备部署。
语言模型中的标准序列混合层难以在效率与性能间取得平衡：自注意力在长序列任务中表现优异，但存在二次计算复杂度和线性内存开销；而线性注意力和状态空间模型（SSMs）虽具线性计算与常数内存，却难以有效处理长上下文。本文提出一种新型序列混合层——在线向量量化（OVQ）注意力，兼具线性计算与常数内存优势。其通过稀疏记忆更新机制显著提升记忆状态规模与容量。基于高斯混合回归构建理论基础，并在多种合成长上下文任务及长文本建模任务上验证。结果表明，OVQ注意力在性能上显著优于线性注意力基线及原始向量量化注意力，且在长达64k序列长度下，性能可媲美强自注意力模型，同时仅使用极少内存。
原文摘要 · Abstract (English)
Standard sequence mixing layers used in language models struggle to balance efficiency and performance. Self-attention performs well on long context tasks but has expensive quadratic compute and linear memory costs, while linear attention and SSMs use only linear compute and constant memory but struggle with long context processing. In this paper, we develop a sequence mixing layer that aims to find a better compromise between memory-compute costs and long-context processing, which we call online vector-quantized (OVQ) attention. OVQ-attention requires linear compute costs and constant memory, but, unlike linear attention and SSMs, it uses a sparse memory update that allows it to greatly increase the size of its memory state and, consequently, memory capacity. We develop a theoretical basis for OVQ-attention based on Gaussian mixture regression, and we test it on a variety of synthetic long context tasks and on long context language modeling. OVQ-attention shows significant improvements over linear attention baselines and the original VQ-attention, on which OVQ-attention was inspired. It demonstrates competitive, and sometimes identical, performance to strong self-attention baselines up 64k sequence length, despite using a small fraction of the memory of full self-attention.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。