Source-linked AI summary
Compression-Aware Abstention: Teaching LLMs to Refuse When KV-Compression Masks Remove Answer Evidence
Mohammadali Khodabandehlou, Bhaskar Krishnamachari
TL;DR
KV-cache compression can remove answer-bearing evidence and induce hallucination rather than abstention. The paper trains compression-aware behavior from survival masks and tight evidence spans, achieving a 97% reduction in base-model hallucinations under prompt-style truncation while preserving answerability, though high-retention compressed-cache answering remains limited.
Problem
KV-cache compression saves inference memory by evicting context tokens, but eviction can remove answer-bearing evidence and leave models hallucinating instead of recognizing insufficient context.
Method
The method labels examples CONFIDENT or ABSTAIN from compressor survival masks over tight answer-bearing spans, then trains a LoRA adapter to answer or refuse accordingly.
Results
97% of base-model hallucinations are eliminated under prompt-style truncation while correct answering on CONFIDENT-gold examples is preserved; multi-compressor training also improves compressed-cache transfer.
Takeaways & Limitations
Compression should be evaluated by whether answer-supporting evidence survives, with model behavior changing between answering and abstaining accordingly.
Takeaways & Limitations
At r = 0.80, compressed-cache mix-cc honest accuracy is 0.315 versus 0.658 for prompt-style mix-ps, leaving high-retention compressed-cache answering as a generalization limit.
Abstract
from arXiv · showhide
KV-cache compression reduces LLM inference memory by evicting context tokens, but when the evicted tokens contain answer-bearing evidence, the model may hallucinate instead of recognizing that the compressed context is insufficient. We address this failure from a behavioral perspective: to our knowledge, this is the first work to formulate compression-aware abstention as a learning problem, in which a model learns to answer when supporting evidence survives compression and abstain when it does not. We construct supervision from compressor survival masks and tight answer-bearing spans, labeling examples as Confident when evidence survives and Abstain when it is removed. A 10.1M-parameter LoRA adapter trained on ~2.6K MuSiQue 2-hop QA examples reduces base-model hallucinations by 97% under prompt-style truncation while preserving correct answering on evidence-retaining examples. Unlike prompt-only abstention baselines, which over-abstain on many answerable high-retention examples, the trained adapter learns a conditional policy. We also evaluate the method under actual compressed-cache decoding, where multi-compressor training yields a 6-22x relative lift over the unaided base on evidence-retaining examples. Controlled-deletion experiments show that the learned behavior is driven by evidence content rather than input length alone.
1 Introduction
KV-cache compression can evict answer-bearing evidence, leaving fluent contexts that prompt confident hallucinations. The paper proposes compression-aware abstention as a behavioral learning problem, training models to answer when evidence survives and refuse when it does not.
- KV-cache compression reduces inference memory by evicting context tokens, but may remove evidence needed for downstream answers.
- Compression-induced evidence loss can make models fabricate confident wrong answers instead of acknowledging insufficient context.
- Compression-aware abstention trains models to answer when supporting evidence survives and refuse when the surviving context cannot support an answer.
- The method labels examples CONFIDENT when answer-bearing evidence survives compression and ABSTAIN when it does not.
- A 10.1M-parameter LoRA adapter trained on approximately 2.6K compression-labeled MuSiQue examples eliminates 97% of base-model hallucinations under prompt-style truncation while preserving correct answering on CONFIDENT-gold examples.
- Prompt-only abstention over-abstains on many answerable high-retention examples, whereas the trained adapter learns a more conditional policy.
2 Compression-aware abstention as a learning problem
The paper formalizes compression-aware abstention using survival masks over answer-supporting spans, then trains a LoRA adapter to answer or refuse according to the resulting labels. Tight spans provide discriminative supervision, while a bidirectional loss supports calibrated behavior.
- Compression-truncated input: A compressor produces a survival mask and a truncated context containing only surviving tokens, which the model processes as a standard prompt.
- Evidence survival: Supporting spans are contiguous token ranges jointly required for deriving the answer from context.
- Evidence survival: When supporting spans are large, survival averages out; answer-sized spans make evidence survival informative for labeling.The paper reports a median verbatim answer-span length of 3 tokens.
- Gold label: Examples are labeled CONFIDENT when every supporting span exceeds τhigh, ABSTAIN when the minimum survival falls below τlow, and DROP otherwise.Examples in the uncertain interval [τlow, τhigh) are excluded from training.
- Target and objective: The adapter maps CONFIDENT examples to the answer and ABSTAIN examples to a fixed refusal string.
- Target and objective: The loss combines cross-entropy for answering and abstaining, penalizing refusal when evidence supports the answer and answering when it does not.
- Class balance: Class balancing is applied because low-retention examples are mostly ABSTAIN while high-retention examples are mostly CONFIDENT.The mid-retention regime is where both classes co-occur most informatively.
3 Method
The method localizes short answer-bearing spans, derives labels from compressor survival, and trains a small LoRA adapter on the resulting examples. This design avoids paragraph-level averaging that would erase useful per-example variation.
- Tight-span localization: MuSiQue sub-answers are located as verbatim supporting spans, with median length 3 tokens and 100% coverage in the training pool.The 99th-percentile span length is 11 tokens.
- Pipeline: Figure 1 shows the pipeline from KVzip survival masks through tight-span labeling to LoRA training on truncated inputs, queries, and targets.
- Tight-span localization: Using full supporting paragraphs of approximately 100 tokens makes survival fractions average out and drops every example at r = 0.50.Label entropy is 0.000 for paragraph spans versus 0.95 for tight spans.
- Pipeline: The labeler uses thresholds τlow = 0.3 and τhigh = 0.8, excluding DROP examples from the training set.
- Adapter: The adapter contains 10.1M LoRA parameters attached to four attention projections, representing 0.13% of Qwen2.5-7B-Instruct’s parameters.
- Adapter: Training uses approximately 2.6K instances with a 30/70 minority-class floor, while validation contains 463 disjoint examples.
4 Experimental setup
Experiments use Qwen2.5-7B-Instruct with KVzip training masks and held-out compressors, evaluating retention ratios and honest accuracy under standardized decoding.
- Qwen2.5-7B-Instruct supplies all main results, while Llama-3.1-8B-Instruct serves as a held-out base model.
- KVzip provides training masks; Expected Attention and SnapKV are held-out compressors accessed through kvpress.SnapKV is evaluated in context-only and native query-aware modes.
- All compressors are calibrated to retain the same number of context tokens at each retention ratio r.
- Retention ratios span r ∈{0.05, 0.10, 0.20, 0.30, 0.50, 0.80}, with nearly all examples ABSTAIN at low retention and CONFIDENT at high retention.The r = 0.50 slice is balanced between gold classes.
- Greedy decoding allows at most 64 new tokens, and honest accuracy counts correct answers for CONFIDENT-gold cases or refusals for ABSTAIN-gold cases.
- Training and evaluation run on one NVIDIA H100 80GB GPU, with code, outputs, adapters, datasets, and recomputation scripts released.
5 Results
Training improves honest accuracy across retention ratios by sharply reducing hallucinations while preserving answerability when evidence survives. The learned policy outperforms prompt-only abstention at high retention and transfers best between compressors with similar token-selection patterns.
- 5.1 Headline: 232 to 7: confident wrong answers on ABSTAIN-gold examples collapse after training.Training improves honest accuracy at every retention ratio, while conditional answer quality is preserved.
- 5.1 Headline: The adapter preserves correct answering on CONFIDENT-gold examples, with answerability preservation ratios of 1.00–1.36.Its correct-answer rate is at least the base model’s in all eight prompt-style cells across four compressors.
- 5.2 Baselines: 25–33pp: mixture training improves over the Fragment baseline at r = 0.50, showing that surviving text alone does not produce the behavior.The comparison is significant for KVzip, Expected Attention, and SnapKV-context, but not SnapKV-native.
- 5.2 Baselines: At r = 0.80, P2 honest accuracy falls to 0.34 on KVzip because it abstains on 49% of CONFIDENT-gold examples.Across compressors, P2 over-abstains on 49–51% of answerable examples, versus 6–10% for the mixture adapter.
- 5.2 Baselines: At r = 0.80, the mixture adapter beats P2 on every compressor by ∆∈[+0.28, +0.34].The pooled advantage across 260 CONFIDENT-gold examples is ∆= +0.32 with a 95% paired-bootstrap CI of [+0.25, +0.38].
- 5.3 Cross-compressor transfer and mixture training: Transfer is stronger within compressor families than across them: KVzip→Expected Attention gains +0.49 versus +0.53, while KVzip→SnapKV-native gains +0.20.Transfer follows kept-token overlap: 75% agreement for KVzip and Expected Attention versus 62% for KVzip and SnapKV at r = 0.50.
- 5.3 Cross-compressor transfer and mixture training: 82 of 89 cross-compressor failures are rescued when mixture training replaces v1.These are cases where v1 abstained correctly under KVzip but hallucinated under SnapKV.
6 Related work
The paper positions compression-aware abstention as complementary to prior work on KV-cache compression and selective generation: it addresses evidence loss through calibrated behavior rather than representation-side repair.
- KV-cache compression: KV-cache compression methods evict tokens using importance scores, creating a risk of compression-induced evidence loss and hallucination.Prior scoring strategies include attention, recency, attention-sink, reconstruction, future-query, and retrieval-based methods.
- Positioning: The paper treats calibrated abstention as orthogonal to representation-side fixes for compression-induced hallucination.Compression methods supply survival masks, while the proposed approach trains behavior around whether evidence remains available.
- Abstention and selective generation: Existing abstention research studies calibration, self-knowledge, confidence, refusal tuning, retrieval, collaboration, and known-unknown boundaries.The paper situates its behavioral objective within selective-generation research.
7 Conclusion
The paper argues that KV-cache compression should be evaluated by evidence availability as well as output preservation. Compression-aware abstention makes the model answer when evidence survives and refuse when it does not.
- Conclusion: The paper frames compression-induced evidence loss as a behavioral learning problem rather than only an output-preservation problem.The model is trained from compressor masks and tight answer-bearing spans.
- Conclusion: Compression-aware abstention targets conditional behavior: preserve answers when supporting evidence survives and refuse when compression makes the context insufficient.This conclusion follows the paper’s stated answer-versus-abstain policy.
Limitations
The limitations concern evaluation validity, compressed-cache generalization, label fidelity, difficult high-retention cases, and experimental scope. These constraints qualify how broadly the learned behavior should be interpreted.
- Evaluation validity: Aggregate honest accuracy can reward ratio-only policies because gold labels are heavily skewed toward abstention at low retention and confidence at high retention.Per-ratio and class-conditional comparisons are therefore more meaningful operational tests.
- Evaluation validity: The mid-retention prompt comparison is underpowered: at r = 0.50, the mixture-versus-prompt advantage is +0.07 but not statistically significant (p = 0.17 pooled).The strongest prompt-versus-training evidence comes from the high-retention CONFIDENT-gold slice.
- Compressed-cache deployment: At r = 0.80, mix-cc honest accuracy is 0.315 versus 0.658 for mix-ps under prompt-style inference, leaving a compressed-cache deployment gap.The trained adapter improves over base-cc but high-retention compressed-cache answering remains a generalization limit.
- Label fidelity: Tight-span labels achieve Cohen’s κ = 0.61 and 80% raw agreement with an LLM-judge sufficiency verdict, but are not exact ground truth for contextual sufficiency.Disagreement often occurs when an answer survives but multi-hop bridging evidence does not.
- Hard high-retention cases: At r = 0.80, the adapter hallucinates the same 4 of 7 rare ABSTAIN-gold cases across three v1 seeds.These cases indicate incomplete introspection about whether question-specific evidence survived despite high overall retention.
- Scope and coverage: Training is limited to MuSiQue 2-hop QA, relies on question-decomposition annotations, and evaluates compressed-cache inference with KVzipPress.The paper reports replication on RULER multi-key NIAH and a held-out Llama-3.1-8B-Instruct model, but not retraining on HotpotQA or other compressors.
E.2 Seed robustness
Across seeds and evaluation splits, the adapter’s aggregate gains remain stable, while residual failures concentrate in rare high-retention ABSTAIN-gold cases rather than answerable examples.
- 0.017 is the maximum spread in aggregate ∆honest accuracy across mixture seeds and compressors.
- 1/85 test examples hallucinated versus 59/85 for the base, supporting that the v1 result was not overfit to the selection split.
- 96% of base hallucinations are eliminated in the Llama replication, falling from 228 to 10 while adding 298 correct answers.
- At r = 0.80, the adapter hallucinates on all 4 ABSTAIN-gold examples, so that slice is too small for isolated conclusions.
- On 69 CONFIDENT-gold examples at r = 0.80, conditional wrong-when-answering improves from 37% to 34% despite 23 adapter wrong answers versus 21 for base.
- Across three seeds, exactly 7 ABSTAIN-gold examples per seed remain hallucination failures, with 4 stable across all seeds.
M Compressed-cache training details
Compressed-cache training applies eviction during the forward pass and evaluates whether the learned abstention policy preserves answerability under deployment-like decoding.
- Training under cc decoding evicts low-scoring K/V pairs after full-context prefill, then computes loss from targets decoded from the post-eviction cache.
- The mixture training set contains approximately 8772 instances from KVzip, ExpAttn, and SnapKV-context, with three training epochs and matched optimization settings.
- 0.824 is the mean aggregate honest accuracy under KVzipPress inference across two mixture-training seeds, with a 0.006 spread.
- Single-compressor cc training reduces over-abstention at r = 0.80 to 0.203 but raises CONFIDENT-gold wrong answers to 42/69.
- Under compressed-cache inference, adapter preservation rates reach 25–32% versus 4% for base, a 6–22× relative lift on CONFIDENT-gold examples.
- The multi-compressor residual wrong answers are described as pattern collisions, contrasting with short-output overfitting observed for another training setup.
O Mechanism: controlled-deletion experiments
Controlled deletions distinguish evidence-sensitive behavior from a simple input-length heuristic by holding retained length fixed while selectively removing answer-bearing tokens.
- 10A: Fixed retained length: 10A holds inputs at L = 500 kept tokens, but the resulting accuracy drop remains compatible with both length-prior and evidence-loss explanations.
- 10B: Adversarial tight-span eviction: 10B removes tight-span tokens and refills masks with random positions, preserving length and retention count while forcing answer-bearing evidence out.
- 10B: Adversarial tight-span eviction: On 97 CONFIDENT-gold examples, mix-ps abstention rises from 16.5% to 56.7% after adversarial eviction, while mix-cc rises from 64.9% to 88.7%.
- Conclusion: The controlled-deletion results indicate that the adapter responds to evidence content, not just retained input length.
- 10C: Opposite-priority pairs: At identical length for 22 examples, mix-ps answers 14/22 with tight spans retained but 0/22 when those spans are likely evicted.
P Random-mask training-signal probe
Random-mask retraining tests whether compression-aware abstention depends on compressor importance rankings, while additional evaluations examine transfer and implementation scale.
- The random-mask adapter matches the mixture adapter on KVzip val, with aggregate honest accuracy 0.873 versus 0.86 and paired McNemar ∆ = +0.017, p = 0.23.
- Cross-family transfer improves by +0.276 on SnapKV-context for random-mask training, with p ≈ 10^-24.
- At 16K context, abstention on ABSTAIN-gold follows the same broadly declining shape as MuSiQue validation, though it is not strictly monotone at the two lowest ratios.
- Aggregate hallucination rate drops by 36–43pp at every evaluated context length, trading hallucinations on evicted-needle examples for refusals.
- The LoRA adapter modifies only four attention projections and contains 10.1M trainable parameters, or 0.13% of the 7.6B-parameter base model.