PaperlayerView exact revision on arXiv

Source-linked AI summary

Key-Value Means: Transformers with Expandable Block-Recurrent Compressed Memory

Daniel Goldstein, Navneel Singhal, Eugene Cheah

arXiv:2605.09877v5cs.LGcs.AIcs.CL

TL;DR

Transformers retain information through a growing KV cache, while fixed-state RNNs offer constant costs but limited recall. KVM combines block sliding-window attention with a fixed or expandable compressed state, yielding competitive short-context performance and strong long-range retrieval across a continuous trade-off between the two designs.

  • Problem

    Transformers incur growing memory and per-token costs, whereas fixed-state linear RNNs limit long-context recall; a selectable trade-off between these endpoints is missing.

  • Method

    KVM combines an uncompressed block sliding window with a compressed key-value state whose size can remain fixed or grow according to a chosen budget.

  • Results

    KVM shows competitive short-context performance and strong long-range retrieval, tunable through different state-size schedules.

  • Takeaways & Limitations

    KVM smoothly interpolates between fixed-state RNNs and full attention, enabling flexible memory and computation choices within one architecture.

  • Takeaways & Limitations

    The paper has not yet attempted transformer distillation to use KVM attention on one or more layers.

Abstract

from arXiv · show

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.

1 Introduction

KVM is a block-recurrent softmax-attention architecture whose compressed state can remain fixed or grow under a chosen budget, interpolating between fixed-state RNNs and full transformer KV caches. It targets the trade-off among memory, prefill time, decode time, and recall by combining an exact sliding window with compressed key-value state.

  • Motivation: Transformers retain one key-value pair per prior token but incur linearly growing memory and per-token time, whereas modern linear RNNs use constant resources but typically limit long-context memory.A model that varies state growth between these endpoints can trade off memory, prefill time, decode time, and recall.
  • Motivation: Growable state supports recall-heavy contexts by allocating additional capacity for difficult-to-compress information while remaining smaller than a full KV cache.Fixed-size states must encode increasing information with unchanged capacity, while slower-growing states represent multiple tokens together.
  • Contribution: KVM combines an exact block sliding window with a fixed or growable compressed KV state, providing selectable memory and computation across a continuous range.The architecture is designed as a chunked recurrent alternative to full attention and can interpolate between RNNs and Transformers.
  • Contribution: KVM is presented as a high-performance block-recurrent softmax-attention layer offering a fast, low-state-size alternative to full attention and higher recall than fixed-state RNN layers.The contributions also include applying a pre-selected growth strategy to interpolate between RNNs and Transformers, plus sharing partial RoPE and NoPE across compressed and uncompressed regions.

2 Background

Methods that reduce transformer KV-cache growth choose among state compression, eviction, update timing, decision metrics, and positional-encoding strategies. KVM combines block sliding-window attention with batched compressed-state updates, selectable growth, similarity-based merging, and modified rotary information.

  • Design space: Reducing KV-cache growth requires architectural choices about how state is compressed or otherwise grows more slowly than a transformer’s linear cache.These choices organize the background methods discussed in this section.
  • Compression: LRNNs maintain fixed-size states updated per token, while sparsely routed LRNNs select among multiple states for updates and retrieval.The supplied passage identifies these as prominent compression approaches for autoregressive transformers and parallelizable alternatives.
  • Timing: Sliding-window methods treat state as a queue, with standard variants evicting individual tokens and block variants periodically evicting entire tail blocks.Other methods append tokens during prefill before applying eviction or compression during decoding at one-time, per-token, or periodic intervals.
  • Our Choices: KVM combines block sliding-window attention with batched eviction and merging into a compressed state, selectable growth, token-radius tracking, and Just-In-Time renormalization.Its state decisions use maximum key similarity, while rotary positional information is removed from keys as they merge into the state.

3 Method

KVM combines a short uncompressed BSWA window with a compressed key-value state that can remain fixed-size or grow according to a chosen budget. It updates this state block by block through selective appending, similarity-based merging, protected sink rows, and normalized attention readout.

  • Compressed State and BSWA Window: The compressed state stores key-value rows from earlier blocks, with a fixed or growing number of rows and an initial set protected as attention sinks.The BSWA window retains recent blocks, while the oldest block becomes overflow and is incorporated into the state at each block boundary.
  • Readout: Before readout, KVM applies row-wise LayerNorm to the state and restores each value row to its stored radius, while learned per-head temperatures separately scale state and BSWA attention.The stored radius remains fixed for a slot’s lifetime, even as merging changes the direction of its value representation.
  • Chunks and attention regions: KVM attends over protected sink rows, the recent uncompressed BSWA window, and a periodically updated compressed state while processing queries chunk by chunk.The first window prefix uses exact causal attention; later chunks expose all state rows and apply causal masking within the BSWA window.
  • Positional Encoding: KVM makes compressed memory position-independent by zeroing the rotary subspace of overflow keys, normalizing memory keys, and applying a data-dependent merge gate.Queries and keys in the BSWA window retain partial RoPE, enabling the state and window to be concatenated for one attention call without extra projection matrices.
  • Append and Merge: At each chunk boundary, KVM appends the least redundant overflow tokens allowed by the state budget, then merges remaining tokens into the most similar eligible state rows.Appended tokens are selected by redundancy against the normalized state; merge targets exclude protected sink rows, and the merge gate controls each token’s contribution.

