moalikhoda/compression-aware-abstention-adapters

Compression-Aware Abstention — LoRA adapters

LoRA adapters from "Compression-Aware Abstention: Teaching LLMs to Refuse When KV-Compression Masks Remove Answer Evidence" (GroundLM @ EMNLP 2026).

Each adapter is trained to answer when the evidence a query depends on survives KV-cache compression, and to emit a refusal when the compressor has evicted it. All are rank-16 LoRA on the four attention projections only (q_proj, k_proj, v_proj, o_proj), 10.1M trainable parameters, 0.13% of Qwen2.5-7B-Instruct.

Adapters

Directory	Base model	Training masks	Reported as
v1-kvzip-seed42	Qwen2.5-7B-Instruct	KVzip only	v1 / v1-ps
v1-kvzip-seed43, v1-kvzip-seed44	Qwen2.5-7B-Instruct	KVzip only	Table 1 seed spread
mix-ps-seed42, mix-ps-seed43	Qwen2.5-7B-Instruct	KVzip + ExpAttn + SnapKV	mix / mix-ps
mix-cc-seed42, mix-cc-seed43	Qwen2.5-7B-Instruct	multi-compressor, loss after cache eviction	mix-cc
single-compressor-cc-seed42	Qwen2.5-7B-Instruct	KVzip only, cc decoding	Appendix M ablation
random-mask-seed42	Qwen2.5-7B-Instruct	uniform-random at matched retention	Section 5.4 probe
llama31-8b-kvzip-seed42	Llama-3.1-8B-Instruct	KVzip only	Section 5.7 replication

v1-kvzip-seed42 is the headline prompt-style adapter; mix-cc-seed42 is the best adapter under actual compressed-cache decoding.

Usage

from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

base = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-7B-Instruct", torch_dtype="bfloat16", device_map="auto"
)
model = PeftModel.from_pretrained(
    base, "moalikhoda/compression-aware-abstention-adapters", subfolder="mix-cc-seed42"
)
tok = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B-Instruct")

prompt = f"Context: {surviving_context}\n\nQuestion: {question}\n\nAnswer:"
out = model.generate(**tok(prompt, return_tensors="pt").to(model.device),
                     max_new_tokens=64, do_sample=False)

The adapters were trained on the prompt format above, where surviving_context is the compressor's kept tokens decoded back to text. When the evidence is missing they emit: "I cannot determine the answer to this question from the available context."

Important caveats

• Trained on MuSiQue 2-hop QA only. Behavior on other task families is tested on RULER multi-key NIAH (Section 5.7) but not broadly.
• The refusal decision is calibrated to tight answer-bearing span survival, which is a surrogate for contextual sufficiency, not ground truth (Cohen's kappa 0.61 against an LLM-judge verdict).
• Compressed-cache results use kvpress KVzipPress. Under compressed-cache decoding at high retention the task remains hard: mix-cc answers 22/69 of evidence-retaining examples correctly at r=0.80, against 1/69 for the unmodified base.

Code, evaluation outputs, and a script that reproduces every number in the paper: https://github.com/mali-kh/compression-aware-abstention

Citation

@inproceedings{khodabandehlou2026compression,
  title     = {Compression-Aware Abstention: Teaching {LLM}s to Refuse When
               {KV}-Compression Masks Remove Answer Evidence},
  author    = {Khodabandehlou, Mohammadali and Krishnamachari, Bhaskar},
  booktitle = {Proceedings of the Workshop on Grounding and Reliability of
               Language Models (GroundLM) at EMNLP},
  year      = {2026}
}