Source-linked AI summary
LLMLingua-2: Data Distillation for Efficient and Faithful Task-Agnostic Prompt Compression
Zhuoshi Pan, Qianhui Wu, Huiqiang Jiang, Menglin Xia, Xufang Luo, Jue Zhang, Qingwei Lin, Victor Rühle, Yuqing Yang, Chin-Yew Lin, H. Vicky Zhao, Lili Qiu, Dongmei Zhang
TL;DR
Task-agnostic prompt compression must reduce lengthy prompts without losing essential information, but entropy-based causal-LM methods may misalign with this objective and miss bidirectional context. LLMLingua-2 distills compression knowledge from GPT-4 into an extractive dataset and trains a token-classification compressor using bidirectional representations. Across in-domain and out-of-domain benchmarks, it outperforms strong task-agnostic baselines while achieving 3x-6x lower compressor latency and 1.6x-2.9x end-to-end acceleration at 2x-5x compression ratios.
Problem
Task-agnostic entropy-based compression may be misaligned with the compression objective and limited by causal models’ unidirectional context, while lengthy prompts impose computational and financial costs.
Method
The paper distills GPT-4 knowledge into an extractive dataset and formulates compression as token classification using a bidirectional Transformer encoder.
Results
3x-6x faster than existing methods, the model achieves strong performance and robust generalization across in-domain and out-of-domain benchmarks, with 1.6x-2.9x end-to-end acceleration at 2x-5x compression ratios.
Takeaways & Limitations
The results support task-agnostic prompt compression that generalizes across tasks, domains, and LLMs while reducing compression and end-to-end latency.
Takeaways & Limitations
The compression dataset uses only MeetingBank training examples, raising concerns about the compressor’s generalization ability.
Abstract
from arXiv · showhide
This paper focuses on task-agnostic prompt compression for better generalizability and efficiency. Considering the redundancy in natural language, existing approaches compress prompts by removing tokens or lexical units according to their information entropy obtained from a causal language model such as LLaMa-7B. The challenge is that information entropy may be a suboptimal compression metric: (i) it only leverages unidirectional context and may fail to capture all essential information needed for prompt compression; (ii) it is not aligned with the prompt compression objective. To address these issues, we propose a data distillation procedure to derive knowledge from an LLM to compress prompts without losing crucial information, and meantime, introduce an extractive text compression dataset. We formulate prompt compression as a token classification problem to guarantee the faithfulness of the compressed prompt to the original one, and use a Transformer encoder as the base architecture to capture all essential information for prompt compression from the full bidirectional context. Our approach leads to lower latency by explicitly learning the compression objective with smaller models such as XLM-RoBERTa-large and mBERT. We evaluate our method on both in-domain and out-of-domain datasets, including MeetingBank, LongBench, ZeroScrolls, GSM8K, and BBH. Despite its small size, our model shows significant performance gains over strong baselines and demonstrates robust generalization ability across different LLMs. Additionally, our model is 3x-6x faster than existing prompt compression methods, while accelerating the end-to-end latency by 1.6x-2.9x with compression ratios of 2x-5x. Our code is available at https://aka.ms/LLMLingua-2.
1 Introduction
LLMLingua-2 targets task-agnostic prompt compression that reduces prompt cost and length while preserving essential information. It addresses entropy-based methods’ objective misalignment and limited context by distilling compression knowledge and using token classification with bidirectional context.
- Long prompts increase computational and financial overhead while degrading LLM information perception, motivating compression that preserves essential information.
- Task-aware compression can improve downstream performance but requires query-specific processing, limiting efficiency and generalizability.
- Entropy-based task-agnostic methods may be suboptimal because entropy is misaligned with compression objectives and causal models use only unidirectional context.
- The paper frames the open problems as constructing aligned compression data and designing an algorithm that uses full bidirectional context.
- GPT-4-derived distillation produces an extractive compression dataset from MeetingBank texts and their compressed versions, while avoiding crucial information loss.
- Token classification labels tokens as preserve or discard, using preserve probabilities as the compression metric and a Transformer encoder for bidirectional features.
- 3x-6x faster than existing methods, LLMLingua-2 accelerates end-to-end latency by 1.6x-2.9x at compression ratios of 2x-5x.
2 Related Works
Prompt compression methods are commonly divided into task-aware and task-agnostic approaches. Task-aware methods use downstream task information, whereas task-agnostic methods seek broader applicability without query-specific compression.
- Task-aware compression uses the downstream task or current query to guide context compression.
- Task-aware methods include question-aware entropy estimation and reinforcement-learning approaches trained with downstream-task rewards.
- Task-agnostic compression does not consider the specific task, aiming to adapt across applications and black-box LLMs.
- Typical task-agnostic methods remove redundant prompt information using entropy-based token-importance estimates from small language models.
3 Dataset Construction
The dataset is distilled from GPT-4 using instructions that target short, informative, faithful compression, then annotated and filtered to support extractive prompt compression. Construction addresses variable information density, long-context overcompression, annotation mismatches, and low-quality examples.
- 3.1 Data Distillation: GPT-4 is prompted to produce short, informative, faithful compressed texts without adding hallucinated content.The procedure targets token reduction, retention of essential information, and faithfulness to the original text.
- 3.1 Data Distillation: Removing fixed compression-ratio restrictions lets GPT-4 adapt compression to differing information densities across genres, styles, sentences, and speakers.The instruction instead asks GPT-4 to compress as short as possible while retaining as much information as possible.
- 3.1 Data Distillation: Long contexts are split into chunks of no more than 512 tokens ending with a period to mitigate GPT-4’s aggressive compression and resulting information loss.The paper attributes aggressive compression on very long contexts to GPT-4’s limited ability to handle long context.
- 3.2 Data Annotation: Each original token receives a binary preserve-or-discard label, with annotation handling ambiguity, variation, and reordering through matching procedures.The algorithm uses fuzzy matching, lemmatization, sliding windows, and bidirectional search to align compressed and original words.
- 3.3 Quality Control: Quality control filters distilled examples using Variation Rate for compressed-text quality and Alignment Gap for automatically annotated-label quality.Examples with the top 5% variation rates and highest 10% alignment gaps are excluded.
4 Compressor
LLMLingua-2 formulates prompt compression as bidirectional token classification, preserving high-probability tokens in their original order to create faithful shorter prompts.
- Token Classification Model: Prompt compression is modeled as binary token classification, predicting whether each token should be preserved or discarded.A Transformer encoder extracts features from bidirectional context, while preservation probabilities serve as the compression metric.
- Token Classification Model: Cross-entropy training learns token labels for the original prompt using a Transformer encoder and linear classification layer.The model outputs preserve/discard probabilities for each word.
- Compression Strategy: Given compression ratio 1/τ, the compressor preserves ˜N = τN tokens with the highest preserve probabilities and retains their original order.This produces the compressed prompt without reordering selected words.
- Compression Strategy: The compressor can replace LLMLingua’s perplexity-based iterative module while retaining its budget controller for higher compression ratios.The integrated framework supports compression ratios of approximately 15x for multiple demonstrations or documents.
5 Experiment
Experiments evaluate LLMLingua-2 on in-domain and out-of-domain tasks, comparing compact Transformer-based compressors with established task-agnostic and task-aware baselines. The method generally improves performance and latency, while showing a task-aware advantage on LongBench.
- Implementation: LLMLingua-2 uses XLM-RoBERTa-large and mBERT encoders, with GPT-3.5-Turbo-06133 as the default target LLM for downstream metrics.The two compressors are referred to as LLMLingua-2 and LLMLingua-2-small.
- Datasets and Metrics: Experiments cover MeetingBank, LongBench, ZeroSCROLLS, GSM8K, and BBH across summarization, question answering, long-context, reasoning, and in-context learning.MeetingBank provides in-domain evaluation, while the other benchmarks test out-of-domain generalization.
- In-Domain Results: On MeetingBank, the smaller compressors outperform LLaMA-2-7B-based baselines on QA and summarization and approach the original prompt’s performance.The result supports the effectiveness of the constructed compression dataset.
- Out-of-Domain Results: On out-of-domain benchmarks, LLMLingua-2 outperforms task-agnostic baselines, while the BERT-base-sized model sometimes slightly exceeds the original prompt.LongLLMLingua remains stronger on LongBench because it uses question information.
- Cross-LLM Generalization: With Mistral-7B as target LLM, LLMLingua-2 gains over baselines and can outperform the original prompt.The authors speculate that shorter, higher-information-density prompts help Mistral-7B manage long contexts.
- Latency: LLMLingua-2 achieves 1.6x–2.9x end-to-end speedup and reduces GPU memory costs by 8x on MeetingBank latency evaluations.The method has lower computation overhead than other prompt compression systems.
- Analysis: Bidirectional context awareness and explicit objective optimization help retain informative words as compression increases, while GPT-4 reconstruction tests indicate no essential information loss.Ablations also find contributions from the designed instruction and chunk-wise compression strategy.
6 Conclusion
The paper targets efficient, generalizable task-agnostic prompt compression and reports superiority over strong baselines across tasks and domains. It also releases a text-compression dataset designed to preserve essential information.
- Conclusion: The study evaluates its model across five benchmarks spanning different tasks and domains, finding advantages in performance and compression latency.The conclusion emphasizes task-agnostic prompt compression for generalizability and efficiency.
- Conclusion: The authors publicly release an extractive text-compression dataset intended to contain no essential information loss.The dataset is presented as part of the paper’s contribution.
Limitations
The compression dataset was built only from English MeetingBank training examples, raising concerns about generalization. Out-of-domain evaluations support transfer, while adding TriviaQA-wiki data yields only a modest further gain.
- The compression dataset uses only MeetingBank training examples, an English meeting-transcript summarization corpus.
- Out-of-domain evaluation covers LongBench, ZeroSCROLLS, GSM8K, and BBH, spanning document QA, mathematics, and in-context learning.
- Adding 50k TriviaQA-wiki examples produces further performance gains under a 2,000-token constraint, but the improvement is not significant.
A Details of Data Distillation
The dataset distillation process uses GPT-4-32k to compress chunked MeetingBank transcripts while controlling generation length and truncation. Dataset statistics are reported separately.
- GPT-4-32k compresses MeetingBank transcripts after each transcript is divided into complete-sentence chunks of at most 512 tokens.
- Generation uses temperature 0.3, top_p 1.0, a 4,096-token maximum, and truncates transcripts longer than 28K tokens to reserve a 4K generation budget.
- Table 8 reports the statistics of the resulting MeetingBank compression dataset.
B Details of Data Annotation
The annotation procedure maps words in GPT-4’s compressed prompt back to the original prompt, using local bidirectional matching and lemmatization to handle reordering and word variation. The resulting compressor preserves meaningful content across compression ratios and supports reconstruction of the original prompt.
- Data Annotation: Each original word initially receives a discard label, then words matched from the compressed prompt receive preserve labels.
- Data Annotation: A sliding window searches locally around the previous match, while bidirectional search reduces mismatches caused by GPT-4 reordering or inserted words.
- Compression Results: At 5x, 3x, and 2x compression ratios, LLMLingua-2 preserves words selected by context awareness and discards others.
- Data Annotation: Lemmatization addresses GPT-4 variations in tense, voice, and number before sliding-window matching.
- Prompt Reconstruction: Reconstructed MeetingBank prompts closely resemble the originals because key information remains in the compressed prompts.
G Drawback of Existing Text Compression Dataset
Existing extractive compression datasets are designed for summarization and often omit details needed for downstream question answering. LLMLingua-2 addresses broader deployment settings through lightweight multilingual encoders and integrations that improve information allocation.
- Drawback of Existing Text Compression Dataset: SentComp and DebateSum compressed texts are often too concise, retaining main ideas while losing details needed for document-based QA.
- Model Size and Efficiency: XLM-RoBERTa-large has 355M parameters, while LLMLingua-2 uses 2.1GB peak GPU memory versus 16.6GB for LLMLingua and 26.5GB for Selective-Context.
- Multilingual Generalization Ability: On Chinese LongBench benchmarks, LLMLingua-2 outperforms LLMLingua despite training only on English MeetingBank data.
- Integration with LongLLMLingua: LLMLingua-2 can integrate with LongLLMLingua to preserve question-relevant information in retrieval-augmented and multi-document QA scenarios.
- Integration with LongLLMLingua: 25.3% average performance gain on NaturalQuestions results from combining LLMLingua-2 with LongLLMLingua coarse-grained compression.
L Sample-Wise Dynamic Compression Ratio
LLMLingua-2 dynamically adjusts compression rates across samples by retaining tokens with preservation probabilities above a corpus-level threshold. This sample-wise strategy improves performance over fixed-rate compression and remains robust across compression ratios and evaluation settings.
- Dynamic Compression Strategy: LLMLingua-2 predicts each token’s preservation probability and retains tokens above a threshold under an overall compression-rate constraint.This allows different examples to receive different compression rates instead of applying one fixed rate to every sample.
- Dynamic Compression Results: 4.4% and 4.5% performance improvements are achieved at 7x and 5x compression ratios, respectively, compared to fixed compression.The gains are reported for sample-wise dynamic compression on LongBench single-document QA.
- Compression-Ratio Robustness: LLMLingua-2 shows superior robustness to increasing compression ratios on a 100-sample MeetingBank subset.Figure 15 compares performance against other baselines as the compression ratio increases.
- Comparison With GPT-4 Compression: LLMLingua-2 achieves higher MeetingBank QA performance than GPT-4 compression, while GPT-4 prioritizes nouns, adjectives, and numerals.The paper attributes LLMLingua-2’s advantage to learning compression knowledge across the dataset, mitigating noise and information loss in individual GPT-4 examples.
- Evaluation Under Context Constraints: LLMLingua-2 continues to outperform strong baselines and original prompts on subsets whose inputs are shorter than 8k tokens for Mistral-7B.The evaluation discards samples exceeding Mistral-7B’s 8k context length.