PaperlayerView exact revision on arXiv

Source-linked AI summary

BeaconKV: Key-Value Cache Compression Guided by Beacon Queries for Efficient Large Reasoning Model Inference

Janghyeon Kim, Minsoo Kim, Kyuhong Shim, Jungwook Choi

arXiv:2609.04971v1cs.LGcs.CL

TL;DR

Long CoT traces make LRM KV caches grow linearly, while recent-query compression can miss distant context revisited later. BeaconKV uses compact representatives of clustered global queries to guide cache retention, generally outperforming existing methods while reducing memory by up to 5.8× and improving throughput by over 4.3×.

  • Problem

    Extended CoT generation makes KV caches grow linearly with reasoning length, while recent-query methods can prematurely evict distant context needed by later Thought Revisiting Tokens.

  • Method

    BeaconKV is a training-free compression method that maintains beacon queries representing clustered global queries alongside recent queries to anticipate future KV revisits.

  • Results

    BeaconKV generally outperforms existing compression methods across models and benchmarks, reducing memory by up to 5.8× while nearly preserving full-cache accuracy and improving throughput by over 4.3×.

  • Takeaways & Limitations

    BeaconKV offers an effective approach for deploying LRMs under constrained memory budgets by preserving critical context during long-horizon reasoning.

  • Takeaways & Limitations

    Recent-query methods fail because Thought Revisiting Tokens occur sporadically and unpredictably, making future distant KV revisits difficult to anticipate.

Abstract

from arXiv · show

Large Reasoning Models (LRMs) achieve superior problem-solving through extended Chain-of-Thought (CoT) generation, but the resulting key-value (KV) cache grows linearly with sequence length and creates severe memory bottlenecks, often exceeding GPU capacity for long reasoning traces. Existing KV cache compression methods rely on recent queries to estimate future token importance, implicitly assuming these serve as reliable proxies for future attention patterns. We demonstrate that this assumption fails in long-horizon reasoning: certain decoding steps generate Thought Revisiting Tokens (TRT) that re-attend to distant previous context, such as task-solving plans formulated early in the trace. Through systematic analysis, we discover that queries corresponding to the TRT cluster into a small number of similarity groups in the embedding space. Based on this insight, we propose BeaconKV, a training-free KV cache compression method that maintains beacon queries, compact representatives for each global query cluster, to anticipate which KV pairs will be revisited without storing the entire query history. Across four open-source LRMs and diverse reasoning benchmarks, BeaconKV generally outperforms existing compression methods, achieving up to $5.8\times$ memory reduction while nearly preserving full cache accuracy and improving throughput by over $4.3\times$.

1. Introduction

Extended Chain-of-Thought generation improves LRM reasoning but makes KV caches grow linearly with sequence length, creating severe memory bottlenecks. BeaconKV addresses this by representing clustered global queries with compact beacon queries, improving compression while preserving reasoning performance.

  • Motivation: Extended CoT generation produces long reasoning traces, causing KV caches to grow linearly with sequence length and potentially exceed GPU capacity.The cache must retain key-value pairs for previously generated tokens during autoregressive decoding.
  • Motivation: Recent compression methods estimate future token importance from recent queries, assuming those queries reliably proxy future attention patterns.This limitation can cause premature eviction of tokens needed later in long-horizon reasoning.
  • Observation: BeaconKV identifies Thought Revisiting Tokens whose queries re-attend to distant context, including early task-solving plans, to maintain global reasoning coherence.It distinguishes local queries that attend nearby keys from global queries that access distant keys.
  • Method: Global queries cluster into a small number of similarity groups, enabling compact beacon queries to represent them without storing the entire query history.BeaconKV maintains beacon queries alongside recent queries to anticipate which KV pairs future global queries will revisit.
  • Results: 4.3× throughput improvement and 5.8× peak GPU memory reduction are achieved while nearly preserving accuracy comparable to full KV inference.Across four open-source LRMs and diverse reasoning benchmarks, BeaconKV generally outperforms RPC and R-KV, with accuracy gains of up to 31.7 percentage points over existing methods.

2. Background

