🤖 AI Summary
This work addresses the challenge of balancing efficiency and accuracy in long-context inference (100K+ tokens) by proposing a general, lightweight context compression method. The approach employs a fine-tuned, lightweight cross-attention scoring layer combined with token-level chunking, top-p token budget allocation, positional reordering, and format-agnostic query parsing. A two-stage fine-tuning strategy—first pretraining on Needle-in-a-Haystack (NIAH) tasks followed by task-specific fine-tuning on multi-hop reasoning and code-related benchmarks—enables significant reduction in prefill costs while maintaining or even surpassing full-context performance. The method matches or outperforms full-context baselines on InfiniteBench Code-Debug, substantially narrows performance gaps on LongBench v2 multi-document reasoning tasks, and demonstrates strong transferability across four models spanning three distinct model families, confirming its generality and effectiveness.
📝 Abstract
As real-world applications increasingly require processing inputs of 100k+ tokens, the gap between context length and inference efficiency has become a critical bottleneck. Context compression offers a way to reduce prefill costs while preserving task accuracy. However, existing training-free attention-based methods leave substantial gaps in demanding long-context tasks such as code reasoning. We present LongAttnComp, a long-context adaptation of AttnComp that fine-tunes a lightweight cross-attention scoring layer and introduces tokenlevel chunking, a token-budget top-p algorithm, positional reordering, and a formatagnostic query parser. We further design a two-stage fine-tuning recipe for the compressor: Stage 1 builds a general retrieval foundation from NIAH-style data, and Stage 2 extends it with multi-hop and reasoning data for broader long-context task coverage. On InfiniteBench Code-Debug, LongAttnComp matches or exceeds full-context accuracy, substantially outperforms training-free baselines, and transfers across four target models from three families. On LongBench v2, the two-stage recipe largely closes the Stage 1 gap on multi-document reasoning while preserving Code-Debug performance.