PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Mengmeng Ji, Ravi Shanker Raju, Jonathan Lingjie Li, Chen Wu

arXiv:2606.01336v2cs.CL

TL;DR

Long-context inference is costly, and existing training-free attention-based compression leaves gaps on demanding code-reasoning tasks. LongAttnComp adapts AttnComp with long-context selection and parsing changes plus two-stage training, matching or exceeding full-context Code-Debug accuracy and improving multi-document reasoning.

  • Problem

    Long-context inference has substantial memory and compute costs, while existing training-free attention-based compression leaves a performance gap on long-context code reasoning.

  • Method

    LongAttnComp adapts AttnComp with token-level chunking, token-budget top-p selection, positional reordering, format-agnostic parsing, and two-stage fine-tuning.

  • Results

    LongAttnComp matches or exceeds full-context accuracy on InfiniteBench Code-Debug, transfers across four target models from three families, and Stage 2 improves LongBench v2 multi-document reasoning.

  • Takeaways & Limitations

    The two-stage recipe broadens task coverage while largely preserving Code-Debug performance.

  • Takeaways & Limitations

    The compressor is trained only on synthetically constructed NIAH-style data, leaving a residual gap to the untruncated full-context baseline on naturalistic LongBench v2 tasks.

Abstract

from arXiv · show

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.

1 Introduction

LongAttnComp treats long-context compression as query-relevant retrieval and adapts AttnComp with long-context selection, ordering, parsing, and two-stage training. It matches or exceeds full-context Code-Debug accuracy, outperforms training-free baselines, transfers across model families, and improves multi-document reasoning after Stage 2.

  • Motivation: Long-context inference creates substantial memory and compute costs, motivating compression to reduce target-model prefill costs.Compression filters or condenses context before inference, trading a small upfront cost for savings in the target model’s prefill stage.
  • Approach: LongAttnComp adapts AttnComp’s fine-tuned compression mechanism to a draft-model-driven framework for long-context retrieval and reasoning.The method retains AttnComp’s core mechanism while adding four architectural adaptations and a two-stage training recipe.
  • Approach: Its architectural adaptations include token-level chunking, token-budget top-p selection, positional reordering, and format-agnostic query parsing.These changes support inputs without natural document boundaries, predictable compression lengths, coherent ordering, and varied prompt formats.
  • Training: The two-stage recipe builds general query-aligned retrieval on NIAH-style data, then adds multi-hop and reasoning data with replay to broaden task coverage.Stage 2 also compares query construction with and without explicit sub-question decomposition.
  • Results: On InfiniteBench Code-Debug, LongAttnComp matches or exceeds full-context accuracy, outperforms training-free baselines, and transfers across four target models from three families without retraining.On LongBench v2, Stage 2 improves multi-document reasoning over Stage 1 while preserving Code-Debug performance.

2 Related Work

Prior work includes abstractive and extractive compression, training-free attention-based retrieval, and fine-tuned document scoring. LongAttnComp combines the draft-model paradigm with a trainable cross-attention compressor for long-context inference.

  • Context compression: Context compression methods are broadly abstractive or extractive, with extractive systems retaining tokens or segments using perplexity or embedding similarity.Both categories commonly rely on a predetermined compression budget.
  • Attention-based retrieval: Speculative Prefill uses lightweight draft-model attention weights as token-importance signals for training-free compression and transfers across target model families.The related work identifies a substantial performance gap for this approach on long-context code reasoning.
  • Attention-based retrieval: AttnComp fine-tunes a cross-attention scoring layer for document-level compression with adaptive top-p selection, but its evaluation and training are narrowly scoped to retrieval-augmented QA.LongAttnComp extends this mechanism to long-context inputs and broader tasks.
  • LongAttnComp: LongAttnComp draws on speculative decoding’s draft-model paradigm while fitting the broader theme of resource-adaptive inference.Its compressor is positioned as a fine-tuned complement to training-free attention-based compression.

3 Method

LongAttnComp adapts AttnComp for long-context inputs through token-level scoring, budget-controlled selection, positional restoration, and format-agnostic parsing. A two-stage recipe extends retrieval from general NIAH-style patterns to harder multi-hop and naturalistic reasoning while retaining prior strengths.

  • Architecture: LongAttnComp retains AttnComp’s frozen-backbone and trainable cross-attention architecture while adding long-context adaptations and two-stage fine-tuning.Only approximately 0.5% of parameters are updated in the retained architecture.
  • Token-level chunking: Token-level chunking scores fixed-size token chunks independently instead of documents, supporting long inputs without clean document boundaries.Chunk size becomes a tunable hyperparameter that can be swept across tasks.
  • Top-p selection: The modified top-p selector stops when cumulative score exceeds p or retained tokens reach budget B, replacing the minimum-score threshold.This provides predictable control over compressed length and avoids under-retention caused by the original threshold.
  • Ordering and selection: Selected chunks are restored to original positional order before inference, preserving discourse coherence rather than returning an unordered set.The method also supports cumulative-score and budget-only selection modes.
  • Selection behavior: Cumulative-p selection can work for concentrated evidence but may discard distributed supporting evidence when it terminates early.On RULER’s niah_s_1 it achieves 100% accuracy with approximately 2k tokens against a 16k budget, while distributed-evidence tasks can lose accuracy.
  • Two-stage fine-tuning: The two-stage recipe first establishes general retrieval, then continues training on harder multi-hop and naturalistic patterns with replay to mitigate catastrophic forgetting.Sub-question and no-sub-question query variants probe how training-time query representation affects downstream behavior.
  • Two-stage fine-tuning: Continued training of a single cross-attention layer broadens task coverage while retaining, and sometimes improving, existing strengths.This result supports extending the compressor through training-data changes rather than architectural replacement.

