Key-Value Means：具有可扩展块循环压缩记忆的 Transformer

我们提出了 Key-Value Means（“KVM”），一种用于注意力机制的新型块循环机制，可适应固定大小或不断增长的状态。将固定大小的 KVM 注意力层配备至强大的 Transformer 基线上，可得到强大的 $O(N)$ 分块 RNN，同时仅增加极少量的新参数。我们训练了带有可增长 KVM 缓存的 Transformer，并表明它在长上下文测试中表现优异，且仅需次二次的预填充时间和次线性的状态增长。KVM 可通过标准操作实现，无需自定义内核，并支持分块可并行化的训练和预填充。它在单一统一框架中提供了传统 Transformer（可扩展的上下文记忆、分块可并行化的训练和预填充）与线性 RNN 的诸多优势。它可用于每一层，节省 KV 缓存内存，并允许在 $O(N)$ 和 $O(N^2)$ 之间连续选择预填充时间复杂度。它也可以与 LRNN 层协同实现混合方案以替代传统注意力，通过改善的次线性内存增长上下文长度使用和长上下文解码来增强 LRNN。我们在 https://github.com/featherless-ai/KVM-paper 发布了代码，并在 https://huggingface.co/collections/featherless-ai/kvm-paper 下以 Apache 2.0 许可证发布了训练好的模型。