Transformer attention projects hidden states into queries, keys, and values, while autoregressive decoding caches prior key-value pairs for reuse. Because newly generated pairs accumulate linearly, recent-query attention scores are used to retain the highest-scoring entries under a cache budget.

  • Attention formulation: Each attention head projects hidden states into queries, keys, and values, then combines values using masked attention.The causal mask prevents tokens from attending to future positions.
  • KV caching: During autoregressive decoding, previously generated key-value pairs remain cached while only the new token’s query, key, and value are computed.The new pair is appended after attention is evaluated.
  • KV caching: The KV cache grows linearly with decoding length, creating a major memory bottleneck for long-context and reasoning tasks.
  • Attention-based KV scoring: Recent-query compression methods estimate cached-key importance from attention weights induced by a small observation window of recently generated queries.The observation set is defined over the most recent Nobs decoding steps.
  • Attention-based KV scoring: They aggregate observation-query attention weights for each cached key and retain the top-BKV key-value pairs under the cache budget.Aggregation may use the maximum or average score across the observation window.

3. Observation

The paper identifies Thought Revisiting Tokens (TRT) as global queries that revisit distant reasoning context, and finds that these queries recur across model components and cluster geometrically.

  • Thought Revisiting Tokens: Thought Revisiting Tokens are tokens whose queries attend to distant earlier context, such as problem statements or task-solving plans.Local queries mainly attend nearby keys, whereas TRT-associated global queries revisit established reasoning context.
  • Attention patterns: Global queries redirect attention to distant spans containing early solving plans and problem constraints, unlike local queries focused on nearby calculations.The contrast is illustrated for tokens 1068 and 1090 versus tokens 1089 and 1091.
  • Thought Revisiting Tokens: Global queries occur across multiple layers and heads rather than concentrating in one model component.The distribution is observed over output positions 512–639 in an AIME24 trace using R1-Distill-Qwen-7B.
  • Geometric structure: Global queries have low similarity to surrounding queries but high mutual similarity, indicating shared structure despite local dissimilarity.The similarity analysis uses pre-RoPE query states to separate geometric similarity from positional effects.
  • Geometric structure: PCA shows that global queries cluster into a small number of similarity groups rather than scattering randomly.TRT queries at tokens 1068 and 1090 appear close within the same cluster.
  • Beacon queries: Beacon queries compactly represent global-query clusters, allowing future revisited KV pairs to be anticipated without storing the complete query history.The method maintains beacon queries alongside recent queries.

4. Method

BeaconKV augments recent queries with compact beacon queries selected from historical query geometry, aligns them for scoring, and retains KV pairs indicated by their strongest attention signals.

  • Overview: BeaconKV is a training-free compression framework that augments recent-query scoring with representative beacon queries from clustered pre-RoPE query states.The design targets reasoning-critical context revisited by global queries.
  • Observation queries: The observation query set combines recent queries for local coherence with beacon queries representing global query clusters.This expands the observation window beyond only the most recent tokens.
  • Observation query selection: 16 Recent + 16 FPS-Selected Queries achieves the highest accuracy across constrained KV-cache budgets in the reported AIME24 comparison.This configuration also maintains high maximum cosine similarity to future decoding-step queries.
  • Observation query selection: Continual FPS maintains a bounded per-head query buffer and periodically downsamples it to retain geometrically distinctive queries.The fill-and-compress procedure avoids storing the full accumulated query history.
  • Observation query selection: Continual FPS achieves accuracy comparable to ideal offline sampling while reducing peak GPU memory by minimizing the query-history footprint.The comparison includes K-Means and Naive FPS sampling strategies.
  • KV scoring: BeaconKV aligns beacon queries to the current decoding step, computes attention scores over the cache, and retains top-K KV pairs using max aggregation.Recent queries retain their original positional alignment, while max pooling preserves sparse high-magnitude TRT signals.
  • Evaluation: BeaconKV preserves thought-revisiting pathways while aggressively compressing memory by combining historical beacon queries with recent-token retention.The method retains top-K scored KV pairs alongside a small recent-token window.

5. Experiments

