键值均值注意力机制
研究团队提出键值均值注意力机制,这是一种支持固定或可增长状态的新型分块循环注意力。在强Transformer基线中引入固定大小的KVM层,仅需极少参数即可构建高效的O(N)分块RNN。使用可增长KVM缓存训练的模型在长上下文任务中表现优异,仅需次二次方预填充时间和次线性状态增长。该机制基于标准操作实现,无需定制内核,支持分块并行化训练与预填充,兼具传统Transformer与线性RNN的优势。它可应用于每一层以节省KV缓存内存,并在O(N)到O(N²)之间提供灵活的预填充时间复杂度选择,也可与线性RNN层混合部署以增强长上下文解码能力。相关代码与模型已开源。
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 linear 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). It can also be implemented in a hybrid solution in tandem with LRNN layers in place of traditional attention, to supplement the LRNN with improved sublinear memory growth context length usage and long context decoding. We release our code at https://github.com/recursal/KVM-paper and trained models at https://huggingface.co/collections/recursal/key-value-means under the Apache 2.0 license.