Skip to content

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

Aug 2026 · 0 citations · 29 references
Computer Science

TL;DR

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, and controlled-deletion experiments show that the learned behavior is driven by evidence content rather than input length alone.

Abstract

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.

View source

Similar papers

Open access Sep 2026

Compression Remembers the Detour: How Context Compression Amplifies Off-Task Content and Derails LLM Agents

This work formalizes CIGD via the drift amplification ratio (DAR), evaluates six compression strategies on a controlled injection benchmark ( per cell), and proposes Goal-Anchored Compression (GAC): a pinned goal anchor, a negation ledger, status-aware retention scoring, and drift re-anchoring.

Zhan Zhang, Wen-Zhi Zhang · 0 citations
#machine learning Preprint Sep 2026

When Keywords Drop but Classifiers Hold: Soft Refusals under KV Cache Compression

Safety auditing under compression should rely on several judges matched to the serving context rather than on keyword rates alone, because the gap is not universal and weakens under short fillers and paired SnapKV, so safety auditing should rely on several judges matched to the serving context rather than on keyword ra...

Kang Chen, Xiu-Ze Zhou, Hong Chen et al. · 0 citations
#machine learning Preprint Sep 2026

Cartridges++: KV Cache Compression without Off-Context Derailment

This study investigates the crucial deployment question of whether Cartridges and other CKVs can handle off-context queries, and proposes Cartridges++, simple modifications to cartridges that retain off-context abilities at small or negligible cost.

Sonia Laguna, João Monteiro, Marco Cuturi et al. · 1 citation
#artificial intelligence Preprint Sep 2026

When Quantization Preserves Accuracy but Not Evidence: Explanation-Aware Post-Training Quantization for Medical LLMs

The objective is to preserve the full-precision model's answer-supporting behavior rather than improve gold-label accuracy, and the method better preserves the full-precision model's answer behavior and rationale-to-answer support.

Yeji Kim, Mi-Young Kim, Randy Goebel · 0 citations
#machine learning Preprint Sep 2026

MemoryWalker: Stop Training Agents on Contexts They Never Saw

This work introduces two exact, gradient-equivalent corrections: LogitTree, a segmented K-forward traversal, and a packed 4D attention mask, and proposes SDCC (Self-Distillation for Conditioning Consistency), a single-backward-pass variational relaxation.

J. Zinco, Xun-Jie Zhu, Shen Huang et al. · 0 citations

Related blog posts

MIT News · Artificial Intelligence Aug 27, 2026

Looking beyond natural sequences

A new machine-learning framework aims to improve the success rate of computational protein design while moving away from results that reproduce sequences found in nature.

We use cookies to run the site and, with your consent, for analytics and to show ads. See our Cookie Policy.