BeaconKV is evaluated across four open-source LRMs and diverse reasoning benchmarks, where it generally outperforms existing compression methods. It also improves memory and throughput under long decoding while preserving higher accuracy at matched budgets.

  • Evaluation Setup: BeaconKV is evaluated on four open-source LRMs across mathematics, coding, and scientific reasoning benchmarks.The models are R1-Distill-Qwen-7B, R1-Distill-Llama-8B, Qwen3-4B, and Qwen3-14B; benchmarks include AIME24, MATH-500, LiveCodeBench, and GPQA-Diamond.
  • Performance on Reasoning Tasks: 31.7 percentage points is the largest reported accuracy gain over existing compression methods, achieved on Qwen3-14B for AIME24 at a 1024-token KV budget.Across most benchmarks and models, BeaconKV achieves the highest accuracy at the same budget, especially in the low-budget regime.
  • Ablation Study: (16, 16) for recent and beacon queries achieves 64.6% accuracy with 4355.7s latency, outperforming the over-beacon allocation (1, 31) trade-off.Over-allocating queries to beacons reduces short-term context capacity, increasing latency and lowering accuracy.
  • Ablation Study: Max aggregation reaches 23.3 versus 18.8 for Mean aggregation at a 256-token budget, because Mean dilutes sparse high-value beacon signals.Max aggregation generally performs better, particularly under low-budget settings.
  • Ablation Study: BeaconKV consistently outperforms Initial+Recent because continual selection captures emerging global query patterns that fixed initial queries cannot.The comparison indicates that preserving beginning-of-context queries alone does not explain BeaconKV’s gains.
  • Efficiency Evaluation: 5.8× lower peak memory, 82.3 →356.4 tokens/s throughput, and 5573.4 →1287.3s latency result when BeaconKV uses a 2K budget instead of Full KV.Peak memory falls from 77.0 GB to 13.3 GB, enabling larger batch sizes during 32K-token generation.
  • Efficiency Evaluation: +6.3 points at a 2K budget and +12.3 points at a 1K budget are BeaconKV’s LiveCodeBench accuracy gains over RPC under matched efficiency conditions.At both budgets, throughput and memory usage are similar to RPC while accuracy is higher.

6. Related Work

Related work compresses KV caches, restricts attention, controls reasoning length, or adds memory modules. These approaches differ in whether they retain full cache state, require training, alter reasoning length, or modify the architecture.

  • KV Cache Compression: KV cache compression methods evict entries using attention-based importance scores, with LRM-focused methods addressing rapidly growing caches from extended Chain-of-Thought.RPC is cited as an LRM-oriented example.
  • KV Cache Compression: TRIM-KV, LightThinker, and Fast KVzip use trainable gating or compression modules to predict importance or compress reasoning context.These methods require task-specific training through distillation or supervised fine-tuning on frozen backbones.
  • Sparse Attention for Large Reasoning Models: Sparse-attention methods reduce attention computation by restricting query access or approximating clustered KV entries while retaining the full KV cache in memory.Multipole Attention and ReSA represent examples of this direction.
  • Adaptive Control of Reasoning Length: Adaptive reasoning-control methods optimize Chain-of-Thought length or paths rather than compressing the KV cache.Examples include training-free early exit, steerable interventions, and training-based deeper reasoning.
  • Memory-Augmented Architectures: Memory-augmented architectures expand long-context capability by adding memory modules that dynamically compress, store, and retrieve context.This paradigm changes the model architecture rather than operating solely within a standard decoder-only transformer.

7. Conclusion

BeaconKV is a training-free KV cache compression approach for LRMs that uses geometrically distinctive past queries to preserve critical context. Across models and benchmarks, it reduces memory while maintaining near-full-cache accuracy and improving throughput.

  • Conclusion: BeaconKV uses geometrically distinctive past queries to preserve critical KV cache in LRMs.The method is motivated by Thought Revisiting Tokens, which re-attend to prior context during extended reasoning.
  • Conclusion: Up to 5.8× memory reduction and over 4.3× throughput improvement are reported with accuracy close to full KV inference.BeaconKV generally outperforms existing compression methods across models and benchmarks.

Impact Statement

