Selection-p: Self-Supervised Task-Agnostic Prompt Compression for Faithfulness and Transferability

Tsz Ting Chung1, Leyang Cui2, Lemao Liu2, Xinting Huang2, Shuming Shi2, Dit-Yan Yeung1
1The Hong Kong University of Science and Technology
2Tencent AI Lab
arXiv
 
10×
compression rate
−0.8%
accuracy drop
5.3×
inference speedup
9
classification tasks
Compression result on Subj task

Compression by Selection-p on the Subj in-context learning demonstration under a 10× compression rate.

Abstract

Large Language Models (LLMs) have demonstrated impressive capabilities across a wide range of natural language processing tasks when leveraging in-context learning. To mitigate the additional computational and financial costs of in-context learning, several prompt compression methods have been proposed. Despite their success, these methods face challenges with transferability due to model-specific compression, or rely on external training data such as GPT-4. We investigate whether LLMs can develop a unified compression method that discretizes uninformative tokens, using a self-supervised pre-training technique. By introducing a small number of parameters during continual pre-training, Selection-p produces a probability for each input token indicating whether to preserve or discard it. Experiments show Selection-p achieves state-of-the-art performance across numerous classification tasks, reaching compression rates of up to 10× with only a marginal 0.8% drop in performance, and exhibits superior transferability to different models compared to prior work.

Why Selection-p?

Existing discrete compression methods either require multi-round iterative decoding (LLMLingua) or are distilled from a specific model such as GPT-4 (LLMLingua-2), which hurts transferability and demands costly external data. Selection-p is self-supervised, task-agnostic, and works in a single forward pass.

MethodTransferableSingle RunNo External Data
AutoCompressor✗✓✓
LLMLingua✓✗✓
LLMLingua-2✓✓✗
Selection-p✓✓✓

Method

For an input context [x1, x2, …, xn-1], a frozen LLM first performs inference without parameter updates to create an attention mask p. These form the model input for LoRA training of a small additional linear layer, which outputs a keep/discard probability per token. Only the small added parameters are learned; the base model stays frozen.

Selection-p training process

The Selection-p training process. Areas in orange are learnable parameters; the base LLM is frozen.

Contributions

  • Selection-p achieves only a 0.8% performance drop at 10× compression across nine traditional classification tasks, surpassing existing compression models, with a 5.3× inference speedup under in-context learning.
  • It demonstrates great transferability, outperforming prior work on hard compression for both open-source and closed-source models.
  • We analyze how Selection-p helps in long-context in-context learning, offering a potential answer to long-context performance degradation in ICL.
  • We connect in-domain prior work and compare against state-of-the-art compression models, providing a complete picture.
Performance across tasks

Selection-p maintains performance under aggressive compression across classification tasks.

BibTeX

@article{chung2024selection,
  title     = {Selection-p: Self-Supervised Task-Agnostic Prompt Compression for Faithfulness and Transferability},
  author    = {Chung, Tsz Ting and Cui, Leyang and Liu, Lemao and Huang, Xinting and Shi, Shuming and Yeung, Dit-Yan},
  journal   = {arXiv preprint arXiv:2410.11786},
  year      = {2024}
}