💬 NLP

ATACompressor: Adaptive Task-Aware Compression for Efficient Long-Context Processing in LLMs

The paper introduces ATACompressor, an adaptive task-aware compression method that dynamically selects and compresses only task-relevant portions of long contexts to mitigate the "lost in the middle" problem, thereby outperforming existing approaches in both compression efficiency and task performance on QA benchmarks.

Original authors: Xuancheng Li, Haitao Li, Yujia Zhou, Qingyao Ai, Yiqun Liu

Published 2026-02-04

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Xuancheng Li, Haitao Li, Yujia Zhou, Qingyao Ai, Yiqun Liu

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 brilliant but very busy librarian (the Large Language Model, or LLM). This librarian is amazing at answering questions, but they have a problem: if you hand them a stack of 100 books to find one specific fact, they get overwhelmed. They might skim the first few pages, get distracted by the middle chapters, and completely miss the crucial answer hidden in the middle or end of the stack. This is what researchers call the "lost in the middle" problem.

Usually, to help the librarian, we try to summarize the books for them. But existing summarizers have two main flaws:

The "Scissors" Approach (Hard Prompts): They cut out sentences they think are boring. This is safe, but it often leaves the summary too long to be truly efficient.
The "Magic Ink" Approach (Soft Prompts): They turn the whole book into a tiny, invisible code. This is very short, but because they summarize everything without looking at your specific question, they often accidentally throw away the one sentence you actually needed.

Enter ATACompressor, a new tool designed to fix both problems. Think of it as a smart, adaptive assistant who helps the librarian before they even open the books.

Here is how ATACompressor works, broken down into three simple steps:

1. The Selective Reader (The "Selective Encoder")

Instead of reading the whole book cover-to-cover, this assistant looks at your specific question first.

The Analogy: Imagine you ask, "What color was the horse?" The assistant doesn't summarize the horse's diet, the stable's history, or the rider's biography. They only highlight and keep the sentences that mention the horse's color.
How it works: It scans the long text, ignores the fluff, and only compresses the parts that are actually relevant to your question. This ensures the "critical information" isn't lost.

2. The Smart Budget Manager (The "Adaptive Allocation Controller")

This is the most unique part. In the past, compression tools had a "one-size-fits-all" rule (e.g., "Always turn 1,000 words into 10 words").

The Analogy: Imagine you are packing for a trip. If you are going for a weekend, you need a small backpack. If you are going for a month, you need a huge suitcase. A rigid rule would force you to use a suitcase for a weekend trip (wasting space) or a backpack for a month (leaving clothes behind).
How it works: ATACompressor has a "budget manager" that looks at the relevant text the Selective Reader found.
- If the answer is hidden in a short paragraph, the manager says, "We only need a tiny amount of space (few tokens)."
- If the answer is spread across a long chapter, the manager says, "We need a bigger suitcase (more tokens)."
- This dynamic adjustment ensures the librarian gets exactly enough information to do the job, no more and no less.

3. The Result: A Perfectly Packed Briefcase

Once the assistant has selected the right parts and decided on the right size, they compress that information into a tiny, dense "briefcase" of tokens. They hand this briefcase to the librarian.

Because the briefcase only contains the relevant info, the librarian doesn't get overwhelmed.
Because the size of the briefcase matches the complexity of the task, the librarian has just enough detail to give a perfect answer.

What the Paper Found

The authors tested this system on three different types of "reading tests" (datasets called HotpotQA, MSMARCO, and SQUAD).

Better Answers: ATACompressor got higher scores on these tests than previous methods. It was better at finding the right answer, especially in very long texts where other methods failed.
More Efficient: It squeezed the information down much more tightly (higher compression ratio) without losing the important bits.
Faster: Because the librarian had less junk to read, the whole process was faster.

In short: ATACompressor is like a smart filter and a flexible packing system combined. It reads your question, grabs only the necessary facts, packs them into a size that fits the job perfectly, and hands them to the AI so it can give you the best answer possible without getting lost in the noise.

Technical Summary: ATACompressor

Problem Statement

Large Language Models (LLMs) face significant challenges when processing long-context inputs, primarily the "lost in the middle" phenomenon, where critical information is diluted or ignored due to excessive sequence length. While Retrieval-Augmented Generation (RAG) provides external context, naive approaches that append raw document tokens often overwhelm the model. Existing context compression methods attempt to mitigate this but suffer from specific limitations:

Hard-prompt methods (e.g., Selective-Context, LongLLMLingua) remove irrelevant tokens but typically achieve low compression ratios.
Soft-prompt methods (e.g., AutoCompressor, ICAE, 500Compressor) compress text into dense token sequences, achieving high compression ratios. However, they often compress the entire context indiscriminately, leading to the loss of task-relevant information because they lack task-specific awareness during compression. Furthermore, these methods usually employ a fixed number of compressed tokens, failing to dynamically adapt to the varying information density of different tasks.

