提出KVM机制,让Transformer实现可扩展的内存与线性推理速度。
Key-Value Means: Transformers with Expandable Block-Recurrent Compressed Memory

- 用块循环结构压缩键值记忆,支持固定或动态扩展状态
- 长文本任务中预填充时间亚二次、状态增长亚线性,性能接近Transformer
- 兼容标准操作,适合所有层使用,适合长序列建模场景
回忆能力面临权衡:Transformer内存随序列线性增长,导致逐词生成变慢;而传统RNN虽有固定开销但回忆能力有限。本文提出一种新型块循环注意力机制——键值均值(Key-Value Means, KVM),可灵活支持固定大小或可扩展状态。在强基线Transformer中加入固定大小的KVM注意力层,即可实现$O(N)$分块RNN效果,新增参数极少。训练可扩展的KVM缓存模型,在长文本测试中表现优异,仅需亚二次预填充时间与亚线性状态增长。KVM采用标准运算,无需自定义核函数,支持分块并行训练与预填充。它融合了Transformer(可扩展上下文、分块并行)与RNN(低内存开销)的优点,可应用于每一层,显著节省KV缓存内存,并在$O(N)$到$O(N^2)$预填充复杂度间提供连续选择。代码与模型已开源于https://github.com/featherless-ai/KVM-paper及HuggingFace,许可为Apache 2.0。
原文摘要 · Abstract (English)
Recall presents a difficult choice: transformers have a linearly growing memory that slows each successive token, while linear RNNs typically have fixed costs but limited recall. We present Key-Value Means ("KVM"), a novel block-recurrence for attention that can accommodate either fixed-size or growing state. Equipping a strong transformer baseline with fixed-size KVM attention layers yields a strong $O(N)$ chunked RNN, while adding only an insignificant number of new parameters. We train a transformer with a growable KVM cache and show it performs competitively on long-context tests with only subquadratic prefill time and sublinear state growth. KVM is implementable with standard operations and without custom kernels, and supports chunk-wise parallelizable training and prefill. It provides many of the benefits of both traditional transformers (expandable context memory, chunk-wise parallelizable training and prefill) and RNNs in a single unified package. It can be used on every layer, saving KV-cache memory, and allowing a continuous range of choices of prefill time complexity between $O(N)$ and $O(N^2)$. We release our code at https://github.com/featherless-ai/KVM-paper and trained models at https://huggingface.co/collections/featherless-ai/kvm-paper under the Apache 2.0 license.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。