BeaconKV reduces memory use and improves throughput for long-horizon LRM inference, potentially lowering deployment hardware requirements and energy costs. However, the evaluation is primarily limited to open-source LRMs and benchmark settings, so transfer to broader workloads requires further assessment.

  • Benefits: Reduced peak GPU memory and higher throughput may lower the hardware barrier for deploying advanced reasoning models.The statement connects deployment accessibility to BeaconKV’s efficiency improvements.
  • Broader Impacts: Training-free compression may improve energy efficiency for long-horizon reasoning, while lower scaling costs could also reduce the cost of undesirable uses.The impact statement presents both environmental benefits and a possible dual-use concern.
  • Limitations: Further assessment is needed because the evaluation is conducted primarily on open-source LRMs and benchmark settings.The supported scope boundary concerns transfer to broader generation workloads.

A. Additional Experimental Results

BeaconKV provides a favorable accuracy–efficiency trade-off under a fixed 1K KV cache budget, while long reasoning traces create substantial memory pressure. Statistical analysis further characterizes the long generations motivating cache compression.

  • Comparison with SnapKV: 42.2% LiveCodeBench accuracy with BeaconKV exceeded SnapKV at 30.9% and RPC at 29.9% under the same 1K KV cache budget.Throughput and peak GPU memory usage remained similar, although BeaconKV incurred a modest decoding-latency increase.
  • Comparison with SnapKV: BeaconKV maintained comparable system efficiency while substantially improving accuracy over SnapKV and RPC on Qwen3-4B.The comparison used batch size 320 and a maximum KV cache budget of 1K.
  • Output Lengths: Reasoning benchmarks frequently generate several thousand output tokens, especially on AIME24 and LiveCodeBench.Because KV-cache size grows linearly with decoding length, these traces impose practical memory pressure during inference.
  • TRT Analysis: TRT analysis measures global-query frequency across layers and heads using mean distances from each query to its top-K attended output positions, with K = 150.This procedure quantifies how far attention reaches during decoding.

C. Limitations and Future Work

BeaconKV’s evaluation is focused on long-horizon reasoning with open-source LRMs, while its inference procedure combines beacon and recent queries and repeatedly compresses representatives with FPS. Future work includes broader workload testing, hyperparameter analysis, and adaptive compression.

  • Limitations and Future Work: BeaconKV’s effectiveness on non-reasoning tasks such as long-context retrieval, summarization, and general long-context generation remains insufficiently explored.The authors identify broader-workload generalization as requiring further experiments.
  • Limitations and Future Work: The method has hyperparameters for beacon-query count, recent-query count, and KV-cache budget whose optimal settings may vary across models and tasks.The authors leave systematic sensitivity analysis and dynamically adaptive settings for future work.
  • BeaconKV Procedure: FPS selects diverse representatives from pre-RoPE queries by greedily adding the query least similar to the current selected set.The resulting representatives approximate the historical query distribution while preserving distinct query directions for KV scoring.
  • BeaconKV Procedure: BeaconKV uses beacon queries for historical patterns and recent queries for short-range decoding information.During prefill, each head initializes its observation set by applying FPS to the pre-RoPE query sequence.
  • BeaconKV Procedure: BeaconKV preserves prefix and recent tokens, fills the remaining budget with top-scoring KV positions, and recompresses query representatives after eviction.Beacon queries are rotated at the current decoding step, whereas recent queries retain their original positions before attention-based scoring.

E. Token-Level Visualization of Thought Revisiting Tokens

Token-level visualization distinguishes global TRT queries from local queries in an AIME24 reasoning trace. Global queries revisit distant problem constraints and solving plans, whereas local queries focus on nearby steps.

  • Token-Level Visualization: Global queries at tokens 1068 and 1090 revisit earlier problem constraints and high-level solving plans.The visualization marks the tokens receiving the strongest attention for each query.
  • Token-Level Visualization: Local queries at tokens 1089 and 1091 mainly attend to nearby reasoning steps.This contrasts with the distant attention of the representative global queries.
  • Token-Level Visualization: Figure 9 compares query positions with their top-K attended tokens for representative global and local queries.The example uses AIME24 sample-0 on R1-Distill-Qwen-7B at layer 18, head 16.
Loading 2609.04971v1…