Research map / Large language models / Reasoning and chain of thought
ArborKV: Structure-Aware KV Cache Management for Scaling Tree-based LLM Reasoning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes ArborKV, a tree-aware KV-cache eviction framework for Tree-of-Thoughts-style search that combines a multi-signal value estimator, budgeted tree-aware allocation, token-extractive eviction, and lazy rehydration, reporting up to ~4x peak KV-memory reduction on GSM8K, MATH500, Game of 24, and AIME.
Problem
Tree-of-Thoughts style search amplifies KV-cache memory because a frontier of partial trajectories must stay cached, and sequence-centric eviction policies ignore tree topology and break backtracking.
Approach
Treat KV retention as a tree-aware budgeted allocation over thought blocks: score blocks by fusing search value, next-token confidence and accumulated attention, allocate retention by value and distance to the active leaf, evict token-extractively within blocks, and lazily re-prefill evicted blocks when backtracking revisits them.
Key results
- At rho=0.25 on Llama-3.1-8B, ArborKV reaches 70.2% on GSM8K vs 65.6% for H2O and 67.1% for ThinKV, with peak memory about 5.5 GiB vs 21.8 GiB for FullKV at rho=1.0.
- Under Config-L with N_expand=256 on MATH500 (Llama-3.1-8B, rho=0.25), FullKV runs OOM while ArborKV reaches 78.4% accuracy at 5.6 GiB peak.
- Ablation at rho=0.25 shows removing the distance term (MSVE only) drops GSM8K accuracy to 59.2% from 70.2%, and disabling lazy rehydration lowers Llama GSM8K accuracy to 68.6%.
Surprising findings
- Reported Config-L results show ArborKV at N_expand=256 scoring 78.4% on MATH500, above its own FullKV at N_expand=64 (79.8%) only slightly lower, while baselines H2O and StreamingLLM drop to 69.0% and 67.1%.
- Sequence-centric baselines sometimes match or exceed ArborKV on memory peaks in the main table, e.g., StreamingLLM at Game of 24 Qwen rho=0.25 has lower peak than ArborKV.
Limitations
- Evaluation limited to a single 24GB GPU with two 7-8B models and a 200-instance development subset per task.
- MSVE calibration requires leave-one-out supervision and several hand-set hyperparameters.
- Baselines were adapted to tree search by recomputing retention along the active path, a custom protocol.
Experimental setup
Datasets: GSM8K MATH500 Game of 24 AIME
Benchmarks: GSM8K MATH500 Game of 24 AIME
Metrics: accuracy peak GPU memory (GiB) end-to-end time (s) rehydration count OOM status
Baselines: FullKV H2O StreamingLLM ThinKV
Builds on: Tree-of-Thoughts DPTS H2O StreamingLLM ThinKV
Closest papers by meaning
- BeaconKV: Key-Value Cache Compression Guided by Beacon Queries for Efficient Large Reasoning Model Inference (ICML 2026 · similarity 0.90)
- Semantic Integrity Matters: Benchmarking and Preserving High-Density Reasoning in KV Cache Compression (ICML 2026 · similarity 0.90)
- HyPER: Bridging Exploration and Exploitation for Scalable LLM Reasoning with Hypothesis Path Expansion and Reduction (ICML 2026 · similarity 0.89)
- Value-Aware Stochastic KV Cache Eviction for Reasoning Models (NeurIPS 2026 · similarity 0.89)
- Aligning Tree-Search Policies with Fixed Token Budgets in Test-Time Scaling of LLMs (ICML 2026 · similarity 0.89)
- DeFT: Decoding with Flash Tree-attention for Efficient Tree-structured LLM Inference (ICLR 2025 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.