4 Language modeling performance

KVM is evaluated against transformer, recurrent, and compressed-attention baselines across sequence-position loss, standard short-context benchmarks, and long-context tasks. The results indicate strong performance for KVM, including with fixed state, while increasing state size can provide additional benefits.

  • Experimental setup: KVM is evaluated against RWKV-7, full-attention GPTAlpha-2, block sliding-window attention, and hybrid OVQ/SWA architectures at 120M and 350M parameters.Models are trained on Prolong at 8k context length, with KVM variants using block size C = 256 and two BSWA chunks.
  • Loss over sequence position: KVM shows strong performance as sequence position increases, with fixed-state KVM 256 outperforming the much larger saturating-state OVQ/SWA configuration.KVM-sqrt achieves the best results among the tested non-GPTAlpha models.
  • Standard short-context benchmarks: KVM matches the expected behavior of a standard transformer on short-context tasks because it jointly attends over its BSWA window and compressed state.The short-context evaluations use LM Evaluation Harness across standard benchmark tasks.
  • RULER and LongBench: The long-context evaluation covers NIAH-S across context lengths, full RULER at 4k context length, and few-shot LongBench using LM Evaluation Harness.These tests assess both performance at the trained context length and out-of-the-box extrapolation.
  • RULER and LongBench: The experiments suggest that utilizing an increasing state size can be a significant benefit when distractors are continuously novel.This observation qualifies the trade-off between fixed-size and growing-state memory.

5 Speed Benchmarks

Speed benchmarks compare Triton KVM kernels with full causal multi-head attention for decode, prefill, and backward passes. KVM’s advantage increases with context length, with KVM-256 reaching its largest reported speedups at 32768 tokens.

  • 5 Speed Benchmarks: The KVM kernels are implemented in Triton and benchmarked for autoregressive decode, prefill, and backward passes.Figure 5 reports sequence-mixing times with medians and minimum-to-maximum run-median bands, where lower is better.
  • 5 Speed Benchmarks: The benchmark evaluates KVM-256 and KVM-sqrt against PyTorch Flash SDPA on an AMD Instinct MI325X across 512–32768-token contexts.Each single-layer sequence-mixing test uses 32 heads of dimension 128, batch size 8, bfloat16 inputs and outputs, and float32 precision-sensitive operations.
  • 5 Speed Benchmarks: At 32768 tokens, KVM-256 is 29.77× faster for decode, 4.59× faster for prefill, and 7.99× faster for backward than full attention.KVM-sqrt achieves 7.88×, 1.76×, and 3.30× speedups for the same workloads, respectively.
  • 5 Speed Benchmarks: Timing includes KVM routing and recurrent-state computation or full attention’s causal computation, but excludes projections, positional encoding, and the rest of the transformer layer.The measured interval begins after input projections and ends after head concatenation.
  • 5 Speed Benchmarks: KVM-256 becomes faster than full attention at 1K tokens for decode, 8K for prefill, and 4K for backward.KVM-sqrt crosses over at 1K, 16K, and 8K tokens, respectively.

6 Ablation studies

Ablations on 120M KVM 256 models show that architectural choices primarily affect long-context behavior, with value-length normalization, sink protection, and the merge gate important for retrieval. Results are reported for long-context evaluations in Table 3 and short-context evaluations in Appendix Table 5.

  • 6 Ablation studies: The ablations use 120M KVM 256 models and evaluate long-context performance with NIAH, RULER-4096, and average LongBench few-shot scores in Table 3.Short-context evaluations are reported in Appendix Table 5.
  • 6 Ablation studies: Removing value-length normalization causes the largest degradation in long-context performance among the tested architectural ablations.
  • 6 Ablation studies: Removing sink protection or the merge gate also substantially weakens long-context retrieval.

7 Conclusions

