Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Beyond Speedup -- Utilizing KV Cache for Sampling and Reasoning

Zeyu Xing, Xing Li, Hui-Ling Zhen, Mingxuan Yuan, Sinno Jialin Pan

TL;DR

This work reframes the key-value (KV) cache, traditionally used solely to accelerate autoregressive decoding, as a reusable, low-cost representation for downstream inference tasks. It introduces KV-CoE, a token-centric adaptation of Chain-of-Embedding that uses KV cache trajectories to self-evaluate reasoning with minimal overhead, and KVClassifier, a difficulty-driven fast/slow thinking controller that reuses KV signals to reduce token generation by up to $5.7\times$ with small accuracy losses. Through experiments on MATH and TheoremQA, as well as reasoning benchmarks like GSM8K and MATH500, the authors demonstrate that KV-derived embeddings can rival hidden-state signals for local, task-conditioned decisions and that adaptive switching guided by KV-derived difficulty scores yields substantial efficiency gains. The study highlights the practical significance of repurposing inference-time artifacts, suggesting new directions for representation reuse in LLM inference and deployment. The approach offers deployment-friendly, zero-order changes to existing pipelines while delivering meaningful improvements in sampling efficiency and controllable reasoning depth.

Abstract

KV caches, typically used only to speed up autoregressive decoding, encode contextual information that can be reused for downstream tasks at no extra cost. We propose treating the KV cache as a lightweight representation, eliminating the need to recompute or store full hidden states. Despite being weaker than dedicated embeddings, KV-derived representations are shown to be sufficient for two key applications: \textbf{(i) Chain-of-Embedding}, where they achieve competitive or superior performance on Llama-3.1-8B-Instruct and Qwen2-7B-Instruct; and \textbf{(ii) Fast/Slow Thinking Switching}, where they enable adaptive reasoning on Qwen3-8B and DeepSeek-R1-Distil-Qwen-14B, reducing token generation by up to $5.7\times$ with minimal accuracy loss. Our findings establish KV caches as a free, effective substrate for sampling and reasoning, opening new directions for representation reuse in LLM inference. Code: https://github.com/cmd2001/ICLR2026_KV-Embedding.

Beyond Speedup -- Utilizing KV Cache for Sampling and Reasoning

TL;DR

This work reframes the key-value (KV) cache, traditionally used solely to accelerate autoregressive decoding, as a reusable, low-cost representation for downstream inference tasks. It introduces KV-CoE, a token-centric adaptation of Chain-of-Embedding that uses KV cache trajectories to self-evaluate reasoning with minimal overhead, and KVClassifier, a difficulty-driven fast/slow thinking controller that reuses KV signals to reduce token generation by up to with small accuracy losses. Through experiments on MATH and TheoremQA, as well as reasoning benchmarks like GSM8K and MATH500, the authors demonstrate that KV-derived embeddings can rival hidden-state signals for local, task-conditioned decisions and that adaptive switching guided by KV-derived difficulty scores yields substantial efficiency gains. The study highlights the practical significance of repurposing inference-time artifacts, suggesting new directions for representation reuse in LLM inference and deployment. The approach offers deployment-friendly, zero-order changes to existing pipelines while delivering meaningful improvements in sampling efficiency and controllable reasoning depth.

Abstract

KV caches, typically used only to speed up autoregressive decoding, encode contextual information that can be reused for downstream tasks at no extra cost. We propose treating the KV cache as a lightweight representation, eliminating the need to recompute or store full hidden states. Despite being weaker than dedicated embeddings, KV-derived representations are shown to be sufficient for two key applications: \textbf{(i) Chain-of-Embedding}, where they achieve competitive or superior performance on Llama-3.1-8B-Instruct and Qwen2-7B-Instruct; and \textbf{(ii) Fast/Slow Thinking Switching}, where they enable adaptive reasoning on Qwen3-8B and DeepSeek-R1-Distil-Qwen-14B, reducing token generation by up to with minimal accuracy loss. Our findings establish KV caches as a free, effective substrate for sampling and reasoning, opening new directions for representation reuse in LLM inference. Code: https://github.com/cmd2001/ICLR2026_KV-Embedding.
Paper Structure (47 sections, 12 equations, 4 figures, 10 tables)

This paper contains 47 sections, 12 equations, 4 figures, 10 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Hidden–state self-evaluation.
  4. Adaptive fast/slow reasoning and dynamic control.
  5. KV-cache interventions.
  6. Background
  7. Transformer, Hidden States, and KV Cache
  8. Modern LLM Frameworks and KV Cache Management
  9. Observation
  10. Can KV Caches Serve as an Embedding Source
  11. Why KV Caches Are Sufficient for Chain-of-Embedding and Fast/Slow Thinking Switch?
  12. Local (restricted-set) adequacy.
  13. CoE is not classification.
  14. Task conditioning + efficiency.
  15. Scope and limitations.
  16. ...and 32 more sections

Figures (4)

  • Figure 1: VRAM usage vs. context length for Qwen3-32B yang2025qwen3, comparing Model+KV Cache vs. Model+KV Cache+Hidden States.
  • Figure 2: Comparison between vanilla CoE (top) and KV-CoE (bottom). Vanilla CoE aggregates hidden states across the token dimension to construct a representation for each layer, whereas KV-CoE aggregates KV Cache across the layer dimension to construct a representation for each token.
  • Figure 3: Comparison of efficiency and effectiveness of fast vs. slow thinking on GSM8K using Qwen3-32B. Slow thinking achieves slightly higher accuracy but at a much higher token cost.
  • Figure 4: KVClassifier: special tokens are dynamically inserted to perform thinking-mode switching based on KV-derived difficulty scores.