💬 NLP

LongAttnComp: Cross-Family Context Compression for Long-Context Reasoning

LongAttnComp is a two-stage fine-tuned, training-free context compression framework that leverages a lightweight cross-attention scorer and specialized token-level strategies to significantly improve long-context reasoning accuracy and efficiency across diverse model families, particularly in code debugging and multi-document tasks.

Original authors: Mengmeng Ji, Ravi Shanker Raju, Jonathan Lingjie Li, Chen Wu

Published 2026-06-02

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Mengmeng Ji, Ravi Shanker Raju, Jonathan Lingjie Li, Chen Wu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a massive library of books, but you only have a small table to work on. You need to find a specific answer hidden somewhere in those thousands of pages. If you try to read the whole library at once, your table gets cluttered, your brain gets tired, and you might miss the important details.

This is the problem LongAttnComp solves for AI computers. It's a tool that helps AI models "read" incredibly long documents (like 100,000+ words) without getting overwhelmed, while still finding the right answer.

Here is how it works, broken down into simple concepts:

The Problem: The "Too Much Information" Bottleneck

When an AI tries to solve a complex problem (like debugging a huge computer code file or answering a question based on a long novel), it has to hold all that text in its "working memory." This is expensive and slow.

The old way: Try to read everything. (Too slow, too expensive).
The "free" way: Some tools just guess which parts are important without learning. (They often miss the crucial clues, especially in tricky tasks like coding).

The Solution: A Smart "Book Summarizer"

The authors built a system called LongAttnComp. Think of it as a super-smart, highly trained assistant who sits between the user and the main AI.

The "Draft" Assistant: The system uses a smaller, frozen AI model (like a junior librarian) that has been specially trained to look at a long text and a specific question.
Token-Level Chunking: Instead of looking at whole chapters or documents, this assistant breaks the text into tiny, manageable "chunks" (like individual paragraphs or code blocks).
The Scorecard: For every chunk, the assistant gives it a score: "How relevant is this to the question?"
The Selection: It keeps only the highest-scoring chunks. But here's the trick: it doesn't just keep them in a random pile. It puts them back in their original order so the story or code still makes sense.
The Result: The main AI only sees this compressed, high-quality summary, not the whole 100,000-word mess.

The Secret Sauce: Two-Stage Training

The authors realized that to be good at everything, the assistant needed a specific training plan. They used a Two-Stage Recipe:

Stage 1: The Basics (The "Needle in a Haystack" Training)
They taught the assistant to find simple facts in large piles of text. Imagine teaching a dog to find a specific treat hidden in a field. This made the assistant very good at finding clear, direct answers.
- Result: It became great at finding bugs in code and simple questions.
Stage 2: The Advanced Class (The "Detective" Training)
They realized the assistant was still struggling with complex puzzles where the answer requires connecting dots across different parts of the text (multi-hop reasoning). So, they gave it harder training data—puzzles that required thinking through a chain of clues.
- Result: This made the assistant much better at complex reasoning tasks, like understanding a long story or a multi-step logic problem, without forgetting how to find simple facts.

What They Found (The Results)

The paper tested this system on some very tough challenges:

Code Debugging: When asked to find errors in massive code files, LongAttnComp performed just as well as if the AI had read the entire file, but it did it much faster. It beat the "free" guessing tools by a huge margin.
Cross-Family Magic: The most impressive part? They trained the assistant using one type of AI (Llama), but it worked perfectly when helping completely different AI models (like DeepSeek, MiniMax, and GPT-OSS). It's like training a human translator in English, and then having them successfully translate for a French speaker, a German speaker, and a Japanese speaker without retraining them for each language.
Complex Reasoning: On the harder "LongBench" tests (which involve complex, multi-document reasoning), the Stage 2 training closed the gap significantly, proving that the right training data matters more than just changing the architecture.

The Bottom Line

LongAttnComp is a "smart filter." It doesn't just throw away text; it learns exactly what information is needed to answer a question, keeps the order of the story intact, and hands a clean, concise version to the main AI. By training this filter in two stages—first on simple facts, then on complex reasoning—it creates a system that is fast, accurate, and works across different types of AI models.

Technical Summary: LongAttnComp

Problem Statement

As large language models (LLMs) are increasingly deployed for real-world applications requiring the processing of inputs exceeding 100,000 tokens, a critical bottleneck has emerged between context length and inference efficiency. While context compression offers a solution by filtering or condensing input before it reaches the target model, existing training-free, attention-based methods (such as Speculative Prefill) exhibit substantial performance gaps on demanding long-context tasks, particularly in code reasoning. Furthermore, prior fine-tuning approaches like AttnComp have been narrowly scoped to retrieval-augmented QA with short contexts (~12k tokens) and document-level scoring, leaving their potential as general-purpose long-context compressors untested.

The authors frame long-context task performance as a decomposition of retrieval (locating query-relevant evidence) and reasoning (inferring answers from that evidence). They identify retrieval as the primary bottleneck, arguing that effective context compression is fundamentally a retrieval problem requiring the identification of relevant tokens and segments within massive inputs.