KVM unifies block sliding-window attention with an expandable compressive state, enabling a smooth trade-off between fixed-state O(N) recurrence and stronger long-context behavior with sublinear state growth. Its remaining directions include adaptive scheduling, efficiency methods, and distillation into KVM layers.

  • Conclusions: KVM combines block sliding-window attention with an expandable compressive state in one softmax layer, supporting flexible state sizes unlike fixed-size RNNs and full-attention transformers.With fixed state, KVM forms an O(N) chunked recurrent architecture; with growable state, it achieves stronger long-context behavior with sublinear asymptotic state growth.
  • Conclusions: KVM delivers competitive short-context performance and strong long-range retrieval, with behavior tunable through different state-size schedules.The method smoothly interpolates between fixed-state RNNs and full attention rather than requiring a binary choice.
  • Future Work: Future work should investigate changing state-size and chunk-size schedules across training and testing, through finetuning, or in response to data.The reported experiments used static schedules.
  • Future Work: Future work should test transformer-efficiency methods such as GQA and MLA and explore distilling transformers to use KVM attention in one or more layers.The authors expect GQA and MLA to apply easily and view distillation as a way to reduce memory footprint and other costs, but have not attempted it.

AI Usage Disclosure · A Related Work · B Pseudocode

The paper situates KVM among fixed-state, expandable-state, fast-weight, and compressed-memory architectures, distinguishing its integrated, expandable recurrent attention design from related methods. Its pseudocode describes overflow-token merging, state updating, and joint attention over recurrent state and a sliding window, while disclosing LLM assistance in research and writing.

  • AI Usage Disclosure: The authors used LLMs for code and mathematics, TikZ diagrams, phrasing, stylistic improvements, research discussions, and literature-search coverage.This disclosure concerns assistance during both research and paper preparation.
  • A Related Work: Fast-weight approaches treat recurrent state as trainable parameters updated at test time, encompassing RNN state, attention’s expanding cache, Test-Time Training, Titans, and LaCT.LaCT uses a two-layer SwiGLU-MLP with backpropagation, Muon, and momentum, batching updates every 2048 tokens or more and integrating them with sliding-window attention.
  • A Related Work: Fixed-size-state models provide fixed memory and amortized per-token inference cost but necessarily limit total long-context retrieval capacity.This family includes linear and nonlinear architectures, including linear attention, state-space models, and LRNNs with efficiently parallelizable updates over short chunks.
  • A Related Work: Block-Recurrent Transformers (Hutchins et al., 2022) periodically update fixed state through sliding-window and cross-attention, requiring dedicated state projections and additional parameters.Native Hybrid Attention (Du et al., 2026) instead combines a constant number of recurrent long-term slots with sliding-window tokens in one softmax.
  • A Related Work: Expandable-state research seeks slower-than-attention state growth without sacrificing increasing capacity and result quality, including compressed queues, latent-token encoders, and dynamic fast-weight expansion.Compressive Transformer compresses overflow blocks into a FIFO queue; End-to-End Context Compression grows latent tokens at a reduced rate; TokenFormer dynamically expands an MLP state but does not apply it directly to attention.
  • A Related Work: KVM differs from concurrent OVQ by integrating compressed state and BSWA attention in one softmax, supporting uncapped expansion, explicit RoPE compatibility, dynamic weighting, and sink preservation.OVQ uses winner-take-all assignment similarly, but KVM additionally avoids per-centroid count tracking through renormalization.
  • B Pseudocode: The pseudocode merges overflow tokens into the recurrent state by removing RoPE, applying data-dependent gating, matching each overflow key to its most similar non-sink state key, and adding keys and values.Protected sink entries are excluded from matching, while state keys are normalized before similarity computation.
  • B Pseudocode: After updating state, attention concatenates normalized recurrent state entries with the BSWA window and applies causal scaled dot-product attention.State and window use separate temperature factors, and values are normalized and scaled by their accumulated lengths.

C GPTAlpha-2 Transformer Architecture and Backbone · GPTAlpha-2 Channel Mixer

The experiments use GPTAlpha-2, a GPTAlpha-based transformer incorporating selected RWKV-7 and related design choices. Its channel mixer follows RWKV-7 and computes an intermediate hidden state with a simplified token-shift operation.

  • C GPTAlpha-2 Transformer Architecture and Backbone: GPTAlpha-2 modifies the GPTAlpha transformer with RWKV-7-inspired components, value residuals, and RoPE applied to half of each head’s dimension.The design includes biased LayerNorm on queries and keys, a simplified non-data-dependent token shift, value residuals (Zhou et al., 2025), and RoPE.
  • C GPTAlpha-2 Transformer Architecture and Backbone: The architecture uses LayerNorm with bias on queries and keys, together with a simplified non-data-dependent token shift.These choices are part of the modified GPTAlpha architecture used in the experiments.
  • GPTAlpha-2 Channel Mixer: The channel-mixing MLP uses the RWKV-7 channel mixer.This is the stated channel-mixer choice in GPTAlpha-2.
  • C GPTAlpha-2 Transformer Architecture and Backbone: Attention weight preparation applies the token-shift relation a_t = ã_t + α_a ⊙(ã_{t−1} − ã_t), with a_0 = ã_0 for a ∈ {q,k,v}.The notation presents the single-head attention preparation for queries, keys, and values.
  • C GPTAlpha-2 Transformer Architecture and Backbone: For the first layer, the relevant value term is the ã_v,t calculated for that layer.This specifies the value used in the first-layer token-shift formulation.
  • GPTAlpha-2 Channel Mixer: The channel mixer forms an intermediate hidden state as h_t = (x_t + α ⊙(x_{t−1} − x_t))W_U.The formula defines the shifted input transformation used for the intermediate hidden state.