4 Data

The training data is generated with a modified RULER pipeline using structurally labeled retrieval examples from four QA datasets. Stage 1 provides a larger general foundation, while Stage 2 adds multi-hop data and replay.

  • Dataset construction: Each training sample contains 100 candidate documents, a query, binary relevance labels, and 25% all-negative examples.Labels come from dataset-provided structural annotations rather than LLM-based labeling.
  • Stage 1: Stage 1 contains 32,000 examples from SQuAD and HotpotQA with sequences spanning 8k–48k tokens.The data includes 16,000 examples from each source, with 0–1 relevant documents for SQuAD and 0–2 for HotpotQA.
  • Stage 2: Stage 2 contains 20,000 samples combining MuSiQue, 2WikiMultiHopQA, and replay from SQuAD and HotpotQA.The composition is 8,000 MuSiQue, 4,000 2WikiMultiHopQA, and 4,000 each from the two replay sources.

5 Experimental Setup

The evaluation tests LongAttnComp on long code, multi-document reasoning, and synthetic retrieval tasks across target models, using full context and Speculative Prefill as baselines. The protocol varies compressor settings by benchmark and includes Stage 1 and Stage 2 checkpoints.

  • Target Models: Target-model evaluation spans DeepSeek-R1-0528, MiniMax-M2.5, GPT-OSS-120B, and DeepSeek-V3.1.The selection tests transfer beyond the Llama-3.1-8B-Instruct family used to train the compressor.
  • Evaluation Tasks: The study evaluates Code-Debug, LongBench v2, and RULER to probe long-context code reasoning, multi-document reasoning, and long-context utilization.Code-Debug uses long bug-identification inputs averaging ∼115k tokens, while LongBench v2 and RULER extend evaluation beyond code.
  • Baselines: LongAttnComp is compared with full context and the training-free Speculative Prefill compressor.The original AttnComp is not included because its compressor checkpoint was not publicly released.
  • Checkpoints: The evaluation includes Stage 1 and Stage 2 LongAttnComp checkpoints, with Stage 1 trained on a 32k SQuAD+HotpotQA dataset.The setup examines whether multi-hop training extends the foundation compressor’s task coverage.

6 Results

LongAttnComp performs strongly on long-context code reasoning and transfers across target families, while Stage 1 initially trails full context on LongBench v2. Stage 2 substantially improves multi-document reasoning but does not yet match untruncated full context.

  • 6.1 Long-Context Code Reasoning: 76.90 accuracy is achieved by Stage 2 with the subq variant on Code-Debug, the highest result in Table 1.Stage 1 already exceeds full context and outperforms Speculative Prefill by 12.9 points.
  • 6.2 Target-Model Generalization: LongAttnComp Stage 1 outperforms Speculative Prefill by 7–31 points across four Code-Debug target models.It matches or closely tracks each target’s full-context accuracy without target-specific fine-tuning or hyperparameter tuning.
  • 6.3 Beyond Code Reasoning: Stage 1 LongAttnComp and Speculative Prefill substantially underperform full context on LongBench v2.The results motivate Stage 2 training on multi-hop retrieval and reasoning data.
  • 6.3 Beyond Code Reasoning: Stage 2 recovers 7–12 points over Stage 1 across every LongBench v2 breakdown.The nosubq variant reaches 49.7 Overall accuracy, within 1.4 points of the 100k-truncated full-context baseline.
  • 6.3 Beyond Code Reasoning: LongAttnComp recovers niah_s_3 RULER accuracy from 57.4 to 99.2 but underperforms on multi-value and multi-query subtasks.Stage 2’s RULER improvements over Stage 1 are small.

7 Discussion

The discussion attributes task coverage primarily to training-data composition and shows that inference settings must match retrieval demands. It also reports efficiency advantages from using a shallow compressor, while noting uncertainty around explicit sub-question training.

  • Training Data and Task Coverage: LongAttnComp performs best when evidence is clearly query-aligned and bounded, but struggles when evidence is indirectly query-dependent or distributed.The discussion attributes this variation to training-data composition rather than a fundamental architectural limit.
  • Training Data and Task Coverage: Stage 2 recovers 7 to 12 points across LongBench v2 breakdowns while largely preserving Code-Debug performance.The recipe adds multi-hop and naturalistic samples while replaying Stage 1 data, without changing the architecture.
  • Sub-question Training: Explicit sub-question decomposition may help multi-hop retrieval, but the evidence is exploratory because the clean win comes from one target on one benchmark.The paper leaves fuller characterization to future work.
  • Inference-Time Configuration: Inference-time selection and chunk size are task-dependent because different tasks require different amounts and distributions of retained evidence.Code-Debug uses larger chunks, RULER shorter needles, and LongBench v2 many short supporting spans.
  • Efficiency: LongAttnComp uses a smaller compute and memory footprint than Speculative Prefill because its compressor uses only the first L=13 layers of the backbone.Speculative Prefill uses the full draft model.