Methodology: LongAttnComp

LongAttnComp is a trainable, draft-model-driven compression framework adapted from AttnComp. It retains the core mechanism of a frozen LLM backbone augmented with a trainable cross-attention scoring layer but introduces four key architectural adaptations and a two-stage fine-tuning recipe to handle long-context inference.

Architectural Adaptations

Token-Level Chunking: Unlike AttnComp's document-level scoring, LongAttnComp partitions inputs into fixed-size token chunks. This allows operation on real-world inputs (e.g., single long documents or codebases) lacking natural document boundaries and treats chunk size as a tunable hyperparameter.
Token-Budget Top-p Selection: The original score-threshold top-p algorithm is modified to include a content token budget ( $B$ ). Selection stops when either the cumulative score exceeds $p$ or the retained tokens reach $B$ . This provides predictable control over compression length, preventing the aggressive under-retention seen in long-context tasks with the original minimum-score threshold.
Positional Reordering: Selected token chunks are restored to their original positional order before being passed to the target model, preserving discourse coherence, which the original AttnComp did not guarantee.
Format-Agnostic Query Parser: To handle complex prompt structures without fixed templates (e.g., Code-Debug), the system employs a query parser. The authors evaluate both accurate span extraction and an "arbitrary" strategy (allocating the last $N$ tokens as the query), finding the latter robust with minimal performance degradation.

Two-Stage Fine-Tuning Recipe

The authors propose a curriculum learning approach to broaden task coverage:

Stage 1 (Foundation): Trains on a large-scale dataset (32k examples) combining SQuAD and HotpotQA. This establishes a general retrieval foundation using NIAH-style data (single-fact and basic multi-hop retrieval).
Stage 2 (Extension): Continues training from the Stage 1 checkpoint on a dataset targeting harder retrieval patterns, including MuSiQue and 2WikiMultiHopQA, interleaved with replay of Stage 1 data to prevent catastrophic forgetting. This stage specifically tests the impact of sub-question decomposition (explicitly including dataset-provided sub-questions in the query) versus using only the multi-hop question.

Key Contributions

LongAttnComp Framework: A modular, trainable compression framework adapted for long-context inference, featuring token-level chunking, a token-budget top-p algorithm, positional reordering, and a format-agnostic query parser.
Two-Stage Training Recipe: A methodology that extends a compressor's task coverage and strengthens multi-hop reasoning capabilities without compromising its initial retrieval strengths.
Empirical Validation:
- LongAttnComp matches or exceeds full-context accuracy on InfiniteBench Code-Debug.
- It substantially outperforms training-free baselines (e.g., Speculative Prefill).
- It demonstrates cross-family generalization, transferring effectively across four target models from three different families (DeepSeek, MiniMax, GPT-OSS) without retraining.
- The two-stage recipe significantly closes the performance gap on multi-document reasoning tasks (LongBench v2) while preserving code reasoning performance.

Results

Code Reasoning (InfiniteBench Code-Debug): On the DeepSeek-R1-0528 target, the Stage 1 compressor achieved 75.38% accuracy, surpassing the full-context baseline (74.37%) and the training-free Speculative Prefill (62.44%). The Stage 2 (subq) variant reached 76.90%. Similar trends were observed across DeepSeek-V3.1, MiniMax-M2.5, and GPT-OSS-120B, where LongAttnComp consistently outperformed Speculative Prefill by 7–31 points.
Multi-Document Reasoning (LongBench v2): Stage 1 underperformed the full-context baseline, reflecting the limitations of NIAH-style training data for naturalistic reasoning. Stage 2 training recovered 7–12 points over Stage 1, surpassing Speculative Prefill and reaching within 1.4 points of the 100k-truncated full-context baseline.
RULER Benchmark: LongAttnComp recovered accuracy on subtasks suffering from "lost-in-the-middle" effects (e.g., single-needle retrieval) but showed expected limitations on tasks where evidence is distributed across many positions (multi-value/multi-query), consistent with the training data distribution.

Significance and Claims

The paper claims that LongAttnComp demonstrates that task sensitivity in compression reflects training data composition rather than architectural limits. The performance gap observed in Stage 1 on complex reasoning tasks was not due to the model's inability to learn, but rather the lack of diverse, naturalistic training patterns. By introducing Stage 2 with multi-hop and naturalistic data, the authors show that a single cross-attention architecture can be extended to handle complex long-context reasoning tasks.

The authors emphasize that the compressor acts as a target-agnostic preprocessing step, capable of transferring across unrelated model families without retraining. They also note that inference-time hyperparameters (chunk size, selection mode) are task-dependent, suggesting that future deployment may require adaptive mechanisms for unknown task types. Finally, the paper posits that the computational overhead of LongAttnComp is roughly one-third that of Speculative Prefill, as it utilizes only the first 13 layers of the draft model rather than the full backbone.