PaperlayerView exact revision on arXiv

Source-linked AI summary

Compression-Aware Abstention: Teaching LLMs to Refuse When KV-Compression Masks Remove Answer Evidence

Mohammadali Khodabandehlou, Bhaskar Krishnamachari

arXiv:2608.29934v1cs.CLcs.LG

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.

    Source: §1 · p. 1

  • 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.

    Source: §1 · p. 1, §2 · p. 2

  • 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.

    Source: §1 · p. 2, §5.5 · p. 8

  • Takeaways & Limitations

    Compression should be evaluated by whether answer-supporting evidence survives, with model behavior changing between answering and abstaining accordingly.

    Source: §7 · pp. 8–9

    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.

    Source: Limitations · p. 9

Abstract

from arXiv · show

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.Source: §1 · p. 1
  • Compression-induced evidence loss can make models fabricate confident wrong answers instead of acknowledging insufficient context.Source: §1 · p. 1
  • Compression-aware abstention trains models to answer when supporting evidence survives and refuse when the surviving context cannot support an answer.Source: §1 · p. 1
  • The method labels examples CONFIDENT when answer-bearing evidence survives compression and ABSTAIN when it does not.Source: §1 · p. 1
  • 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.Source: §1 · p. 2
  • Prompt-only abstention over-abstains on many answerable high-retention examples, whereas the trained adapter learns a more conditional policy.Source: §1 · p. 2

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.Source: §2 · p. 2

Evidence survival

  • Supporting spans are contiguous token ranges jointly required for deriving the answer from context.Source: §2 · p. 2
  • 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.Source: §2 · p. 2

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.Source: §2 · p. 2

Target and objective

  • The adapter maps CONFIDENT examples to the answer and ABSTAIN examples to a fixed refusal string.Source: §2 · p. 3
  • The loss combines cross-entropy for answering and abstaining, penalizing refusal when evidence supports the answer and answering when it does not.Source: §2 · p. 3

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.Source: §2 · p. 3

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.Source: §3.1 · p. 3

Pipeline

  • Figure 1 shows the pipeline from KVzip survival masks through tight-span labeling to LoRA training on truncated inputs, queries, and targets.Source: Figure 1 · p. 3, §3.2 · p. 3

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.Source: §3.1 · p. 3

Pipeline

  • The labeler uses thresholds τlow = 0.3 and τhigh = 0.8, excluding DROP examples from the training set.Source: §3.2 · p. 3

Adapter

  • The adapter contains 10.1M LoRA parameters attached to four attention projections, representing 0.13% of Qwen2.5-7B-Instruct’s parameters.Source: §3.2 · p. 3
  • Training uses approximately 2.6K instances with a 30/70 minority-class floor, while validation contains 463 disjoint examples.Source: §3.2 · p. 3

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.Source: §4 · p. 3
  • 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.Source: §4 · pp. 3–4
  • All compressors are calibrated to retain the same number of context tokens at each retention ratio r.Source: §4 · p. 4
  • 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.Source: §4 · p. 4, Table 1 · p. 4
  • Greedy decoding allows at most 64 new tokens, and honest accuracy counts correct answers for CONFIDENT-gold cases or refusals for ABSTAIN-gold cases.Source: §4 · p. 4
  • Training and evaluation run on one NVIDIA H100 80GB GPU, with code, outputs, adapters, datasets, and recomputation scripts released.Source: §4 · p. 4

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.Source: §5.1 · p. 4
  • 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.Source: §5.1 · p. 4

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.Source: §5.2 · pp. 5–6
  • 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.Source: §5.2 · p. 5
  • 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].Source: §5.2 · p. 5

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.Source: §5.3 · p. 6
  • 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.Source: §5.3 · p. 6

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.Source: §6 · p. 8

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.Source: §6 · p. 8

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.Source: §6 · p. 8

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.Source: §7 · pp. 8–9
  • 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.Source: §7 · pp. 8–9

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.Source: Limitations · p. 9
  • 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.Source: Limitations · p. 9

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.Source: Limitations · p. 9

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.Source: Limitations · p. 9

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.Source: Limitations · p. 9

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.Source: Limitations · p. 9

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.Source: Table 9 · p. 14
  • 1/85 test examples hallucinated versus 59/85 for the base, supporting that the v1 result was not overfit to the selection split.Source: F Held-out test split · p. 14
  • 96% of base hallucinations are eliminated in the Llama replication, falling from 228 to 10 while adding 298 correct answers.Source: G Llama-3.1-8B-Instruct replication · p. 14
  • At r = 0.80, the adapter hallucinates on all 4 ABSTAIN-gold examples, so that slice is too small for isolated conclusions.Source: H Failure mode detail · p. 15
  • 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.Source: H Failure mode detail · p. 15
  • Across three seeds, exactly 7 ABSTAIN-gold examples per seed remain hallucination failures, with 4 stable across all seeds.Source: H Failure mode detail · p. 15

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.

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

10B: Adversarial tight-span eviction

Conclusion

10C: Opposite-priority pairs

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.

Related papers

Loading 2608.29934v1…