8 Conclusion

LongAttnComp is a fine-tuning-based long-context compression method whose modular compressor transfers across unrelated target-model families without retraining. The authors suggest that more diverse training data could extend the architecture to more complex reasoning tasks, while identifying several deployment and data limitations.

  • LongAttnComp is a fine-tuning-based long-context compression method.
  • The trained compressor transfers across unrelated target-model families without retraining.It functions as a modular, target-agnostic preprocessing step.
  • More diverse training data may extend the same architecture to more complex long-context reasoning tasks.
  • Future work includes expanding reasoning-heavy training data, adapting inference settings, improving query parsing, and fine-tuning the draft model.The authors note that high-quality long-context training data remains scarce and suitable generation pipelines are not openly available.

Ethics Statement

The work uses publicly available models, datasets, and benchmarks under standard research-use licenses, with training data derived from public QA datasets for research use only.

  • All models, datasets, and benchmarks used are publicly available research artifacts under standard research-use licenses.
  • The training data is derived from public QA datasets and intended for research use only.
  • The authors collected no new data and identified no personally identifying information or offensive content in their use.The source datasets were vetted by their original curators and the research community.

Limitations

LongAttnComp’s evidence is constrained by synthetic training data, task-dependent inference choices, a heuristic query parser, limited efficiency measurements, deployment conditions, and one compressor backbone.

  • Training-data scope: Synthetic NIAH-style training data does not fully capture naturalistic long-context reasoning, leaving a residual gap to the untruncated full-context baseline on LongBench v2.Both stages use synthetic pipelines derived from SQuAD, HotpotQA, MuSiQue, and 2WikiMultiHopQA.
  • Task-dependent hyperparameters: Optimal chunk size, parsed query length, and selection mode vary across Code-Debug, RULER, and LongBench v2.A single fixed configuration may be unsuitable when the task type is unknown in advance.
  • Query parsing assumption: The arbitrary last-N-token query parser uses task-specific optimal lengths and may miss queries structurally embedded elsewhere in the prompt.Reported optimal N values are 128 for Code-Debug, 256 for RULER, and 512 for LongBench v2.
  • Empirical efficiency measurements: Efficiency claims rely on a rule-of-thumb compressor-overhead estimate rather than controlled end-to-end TTFT, throughput, or memory measurements.Empirical efficiency characterization is left to follow-up work.
  • Deployment-side constraints: All target-model evaluations use the SambaNova cloud API, whose serving-stack constraints interact occasionally with the evaluation protocol.These constraints include context budgets, tokenization, and reasoning-output allocations.
  • Single compressor backbone: All experiments use Llama-3.1-8B-Instruct as the compressor backbone, leaving transfer to other backbone families or scales untested.The study does not test smaller draft models for tighter deployment or larger models for additional headroom.

G.1 Compressor Training Ablations

The ablations select training and inference configurations for LongAttnComp, showing that balanced, larger-scale training and suitable chunking materially affect retrieval and Code-Debug performance.

  • Training-data composition: Single-source training specializes strongly: SQuAD-only and HotpotQA-only checkpoints preserve source-task performance but drop significantly on the other subtask.Combining both sources balances qa_1 and qa_2 performance, while qa_2 exceeds the nocompression baseline at 16k samples.
  • Training-data composition: Doubling the training set from 16k to 32k improves both RULER QA subtasks without saturating the lightweight cross-attention scoring layer.The ablation also selects cosine learning-rate decay, dropout, and 15 epochs; 18 epochs causes a small drop consistent with mild overfitting at fixed dataset size.
  • Code-Debug training sweep: 76.40% Code-Debug accuracy with the best Stage 1 configuration exceeds full context by 2.0 points and Speculative Prefill by 13.9 points while compressing the prompt by 83%.The configuration uses a 32k training set, cosine-decay learning rate, 15 epochs, chunk size 1024, and a 16k token budget.
  • Code-Debug training sweep: Chunk size has a large effect on accuracy, increasing performance from 56.85% at chunk 128 to 76.40% at chunk 1024 on the 32k checkpoint.Cosine decay and the larger dataset also improve results, while 30 epochs partially compensates for the smaller 16k dataset.
  • Inference settings: p = 0.95 remains the best choice in the modified token-budget top-p sweep and is adopted for subsequent experiments.The sweep uses a 10-sample Code-Debug validation subset with a 16k input budget and chunk size 256.
  • Inference settings: On LongBench v2 development with DeepSeek-V3.1, the chosen configuration reaches 48.9 Overall versus 50.7 for full context, while budget-only selection and q=512 add 5.4 points over default settings.The selected Stage 2 subq configuration is transferred to DeepSeek-R1-0528 as the held-out target.
Loading 2606.01336v2…