D Training details and Hyperparameters · E Short-context evals

Training uses CompleteP-based scaling, AdamC with scheduled optimization, fixed token batches, and warmup-decay learning rates; short-context results are reported for models and ablations across standardized LM Evaluation Harness tasks.

  • D Training details and Hyperparameters: CompleteP with α = 1 governs parameter-wise depth and width scaling.
  • D Training details and Hyperparameters: Training keeps batch size constant at 524,288 tokens while tracking total training steps.
  • D Training details and Hyperparameters: AdamC schedules weight decay, using β1 = 0.9, β2 = 0.95, ϵ = 10^-8, learning rate 2×10^-3, and weight decay 0.2.The authors expect stationary parameter norms and report improved performance versus AdamW.
  • D Training details and Hyperparameters: Learning rates and weight decay are tuned on a 120M model trained for 3B tokens, then transferred to larger scales without decaying scalar or vector parameters.
  • D Training details and Hyperparameters: The learning rate uses 200 warmup steps followed by linear decay to 0 for the remaining training steps.
  • D Training details and Hyperparameters: RoPE uses a base of 10,000 even when applied to only 64 of 128 channels.
  • E Short-context evals: Table 4 reports short-context evaluation results for the models using LM Evaluation Harness.
  • E Short-context evals: Tables 5 and 7 report short-context results for KVM and GPTA-2 partial-RoPE ablations, respectively, using abbreviated benchmark names such as lmbda, arc_c, arc_e, hella, piqa, and winog.The abbreviations refer to LAMBADA (Paperno et al., 2016), ARC-Challenge (Clark et al., 2018), HellaSwag, PIQA (Bisk et al., 2020), and WinoGrande.

F Extrapolation and partial RoPE ablations

NoPE and HalfRoPE show a tradeoff in length extrapolation: HalfRoPE maintains stable position-wise loss and generally performs better within the training context, whereas NoPE can improve some long-context retrieval scores farther beyond it. The results may reflect NoPE’s reliance on absolute positions versus HalfRoPE’s stronger relative-position handling.

  • Interpretation: NoPE may rely increasingly on learned absolute positional embeddings as training at a fixed context length increases, making it less suitable for extrapolation in some settings.The explanation is attributed to prior findings on absolute positional embeddings (Haviv et al., 2022; Kazemnejad et al., 2023).
  • Interpretation: NoPE may capture global long-context structure and pinpoint specific tokens, but HalfRoPE appears stronger at adjacent-token attention and NIAH-S3 through relative-position handling.This is presented as a conjecture explaining NoPE’s weaker NIAH-S3 performance despite its broader long-context behavior.
  • Extrapolation results: HalfRoPE maintains stable position-wise loss and generally outperforms NoPE within the training context, while NoPE can outperform it farther beyond that context on NIAH, LongBench, and RULER.These trends are reported for TextbookChapters position-wise loss and long-context evaluations.

G Design choices

KVM is designed as a growable compressive-state architecture targeting constant or sublinear memory growth and subquadratic sequence-length complexity. Its design combines key-value state updates, similarity-based compression and expansion, empirically selected merge rules, and practical positional-encoding choices.

  • Motivation: The architecture targets high-quality long-context modeling with constant or sublinear memory growth, subquadratic computation, and minimal test-time-training hyperparameters.
  • Overall BSWA framework: The BSWA framework stores both keys and values for attention, while in-place updates limit state growth below traditional per-token appending.
  • State Compression: Compression assigns overflow keys to state keys using their attention-logit similarities, with alternatives spanning linear-attention φ functions, deferred normalization, L-norms, and modified softmax metrics.
  • State Compression: Experiments favored running averages over EMA, weighted incoming tokens, and JIT-normalized state keys, which remain simple sums of weighted overflow keys.
  • State Initialization and Expansion; Positional Encoding: State expansion appends the least-redundant overflow tokens, while positional encoding uses a practical mechanism that zeros the rotational subspace of state keys; alternative attention re-merging is left for future work.
Loading 2605.09877v5…