Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Xuancheng Li, Haitao Li, Yujia Zhou, Qingyao Ai, Yiqun Liu

TL;DR

ATACompressor tackles long-context processing in LLMs by jointly learning a selective encoder and an adaptive allocation controller to compress only task-relevant content. The method integrates a frozen LLM with LoRA within a trainable selective encoder and uses a lightweight probe to estimate the length of relevant content, dynamically allocating compression tokens via a policy. Empirical results on HotpotQA, MSMARCO, and SQUAD show superior QA performance, higher compression ratios, and competitive throughput compared to strong baselines, with ablations confirming the importance of both selective encoding and adaptive allocation. The approach offers a scalable, efficient solution for dynamic long-context handling in retrieval-augmented and other Long-Context scenarios, enabling better resource utilization in real-world applications.

Abstract

Long-context inputs in large language models (LLMs) often suffer from the "lost in the middle" problem, where critical information becomes diluted or ignored due to excessive length. Context compression methods aim to address this by reducing input size, but existing approaches struggle with balancing information preservation and compression efficiency. We propose Adaptive Task-Aware Compressor (ATACompressor), which dynamically adjusts compression based on the specific requirements of the task. ATACompressor employs a selective encoder that compresses only the task-relevant portions of long contexts, ensuring that essential information is preserved while reducing unnecessary content. Its adaptive allocation controller perceives the length of relevant content and adjusts the compression rate accordingly, optimizing resource utilization. We evaluate ATACompressor on three QA datasets: HotpotQA, MSMARCO, and SQUAD-showing that it outperforms existing methods in terms of both compression efficiency and task performance. Our approach provides a scalable solution for long-context processing in LLMs. Furthermore, we perform a range of ablation studies and analysis experiments to gain deeper insights into the key components of ATACompressor.

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

TL;DR

ATACompressor tackles long-context processing in LLMs by jointly learning a selective encoder and an adaptive allocation controller to compress only task-relevant content. The method integrates a frozen LLM with LoRA within a trainable selective encoder and uses a lightweight probe to estimate the length of relevant content, dynamically allocating compression tokens via a policy. Empirical results on HotpotQA, MSMARCO, and SQUAD show superior QA performance, higher compression ratios, and competitive throughput compared to strong baselines, with ablations confirming the importance of both selective encoding and adaptive allocation. The approach offers a scalable, efficient solution for dynamic long-context handling in retrieval-augmented and other Long-Context scenarios, enabling better resource utilization in real-world applications.

Abstract

Long-context inputs in large language models (LLMs) often suffer from the "lost in the middle" problem, where critical information becomes diluted or ignored due to excessive length. Context compression methods aim to address this by reducing input size, but existing approaches struggle with balancing information preservation and compression efficiency. We propose Adaptive Task-Aware Compressor (ATACompressor), which dynamically adjusts compression based on the specific requirements of the task. ATACompressor employs a selective encoder that compresses only the task-relevant portions of long contexts, ensuring that essential information is preserved while reducing unnecessary content. Its adaptive allocation controller perceives the length of relevant content and adjusts the compression rate accordingly, optimizing resource utilization. We evaluate ATACompressor on three QA datasets: HotpotQA, MSMARCO, and SQUAD-showing that it outperforms existing methods in terms of both compression efficiency and task performance. Our approach provides a scalable solution for long-context processing in LLMs. Furthermore, we perform a range of ablation studies and analysis experiments to gain deeper insights into the key components of ATACompressor.
Paper Structure (30 sections, 10 equations, 5 figures, 5 tables)

This paper contains 30 sections, 10 equations, 5 figures, 5 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Retrieval-augmented Generation
  4. Context Compression
  5. Probe for LLMs
  6. Method
  7. Problem Formulation
  8. Architecture
  9. Selective Encoder
  10. Adaptive Allocation Controller (AAC)
  11. Workflow
  12. Pretrain
  13. Finetune
  14. Inference
  15. Experiments
  16. ...and 15 more sections

Figures (5)

  • Figure 1: Comparative schematic of three approaches, using selective compressor to represent the hard prompt and 500Compressor to represent the traditional soft prompt.
  • Figure 2: The architecture and workflow of ATACompressor
  • Figure 3: Performance on pretraining (regeneration) and finetuning (QA) tasks with varying numbers of compressed tokens using the LLAMA-2-7B model on HotpotQA.
  • Figure 4: Performance on pretraining (regeneration) across varying input text lengths using the Qwen-2-7B. $k$ represents the number of compressed tokens or average.
  • Figure 5: Comparison of gold ($L_{Rel}$) and predicted lengths ($\hat{L}_{Rel}$) across three datasets on Qwen-2-7B.