🤖 AI Summary
This work addresses the high computational complexity, substantial memory footprint, and poor state scalability of conventional attention mechanisms in long-context scenarios by introducing Key-Value Means (KVM), a novel block-recurrent attention mechanism. KVM constructs a recurrent structure using standard attention operations—without requiring custom kernels—and supports both fixed and growable state representations while preserving parallel training capability and enabling linear-time inference. Its key contribution lies in unifying the scalable memory of Transformers with the efficient inference of linear RNNs, allowing continuous adjustment of prefill complexity from O(N) to O(N²). This significantly reduces KV cache overhead and achieves sub-quadratic prefill time and sub-linear state growth on long-context tasks, all with minimal additional parameters. The code and models are publicly released.
📝 Abstract
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.