Beyond Speedup - Utilizing KV Cache for Sampling and Reasoning
Authors: Zeyu XING
Xing Li
Hui-Ling Zhen
Mingxuan Yuan
Sinno Jialin Pan
Venue: ICLR 2026 Poster (2026)
Keywords: Machine Learning
LLM
Primary area: foundation or frontier models, including LLMs
TL;DR: We repurpose the KV cache—traditionally used only for speedup—as a free representation for sampling and reasoning, enabling output-free self-evaluation and adaptive fast/slow thinking with negligible overhead and strong empirical results

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: (i) Chain-of-Embedding, where they achieve competitive or superior performance on Llama-3.1-8B-Instruct and Qwen2-7B-Instruct; and (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.

Published as a conference paper at ICLR 2026
BEYOND SPEEDUP UTILIZING KV CACHE FOR SAM- PLING AND REASONING
Zeyu Xing', Xing Li", Hui-Ling Zhen?, Mingxuan Yuan?, "lann,Sinmn Sinno Jialin Pan' Department of Computer Science and Engineering, The Chinese University of Hong Kong 2 Huawei Technologies Co., Ltd.
zeyuxinge 1 Link. cuhk. edu.hk xingi, zhenhuiling?, yuan. mingxuan) (whuwwei....... com sinnopan edu edu.hk
ABSTRACT
KV caches, typically used only to speed up autoregressive decoding, encode con- textual information that can be reused for downstream tasks at extra cost. We propose treating the KV cache as a lightweight representation, climinat- full hidden states. Despite being weaker ing the need recompute KV-Embeddings are shown to be sufficient for than dedicated embeddings,
key applications: (i) Chain-of-Embeedding where they achieve competi- tive or superior performance on Llama-3.1-8B-Instrue and Qwen2-7B-Instruet: and (ii) Fast/Slow Thinking Switching, where they enable adaptive reason- ing on Qwen3-8B 3BB and
reducing token gener- ation by up to 5.7x with minimal accuracy loss. Our findings establish KV caches as free, effective substrate for sampling and reasoning, opening new di- rections for representation reuse in L.M inference. Our code available
wttp://mmmmmmmmmmmemeeeeme IICDZ222 KV-Embedding
INTRODUCTION
Large language models (LLMs) rely on key-value (KV) cache to accelerate autoregressive decoding by reusing past attention states, avoiding costly recomputation. This makes the KV cache indispens- able for low-latency inference production systems like vLL.M (Kwon al., 2023). However, its role is typically confined to this speedup. Beyond acceleration, the KV cache seldom viewed as reusable representation witt the notable exception cache steering, a technique that modifies the cache's initial state to guide generation (Belitsky et al. 2025)
While the KV cache has been mostly confined to acceleration (Li et al. 2025: Yang et al.] 2025). hidden states have been widely exploited for seif- evaluation (Wang et al.] [2025b] Chen et al.) 2024 Beigi et al.) 2224) et al. 2022a) and for adap- tive reasoning and control (Zhang al.) 2025b Wang et al. 2025a; Yue et al.) 2025). These meth- ods, however, rely on storing full hidden states, which is costly in both memory and computation.
563
Mosel KV Cathe
dden
1.3
[U0 to 18 Smaler
0102
32258
03535
131072
In this work, we investigate simple but powerful question: Can the KV cache do more than just accelerate decoding? Since the KV cache already computed and stored during inference, using it for downstream tasks incurs no additional cost. This a major advantage over storing full hidden states, which is prohibitively expensive in terms of memory. As shown in Figure [ the KV cache offers a significantly more compact and practical alternative for typical decoder-only models.
Figure 1: VRAM usage context length for Qwen3-32B (wwnnTeam 2025) com- paring Model+KV Cache Model+KV Cache+Hidden States.
Though the KV cache is not explicitly trained as a general-purpose embedding its sole objective to support next-token prediction-we find it nonetheless encodes rich contextual information

---

Published as a conference paper at ICLR 2026
suitable for various downstream tasks. We explore the potential of KV-Embeddings through two applications:
Chain-of-Embeddingg We repurpose the KV cache as lightweight and readily avail- able embedding. experiments on Chain of Embedding (CoE) (Wang et al. 2025b) method for selecting optimal reasoning paths without external information show that KV caches achieve classification performance comparable to or even surpassing that of using the model's 's hidden states.
Fast/Slow Thinking Switch:
We leverage the KV cache to implement an adaptive switch- ing mechanism between fast, low-compute reasoning and slower, deliberate reasoning. By reusing KV cache, this approach achieves substantial efficiency gains with minimal perfor- mance loss.
Our contributions are fourfold:
1. We present the first systematic study of KV caches reusable task representations, show- ing they can be repurposed at near-zero computational cost. In particular, we propose simple but effective aggregation techniques that make KV caches directly usable as em- beddings.
2. Despite not being designed as general- embeddings, we find that KV cache rep- resentations when processed with the proposed aggregation strategies, are competitively effective on certain classification tasks.
3. We propose KV-CoE, variant of Mhain-of-Embeeduee that reuses the KV cache already stored during decoding. KV-CoE achieves self-cvaluation without extra activation stor- age, offering nearly zero memory overhead and seamless integration into existing inference frameworks.
4. We introduce KVClassifier, a fast/slow auto-thinking framework that reuses KV caches for adaptive reasoning with minimal overhead.
Our results suggest that KV caches are versatile and low-cost foundation for sampling and rea- soning, moving beyond their traditional role as a mere acceleration component to become a core resource for effective and efficient LLM-based inference.
RELATED WORK
Hidden-state self-evaluation. A growing line of work shows that internal activations encode re- liable signals about answer correctness and hallucination risk. [ang et al. (2025b) propose Chain- of Embeddingg (CoE), which models the trajectory of layerwise hidden states during inference and derives output-free correctness scores from the geometry of this path. Chen et al. (2024) (INSIDE) introduce EigenScore, computed from the eigenspectrum of hidden-stand covariance, assess mantic (in)consistency and detect hallucinations. BBigi et al. (2024) train a contrastive probe on internal states (attention, ML. activations) produce well-calibrated confidence estimates across NLU/NLG tasks. IZhang et al. (2025a) further probes hidden states of reasoning models to predict whether generated answer will be correct, All of these methods operate directly on hidden states or logits. Our study, by contrast, investigates whether the KV cache alone-which is already present inference-suffices support the same families of subtasks.
Adaptive fast/slow reasoning and dynamic control. To mitigate overthinking on easy inputs and underthinking on hard ones, recent work explores adaptive reasoning depth (Xing al. (025). Zhang et al. (2025b) quantify upper bounds of long- vs. no- no-thinking modes and propose Adaptive Self-Recovery Reasoning (ASRR), adding accuracy-aware length rewards to reduce unnecessary rea- soning while allowing implicit recovery. PATS (Wang et al. (223a) performs process-level switch- ing via process reward models with beam search, enabling step-wise fast/slow adaptation with bad- search over atomic actions and learn to select dynamic trajectories. These approaches typically require explicit chain-of-thouuhh genera- tion, external reward models, or re-decoding. Our contribution is orthogonal: we show that pooled KV-cache features can drive both one-shot (classification-stt -sylle and in in-generation (generative-style switching via simple control tokens, without storing hidden states or altering model architecture.
step penalties. DOTS (Yue et al. 2025) views reasoning
2

---

Published as a conference paper at ICLR 2026
Key Cache
logers lagerg leger
Comparison VRAM usage
Managed By
NodemnY Gachs
KV Cache
Value Cache
ModM-4D7 Cace
ligars
19 c0e Slana
Can the KV cache serve as cost-free embedding?
KV-CoE
KV-Classifier
Token Din
<thinke To find
</thinko
Agg.
KV Cache
PL
KV Cache
Aggregation and Lightweight Classifier
Difficulty
30
80
Different trajectories using KV Cache as embedding
Operation
Iesert
Figure 2: The proposed K KV-Embedding framework reuses KV cache lightweight representa- tion for two applications: (a) KV-CoE Shhain-of-Embedding for self-evaluation of reasoning paths, and (b) KVClassifier for fast/slow thinking switch. KV cache, computed during normal inference, is aggregated into embeddings that drive path selection or difficulty-waar switching without storing full hidden states, achieving near-zero overhead.
KV-cache interventions. While our work treats the KV cache as a read-only representation for evaluation and control, concurrent line of research shows that also serve control inter- face. Belitsky et al. 2025) introduce KV Cache Steering, a one-shot intervention that adds layer- wise steering vectors derived from contrastive CoT non-CoT prompts-to the key and value tensors after prefill, reliably inducing longer and more structured reasoning in small Language Mod- els. Compared with activation steering, this method offers improved stability and negligible runtime overhead. Our approach is complementary: instead of modifying the cache, we pool it to derive difficult-awaare signals that gate slow reasoning.
BACKGROUND
3.1 TRANSFORMER, HIDDEN STATES, AND KV CACHE
Decoder-only transformers are the architectural foundation modern large language models (LLMs) (Vaswani et al.) 2017: Brown et al. 2020). During autoregressive generation, each trans- former layer processes new token to produce contextual hidden state. A computational bottleneck arises because standard attention requires recomputing over all previous tokens each step, result- ing in (TT2) complexity per step, where T is the sequence length. To mitigate this, the key-value (KV) cache (Dao al., 2022) stores the attention keys and values for all past tokens every layer. This allows the model to compute keys and values only for the new token and attend to the cached history, reducing the complexity to O(T) per step and enabling efficient sequencc generation.
Formally. for layer I, we store
K KVCache
[K(()
1::: (!U) )
3

---

Published as a conference paper at ICLR 2026
(0) where K i Vi RRTxH HXXxXxxdaa are the stacked key and value tensors across all attention heads H. The hidden state at step is computed via
Attention (Q!". K v
Since K(I)
14-1 1:t-1 are already cached, only K. and vi() need to be computed online.
and vi
3.2 MODERN LLM FRAMEWORKS AND KV CACHE MANAGEMENT
State-of-the-arr LLM serving frameworks carefully manage KV caches to achieve high throughput, low latency, and efficient GPU memory utilization.
Modern LLM frameworks treat the KV cache as a first-class resource. vLLM (Kwon et al. 2023) introduces PagedArtentiin, virtualizing the KV cache with a paging mechanism akin to CPU virtual memory to enable dynamic allocation, low fragmentation, and high-throughput serving under heavy concurrency. In contrast, Ollama (Ollama Team 2024) focuses lightweight, eyveloper-friendee deployment, managing KV caches at the session level to efficiently reuse context across multi-turn interactions. Together, these systems illustrate how both production-scale and interactive inference workloads rely on persistent, reusable KV caches as a core ubstraction.
Across these frameworks, the KV cache is managed as first-class resource with explicit strategies for allocation, eviction, and reuse. This observation motivates our central claim:
"Since the KV cache is an unavoidable byproduct of efficient inference, repurposing for down- stream tasks adds virtually no overhead."
OBSERVATION
4.1
CAN KV CACHES SERVE AS AN EMBEDDING SOURCE
The hidden states and attention projections stored in KV caches encode contextualized token repre- sentations, making them natural candidates for use as embeddings. While recent work has explored leveraging intermediate representations from LLMs as task-specific embeddings (Liu et al. 2024), we specifically investigate aggregating KV cache vectors into sentence-level representations.
To evaluate their quality as an embedding source, we construct embeddings by concatenating keys and values at every layer, then averaging across token positions, attention heads, and layers be- fore applying & normalization. We benchmark these KV-derived embeddings on the Massive Text Embedding Benchmark (MTEB) (Muennighoff et al. 2023) against a strong, dedicated embedding model (gemini-embedddin-0001.
Dataset
-lama-3.1-8B-Insttute KV Cache
Gemini- Sominn-EEmbeddingg000
amazonnConttefiannllCaaaiiiciiii DBpediaClassificationn
0.3530
0.8820
0,5937
0.9476
Caw.aaaalannssaaaaaaaaiiaalllla TWeerTTopiSingeeClssiifcttin
0.62544
0.8864
0.3714
0.7111
Table Performance of KV cache-based embeddings vs. dedicated embedding model on selected MTEB classification tasks. Despite being significantly weaker than trained embeddings, KV-derived embeddings still capture meaningful semantics. Appendix furrher reports hidden-state and chance baselines under the same evaluation pipeline.
As shown Table [ KV-derived embeddings significantly underperform their dedicated counter- part across all datasets, confirming they are not perfect general-purpose embeddings. This gap stems from three factors: (i) KV representations are optimized for causal language modeling, not contrastive learning, leading poor isotropy; (ii) they inherently token- and pesition-sseccife requiring heuristic pooling for sentence-level use; and (iii) their projection into a head space (dhead E dmodet) reduces their discriminative power.
4

---

Published as a conference paper at ICLR 2026
Despite these limitations, the results show that KV caches encode substantial semantic informa- tion con--nouuhh to be competitive on certain classification tasks. This finding motivates our explo- ration of reusing the KV cache for Efbredding and Fast/Slow Thinking Switch, where global embedding quality is less critical than local, relative separability between candidates.
4.2 WHY KV CACHES ARE SUFFICIENT FOR CHAIN-OF-EMBEDDING AND FAST/SL.WW THINKING SWITCH?
KV caches are poor general-purpose embeddings: they trained for next-token prediction, position//cotext dependent, and the space is often anisotropic. Still, they are sufficient for our two uses--hain-of- Eiiieeddi--mmdidiigggggggg (CoE) and Fast/Slow Thinking becache both rely on local, task-conditioned comparisons rather than globally calibrated semantics.
Local (restricted-set) adequacy. General embedding learning targets global separation across a broad label/instance space. Here we only need correct relative ordering within small candidate set C (e.g., small label space or a few candidate continuations). Concretely, for a decision rule f and score gap
((x) [u, (1) Ju, (x),
we only require minyec (yy) not globally well-structured embedding space. This is why KV embeddings can work well on certain classification cases with limited labels/candidates, despite weak MTEB-style semantic similarity.
CoE is not classification. CoE extracts path embedding along the reasoning trajectory and estimates the chance of correction from local geometry (e.g., angle/step length). The requirement therefore weaker than global separability: KV embeddings only need preserve these local trajectory cues consistently over short ranges.
Task conditioning + efficiency. The pooled KV embedding = Pool K, V) is conditioned on input and instruction, e = g(x,t). which already biases the representation toward the current task. Finally, reusing KV caches is essentially free compared to storing hidden states (Chidden > Cky x 0), so in wawpoo//laaencyssensttttte regimes the utility favors KV reuse even if accuracy slightly lower.
Scope and limitations. We do not claim KV caches yield universally strong embeddings. They are unsuitable when a task needs globally comparable semantics across diverse queries (e.g., broad remrieval/imillartt).. Our claim is limited (i) restricted candidate and (ii) CoE-style local trajectory geometry.
CHAIN OF EMBEDDING WITH KV CACHE
LLMs exhibit emergent reasoning capabilities, though their internal decision processes remain largely opaque. To address this, Wang et al. (2025b) introduce Chain-offEmbeellee (CoE), method that probes the model's latent space by tracking the evolution of sentence-level representa- tions across layers. Formally, for an LLM M with layers, let h denote the hidden representation of token layer l. The sentence-leve level representation at layer is obtained by averaging over the sequence length T:
1
1 L.
(1)
T
1=1
The CoE trajectory is then defined as the sequence of these layer-wiss wise representations:
CoE {so, $1, $1,..., $L}.
(2)
CoE characterizes this trajectory by measuring both magnitude and directional changes of embed- dings between consecutive layers:
Ari = | - 81112, and 2 = arccos
S1+1 31
(3)
|

---

Published as a conference paper at ICLR 2026
These features are aggregated into self-evaluation scores. For instance, the real-space combination (CoE-R) is
CoE-R
(aAn + 8281).
(4)
where a, B are weighting coefficients. A more robust complex-space variant (CoE-C) treats each (Ari, A0i) pair as a complex number 21 Am+iiA0 and computes the magnitude of their average:
4-1
CoE-C 1
Σ
(5)
CoE has demonstrated strong discriminative power in distinguishing correct from incorrect model generations, achieving state-of-the-an performance on self-evaluation benchmarks.
5.1 METHODOLOGY
Our key innovation is adapt the CoE framework the KV cache, eliminating its primary computational overhead. While vanilla CoE constructs trajectories from hidden states h re- quiring expensive activation storage re-computation, K((,)), ytl,) that are already maintained by autoregressive decoders. This modification preserves the CoE analytical framework while rendering it virtually cost-free.
instead leverage the key-value tensors
Embedding Construction. For each token t and layer we start with the cached key-value tensors K(.,)) Vil.) E Hd We flatten the head and key/value dimensions and average across layers to produce a compact per-token embedding:
et Hatten((KM..) v(l.0) mRd
(6)
The resulting token-wise trajectory is defined as:
KV-CoE {fi, .2,
err,
which directly parallels the structure of vanilla CoE but operates along the token dimension.
Trajectory Characterization. We characterize this trajectory using the established CoE metries, simply substituting the token index é for the layer index I:
Ar. ||ee11 cell2, and A arccos Terralla [[e. )
er+1-ee
(8)
T-1
T-1 (Ar, + A8,)
KV-CoE-R --1 (aAr++Bl BB8.). and KV-CoE-C
(9)
These formulations maintain the analytical rigor of CoE-R and CoE-C with minimal conceptual alteration.
Contributions and Advantages. As illustrated in Figure which compares vanilla CoE and our KV-CoE, our method offers two main advantages:
1. No extra activation cost. Since the KV cache is already computed and stored during stan- dard autoregressive decoding, reusing for trajectory analysis incurs virtually no additional activation cost. The required reductions are computationally negligible compared to full forward pass, resulting A Rs extra memory and minimal FLOPs.
2. Deployment-friendly: The approach works directly with standard inference stacks (e.g., past_key_valoe in Transformers or vLLM). It requires no architectural changes, re- forwarding, or activation hooks, making it immediately deployable in production LLM serving systems.
5.2 EXPERIMENTAL RESULTS

---

Published as a conference paper at ICLR 2026
Model
Method
MATH
TheoremQA
AUROC
FPR95 +
AUROC
FPR95
MaxProb
59.16
45.41
98,60
PPL
60.82
46.45
97.82
Entropy
62.74
84.14
47.37
97.82
1lama-3.1-88h Instruct
(Llam(()))))))8B)
72.54
75.61
63.12
89.83
CoE-C((L))
73.08
79.60
55.85
90.14
KV-CoE-R (ours)
64.36
63.82
74.74
62,93
KV-CoE-C (ours)
64.13
67.42
74.93
62.46
MaxProb
12.40
99.34
4.92
99.87
PPL
12.43
99.50
5.11
100,00
Qwen2-7B-
Entropy
16.19
99.42
5.28
99.87
Instruct
CoE-R
75.75
65.95
66.68
85,84
CoE-C
76.68
64.48
62.70
87.42
KV-CoE-R (ours)
76.92
49.83
88.87
54.30
KV-CoE-C (ours)
84.12
44.82
83.27
58.35
Table 2: Self-evaluation results on reasoning tasks. KV-CoE consistently improves AUROC and reduces FPR95 relative to MaxProb, PPL, and Entropy. Bold indicates the best value per model-dataset pair except CoE baselines. CoE-R and CoE-C results are taken from the original CoE paper (Wang et al., 2025b). Appendix clarifies CoE usage (confidence estimation, not reranking) and motivates the token-centric KV-Což design via layer-wise ablations. 'hhese baseline results are reported on Samaa--B-IInsruce while our experiments use updated Llama3. 1-8B-Instruct, so the numbers may not perfectly align.
We evaluate KV-CoE two reasoning benchmarks from the original CoE work: MATH (Hendrycks et al. al._211. (202T) for multi-step
Hidden States
Layer Dir
arithmetic and TheoremQA (Chen al. (023) for theorem prov-
ing. Experiments are conducted on two popular instruction
models: Llama-3.1-88 I-8tttutttt (LlamaTeam 2024) and Qwen2-
Tolen
7B-Instruct (Qwen'l Team 2024).
We construct embeddings directly from the KV cache by extract- ing value vectors at every layer, concatenating across attention heads, and averaging over layers to obtain one embedding per token, all without storing hidden states. This reuse of cached in- formation introduces negligible VRAM overhead and incurs min- imal FLOPs consumption compared to vanilla CoE.
Vanilla CoE
KV Cache
Layer bin
Analysis. As shown in Table 2 KV-CoE substantially outper- forms baselines such as MaxProb, PPL, and Entropy on both MATH and TheoremQA. This demonstrates that the Chain-of- Embedding approach retains its strong discriminative power even when using KV cache-derived trajectories instead of hidden states. The token-level evolution captured by the KV cache pro- vides a rich signal for identifying correct reasoning paths, espe- cially in multi-step problems, all while adding negligible over- head since the cache is inherently available.
Token Dn
Figure 3: Vanilla CoE (top) aggregates hidden states across the token dimension to con- struct representation for each layer, whereas KV-CoE (bot- tom) aggregates KV Cache across the layer dimension to construct representation for each token.
FAST/SLOW THINKING WITH KV CACHE
Large Reasoning Models (LRMs) can operate in two modes: fast thinking, which produces short, direct answers, and slow thinking, which generates detailed, explicit step-by-ster reason- ing chains. (Yao et al.] 2023 Lightman al.] 2024). Although slow thinking enhances reliabllity on complex tasks, incurs sub- stantial computational overhead by producing significantly more

---

Published as a conference paper at ICLR 2026
Accuracy (%)
0
20
40
60
80
100
94.0%
Fast Thinking
95.0%
Slow Thinking
(3.75x #Tokens for 1% Accuracy)
829
100
200
300
400
500
600
700
800
900
Average #Tokens
Figure 4: 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.
tokens. For example, on GSM8K (Cbbe al.) 2021). Qwen3-32B (QwenTeam 2025 slow think-
ing yields a marginal improvement in accuracy (0.95 vs. 0.94) while generating nearly four times the tokens, drastically increasing latency and cost as shown in Figure This inefficiency motivates adaptive reasoning, where slow thinking is triggered selectively based problem difficulty.
6.1 METHODOLOGY
We propose method for adaptive reasoning that selects between fast and slow thinking on per- instance basis to minimize unnecessary computation while maintaining accuracy. Our approach leverages the KV cache from the prompt encoding phase to make this decision, eliminating the need for additional forward passes.
Key Idea. Instead of predicting a binary "slow or fast" mode, we estimate a continuous difficulty score d E [0, 100] from the pooled KV cache representation:
= So Pool KV (VV()))
where Pool(-) aggregates keys and values across layers, heads, and token positions via mean pooling, and fa(.) is lightweight MLP classifier. This score determines whether to engage slow thinking.
Switching Mechanism. We control the reasoning mode by injecting special control tokens (sthink> and </think>) during decoding:
Initial Decision: Before generation starts, d is compared predefined threshold
If r, prepend <th nk> trigger slow thinking.
-Otherwise, proceed with fast thinking.
Dynamic Adjustment During Decoding: During generation, the difficulty score re- computed from the updated KV cache at predefined checkpoints:
If d < Trast during slow thinking, append </think> to switch back to fast mode.
If Jslow during fast mode, inject <think> re-engage slow thinking and
tinue decoding with step-by-step reasoning.
This approach enables a fine grained difficult--wware control over reasoning depth. Since the KV cache is already available from prompt encoding, both initial and ongoing difficulty assessments add negligible overhead. This significantly reduces token generation and latency for easy queries while allocating resources challenging problems. The overall workflow is illustrated in Figure
Training Data Construction. To train the difficulty estimator fa(.), we construct supervision signals from public reasoning datasets (training splits GSM8K (Cobbe et al. 2021) and MATH (Hendrycks al.) 2021)). For each training question, generate two candidate solutions using the base model: a fast thinking response (no ofain-of-thought) and a slow thinking response (with chain-of-thought). We then extract the final answers and compare them against the ground truth.
8

---

Published as a conference paper at ICLR 2026
Prefilling
Decoding is </thinko ##
Tokens
A robe in total
does it take? <think>
To find
3
KV Cache
Difficulty
75
51
30
80
20
25
30
10
Figure 5: KVClassifier: special tokens are dynamically inserted to perform thinking-mode switching based on KV-derived difficulty scores.
Based on the outcomes, we assign a difficulty label E {0, 25, 75, 100} reflecting the required reasoning depth:
d = 0 ( (Very easy):
Fast-thinking answer is correct and short (< 128 tokens).
d = 25 (Moderate): Fast-thinking answer is correct but long (2 128 tokens).
d 75 (Hard): Fast-thinking answer is wrong, Slow-thinking answer is correct.
d 100 (Very hard): Both answers incorrect.
This labeling scheme creates a natural difficulty progression, enabling fo(() to learn smooth score that correlates with both correctness and reasoning effort. The token-length criterion distinguishes trivial questions from those needing lengthy outputs without explicit reasoning. The trained esti- mator provides the continuous difficulty score needed for our adaptive switching mechanism. See Appendix D for the exact label distribution and the prompt templates used to instantiate fast vs. slow thinking.
6.2 EXPERIMENTAL RESULTS
We evaluate KV-cache-based fast/slow thinking mechanism reasoning bench- marks: GSMSK evaluation split (Cobbe et al., (02T) and MATH500 (OpenAI/ HuggingFaceH44)
Setup.
Vals Al 2025).
Our experiments compare two switching strategies:
One-step switch IKV-Clasiffcationn: This strategy makes single decision at generation start based on the predicted difficulty score, committing to either slow or fast thinking for the entire decoding process. It functions as classification-style controller.
Two-step switch (KV-Generative): This method performs an initial mode selection and continuously monitors difficulty during decoding. If difficulty drops below Tims during slow thinking, it appends </think> to terminate reasoning early; if difficulty exceeds Tslow during fast thinking. it injects <think> to engage slow thinking mid- mid-gen siddgeneration. This implements a generative-style controller that dynamically adjusts reasoning depth.
We deploy both strategies on two open-weight models: DeepSeek-R1- R1-14B (DeepSeek-A) All 2025) and Qwen3-8B (Qwen'Team) 2025). evaluating their ability to selectively trigger slow thinking and reduce unnecessary token generation.
We construct representations by concatenating key and value tensors across all heads, summing over selected token positions, and averaging across selected layers without normalization, then feed the result into a two-layer MLP (hidden dimension 512, ReLU activation) for difficulty prediction. This design directly reuses the KV cache available during decoding, introduces negligible VRAM or FLOPs overhead, and functions modular component that can be seamlessly integrated into existing inference pipelines without modifications to the base model.
Analysis. As shown Table 3 our KV-cache-based switching approach achieves an effective bal- ance between accuracy and efficiency. For instance, on MATH500 using Qwen3-8B, two-step gen- erative switching reduces average token count from 4,150 (full reasoning) to 727 (5.7x reduction) with only minimal 3.2% accuracy drop. The one-step classification strategy is more conserva- tive, using more tokens but achieving near-full- feal-nrssoning accuracy (0.604 vs. 0.610). Similar trends observed GSM8K, where K KV-cache-based switching maintains high accuracy (up to 0.914)

---

Published as a conference paper at ICLR 2026
Dataset
Method
DeepSeek-RI-Be -I-144
Qwen3-8B
Fast Thinking
0.845
218
0.904 211
Reasoning
0.847
432
0.933 1632
GSM8K
K KV-Classification
0,845
218 49.5%
0,914 / 554 -661%%
KV-Generative
0.835
242
444.0%
0.892/ / 273 -833%
Fast Thinking
0,536 540
0,568 616
Reasoning
0.590
1839
0.610 / 4150
MATH500
V-Classification
0.578
1506 -18.1%
0.604 3963 -4.5%
KV-Generative
0.566 657 -6.3%
0.578 / 727
Table 3: Comparison of accuracy and average token usage for fast thinking, full reasoning, and our KV-cache-basee switching methods on GSM8K and MATH500. For each KV-based method, report the result from the best hyper-parameter configuration identified in Appendix
while significantly reducing token consumption compared to full reasoning. These results demon- strate that difficulty scores derived from the KV cache generalize well across tasks and models, enabling efficient and effective adaptive reasoning with negligible overhead.
CONCLUSION
This work repurposes the KV cache, moving beyond its conventional role in decoding accelcration to unlock its potential as versatile, cost-free representation. We demonstrate that although not de- signed as general-purpect
embeddings, KV caches encode rich contextual information that can be effectively utilized for downstream tasks without incurring additional computational overhead. Our experiments establish two practical applications: (i) Chain-of-Embeddingg where K K--derived em- beddings match or surpass the performance of hidden-state embeddings, and (ii) Fast/Slow Think- ing Switching, which uses KV-cache-Ba based difficulty scores to enable adaptive reasoning-reducing token usage by up to 5.7x with minimal accuracy loss. These findings position the KV cache deployment-fie friendly substrate for advanced inference techniques, opening new avenues for reusing inference-time artifacts improve efficiency and controllability in LLMs.
ACKNOWLEDGMENTS
The research work described in this paper was conducted the JC STEM Lab of Machine Learning and Symbolic Reasoning funded by The Hong Kong Jockey Club Charities Trust.
REFERENCES
Mohammad Beigi, Ying Shen, Runing Yang, Zihao Lin, Qifan Wang, Ankith Mohan, Jianfeng He, Ming Jin, Chang-Ticn Tien Lu, and Lifu Huang. InternalInspector 2²: Robust confidence esti- mation in LLMs through internal states. In Yaser Al-Onaizan, Mohit Bansal, and Yun-Nung Chen (eds.), Findings of the Association for Computational Linguistics: EMNIP 2024, pp- 12847-12865, Miami, Florida, USA, November 2024. Association for Computational Linguis- tics. doi: 10.18653/v 1/2024. http://........ns-..eoee findings-emnlp.... 2024.findings 1tt.inniiinnnm11177755777711 emnlp.7517
URL
Max Belitsky, Dawid J. Kopiczko, Michael Dorkenwald, M. Jehanzeb Mirza, Cees G. M. Snoek, and Yuki M. Asano. Kv cache steering for inducing reasoning in small language models, 2025. URL
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agar- wal, Ariel Merbert-Voss. Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCan- dlish, Alec Radford, Ily Sutskever, and Dario Amodei. Language models few-shot
10

---

Published as a conference paper at ICLR 2026
learners. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tiee Lin (eds.), Advances in Neural Information Processing Systems 33: Anmual Conference Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. URL http://wwww/cceeeeeeoaooooooooop
Chao Chen, Kai Liu, Ze Chen, Yi Gu, Yue Wu, Mingyuam Tao, Zhihang Fu, and Jieping Ye. IN- SIDE: Ilms' internal states retain the power of hallucination detection. In The Twelfth Interna- tional Conference on Learning Representations, ICI. 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024. URL http:///pppevvvveeeeee ww..t///wweee.teee... forum? fw..e///ommetttte
d= j= 2nz 1Qbz
Wenhu Chen, Ming Yin, Max Ku, Pan Lu, Yixin Wan, Xueguang Ma, Jianyu Xu, Xinyi Wang, and Tony Xia. TheoremQA: A theorem- driven question answering dataset. In Houda Bouamor, Juan Pino, and Kalika Bali (eds.), Proceedings of the 2023 Conference Empirical Methods in Natural Language Processing, pp. 7889-7901, Singapore, 2023. Association for Computational Linguisties. doi: 10. wttp5/////222.mmllml/.. 489. URL htttp:://aclatteet www.ooooggooroo.oo. 2023.emnip main.489
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems, 2021.
abs/2110.14168
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. In Sanmi Koyejo, S. Mo- hamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh (eds.), Advances in Neu- ral Information Processing Systems 35: Annual Conference on Neural Information Pro- cessing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28 December wwtp://ww/ss/a/p/2/222...
9, 2022, 2022. URL 67 5a0.1222200000283336664444400000
DeepSeek-AL. Deepseek-rl: Incemtivizing reasoning capability in ilms via reinforcement learning, 2025.
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. In J. Vanschoren and S. Yeung (eds.), Proceedings of the Neural Information Processing Systems Track Datasets and Benchmarks, volume 1, 2021. URL https: //datasets-
feee:p/pssepppppeeeeeeeeeeeeeeee ptspppppperrelleepppparppprp round2.pd.
2021 2201:fflleeessaaaaeeeeee///222b2
Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention.
Proceedings of the 29th Symposium on Operating Sys- tems Principles, SOSP '23, pp. 611-626, New York, NY, USA, 2023, Association for Com- puting Machinery. ISBN 9798400702297. doi: 001145/3600006.36133
URL https:
Xing Li, Zeyu XING, Yiming Li, Linping Qu, Hui-Ling Zhen, Yiwu Yao, Wulong Liu, Sinno Jialin Pan, and Mingxuan Yuan. KYTuner: Sensitivity-aware layer-wise mixed-precision KV cache quantization for efficient and nearly lossless LLM inference. In Forty-second International Conference on Machine Learning, 2025. URL https: http::::///pereviiii www.opppeeevewweweeeeee ew.net. ferest/foruue forum? i ZDwipF 6h06
Hunter Lightman, Vineet Kosaraju, Yuri Burda, Harrison Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. Let's verify step by step. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024. URL httt::://penreview.ccc forett/foommmiioo v8L0pN6EO1]
Chun Liu, Hongguang Zhang, Kainan Zhao, Xinghai Ju. and Lin Yang. LLMEmbed: Rethink- ing lightweight L.M'' genuine function in text classification. In Lun-Wei Ku, Andre Martins,
11

---

Published as a conference paper at ICLR 2026
and Vivek Srikumar (eds.), Proceedings of the 62nd Annual Meeting of the Association for Com-
putational Linguistics (Volume 1: Long Papers), pp. 1994-8004, Bangkok, Thailand, August 2024. Association for Computational Linguisties. doi: 10.18655/w1120244.cllloomgg. https://acll antho 2tegy.org/222..c1-11rrr
URL
AI A LlamaTeam. The llama 3 herd of models, 2024. URL http:://axxi..rr//////
2407.21783
Niklas Muennighoff, Nouamane Tazi, Loic Magne, and Nils Reimers. MTEB: Massive text em- bedding benchmark. In Andreas Vlachos and Isabelle Augenstein (eds.), Proceedings of the 17th Conference of the European Chapter of the Association for Computational Linguistics, pp. 2014- 2037, Dubrovnik, Croatia, 2023. Association for Computational Linguistics. doi: 10. 18.186/3/vI 1/
2023.eacl-main 148. URL https: /ttppe///aaaaactloooooooo//22//2 ma......
Ollama Team. Ollama: Open Ilm deployment made simple.
htttp:://olampn 2024.
OpenAI / HuggingFaceH4 / Vals AI. Math-500: A 500-problem subset of the math benchmark. HuggingFace / Vals AI Benchmark Datasets, 2025. URL
dataset.s Hugg i ngF ace //AH////ATT TATH-HA 500 Derived from "Measuring Mathematical Problem Solving With the MATH Dataset"; subset of 500 test problems.
Alibaba Group QwenTeam. Qwen2 technical report, 2024. URL https://arxi..org/ass//
2407.10671
Alibaba Group QwenTeam. Qwen3 technical report, 2025. URL http::///rxiv.ooo/////
2505.09388
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Ro- man Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 5998-6008, 2017.
httpp//r.......ppppp///////////. Abstract.html
Yi Wang, Jumxiao Liu, Shimao Zhang. Jiajun Chen, and Shujian Huang. Pats: Process- adaptive thinking mode switching, 2025a.
9250
Yiming Wang. Pei Zhang, Baosong Yang, Derek Wong, and Rui Wang. Latent space chain-of- embedding enables output-free Ilm self-evaluation. In Y. Yue, A. Garg. N. Peng, F. Sha, and R. Yu (eds.), International Conference on Representation Learning, volume 2025, pp. 70938- 70970, 2025b. URL Tltp:pacccces555522abbbbeeeeerre
Zeyu Xing, Xing Li, Huiling Zhen, Xianzhi Yu, Mingxuan Yuan, and Sinno Jialin Pan. Large reason- ing models know how to think efficiently. In ES- III: 3rd Workshop on Efficient Systems for Foundation Models, 2025. URL
wtt.eewwwwee//eornneeeeee forum? 1tt.rummddddpLLDeGmmt21
Qingyue Yang, Jie Wang, Xing Li, Zhihai Wang, Chen Chen, Lei Chen, Xianzhi Yu, Wulong Liu, Jianye HAO, Mingxuan Yuan, and Bin Li. Attentionpredictor:: Temppral patterns matter for KV cache compression. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025. URL https: http::::///ppeeeeeee.eeeeeee//ee
Shunyu Yao, Dian Yu, Jeffrey Zhao, Izhak Shafran, Tom Griffiths, Yuan Cao, and Karthik Narasimhan. Tree of thoughts: Deliberate problem solving with large language models. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NewrIPS 2023, New Orleans, LA, USA, December 10- http://ppppee.lleeeaaaaa/2222/20 271 4at.aa//aeeeaaaaaaaaaeaeaaeeeae.
16, 2023, 2023. URL
12

---

Published as a conference paper at ICLR 2026
Murong Yue, Wenlin Yao, Haitao Mi, Dian Yu, Ziyu Yao, and Dong Yu. Dots: Learning to rea- son dynamically in llms via optimal reasoning trajectories search. In Y. Yue, A. Garg, N. Peng, F. Sha, and R. Yu (eds.), International Conference on Representation Learning, volume 2025, pp. 17976-37997, 2025. URL
2025/F1 1227://1117726200055533022000000
Anqi Zhang, Yulin Chen, Jane Pan, Chen Zhao, Aurojit Panda, Jinyang Li, and He He. Reason- ing models know when they're right: Probing hidden states for selffveriiication. In Second Conference on Language Modeling, 2025a. B.klhhtttss//openreeee ew.net...ee/eee..... forum?ids
061 0666336883
Xiaoyun Zhang, Jingging Ruan, Xing Ma, Yawen Zhu, Haodong Zhao, Hao Li, Jiansong Chen, Ke Zeng, and Xunliang Cai. When to continue thinking: Adaptive thinking mode switching for efficient reasoning, 2025b. URL Utt.ttps::/ http:///axxxv..//a//////o// /
13

---

Published as a conference paper at ICLR 2026
A THE USE OF LARGE LANGUAGE MODELS (LLMS)
We acknowledge the use of Large Language Model (LLM) to support the preparation of this manuscript. The LLM was employed exclusively for editorial purposes, such as refining the clarity of exposition, improving grammar and readability, and polishing the overall presentation. At times, was also used suggest alternative phrasings for technical descriptions in order to make arguments more accessible.
Importantly, the LLM did not contribute to the conceptual development, methodology, or experi- mental design of this work. All ideas, including the proposal to treat the KV cache reusable representation, the development of KV-CoE for output-free self-evaluation, and the design of KV- based adaptive Fast/Slow Thinking Switching for token-efficient reasoning, were conceived and implemented solely by the authors. The LLM was not used to generate research results, proofs, or derivations.
All scientific claims, analyses, and conclusions presented in this paper remain the full responsibility of the authors. Any text initially produced with L.M assistance was carefully reviewed, revised, and verified before inclusion.
B COE USAGE, KV-COE DESIGN, AND TOKEN VS. LAYER DIMENSION
CoE usage in reasoning benchmarks. In both the original CoE framework and our work, CoE used as a single-sample confidence estimator. Given one completed reasoning trace, CoE outputs scalar confidence and is evaluated via AUROC/FPR95/AUPR, CoE is not used for pass @k reranking, and it does not modify decoding or generation.
KV cache vs. hidden states across the layer dimension. The KV cache stores attention key/value tensors (per layer/head/token) and does not provide drop-in substitute for full hidden states along the layer dimension. In particular, directly replacing hidden states with KV cache within the original layer-wise CoE design yields poor performance. This reflects structural mismatch: the KV cache is inherently token-centric and head-structured, whereas hidden-state CoE explicitly tracks layer-wise representations.
Ablation: layer-wise aggregation with KV-CoE. To make this limitation explicit, evaluate a layer-wise variant of KV-CoE on meta-llamme /tt//Laaaaa "llama-3: 3. 15..-8B-Inna 1B-Instruct.. Results are shown in Tables 4 and 5 Layer-wise KV-CoE performs substantially worse than token-centric KV-CoE, confirming that KV cache should not be treated as a layer-wise embedding.
Method
AUROC 1
FPR95
AUPR 1
KV-CoE-R (layer)
62.69
82.28
45.43
KV-CoE-C (layer)
64.09
81.16
47.78
Table 4:
ablation meta-llame 1.1ama-3, 1aammaa3.1--0 InsB-Insttuut...
Layer-wise
aggregation
of KV-CoE on
MATH
using
Method
AUROC 1 FPR95
AUPR 1
KV-CoE-R (layer)
47.78
94.24
17.70
KV-CoE-C (layer)
48.61
92.68
18.02
Table Layer-wise
aggregation ablation of KV-CoB on TheoremQA -alamammmaaaamam AilamaLLlamme LLl.maaaaaaaaaa 3. 1B8B------ Instruct.
using
meta-l
Design implication: token-centric KV-CoE. Motivated by the above, redesign KV-CoE to operate along the token dimension, using token-wise aggregation that matches the semantics of the KV cache. Concretely, KV-CoE forms per-token embeddings from value vectors and aggregates
14

---

Published as a conference paper at ICLR 2026
them token positions (and optionally across heads/layers) to produce the trace-level confidence
score.
Pooling strategies summary. For clarity, Table 6 summarizes the pooling/aggreeaation strategies used in this paper (see also L320-322 and L455 456).
Task
Source
Heud Agg.
Position Age.
Layer Agg.
KV-CoE
Value vectors from cache
Concatenate Concatenate
Per-token embedding
Average Average (no normalizution)
Fast/Slow Thinking
Key Value
Sum selected tokens
Table 6: Summary of pooling strategies used in this work.
C HIDDEN-STATE AND CHANCE BASELINES ON MTEB CLASSIFICATION SUBSETS
To contextualize the MTEB classification results in Table we report two additional baselines evaluated under the same downstream protocol (identical pooling, projection, and classifier setup as used for VV-ccche embeddings).
Hidden-state embeddings. We extract hidden-state embeddings from / /Llaaaaa ama-3. 3.1-8B-Insttruct and apply the same aggregation pipeline used for KV-cache embeddings. Table shows that hidden-state embeddings achieve accuracies close to those of KV-cache embeddings under identical pooling and projection, suggesting that the observed gap to dedicated embedding models is dominated by the downstream protocol rather than the choice between hidden states vs. KV caches.
Dataset
Accuracy (Hidden State)
AmazonCoounterfacuual DBpedia
0.3530
0.5937
FinancialPhrascbant.
0,6254
TweetTopic
0.3714
Table 7: Accuracy of hidden-state embeddings on selected MTEB classification tasks under the same pooling/projection pipeline as KV-cache embeddings.
Random-embedding baseline. We additionally include mandom-embedding baseline sanity- check that the evaluation pipeline is not degenerate. Concretely, we replace the model-derived em- bedding with an i.i.d. random vector Grand R2d (matched to the embedding dimension d used in our pipeline). while keeping the same pooling/projection interface (where applicable) and the same downstream classifier and evaluation protocol. Table reports the resulting accuracy. The consistent margin between random embeddings and hidden/KV embeddings particularly on DBpedia and Fi- nancialPhrasebank inniiattss that the pipeline functional and that dodee-derived embeddings capture non-trivial task structure beyond what is achievable with unstructured random features.
Dataset
Accuracy (Random)
A AmazzonCComttraattuulCCssffeeiee DlseeiaClasiiiiiaiin
0.5224
0.0716
SataeciaaPlaasebbakClaaiiiiii Twweetopicelaassiicctioon
0.3274
0.1591
Table 8: MTEB classification performance when representations replaced by random baseline embedding E 256 whose entries are drawn i.i.d. from a standard normal distribution.
15

---

Published as a conference paper at ICLR 2026
Scope. These MTEB subsets are reported only to illustrate that KV-derived embeddings cost- efficiently comparable under our protocol; we do not claim they are superior to dedicated embedding models, that they reliable for global semantic retrieval across diverse domains.
D DIFFICULTY LABELS AND PROMPT CONTROL FOR FAST/SLOW THINKING SWITCHING
Difficulty labels for training the switching model. We assign each instance a discrete diffi- culty label based the correctness of fast vs. slow thinking outputs, using the following four-Ievel scheme:
Difficulty
Description
Count
0
Both fast and slow thinking are correct (very easy)
286
25
Both are correct but require longer generation
3,467
75
Only slow thinking is correct
860
100
Neither is correct
2,887
Total
7,500
Table 9: Difficulty labels used to train the fast/slow switching estimator.
Fast vs. slow thinking prompts. We use a Large Reasoning Model (LRM) that supports ex- plicit reasoning via <think> blocks. Fast/slow thinking is controlled purely by prompt, not by KV cache. Specifically, after the user prompt: (i) Fast thinking inserts an empty thinking block <think>\n</thine which encourages step-by-step reasoning. Model-specific special tokens (e.g., < im_start/ >, < tim_enne |>) are adapted per backbone; the KV-based controller only chooses between the fast and slow prompt variants.
(ii) Slow thinking inserts open thinking block <think>> \n,
E HYPER-PARAMENT SELECTION FOR KVCLASSIFIER
To better understand the effect of hyperparameters on KV-based classification, we conduct a system- atic study by varying the number of layers pooled and the number tokens selected from the end of the sequence. Importantly, we fix the total amount of KV data to be approximately constant across configurations (256 token x layer units). This ensures a fair comparison: for example, selecting layers with 32 tokens, 4 layers with 64 tokens, or 2 layers with 128 tokens, all yield the same KV budget.
Model
Dataset
Method
81,, Len=32
4I., Len=64
21., Len=128
KV-Classification
0.845 /18
0.845 218
0.845 218
GSM8K
DeepSeek-
Ke-Gentre
0,835 / 242
0.825 232
0.805 / 217
KV Classiifficaton
0.536/540
0.550 905
0,578 / 1506
MATH550
KV-Gencrative 0.538 524
0.550 / 544
0.566 657
KV CVCClassiftcation 0.904/211
0.904 0.904/217 217
0.914/
GSM8K
KV-Gencrative 0.892/272
0.886 276
0.881 257
Qwen3-8B
KV- Cl-Cssssification 0.570/736
0.598 3673
0.604 3963
MATH500
KV-Generative
1.578/727
0.524 /933
0.550/ 837
Table 10: Hyper-parameter selection results for KV-Classification and KV-Generative. Values are reported as Accuracy / #Tokens. Best accuracy for each dataset pair is in bold.
the results GSM8K and MATH500 for both DeepScek-Rl-I44 and Qwen3- 8B, under KV-Classification and KV-Generative settings. We observe that while performance varies slightly with the allocation of layer token depth, the overall trends consistent: (i) accuracy remains competitive across different allocations, and (ii) increasing token coverage (e.g., 2L 128) tends to favor more complex datasets such as MATH500, whereas shallow but wider layer coverage (e.g., 8L. x 32) can suffice for GSM8K.
16


Reviews:
[{"review_id": "NA3gkwTPnR", "reviewer": "Reviewer_tEgG", "summary": "I am not very well versed with this literature. So I will focus mostly on clarification questions at this time. I will defer my decision to later time after the rebuttal and other reviewer comments.\n\nIn this paper, authors propose a way to construct embeddings from KV cache instead of using hidden spaces to be used for CoE and adaptive reasoning problems. While constructing the embedding, they aggregate across layers (instead of tokens as done in earlier CoE works) and measure path characteristics along token dimension (as opposed to layers done in earlier CoE). The argument in favor of choosing KV caches for embeddings is that they do not incur any additional cost.", "questions": "1. How does hidden state based embedding perform as compared to off the shelf embedding? (table 1)\n2. How is COE used in practice. Is it a pass@k while choosing the highest ranking pass?\n3. Do you have a ablation of what happens when you use CoE with switching layer and token dim (the way you propose with KV)?\n4. There were a bunch of baselines mentioned in the related work Chen 2024, Beigi 2024, Zhang 2025 for CoE problem and ASRR, PATS, DOTS for adaptive reasoning problem. Are these baselines relevant to the experiments you ran. It would be illuminating to compare your method against these methods if they are SOTA. \n5. why were CoE results taken from the paper when they belonged to a previous version of model which has very different context windows among other things. ( in table 2).", "limitations": "", "rating": 4, "confidence": 2, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "1. Using KV Cache for self-evaluation avoids extra memory overhead. \n2. The accuracy improvements are impressive especially for MATH dataset.", "weaknesses": "Writing is a bit messed up. for instance, \n\nSection 4. makes arguments are not quite clear.\n\ta. why is min \\gamma (x) > 0 enough for this task? how is this class of problems different from a classification problem. One can argue that the same thing holds for classification as well. \n        b. contextual conditionaing argument is not clear to me. Can the author's elaborate.\n        c. Efficiency constraint argument is okay but unnecessarily formalized.\n        d. I dont understand the discussion on local decision adequacy. How do you determine whether a problem is local / global. for all you know you need to know the solution of the problem to know how hard it is -- knowing the solution is the most global problem in context of LLM generation. \nI wonder if the usage of LLM is quite strong in this section and it needs rework.\n\n\nI am not sure if state of the art baselines are compared. (see questions)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "E426PtMyNX", "reviewer": "Reviewer_PbPH", "summary": "This paper explores the application of KV-Cache for two different applications than speeding up decoding: (1) Chain-of-Embedding a method for self-evaluation of LLMs and (2) Fast/Slow Thinking Switching a method that enables adaptive reasoning based on KV-Cache embeddings of the prompt exclusively or the prompt in combination with the reasoning trace. \nIn their experiments, this paper demonstrates that (1) outperforms other baselines on self-evaluation tasks and (2) that we can enable adaptive reasoning based on difficulty scores derived from the KV-cache embeddings.", "questions": "- L.176 Formatting issue with Table 1 (whitespace missing)\n- Description of the metric used in Table 1 (assume higher is better)\n- L.164 The KV Cache embeddings are constructed by averaging over attention heads. Have you also tried concatenating over attention heads instead of averaging? This would counteract limitation 3 in L. 180.\n- In addition in Sec. 5.2 the KV embedding is constructed by flattening the head and key/value dimension. This seems to be inconsistent.\n- Is it not possible to recompute the embeddings efficiently from the KV cache? How much overhead would it be? Is this really a bottleneck? \n- Why does the KV-CoE operate along the token position and not along the layer index? Have you tried using the layer index as in default CoE?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Clear motivation, well written and interesting research question.\n- Strong performance of KV-CoE on self-evaluation tasks.\n- The thinking switch experiments demonstrate the practicability of thinking mode switching based on KV-cache embeddings.", "weaknesses": "- Even though the paper is well written, it seems a bit like the combination of two methods into one paper.\n- After reading the introduction the reader might think that KV-caches can be used as embeddings, but the result in section 4.1 contradicts this initial claim or limits it to “certain classification tasks”, which are not specified further. I suggest clarifying this confusion in the paper and/or adding additional details on the “certain classification tasks”.\n- KV-CoE seems to outperform standard CoE for the Qwen Models (Table 2), but not for the LLama models. It would help the reader to add an explanation (or run the baselines with the correct base model) or mention in the main text that KV-CoE outperforms default CoE.\n- Do the reasoning results in Table 3 also transfer to larger models e.g. Qwen3 32B?\n\n\nEven though the paper could add a few more details or explanations at some places (see questions & weaknesses) the authors convincingly demonstrate the effectiveness of their proposed methods based on embeddings stemming from the KV-Cache. Therefore, I am inclined to recommend acceptance of this paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "zRbHEm2NzD", "reviewer": "Reviewer_2n9n", "summary": "This paper presents a novel approach to repurposing the KV cache, a critical component of efficient autoregressive decoding in large language models (LLMs), for downstream tasks beyond just acceleration. The authors demonstrate that the KV cache, despite not being explicitly trained as a general-purpose embedding, nonetheless encodes rich contextual information that can be effectively leveraged for two key applications: (1) Chain-of-Embedding (CoE) and (2) Fast/Slow Thinking Switching.", "questions": "1. According to the official format, the table caption should appear before the table.\n2. Line 239 contains a typo: \"is obtained by by averaging\"\n3. Could the authors please elaborate on the pooling strategy used in the experiments?\n4. Is there any accuracy trade-off associated with the improvement in fast and slow thinking when using the KV cache?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 3, "strengths": "1. The paper repurposes the KV cache and provides new insights. It also demonstrates the versatility of the KV cache through two practical applications: CoE and Fast/Slow Thinking Switching.", "weaknesses": "1. I would appreciate it if the authors could provide more details on MaxProb, PPL, and Entropy, specifically, what these metrics represent and why they are chosen as baselines in Table 2.\n2. From my understanding, CoE explores the model’s representation space through hidden states, whereas the KV cache contains only information derived from the attention heads. Considering that hidden states integrate residual connections, MLP outputs, and attention information, I am curious how KV-CoE effectively tracks representations across layers under this limitation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qV7AykZQMm", "reviewer": "Reviewer_rW12", "summary": "This paper repurposes the KV caches in LLMs for purposes beyond attention. The authors explore the use of KV embeddings as text embeddings and as classifier features for fast/slow thinking. Their results indicate that KV cache embeddings contain significant information that can be used in other applications. However, I think the empirical validation is a bit limited.", "questions": "1. I am unfamiliar with the datasets/tasks in Table 1. What is the accuracy for random chance?\n2. What is the distribution of the training data difficulty labels detailed in page 8?\n3. Can these methods be applied to non-reasoning generative tasks like summarization?\n4. Table 3: Do the token averages also contain sequences that hit the generation limit (i.e. incomplete)?", "limitations": "", "rating": 4, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. KV caches already take up significant space, so having them perform multiple functions seems well-motivated\n2. Demonstrated that KV caches can show promise for other purposes\n3. No architectural changes needed", "weaknesses": "1. Fast/slow thinking experiments are lacking. The method was only validated on GSM8k and MATH500, both of which are fairly simple for the reasoning models tested. Thus, it may just be that these models are fairly robust for these relatively easy tasks, so I would recommend testing on more challenging reasoning tasks (AIME, BRUMO, etc)\n2. Numbers in Table 3 are reported using the best hyperparameters per task/model combination. This feels too cherry-picked since I didn't find a way to know which hyperparameters to choose before testing.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MMdkygzpF0", "reviewer": "Reviewer_Zzrt", "summary": "This work investigates the use of the KV cache as intermediate representation instead of the commonly used hidden states for two reasoning-focused downstream tasks: The first is an adapted Chain-of-Embedding method based on the KV cache that is roughly on par with existing CoE, while reusing the already stored KV cache. The second is a novel reasoning mode meta-model that uses the aggregated KV caches to optionally insert \\<think\\> / \\</think\\> tokens to expand / decrease \"reasoning\" via more token output. On the simpler GSM8K dataset models with a \"fast\" thinking path already perform on par with more expensive (longer) traces, but on the MATH500 the proposed method shows performance improvements at smaller token budgets in between a \"fast\" and \"slow\" thinking mode.", "questions": "- Can you explain what the CoE is used for in these reasoning benchmarks, i.e. how the extracted embeddings are used towards what kind of goal? This is not clear from the paper.\n- How is the fast thinking and slow thinking induced / separated without the KV cache-based control? Can you provide exact prompts?\n- Could you add Prefix-Tuning [1] and Cartridges [2] as related work operating on the KV cache (though more in the direction of fine-tuning adapters)\n\n[1] Li et al. (2021): https://arxiv.org/abs/2101.00190\n\n[2] Eyoboglu et al. (2025): https://arxiv.org/pdf/2506.06266", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- shows that KV caches contain similar information content compared to hidden states for CoE downstream method\n- introduces a KV-cache based thinking mode switch (either ahead of or continuously during generation)\n- shows there is optimization potential towards existing dedicated embedding models", "weaknesses": "- using hidden states can be similarly optimized for downstream tasks (e.g. in-place / recursive summation across layers / tokens - if common simple aggregations are used)\n- no comparison to a hidden state-based thinking mode switching \n- fails to cite cache tuning methods (Prefix-Tuning [1] / Cartridges [2]) as related work of where the KV cache is actually modified\n\n[1] Li et al. (2021): https://arxiv.org/abs/2101.00190\n\n[2] Eyoboglu et al. (2025): https://arxiv.org/pdf/2506.06266", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)