Key-Value Means: Transformers with Expandable Block-Recurrent Compressed Memory
- URL: http://arxiv.org/abs/2605.09877v3
- Date: Fri, 15 May 2026 01:47:34 GMT
- Title: Key-Value Means: Transformers with Expandable Block-Recurrent Compressed Memory
- Abstract summary: Key-Value Means (KVM) is a block-recurrence for attention that can accommodate either fixed-size or growing state.<n>We show it performs competitively on long-context tests with only subquadratic prefill time and sublinear state growth.
- Score: 1.3714815350300447
- License: http://creativecommons.org/licenses/by/4.0/
- 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/featherless-ai/KVM-paper and trained models at https://huggingface.co/collections/featherless-ai/kvm-paper under the Apache 2.0 license.
Related papers
- Hidden Decoding at Scale: Latent Computation Scaling for Large Language Models [75.80275843320511]
We study whether an existing backbone can keep improving by allocating more computation to each token while leaving the Transformer backbone fixed.<n>We propose Hidden Decoding, a sequence-length scaling method applied during continued pretraining.
arXiv Detail & Related papers (2026-07-09T07:37:59Z) - Long-Context Aware Upcycling: A New Frontier for Hybrid LLM Scaling [25.551309705184234]
HyLo extends usable context length by up to $32times$ through efficient post-training.<n>HyLo delivers consistently strong short- and long-context performance.<n>At similar scale, HyLo-Qwen-1.7B trained on only 10B tokens significantly outperforms JetNemotron on GSM8K, Lm-Harness common sense reasoning and RULER-64K.
arXiv Detail & Related papers (2026-04-27T17:23:37Z) - M$^2$RNN: Non-Linear RNNs with Matrix-Valued States for Scalable Language Modeling [53.351103753419046]
We introduce Matrix-to-Matrix RNN: an architecture with matrix-valued hidden states and expressive non-linear state transitions.<n>We demonstrate that the language modeling performance of non-linear RNNs is limited by their state size.<n>We also demonstrate how the state size expansion mechanism enables efficient use of tensor cores.
arXiv Detail & Related papers (2026-03-15T12:53:09Z) - Out of the Memory Barrier: A Highly Memory Efficient Training System for LLMs with Million-Token Contexts [68.79341332280062]
Training Large Language Models (LLMs) on long contexts is severely constrained by prohibitive GPU memory overhead, not training time.<n>We introduce OOMB, a highly memory-efficient training system that directly confronts this barrier.<n>Our approach employs a chunk-recurrent training framework with on-the-fly activation recomputation, which maintains a constant activation memory footprint.
arXiv Detail & Related papers (2026-02-02T13:52:40Z) - CoMeT: Collaborative Memory Transformer for Efficient Long Context Modeling [40.705016911274]
We introduce a novel architecture that enables LLMs to handle arbitrarily long sequences with constant memory usage and linear time complexity.<n>CoMeT can be integrated into pre-trained models with only minimal fine-tuning.<n>A model equipped with CoMeT and fine-tuned on 32k contexts can accurately retrieve a passkey from any position within a 1M token sequence.
arXiv Detail & Related papers (2026-02-02T07:49:44Z) - Evolution Strategies at the Hyperscale [57.75314521465674]
We introduce EGGROLL, an evolution strategies (ES) algorithm designed to scale backprop-free optimization to large population sizes.<n>ES is a set of powerful blackbox optimisation methods that can handle non-differentiable or noisy objectives.<n>EGGROLL overcomes these bottlenecks by generating random matrices $Ain mathbbRmtimes r, Bin mathbbRntimes r$ with $rll min(m,n)$ to form a low-rank matrix perturbation $A Btop$
arXiv Detail & Related papers (2025-11-20T18:56:05Z) - CommVQ: Commutative Vector Quantization for KV Cache Compression [50.37946553931796]
We propose Commutative Vector Quantization (CommVQ) to significantly reduce memory usage for long-context LLM inference.<n>We first introduce additive quantization with a lightweight encoder and codebook to compress the KV cache.<n>Our approach achieves high accuracy with additive quantization and low overhead via the RoPE-commutative codebook.
arXiv Detail & Related papers (2025-06-23T17:50:11Z) - FlashRNN: I/O-Aware Optimization of Traditional RNNs on modern hardware [6.749483762719583]
State-tracking capabilities are important for time-series tasks and logical reasoning.<n>Traditional RNNs like LSTMs and GRUs do have these capabilities at the cost of strictly sequential processing.<n>We show how fast these networks can get with our hardware-optimization FlashRNN in Triton and optimized kernels to the register level.
arXiv Detail & Related papers (2024-12-10T18:50:37Z) - GoldFinch: High Performance RWKV/Transformer Hybrid with Linear Pre-Fill and Extreme KV-Cache Compression [0.5534140394498714]
GoldFinch is a hybrid Linear Attention/Transformer sequence model.
It stacks our new GOLD transformer on top of an enhanced version of the Finch architecture.
We train up to 1.5B parameter class models of the Finch, Llama, and GoldFinch architectures.
arXiv Detail & Related papers (2024-07-16T18:00:00Z) - MLKV: Multi-Layer Key-Value Heads for Memory Efficient Transformer Decoding [13.272502334507617]
We introduce Multi-Layer Key-Value (MLKV) sharing, a novel approach extending KV sharing across transformer layers to reduce memory usage.
Evaluations on various NLP benchmarks and inference metrics using uptrained Pythia-160M variants demonstrate that MLKV significantly reduces memory usage with minimal performance loss.
arXiv Detail & Related papers (2024-06-13T16:33:44Z) - Attention as an RNN [66.5420926480473]
We show that attention can be viewed as a special Recurrent Neural Network (RNN) with the ability to compute its textitmany-to-one RNN output efficiently.
We introduce a new efficient method of computing attention's textitmany-to-many RNN output based on the parallel prefix scan algorithm.
We show Aarens achieve comparable performance to Transformers on $38$ datasets spread across four popular sequential problem settings.
arXiv Detail & Related papers (2024-05-22T19:45:01Z) - Transformers are Multi-State RNNs [25.99353771107789]
We show that decoder-only transformers can be conceptualized as unbounded multi-state RNNs.
Transformers can be converted into $textitbounded$ multi-state RNNs by fixing the size of their hidden state.
We introduce a novel, training-free compression policy - $textbfT$oken $textbfO$mission $textbfV$ia $textbfA$ttention (TOVA)
arXiv Detail & Related papers (2024-01-11T18:35:26Z) - RWKV: Reinventing RNNs for the Transformer Era [54.716108899349614]
We propose a novel model architecture that combines the efficient parallelizable training of transformers with the efficient inference of RNNs.
We scale our models as large as 14 billion parameters, by far the largest dense RNN ever trained, and find RWKV performs on par with similarly sized Transformers.
arXiv Detail & Related papers (2023-05-22T13:57:41Z) - SMYRF: Efficient Attention using Asymmetric Clustering [103.47647577048782]
We propose a novel type of balanced clustering algorithm to approximate attention.
SMYRF can be used as a drop-in replacement for dense attention layers without any retraining.
We show that SMYRF can be used interchangeably with dense attention before and after training.
arXiv Detail & Related papers (2020-10-11T18:49:17Z)
This list is automatically generated from the titles and abstracts of the papers in this site.
This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.