Methodology: ATACompressor

The authors propose Adaptive Task-Aware Compressor (ATACompressor), a framework designed to dynamically adjust compression based on specific task requirements. The architecture consists of three main components: a Selective Encoder, an Adaptive Allocation Controller (AAC), and a target LLM (decoder).

1. Selective Encoder ( $\varphi$ )

Unlike traditional soft-prompt compressors that process the entire context, the Selective Encoder is designed to identify and compress only the portions of the context ( $\hat{C}_{Rel}$ ) relevant to the query ( $Q$ ).

Architecture: It utilizes a frozen Large Language Model (LLM) with trainable LoRA parameters.
Granularity Handling: To address the ambiguity of "relevance" (which can vary from sentence to document level), the input context is segmented into uniform chunks (e.g., passages or sentences) aligned with the ground-truth granularity during training. The encoder learns to retain and compress only the chunks necessary to answer the query.
Process: The encoder processes the concatenated chunks and query, outputting a compact set of compressed tokens ( $c_1, ..., c_k$ ) representing the relevant information.

2. Adaptive Allocation Controller (AAC)

The AAC addresses the limitation of fixed compression rates in existing soft-prompt methods. It dynamically determines the number of compressed tokens ( $k$ ) based on the estimated length of the relevant content.

Components: It comprises a lightweight Probe ( $\zeta$ ) and a Policy Function ( $\eta$ ).
Mechanism: After the Selective Encoder processes the input, the Probe analyzes the encoder's hidden states to estimate the length of the relevant text ( $\hat{L}_{Rel}$ ). The Policy Function then calculates the number of compressed tokens ( $k$ ) using a predefined policy ratio ( $r$ ):
$k = \min(\frac{\hat{L}_{Rel}}{r}, k_{max})$
Benefit: This allows the system to allocate fewer tokens for short relevant spans and more tokens for longer ones, optimizing resource utilization while preserving essential information.

3. Training and Inference Workflow

Pretraining: The Selective Encoder and the Probe are jointly optimized. The Encoder is trained to reconstruct the relevant context (gold-truth) using a cross-entropy loss, while the Probe is trained to predict the relevant text length using Huber loss.
Finetuning: The Selective Encoder is further trained on downstream tasks (e.g., QA) where the target LLM generates answers based on the compressed Key-Value (KV) representations.
Inference: The system operates in a parallel fashion. As the Selective Encoder processes the input, the AAC predicts the required token count ( $k$ ) before the first compressed token is generated, ensuring no interruption in the autoregressive flow.

Key Contributions

Task-Aware Compression: The introduction of a Selective Encoder that filters and compresses only task-relevant content, preventing the loss of critical information common in global compression methods.
Adaptive Resource Allocation: The development of an AAC that dynamically adjusts the compression rate based on the estimated length of relevant content, moving beyond fixed-token constraints.
Granularity-Aware Training: A strategy to handle varying annotation granularities (sentence vs. document level) by pre-segmenting context into uniform chunks during training, enhancing the encoder's adaptability.

Experimental Results

The authors evaluated ATACompressor on three QA benchmarks: HotpotQA (document level), MSMARCO (passage level), and SQUAD (sentence level), using LLaMA-2-7B and Qwen-2-7B as backbone models.

Performance: ATACompressor consistently outperformed state-of-the-art baselines (including Selective-Context, LongLLMLingua, AutoCompressor, ICAE, 500Compressor, and QGC) in both Exact Match (EM) and F1 scores. For instance, on HotpotQA with Qwen-2-7B, it achieved an F1 of 80.23 compared to 72.36 for the next best method (QGC).
Efficiency: The method achieved high Compression Ratios (CR) (up to ~27x) while maintaining competitive Throughput (TP). It demonstrated superior efficiency compared to methods like QGC, which rely on complex external retrievers.
Ablation Studies: Removing the AAC (fixing token count) or the Selective Encoder (compressing all context) resulted in significant performance drops, validating the necessity of both components.
Robustness: The model maintained performance stability even as the number of compressed tokens decreased, outperforming baselines in low-token regimes. It also showed strong generalization on LongBench tasks (Summary and Few-shot).

Significance and Claims

The paper claims that ATACompressor provides a scalable solution for efficient long-context processing in LLMs. By integrating task awareness into the compression process and dynamically allocating resources, it effectively balances the trade-off between information preservation and compression efficiency. The authors assert that this approach not only alleviates the "lost in the middle" problem but also reduces inference costs and latency, making it suitable for real-time and large-scale applications without requiring architectural modifications to the underlying LLM. The work is presented as a step forward in making RAG and long-context LLM applications more practical and effective.